Skip to content

Gaussian

GaussianPrior

Bases: PriorDistribution

A subclass representing a Gaussian prior distribution.

Source code in bionemo/moco/distributions/prior/continuous/gaussian.py 
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
class GaussianPrior(PriorDistribution):
    """A subclass representing a Gaussian prior distribution."""

    def __init__(
        self,
        mean: Float = 0.0,
        std: Float = 1.0,
        center: Bool = False,
        rng_generator: Optional[torch.Generator] = None,
    ) -> None:
        """Gaussian prior distribution.

        Args:
            mean (Float): The mean of the Gaussian distribution. Defaults to 0.0.
            std (Float): The standard deviation of the Gaussian distribution. Defaults to 1.0.
            center (bool): Whether to center the samples around the mean. Defaults to False.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        """
        self.mean = mean
        self.std = std
        self.center = center
        self.rng_generator = rng_generator

    def sample(
        self,
        shape: Tuple,
        mask: Optional[Tensor] = None,
        device: Union[str, torch.device] = "cpu",
        rng_generator: Optional[torch.Generator] = None,
    ) -> Tensor:
        """Generates a specified number of samples from the Gaussian prior distribution.

        Args:
            shape (Tuple): The shape of the samples to generate.
            device (str): cpu or gpu.
            mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            Float: A tensor of samples.
        """
        if rng_generator is None:
            rng_generator = self.rng_generator
        samples = torch.randn(*shape, device=device, generator=rng_generator)
        if self.std != 1:
            samples = samples * self.std
        if self.mean != 0:
            samples = samples + self.mean

        if self.center:
            samples = remove_center_of_mass(samples, mask)
        if mask is not None:
            samples = samples * mask.unsqueeze(-1)
        return samples

__init__(mean=0.0, std=1.0, center=False, rng_generator=None)

Gaussian prior distribution.

Parameters:

Name Type Description Default
mean Float

The mean of the Gaussian distribution. Defaults to 0.0.

 0.0
std Float

The standard deviation of the Gaussian distribution. Defaults to 1.0.

 1.0
center bool

Whether to center the samples around the mean. Defaults to False.

 False
rng_generator Optional[Generator]

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

 None
Source code in bionemo/moco/distributions/prior/continuous/gaussian.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def __init__(
    self,
    mean: Float = 0.0,
    std: Float = 1.0,
    center: Bool = False,
    rng_generator: Optional[torch.Generator] = None,
) -> None:
    """Gaussian prior distribution.

    Args:
        mean (Float): The mean of the Gaussian distribution. Defaults to 0.0.
        std (Float): The standard deviation of the Gaussian distribution. Defaults to 1.0.
        center (bool): Whether to center the samples around the mean. Defaults to False.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """
    self.mean = mean
    self.std = std
    self.center = center
    self.rng_generator = rng_generator

sample(shape, mask=None, device='cpu', rng_generator=None)

Generates a specified number of samples from the Gaussian prior distribution.

Parameters:

Name Type Description Default
shape Tuple

The shape of the samples to generate.

 required
device str

cpu or gpu.

 'cpu'
mask Optional[Tensor]

An optional mask to apply to the samples. Defaults to None.

 None
rng_generator Optional[Generator]

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

 None

Returns:

Name Type Description
Float Tensor

A tensor of samples.
Source code in bionemo/moco/distributions/prior/continuous/gaussian.py 
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
def sample(
    self,
    shape: Tuple,
    mask: Optional[Tensor] = None,
    device: Union[str, torch.device] = "cpu",
    rng_generator: Optional[torch.Generator] = None,
) -> Tensor:
    """Generates a specified number of samples from the Gaussian prior distribution.

    Args:
        shape (Tuple): The shape of the samples to generate.
        device (str): cpu or gpu.
        mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        Float: A tensor of samples.
    """
    if rng_generator is None:
        rng_generator = self.rng_generator
    samples = torch.randn(*shape, device=device, generator=rng_generator)
    if self.std != 1:
        samples = samples * self.std
    if self.mean != 0:
        samples = samples + self.mean

    if self.center:
        samples = remove_center_of_mass(samples, mask)
    if mask is not None:
        samples = samples * mask.unsqueeze(-1)
    return samples