Source code for squid.models.classical_utils.to_range

from math import pi
from typing import Tuple, Union

import torch
from torch import nn


class ToRange(nn.Module):
    """ToRange docs

    .. todo:: TODO: Finish ToRange docs

    :param input_range: [description], defaults to (0, torch.Tensor([pi]))
    :type input_range: Tuple[ Union[torch.Tensor, int, float], Union[torch.Tensor, int, float] ], optional
    """

    def __init__(
        self,
        input_range: Tuple[
            Union[torch.Tensor, int, float], Union[torch.Tensor, int, float]
        ] = (0, torch.Tensor([pi])),
    ):
        super().__init__()
        self.start = input_range[0]
        self.delta = input_range[1] - input_range[0]
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.sigmoid(x)
        x *= self.delta
        x += self.start
        return x