Skip to content

Uniform

SymmetricUniformTimeDistribution

Bases: TimeDistribution

A class representing a uniform time distribution.

Source code in bionemo/moco/distributions/time/uniform.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class SymmetricUniformTimeDistribution(TimeDistribution):
    """A class representing a uniform time distribution."""

    def __init__(
        self,
        min_t: Float = 0.0,
        max_t: Float = 1.0,
        discrete_time: Bool = False,
        nsteps: Optional[int] = None,
        rng_generator: Optional[torch.Generator] = None,
    ):
        """Initializes a UniformTimeDistribution object.

        Args:
            min_t (Float): The minimum time value.
            max_t (Float): The maximum time value.
            discrete_time (Bool): Whether the time is discrete.
            nsteps (Optional[int]): Number of nsteps for discretization.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        """
        super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)

    def sample(
        self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
    ):
        """Generates a specified number of samples from the uniform time distribution.

        Args:
            n_samples (int): The number of samples to generate.
            device (str): cpu or gpu.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            A tensor of samples.
        """
        if rng_generator is None:
            rng_generator = self.rng_generator
        if self.discrete_time:
            if self.nsteps is None:
                raise ValueError("nsteps cannot be None for discrete time sampling")
            time_step = torch.randint(
                0, self.nsteps, size=(n_samples // 2 + 1,), device=device, generator=rng_generator
            )
            time_step = torch.cat([time_step, self.nsteps - time_step - 1], dim=0)[:n_samples]
        else:
            time_step = torch.rand(n_samples // 2 + 1, device=device, generator=rng_generator)
            time_step = torch.cat([time_step, 1 - time_step], dim=0)[:n_samples]
            if self.min_t and self.max_t and self.min_t > 0:
                time_step = time_step * (self.max_t - self.min_t) + self.min_t
        return time_step

__init__(min_t=0.0, max_t=1.0, discrete_time=False, nsteps=None, rng_generator=None)

Initializes a UniformTimeDistribution object.

Parameters:

Name Type Description Default
min_t Float

The minimum time value.

 0.0
max_t Float

The maximum time value.

 1.0
discrete_time Bool

Whether the time is discrete.

 False
nsteps Optional[int]

Number of nsteps for discretization.

 None
rng_generator Optional[Generator]

An optional :class:torch.Generator for reproducible sampling. Defaults to None.

 None
Source code in bionemo/moco/distributions/time/uniform.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def __init__(
    self,
    min_t: Float = 0.0,
    max_t: Float = 1.0,
    discrete_time: Bool = False,
    nsteps: Optional[int] = None,
    rng_generator: Optional[torch.Generator] = None,
):
    """Initializes a UniformTimeDistribution object.

    Args:
        min_t (Float): The minimum time value.
        max_t (Float): The maximum time value.
        discrete_time (Bool): Whether the time is discrete.
        nsteps (Optional[int]): Number of nsteps for discretization.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """
    super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)

sample(n_samples, device='cpu', rng_generator=None)

Generates a specified number of samples from the uniform time distribution.

Parameters:

Name Type Description Default
n_samples int

The number of samples to generate.

 required
device str

cpu or gpu.

 'cpu'
rng_generator Optional[Generator]

An optional :class:torch.Generator for reproducible sampling. Defaults to None.

 None

Returns:

Type Description

A tensor of samples.
Source code in bionemo/moco/distributions/time/uniform.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def sample(
    self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
):
    """Generates a specified number of samples from the uniform time distribution.

    Args:
        n_samples (int): The number of samples to generate.
        device (str): cpu or gpu.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        A tensor of samples.
    """
    if rng_generator is None:
        rng_generator = self.rng_generator
    if self.discrete_time:
        if self.nsteps is None:
            raise ValueError("nsteps cannot be None for discrete time sampling")
        time_step = torch.randint(
            0, self.nsteps, size=(n_samples // 2 + 1,), device=device, generator=rng_generator
        )
        time_step = torch.cat([time_step, self.nsteps - time_step - 1], dim=0)[:n_samples]
    else:
        time_step = torch.rand(n_samples // 2 + 1, device=device, generator=rng_generator)
        time_step = torch.cat([time_step, 1 - time_step], dim=0)[:n_samples]
        if self.min_t and self.max_t and self.min_t > 0:
            time_step = time_step * (self.max_t - self.min_t) + self.min_t
    return time_step

UniformTimeDistribution

Bases: TimeDistribution

A class representing a uniform time distribution.

Source code in bionemo/moco/distributions/time/uniform.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class UniformTimeDistribution(TimeDistribution):
    """A class representing a uniform time distribution."""

    def __init__(
        self,
        min_t: Float = 0.0,
        max_t: Float = 1.0,
        discrete_time: Bool = False,
        nsteps: Optional[int] = None,
        rng_generator: Optional[torch.Generator] = None,
    ):
        """Initializes a UniformTimeDistribution object.

        Args:
            min_t (Float): The minimum time value.
            max_t (Float): The maximum time value.
            discrete_time (Bool): Whether the time is discrete.
            nsteps (Optional[int]): Number of nsteps for discretization.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        """
        super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)

    def sample(
        self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
    ):
        """Generates a specified number of samples from the uniform time distribution.

        Args:
            n_samples (int): The number of samples to generate.
            device (str): cpu or gpu.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            A tensor of samples.
        """
        if rng_generator is None:
            rng_generator = self.rng_generator
        if self.discrete_time:
            if self.nsteps is None:
                raise ValueError("nsteps cannot be None for discrete time sampling")
            time_step = torch.randint(0, self.nsteps, size=(n_samples,), device=device, generator=rng_generator)
        else:
            time_step = torch.rand(n_samples, device=device, generator=rng_generator)
            if self.min_t and self.max_t and self.min_t > 0:
                time_step = time_step * (self.max_t - self.min_t) + self.min_t
        return time_step

__init__(min_t=0.0, max_t=1.0, discrete_time=False, nsteps=None, rng_generator=None)

Initializes a UniformTimeDistribution object.

Parameters:

Name Type Description Default
min_t Float

The minimum time value.

 0.0
max_t Float

The maximum time value.

 1.0
discrete_time Bool

Whether the time is discrete.

 False
nsteps Optional[int]

Number of nsteps for discretization.

 None
rng_generator Optional[Generator]

An optional :class:torch.Generator for reproducible sampling. Defaults to None.

 None
Source code in bionemo/moco/distributions/time/uniform.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    min_t: Float = 0.0,
    max_t: Float = 1.0,
    discrete_time: Bool = False,
    nsteps: Optional[int] = None,
    rng_generator: Optional[torch.Generator] = None,
):
    """Initializes a UniformTimeDistribution object.

    Args:
        min_t (Float): The minimum time value.
        max_t (Float): The maximum time value.
        discrete_time (Bool): Whether the time is discrete.
        nsteps (Optional[int]): Number of nsteps for discretization.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """
    super().__init__(discrete_time, nsteps, min_t, max_t, rng_generator)

sample(n_samples, device='cpu', rng_generator=None)

Generates a specified number of samples from the uniform time distribution.

Parameters:

Name Type Description Default
n_samples int

The number of samples to generate.

 required
device str

cpu or gpu.

 'cpu'
rng_generator Optional[Generator]

An optional :class:torch.Generator for reproducible sampling. Defaults to None.

 None

Returns:

Type Description

A tensor of samples.
Source code in bionemo/moco/distributions/time/uniform.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def sample(
    self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
):
    """Generates a specified number of samples from the uniform time distribution.

    Args:
        n_samples (int): The number of samples to generate.
        device (str): cpu or gpu.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        A tensor of samples.
    """
    if rng_generator is None:
        rng_generator = self.rng_generator
    if self.discrete_time:
        if self.nsteps is None:
            raise ValueError("nsteps cannot be None for discrete time sampling")
        time_step = torch.randint(0, self.nsteps, size=(n_samples,), device=device, generator=rng_generator)
    else:
        time_step = torch.rand(n_samples, device=device, generator=rng_generator)
        if self.min_t and self.max_t and self.min_t > 0:
            time_step = time_step * (self.max_t - self.min_t) + self.min_t
    return time_step