Skip to content

Custom

DiscreteCustomPrior

Bases: DiscretePriorDistribution

A subclass representing a discrete custom prior distribution.

This class allows for the creation of a prior distribution with a custom probability mass function defined by the prior_dist tensor. For example if my data has 4 classes and I want [.3, .2, .4, .1] as the probabilities of the 4 classes.

Source code in bionemo/moco/distributions/prior/discrete/custom.py
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
class DiscreteCustomPrior(DiscretePriorDistribution):
    """A subclass representing a discrete custom prior distribution.

    This class allows for the creation of a prior distribution with a custom
    probability mass function defined by the `prior_dist` tensor. For example if my data has 4 classes and I want [.3, .2, .4, .1] as the probabilities of the 4 classes.
    """

    def __init__(self, prior_dist: Tensor, num_classes: int = 10) -> None:
        """Initializes a DiscreteCustomPrior distribution.

        Args:
            prior_dist: A tensor representing the probability mass function of the prior distribution.
            num_classes: The number of classes in the prior distribution. Defaults to 10.

        Note:
            The `prior_dist` tensor should have a sum close to 1.0, as it represents a probability mass function.
        """
        super().__init__(num_classes, prior_dist)
        if torch.sum(self.prior_dist).item() - 1.0 > 1e-5:
            raise ValueError("Prior distribution probabilities do not sum up to 1.0")

    def sample(
        self,
        shape: Tuple,
        mask: Optional[Tensor] = None,
        device: Union[str, torch.device] = "cpu",
        rng_generator: Optional[torch.Generator] = None,
    ) -> Tensor:
        """Samples from the discrete custom prior distribution.

        Args:
            shape: A tuple specifying the shape of the samples to generate.
            mask: An optional tensor mask to apply to the samples, broadcastable to the sample shape. Defaults to None.
            device: The device on which to generate the samples, specified as a string or a :class:`torch.device`. Defaults to "cpu".
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            A tensor of samples drawn from the prior distribution.
        """
        samples = (
            torch.multinomial(self.prior_dist, math.prod(shape), replacement=True, generator=rng_generator)
            .to(device)
            .reshape(shape)
        )
        if mask is not None:
            samples = samples * mask[(...,) + (None,) * (len(samples.shape) - len(mask.shape))]
        return samples

__init__(prior_dist, num_classes=10)

Initializes a DiscreteCustomPrior distribution.

Parameters:

Name Type Description Default
prior_dist Tensor

A tensor representing the probability mass function of the prior distribution.

required
num_classes int

The number of classes in the prior distribution. Defaults to 10.

10
Note

The prior_dist tensor should have a sum close to 1.0, as it represents a probability mass function.

Source code in bionemo/moco/distributions/prior/discrete/custom.py
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(self, prior_dist: Tensor, num_classes: int = 10) -> None:
    """Initializes a DiscreteCustomPrior distribution.

    Args:
        prior_dist: A tensor representing the probability mass function of the prior distribution.
        num_classes: The number of classes in the prior distribution. Defaults to 10.

    Note:
        The `prior_dist` tensor should have a sum close to 1.0, as it represents a probability mass function.
    """
    super().__init__(num_classes, prior_dist)
    if torch.sum(self.prior_dist).item() - 1.0 > 1e-5:
        raise ValueError("Prior distribution probabilities do not sum up to 1.0")

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

Samples from the discrete custom prior distribution.

Parameters:

Name Type Description Default
shape Tuple

A tuple specifying the shape of the samples to generate.

required
mask Optional[Tensor]

An optional tensor mask to apply to the samples, broadcastable to the sample shape. Defaults to None.

None
device Union[str, device]

The device on which to generate the samples, specified as a string or a :class:torch.device. Defaults to "cpu".

'cpu'
rng_generator Optional[Generator]

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

None

Returns:

Type Description
Tensor

A tensor of samples drawn from the prior distribution.

Source code in bionemo/moco/distributions/prior/discrete/custom.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
def sample(
    self,
    shape: Tuple,
    mask: Optional[Tensor] = None,
    device: Union[str, torch.device] = "cpu",
    rng_generator: Optional[torch.Generator] = None,
) -> Tensor:
    """Samples from the discrete custom prior distribution.

    Args:
        shape: A tuple specifying the shape of the samples to generate.
        mask: An optional tensor mask to apply to the samples, broadcastable to the sample shape. Defaults to None.
        device: The device on which to generate the samples, specified as a string or a :class:`torch.device`. Defaults to "cpu".
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        A tensor of samples drawn from the prior distribution.
    """
    samples = (
        torch.multinomial(self.prior_dist, math.prod(shape), replacement=True, generator=rng_generator)
        .to(device)
        .reshape(shape)
    )
    if mask is not None:
        samples = samples * mask[(...,) + (None,) * (len(samples.shape) - len(mask.shape))]
    return samples