Skip to content

Uniform

SymmetricUniformTimeDistribution

Bases: TimeDistribution

A class representing a uniform time distribution.

Source code in bionemo/moco/distributions/time/uniform.py 
 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
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: Union[int, Tuple[int, ...], torch.Size],
        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 not isinstance(n_samples, int):
            n_samples = n_samples[0]
        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 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
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 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def sample(
    self,
    n_samples: Union[int, Tuple[int, ...], torch.Size],
    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 not isinstance(n_samples, int):
        n_samples = n_samples[0]
    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
71
72
73
74
75
76
77
78
79
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: Union[int, Tuple[int, ...], torch.Size],
        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,) if isinstance(n_samples, int) else 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
71
72
73
74
75
76
77
78
79
def sample(
    self,
    n_samples: Union[int, Tuple[int, ...], torch.Size],
    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,) if isinstance(n_samples, int) else 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