_distributions.py module

class _distributions.NegativeBinomial(total_count: Optional[Tensor] = None, probs: Optional[Tensor] = None, logits: Optional[Tensor] = None, mu: Optional[Tensor] = None, theta: Optional[Tensor] = None, scale: Optional[Tensor] = None, validate_args: bool = False)

Bases: Distribution

Negative binomial distribution.

One of the following parameterizations must be provided:

  1. (total_count, probs) where: - total_count is the number of failures until the experiment is stopped - probs is the success probability

  2. (mu, theta) parameterization, which is the one used by scvi-tools, where: - mu controls the mean of the distribution - theta controls the inverse dispersion of the distribution

In the (mu, theta) parameterization, samples from the negative binomial are generated as follows:

\[ \begin{align}\begin{aligned}w \sim \textrm{Gamma}\left(\theta, \frac{\theta}{\mu}\right)\\x \sim \textrm{Poisson}(w)\end{aligned}\end{align} \]
Parameters:

  • total_count (torch.Tensor, optional) -- Number of failures until the experiment is stopped.

  • probs (torch.Tensor, optional) -- The success probability.

  • mu (torch.Tensor, optional) -- Mean of the distribution.

  • theta (torch.Tensor, optional) -- Inverse dispersion.

  • scale (torch.Tensor, optional) -- Normalized mean expression of the distribution.

  • validate_args (bool, optional) -- Raise ValueError if arguments do not match constraints.

arg_constraints

Dictionary of argument constraints for the distribution parameters:

  • mu must be greater than or equal to 0

  • theta must be greater than or equal to 0

  • scale must be greater than or equal to 0

support

Constraints on the support of the distribution (non-negative integers).

arg_constraints = {'mu': GreaterThanEq(lower_bound=0), 'scale': GreaterThanEq(lower_bound=0), 'theta': GreaterThanEq(lower_bound=0)}
log_prob(value: Tensor) → Tensor

Calculate the log probability of a given value under the Negative Binomial distribution.

This method computes the log likelihood of observing a given count value under the specified Negative Binomial distribution parameters.

Args:

value (torch.Tensor): The observed count values for which the log probability is calculated.

Returns:

torch.Tensor: The log probability of the observed counts.

Raises:

UserWarning: If the provided values are outside the support of the distribution.

mean()

Calculate the mean of the Negative Binomial distribution.

This method computes the expected value (mean) of the distribution, which is defined as mu for the Negative Binomial distribution.

Returns:

torch.Tensor: The mean of the distribution.

sample(sample_shape: Optional[Union[Size, Tuple]] = None) → Tensor

Generate samples from the Negative Binomial distribution.

This method samples from the distribution using the gamma-Poisson mixture: a gamma distribution to generate rate parameters, and a Poisson distribution to generate the final count data.

Args:
sample_shape (Union[torch.Size, Tuple], optional): The shape of the tensor to generate.

Defaults to torch.Size(), which generates a single sample.

Returns:

torch.Tensor: Samples from the Negative Binomial distribution.

support = IntegerGreaterThan(lower_bound=0)
variance()

Calculate the variance of the Negative Binomial distribution.

This method computes the variance of the distribution, which is derived from the formula: mean + (mean^2) / theta.

Returns:

torch.Tensor: The variance of the distribution.

class _distributions.Poisson(rate: Tensor, validate_args: Optional[bool] = None, scale: Optional[Tensor] = None)

Bases: Poisson

Poisson distribution.

The Poisson distribution is a discrete probability distribution that expresses the probability of a given number of events occurring in a fixed interval of time or space, if these events occur with a known constant mean rate and independently of the time since the last event.

The Poisson distribution is characterized by a single parameter \(\lambda\) that represents the average rate of occurrence. It is commonly used to model the number of times a random event occurs in a given amount of time, distance, area, etc.

Parameters:

  • rate (torch.Tensor) -- Rate parameter \(\lambda\) of the Poisson distribution.

  • validate_args (bool, optional) -- Whether to validate input arguments.

  • scale (torch.Tensor, optional) -- Normalized mean expression of the distribution.

class _distributions.ZeroInflatedNegativeBinomial(total_count: Optional[Tensor] = None, probs: Optional[Tensor] = None, logits: Optional[Tensor] = None, mu: Optional[Tensor] = None, theta: Optional[Tensor] = None, zi_logits: Optional[Tensor] = None, scale: Optional[Tensor] = None, validate_args: bool = False)

Bases: NegativeBinomial

Zero-inflated negative binomial distribution.

One of the following parameterizations must be provided:

