Skip to content

Distribution

MixTimeDistribution

An abstract base class representing a mixed time distribution.

uniform_dist = UniformTimeDistribution(min_t=0.0, max_t=1.0, discrete_time=False) beta_dist = BetaTimeDistribution(min_t=0.0, max_t=1.0, discrete_time=False, p1=2.0, p2=1.0) mix_dist = MixTimeDistribution(uniform_dist, beta_dist, mix_fraction=0.5)

Source code in bionemo/moco/distributions/time/distribution.py 
 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
class MixTimeDistribution:
    """An abstract base class representing a mixed time distribution.

    uniform_dist = UniformTimeDistribution(min_t=0.0, max_t=1.0, discrete_time=False)
    beta_dist = BetaTimeDistribution(min_t=0.0, max_t=1.0, discrete_time=False, p1=2.0, p2=1.0)
    mix_dist = MixTimeDistribution(uniform_dist, beta_dist, mix_fraction=0.5)
    """

    def __init__(self, dist1: TimeDistribution, dist2: TimeDistribution, mix_fraction: Float):
        """Initializes a MixTimeDistribution object.

        Args:
            dist1 (TimeDistribution): The first time distribution.
            dist2 (TimeDistribution): The second time distribution.
            mix_fraction (Float): The fraction of samples to draw from dist1. Must be between 0 and 1.
        """
        if not 0 <= mix_fraction <= 1:
            raise ValueError("mix_fraction must be between 0 and 1")
        self.dist1 = dist1
        self.dist2 = dist2
        self.mix_fraction = mix_fraction

    def sample(
        self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
    ) -> Float:
        """Generates a specified number of samples from the mixed 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:
            Float: A list or array of samples.
        """
        samples_dist1 = self.dist1.sample(n_samples, device)
        samples_dist2 = self.dist2.sample(n_samples, device)
        mix = torch.rand(n_samples, device=device, generator=rng_generator)
        return torch.where(mix < self.mix_fraction, samples_dist1, samples_dist2)

__init__(dist1, dist2, mix_fraction)

Initializes a MixTimeDistribution object.

Parameters:

Name Type Description Default
dist1 TimeDistribution

The first time distribution.

 required
dist2 TimeDistribution

The second time distribution.

 required
mix_fraction Float

The fraction of samples to draw from dist1. Must be between 0 and 1.

 required
Source code in bionemo/moco/distributions/time/distribution.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def __init__(self, dist1: TimeDistribution, dist2: TimeDistribution, mix_fraction: Float):
    """Initializes a MixTimeDistribution object.

    Args:
        dist1 (TimeDistribution): The first time distribution.
        dist2 (TimeDistribution): The second time distribution.
        mix_fraction (Float): The fraction of samples to draw from dist1. Must be between 0 and 1.
    """
    if not 0 <= mix_fraction <= 1:
        raise ValueError("mix_fraction must be between 0 and 1")
    self.dist1 = dist1
    self.dist2 = dist2
    self.mix_fraction = mix_fraction

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

Generates a specified number of samples from the mixed 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:

Name Type Description
Float Float

A list or array of samples.
Source code in bionemo/moco/distributions/time/distribution.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def sample(
    self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
) -> Float:
    """Generates a specified number of samples from the mixed 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:
        Float: A list or array of samples.
    """
    samples_dist1 = self.dist1.sample(n_samples, device)
    samples_dist2 = self.dist2.sample(n_samples, device)
    mix = torch.rand(n_samples, device=device, generator=rng_generator)
    return torch.where(mix < self.mix_fraction, samples_dist1, samples_dist2)

TimeDistribution

Bases: ABC

An abstract base class representing a time distribution.

Parameters:

Name Type Description Default
discrete_time Bool

Whether the time is discrete.

 False
nsteps Optional[int]

Number of nsteps for discretization.

 None
min_t Optional[Float]

Min continuous time.

 None
max_t Optional[Float]

Max continuous time.

 None
rng_generator Optional[Generator]

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

 None
Source code in bionemo/moco/distributions/time/distribution.py 
24
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
80
81
82
83
class TimeDistribution(ABC):
    """An abstract base class representing a time distribution.

    Args:
        discrete_time (Bool): Whether the time is discrete.
        nsteps (Optional[int]): Number of nsteps for discretization.
        min_t (Optional[Float]): Min continuous time.
        max_t (Optional[Float]): Max continuous time.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """

    def __init__(
        self,
        discrete_time: Bool = False,
        nsteps: Optional[int] = None,
        min_t: Optional[Float] = None,
        max_t: Optional[Float] = None,
        rng_generator: Optional[torch.Generator] = None,
    ):
        """Initializes a TimeDistribution object."""
        self.discrete_time = discrete_time
        self.nsteps = nsteps
        self.rng_generator = rng_generator
        if discrete_time:
            min_t = 0.0
            max_t = 1.0
            if nsteps is None:
                raise ValueError("nsteps must not be None and must be specified for discrete time")
        if min_t is not None and isinstance(min_t, float):
            if not 0 <= min_t < 1.0:
                raise ValueError("min_t must be greater than or equal to 0 and less than 1.0")
        self.min_t = min_t
        if max_t is not None and isinstance(max_t, float):
            if not 0 < max_t <= 1.0:
                raise ValueError("max_t must be greater than 0 and less than or equal to 1.0")
        self.max_t = max_t
        if (
            self.min_t is not None
            and self.max_t is not None
            and isinstance(self.min_t, float)
            and isinstance(self.max_t, float)
        ):
            if self.min_t >= self.max_t:
                raise ValueError("min_t must be less than max_t")

    @abstractmethod
    def sample(
        self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
    ) -> Float:
        """Generates a specified number of samples from the 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:
            Float: A list or array of samples.
        """
        pass

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

Initializes a TimeDistribution object.

Source code in bionemo/moco/distributions/time/distribution.py 
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
def __init__(
    self,
    discrete_time: Bool = False,
    nsteps: Optional[int] = None,
    min_t: Optional[Float] = None,
    max_t: Optional[Float] = None,
    rng_generator: Optional[torch.Generator] = None,
):
    """Initializes a TimeDistribution object."""
    self.discrete_time = discrete_time
    self.nsteps = nsteps
    self.rng_generator = rng_generator
    if discrete_time:
        min_t = 0.0
        max_t = 1.0
        if nsteps is None:
            raise ValueError("nsteps must not be None and must be specified for discrete time")
    if min_t is not None and isinstance(min_t, float):
        if not 0 <= min_t < 1.0:
            raise ValueError("min_t must be greater than or equal to 0 and less than 1.0")
    self.min_t = min_t
    if max_t is not None and isinstance(max_t, float):
        if not 0 < max_t <= 1.0:
            raise ValueError("max_t must be greater than 0 and less than or equal to 1.0")
    self.max_t = max_t
    if (
        self.min_t is not None
        and self.max_t is not None
        and isinstance(self.min_t, float)
        and isinstance(self.max_t, float)
    ):
        if self.min_t >= self.max_t:
            raise ValueError("min_t must be less than max_t")

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

Generates a specified number of samples from the 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:

Name Type Description
Float Float

A list or array of samples.
Source code in bionemo/moco/distributions/time/distribution.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@abstractmethod
def sample(
    self, n_samples: int, device: Union[str, torch.device] = "cpu", rng_generator: Optional[torch.Generator] = None
) -> Float:
    """Generates a specified number of samples from the 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:
        Float: A list or array of samples.
    """
    pass