class ZeroInflatedNegativeBinomial(ZeroInflatedDistribution):
"""
A Zero Inflated Negative Binomial distribution.
:param total_count: non-negative number of negative Bernoulli trials.
:type total_count: float or torch.Tensor
:param torch.Tensor probs: Event probabilities of success in the half open interval [0, 1).
:param torch.Tensor logits: Event log-odds for probabilities of success.
:param torch.Tensor gate: probability of extra zeros.
:param torch.Tensor gate_logits: logits of extra zeros.
"""
arg_constraints = {
"total_count": constraints.greater_than_eq(0),
"probs": constraints.half_open_interval(0.0, 1.0),
"logits": constraints.real,
"gate": constraints.unit_interval,
"gate_logits": constraints.real,
}
support = constraints.nonnegative_integer
def __init__(self,
total_count,
*,
probs=None,
logits=None,
gate=None,
gate_logits=None,
validate_args=None):
base_dist = NegativeBinomial(
total_count=total_count,
probs=probs,
logits=logits,
validate_args=False,
)
base_dist._validate_args = validate_args
super().__init__(base_dist,
gate=gate,
gate_logits=gate_logits,
validate_args=validate_args)
@property
def total_count(self):
return self.base_dist.total_count
@property
def probs(self):
return self.base_dist.probs
@property
def logits(self):
return self.base_dist.logits
class NegativeBinomial(Distribution):
r"""
Creates a Negative Binomial distribution, i.e. distribution
of the number of successful independent and identical Bernoulli trials
before :attr:`total_count` failures are achieved. The probability
of success of each Bernoulli trial is :attr:`probs`.
Args:
total_count (float or Tensor): non-negative number of negative Bernoulli
trials to stop, although the distribution is still valid for real
valued count
probs (Tensor): Event probabilities of success in the half open interval [0, 1)
logits (Tensor): Event log-odds for probabilities of success
"""
arg_constraints = {'total_count': constraints.greater_than_eq(0),
'probs': constraints.half_open_interval(0., 1.),
'logits': constraints.real}
support = constraints.nonnegative_integer
def __init__(self, total_count, probs=None, logits=None, validate_args=None):
if (probs is None) == (logits is None):
raise ValueError("Either `probs` or `logits` must be specified, but not both.")
if probs is not None:
self.total_count, self.probs, = broadcast_all(total_count, probs)
self.total_count = self.total_count.type_as(self.probs)
else:
self.total_count, self.logits, = broadcast_all(total_count, logits)
self.total_count = self.total_count.type_as(self.logits)
self._param = self.probs if probs is not None else self.logits
batch_shape = self._param.size()
super(NegativeBinomial, self).__init__(batch_shape, validate_args=validate_args)
def expand(self, batch_shape, _instance=None):
new = self._get_checked_instance(NegativeBinomial, _instance)
batch_shape = torch.Size(batch_shape)
new.total_count = self.total_count.expand(batch_shape)
if 'probs' in self.__dict__:
new.probs = self.probs.expand(batch_shape)
new._param = new.probs
if 'logits' in self.__dict__:
new.logits = self.logits.expand(batch_shape)
new._param = new.logits
super(NegativeBinomial, new).__init__(batch_shape, validate_args=False)
new._validate_args = self._validate_args
return new
def _new(self, *args, **kwargs):
return self._param.new(*args, **kwargs)
@property
def mean(self):
return self.total_count * torch.exp(self.logits)
@property
def variance(self):
return self.mean / torch.sigmoid(-self.logits)
@lazy_property
def logits(self):
return probs_to_logits(self.probs, is_binary=True)
@lazy_property
def probs(self):
return logits_to_probs(self.logits, is_binary=True)
@property
def param_shape(self):
return self._param.size()
@lazy_property
def _gamma(self):
# Note we avoid validating because self.total_count can be zero.
return torch.distributions.Gamma(concentration=self.total_count,
rate=torch.exp(-self.logits),
validate_args=False)
def sample(self, sample_shape=torch.Size()):
with torch.no_grad():
rate = self._gamma.sample(sample_shape=sample_shape)
return torch.poisson(rate)
def log_prob(self, value):
if self._validate_args:
self._validate_sample(value)
log_unnormalized_prob = (self.total_count * F.logsigmoid(-self.logits) +
value * F.logsigmoid(self.logits))
log_normalization = (-torch.lgamma(self.total_count + value) + torch.lgamma(1. + value) +
torch.lgamma(self.total_count))
return log_unnormalized_prob - log_normalization
class NegativeBinomialMixture(Distribution):
"""
Negative binomial mixture distribution.
See :class:`~scvi.distributions.NegativeBinomial` for further description
of parameters.
Parameters
----------
mu1
Mean of the component 1 distribution.
mu2
Mean of the component 2 distribution.
theta1
Inverse dispersion for component 1.
mixture_logits
Logits scale probability of belonging to component 1.
theta2
Inverse dispersion for component 1. If `None`, assumed to be equal to `theta1`.
validate_args
Raise ValueError if arguments do not match constraints
"""
arg_constraints = {
"mu1": constraints.greater_than_eq(0),
"mu2": constraints.greater_than_eq(0),
"theta1": constraints.greater_than_eq(0),
"mixture_probs": constraints.half_open_interval(0.0, 1.0),
"mixture_logits": constraints.real,
}
support = constraints.nonnegative_integer
def __init__(
self,
mu1: torch.Tensor,
mu2: torch.Tensor,
theta1: torch.Tensor,
mixture_logits: torch.Tensor,
theta2: Optional[torch.Tensor] = None,
validate_args: bool = False,
):
(
self.mu1,
self.theta1,
self.mu2,
self.mixture_logits,
) = broadcast_all(mu1, theta1, mu2, mixture_logits)
super().__init__(validate_args=validate_args)
if theta2 is not None:
self.theta2 = broadcast_all(mu1, theta2)
else:
self.theta2 = None
@property
def mean(self):
pi = self.mixture_probs
return pi * self.mu1 + (1 - pi) * self.mu2
@lazy_property
def mixture_probs(self) -> torch.Tensor:
return logits_to_probs(self.mixture_logits, is_binary=True)
def sample(
self, sample_shape: Union[torch.Size, Tuple] = torch.Size()
) -> torch.Tensor:
with torch.no_grad():
pi = self.mixture_probs
mixing_sample = torch.distributions.Bernoulli(pi).sample()
mu = self.mu1 * mixing_sample + self.mu2 * (1 - mixing_sample)
if self.theta2 is None:
theta = self.theta1
else:
theta = self.theta1 * mixing_sample + self.theta2 * (1 - mixing_sample)
gamma_d = _gamma(mu, theta)
p_means = gamma_d.sample(sample_shape)
# Clamping as distributions objects can have buggy behaviors when
# their parameters are too high
l_train = torch.clamp(p_means, max=1e8)
counts = Poisson(
l_train
).sample() # Shape : (n_samples, n_cells_batch, n_features)
return counts
def log_prob(self, value: torch.Tensor) -> torch.Tensor:
try:
self._validate_sample(value)
except ValueError:
warnings.warn(
"The value argument must be within the support of the distribution",
UserWarning,
)
return log_mixture_nb(
value,
self.mu1,
self.mu2,
self.theta1,
self.theta2,
self.mixture_logits,
eps=1e-08,
)
class ZeroInflatedNegativeBinomial(NegativeBinomial):
r"""
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:
1. :math:`w \sim \textrm{Gamma}(\underbrace{\theta}_{\text{shape}}, \underbrace{\theta/\mu}_{\text{rate}})`
2. :math:`x \sim \textrm{Poisson}(w)`
Parameters
----------
total_count
Number of failures until the experiment is stopped.
probs
The success probability.
mu
Mean of the distribution.
theta
Inverse dispersion.
zi_logits
Logits scale of zero inflation probability.
validate_args
Raise ValueError if arguments do not match constraints
"""
arg_constraints = {
"mu": constraints.greater_than_eq(0),
"theta": constraints.greater_than_eq(0),
"zi_probs": constraints.half_open_interval(0.0, 1.0),
"zi_logits": constraints.real,
}
support = constraints.nonnegative_integer
def __init__(
self,
total_count: Optional[torch.Tensor] = None,
probs: Optional[torch.Tensor] = None,
logits: Optional[torch.Tensor] = None,
mu: Optional[torch.Tensor] = None,
theta: Optional[torch.Tensor] = None,
zi_logits: Optional[torch.Tensor] = None,
validate_args: bool = False,
):
super().__init__(
total_count=total_count,
probs=probs,
logits=logits,
mu=mu,
theta=theta,
validate_args=validate_args,
)
self.zi_logits, self.mu, self.theta = broadcast_all(
zi_logits, self.mu, self.theta
)
@property
def mean(self):
pi = self.zi_probs
return (1 - pi) * self.mu
@property
def variance(self):
raise NotImplementedError
@lazy_property
def zi_logits(self) -> torch.Tensor:
return probs_to_logits(self.zi_probs, is_binary=True)
@lazy_property
def zi_probs(self) -> torch.Tensor:
return logits_to_probs(self.zi_logits, is_binary=True)
def sample(
self, sample_shape: Union[torch.Size, Tuple] = torch.Size()
) -> torch.Tensor:
with torch.no_grad():
samp = super().sample(sample_shape=sample_shape)
is_zero = torch.rand_like(samp) <= self.zi_probs
samp[is_zero] = 0.0
return samp
def log_prob(self, value: torch.Tensor) -> torch.Tensor:
try:
self._validate_sample(value)
except ValueError:
warnings.warn(
"The value argument must be within the support of the distribution",
UserWarning,
)
return log_zinb_positive(value, self.mu, self.theta, self.zi_logits, eps=1e-08)
class ZeroInflatedNegativeBinomial(NegativeBinomial):
r"""Zero Inflated Negative Binomial distribution.
zi_logits correspond to the zero-inflation logits
mu + mu ** 2 / theta
The negative binomial component parameters can follow two two parameterizations:
- The first one corresponds to the parameterization NB(`total_count`, `probs`)
where `total_count` is the number of failures until the experiment is stopped
and `probs` the success probability.
- The (`mu`, `theta`) parameterization is the one used by scVI. These parameters respectively
control the mean and overdispersion of the distribution.
`_convert_mean_disp_to_counts_logits` and `_convert_counts_logits_to_mean_disp`
provide ways to convert one parameterization to another.
Parameters
----------
Returns
-------
"""
arg_constraints = {
"mu": constraints.greater_than_eq(0),
"theta": constraints.greater_than_eq(0),
"zi_probs": constraints.half_open_interval(0.0, 1.0),
"zi_logits": constraints.real,
}
support = constraints.nonnegative_integer
def __init__(
self,
total_count: torch.Tensor = None,
probs: torch.Tensor = None,
logits: torch.Tensor = None,
mu: torch.Tensor = None,
theta: torch.Tensor = None,
zi_logits: torch.Tensor = None,
validate_args=True,
):
super().__init__(
total_count=total_count,
probs=probs,
logits=logits,
mu=mu,
theta=theta,
validate_args=validate_args,
)
self.zi_logits, self.mu, self.theta = broadcast_all(
zi_logits, self.mu, self.theta)
@lazy_property
def zi_logits(self) -> torch.Tensor:
return probs_to_logits(self.zi_probs, is_binary=True)
@lazy_property
def zi_probs(self) -> torch.Tensor:
return logits_to_probs(self.zi_logits, is_binary=True)
def sample(
self, sample_shape: Union[torch.Size, Tuple] = torch.Size()
) -> torch.Tensor:
with torch.no_grad():
samp = super().sample(sample_shape=sample_shape)
is_zero = torch.rand_like(samp) <= self.zi_probs
samp[is_zero] = 0.0
return samp
def log_prob(self, value: torch.Tensor) -> torch.Tensor:
try:
self._validate_sample(value)
except ValueError:
warnings.warn(
"The value argument must be within the support of the distribution",
UserWarning,
)
return log_zinb_positive(value,
self.mu,
self.theta,
self.zi_logits,
eps=1e-08)
return x
def _constraint_hash(constraint: constraints.Constraint) -> int:
assert isinstance(constraint, constraints.Constraint)
out = hash(type(constraint))
out ^= hash(frozenset(constraint.__dict__.items()))
return out
# useful if the distribution's init has multiple ways of specifying (e.g. both logits or probs)
_distribution_to_param_names = {NegativeBinomial: ['probs', 'total_count']}
# torch.distributions has a whole 'transforms' module but I don't know if they provide a mapping
_constraint_to_ilink = {
_constraint_hash(constraints.positive):
torch.exp,
_constraint_hash(constraints.greater_than(0)):
torch.exp,
_constraint_hash(constraints.unit_interval):
torch.sigmoid,
_constraint_hash(constraints.real):
identity,
# TODO: is there a way to make these work?
_constraint_hash(constraints.greater_than_eq(0)):
torch.exp,
_constraint_hash(constraints.half_open_interval(0, 1)):
torch.sigmoid,
# TODO: constraints.interval
}