(1), (total_count, probs) where total_count is the number of failures until the experiment is stopped and probs the success probability. (2), (mu, theta) parameterization, which is the one used by scvi-tools. These parameters respectively control the mean and inverse dispersion of the distribution.

In the (mu, theta) parameterization, samples from the negative binomial are generated as follows:

\[ \begin{align}\begin{aligned}w \sim \textrm{Gamma}\left(\theta, \frac{\theta}{\mu}\right)\\x \sim \textrm{Poisson}(w)\end{aligned}\end{align} \]

where \(\theta\) is the shape parameter and \(\frac{\theta}{\mu}\) is the rate parameter of the Gamma distribution.

Parameters:

  • total_count (torch.Tensor, optional) -- Number of failures until the experiment is stopped.

  • probs (torch.Tensor, optional) -- The success probability.

  • mu (torch.Tensor, optional) -- Mean of the distribution.

  • theta (torch.Tensor, optional) -- Inverse dispersion.

  • zi_logits (torch.Tensor, optional) -- Logits scale of zero inflation probability.

  • scale (torch.Tensor, optional) -- Normalized mean expression of the distribution.

  • validate_args (bool, optional) -- Raise ValueError if arguments do not match constraints.

arg_constraints = {'mu': GreaterThanEq(lower_bound=0), 'scale': GreaterThanEq(lower_bound=0), 'theta': GreaterThanEq(lower_bound=0), 'zi_logits': Real()}
log_prob(value: Tensor) → Tensor

Calculate the log probability of a given value under the Zero-Inflated Negative Binomial distribution.

Args:

value (torch.Tensor): The observed count values for which the log probability is calculated.

Returns:

torch.Tensor: The log probability of the observed counts.

Raises:

UserWarning: If the provided values are outside the support of the distribution.

mean()

Calculate the mean of the Zero-Inflated Negative Binomial distribution.

Returns:

torch.Tensor: The expected mean of the distribution.

sample(sample_shape: Optional[Union[Size, Tuple]] = None) → Tensor

Generate samples from the Zero-Inflated Negative Binomial distribution.

Args:
sample_shape (Union[torch.Size, Tuple], optional): The shape of the tensor to generate.

Defaults to an empty tuple, which generates a single sample.

Returns:

torch.Tensor: Samples from the Zero-Inflated Negative Binomial distribution.

support = IntegerGreaterThan(lower_bound=0)
variance()

Calculate the variance of the Zero-Inflated Negative Binomial distribution.

Raises:

NotImplementedError: Indicates that the method is not implemented.

zi_logits() → Tensor

ZI logits.

zi_probs() → Tensor
_distributions.log_nb_positive(x: ~torch.Tensor, mu: ~torch.Tensor, theta: ~torch.Tensor, eps: float = 1e-08, log_fn: callable = <built-in method log of type object>, lgamma_fn: callable = <built-in method lgamma of type object>)

Log likelihood (scalar) of a minibatch according to a nb model.

Parameters:

  • x (torch.Tensor) -- The input data tensor.

  • mu (torch.Tensor) -- The mean parameter of the negative binomial distribution (must be positive). Shape: (minibatch, vars)

  • theta (torch.Tensor) -- The inverse dispersion parameter of the negative binomial distribution (must be positive). Shape: (minibatch, vars)

  • eps (float, optional) -- A small constant for numerical stability (default: 1e-8).

  • log_fn (callable, optional) -- The logarithm function to use (default: torch.log).

  • lgamma_fn (callable, optional) -- The log-gamma function to use (default: torch.lgamma).

Returns:

The log likelihood of the input data under the specified negative binomial distribution.

Return type:

torch.Tensor

Returns

torch.Tensor

The log likelihood of the input data under the specified negative binomial distribution.

_distributions.log_zinb_positive(x: Tensor, mu: Tensor, theta: Tensor, pi: Tensor, eps=1e-08)

Log likelihood (scalar) of a minibatch according to a zinb model.

Parameters:

  • x (torch.Tensor) -- The observed count data.

  • mu (torch.Tensor) -- The mean parameter of the negative binomial component (must be positive). Shape: (minibatch, vars)

  • theta (torch.Tensor) -- The inverse dispersion parameter of the negative binomial component (must be positive). Shape: (minibatch, vars)

  • pi (torch.Tensor) -- The logit of the zero-inflation probability parameter (has real support). Shape: (minibatch, vars)

  • eps (float) -- A small constant for numerical stability (default: 1e-8).

Returns:

The log-likelihood of the input data under the ZINB model.

Return type:

torch.Tensor

Note

  • We parametrize the Bernoulli component using logits, hence the appearance of the softplus function.