| from typing import Dict |
| |
| import torch |
| from torch.distributions import Categorical, constraints |
| from torch.distributions.distribution import Distribution |
| |
| __all__ = ["MixtureSameFamily"] |
| |
| |
| class MixtureSameFamily(Distribution): |
| r""" |
| The `MixtureSameFamily` distribution implements a (batch of) mixture |
| distribution where all component are from different parameterizations of |
| the same distribution type. It is parameterized by a `Categorical` |
| "selecting distribution" (over `k` component) and a component |
| distribution, i.e., a `Distribution` with a rightmost batch shape |
| (equal to `[k]`) which indexes each (batch of) component. |
| |
| Examples:: |
| |
| >>> # xdoctest: +SKIP("undefined vars") |
| >>> # Construct Gaussian Mixture Model in 1D consisting of 5 equally |
| >>> # weighted normal distributions |
| >>> mix = D.Categorical(torch.ones(5,)) |
| >>> comp = D.Normal(torch.randn(5,), torch.rand(5,)) |
| >>> gmm = MixtureSameFamily(mix, comp) |
| |
| >>> # Construct Gaussian Mixture Model in 2D consisting of 5 equally |
| >>> # weighted bivariate normal distributions |
| >>> mix = D.Categorical(torch.ones(5,)) |
| >>> comp = D.Independent(D.Normal( |
| ... torch.randn(5,2), torch.rand(5,2)), 1) |
| >>> gmm = MixtureSameFamily(mix, comp) |
| |
| >>> # Construct a batch of 3 Gaussian Mixture Models in 2D each |
| >>> # consisting of 5 random weighted bivariate normal distributions |
| >>> mix = D.Categorical(torch.rand(3,5)) |
| >>> comp = D.Independent(D.Normal( |
| ... torch.randn(3,5,2), torch.rand(3,5,2)), 1) |
| >>> gmm = MixtureSameFamily(mix, comp) |
| |
| Args: |
| mixture_distribution: `torch.distributions.Categorical`-like |
| instance. Manages the probability of selecting component. |
| The number of categories must match the rightmost batch |
| dimension of the `component_distribution`. Must have either |
| scalar `batch_shape` or `batch_shape` matching |
| `component_distribution.batch_shape[:-1]` |
| component_distribution: `torch.distributions.Distribution`-like |
| instance. Right-most batch dimension indexes component. |
| """ |
| arg_constraints: Dict[str, constraints.Constraint] = {} |
| has_rsample = False |
| |
| def __init__( |
| self, mixture_distribution, component_distribution, validate_args=None |
| ): |
| self._mixture_distribution = mixture_distribution |
| self._component_distribution = component_distribution |
| |
| if not isinstance(self._mixture_distribution, Categorical): |
| raise ValueError( |
| " The Mixture distribution needs to be an " |
| " instance of torch.distributions.Categorical" |
| ) |
| |
| if not isinstance(self._component_distribution, Distribution): |
| raise ValueError( |
| "The Component distribution need to be an " |
| "instance of torch.distributions.Distribution" |
| ) |
| |
| # Check that batch size matches |
| mdbs = self._mixture_distribution.batch_shape |
| cdbs = self._component_distribution.batch_shape[:-1] |
| for size1, size2 in zip(reversed(mdbs), reversed(cdbs)): |
| if size1 != 1 and size2 != 1 and size1 != size2: |
| raise ValueError( |
| f"`mixture_distribution.batch_shape` ({mdbs}) is not " |
| "compatible with `component_distribution." |
| f"batch_shape`({cdbs})" |
| ) |
| |
| # Check that the number of mixture component matches |
| km = self._mixture_distribution.logits.shape[-1] |
| kc = self._component_distribution.batch_shape[-1] |
| if km is not None and kc is not None and km != kc: |
| raise ValueError( |
| f"`mixture_distribution component` ({km}) does not" |
| " equal `component_distribution.batch_shape[-1]`" |
| f" ({kc})" |
| ) |
| self._num_component = km |
| |
| event_shape = self._component_distribution.event_shape |
| self._event_ndims = len(event_shape) |
| super().__init__( |
| batch_shape=cdbs, event_shape=event_shape, validate_args=validate_args |
| ) |
| |
| def expand(self, batch_shape, _instance=None): |
| batch_shape = torch.Size(batch_shape) |
| batch_shape_comp = batch_shape + (self._num_component,) |
| new = self._get_checked_instance(MixtureSameFamily, _instance) |
| new._component_distribution = self._component_distribution.expand( |
| batch_shape_comp |
| ) |
| new._mixture_distribution = self._mixture_distribution.expand(batch_shape) |
| new._num_component = self._num_component |
| new._event_ndims = self._event_ndims |
| event_shape = new._component_distribution.event_shape |
| super(MixtureSameFamily, new).__init__( |
| batch_shape=batch_shape, event_shape=event_shape, validate_args=False |
| ) |
| new._validate_args = self._validate_args |
| return new |
| |
| @constraints.dependent_property |
| def support(self): |
| # FIXME this may have the wrong shape when support contains batched |
| # parameters |
| return self._component_distribution.support |
| |
| @property |
| def mixture_distribution(self): |
| return self._mixture_distribution |
| |
| @property |
| def component_distribution(self): |
| return self._component_distribution |
| |
| @property |
| def mean(self): |
| probs = self._pad_mixture_dimensions(self.mixture_distribution.probs) |
| return torch.sum( |
| probs * self.component_distribution.mean, dim=-1 - self._event_ndims |
| ) # [B, E] |
| |
| @property |
| def variance(self): |
| # Law of total variance: Var(Y) = E[Var(Y|X)] + Var(E[Y|X]) |
| probs = self._pad_mixture_dimensions(self.mixture_distribution.probs) |
| mean_cond_var = torch.sum( |
| probs * self.component_distribution.variance, dim=-1 - self._event_ndims |
| ) |
| var_cond_mean = torch.sum( |
| probs * (self.component_distribution.mean - self._pad(self.mean)).pow(2.0), |
| dim=-1 - self._event_ndims, |
| ) |
| return mean_cond_var + var_cond_mean |
| |
| def cdf(self, x): |
| x = self._pad(x) |
| cdf_x = self.component_distribution.cdf(x) |
| mix_prob = self.mixture_distribution.probs |
| |
| return torch.sum(cdf_x * mix_prob, dim=-1) |
| |
| def log_prob(self, x): |
| if self._validate_args: |
| self._validate_sample(x) |
| x = self._pad(x) |
| log_prob_x = self.component_distribution.log_prob(x) # [S, B, k] |
| log_mix_prob = torch.log_softmax( |
| self.mixture_distribution.logits, dim=-1 |
| ) # [B, k] |
| return torch.logsumexp(log_prob_x + log_mix_prob, dim=-1) # [S, B] |
| |
| def sample(self, sample_shape=torch.Size()): |
| with torch.no_grad(): |
| sample_len = len(sample_shape) |
| batch_len = len(self.batch_shape) |
| gather_dim = sample_len + batch_len |
| es = self.event_shape |
| |
| # mixture samples [n, B] |
| mix_sample = self.mixture_distribution.sample(sample_shape) |
| mix_shape = mix_sample.shape |
| |
| # component samples [n, B, k, E] |
| comp_samples = self.component_distribution.sample(sample_shape) |
| |
| # Gather along the k dimension |
| mix_sample_r = mix_sample.reshape( |
| mix_shape + torch.Size([1] * (len(es) + 1)) |
| ) |
| mix_sample_r = mix_sample_r.repeat( |
| torch.Size([1] * len(mix_shape)) + torch.Size([1]) + es |
| ) |
| |
| samples = torch.gather(comp_samples, gather_dim, mix_sample_r) |
| return samples.squeeze(gather_dim) |
| |
| def _pad(self, x): |
| return x.unsqueeze(-1 - self._event_ndims) |
| |
| def _pad_mixture_dimensions(self, x): |
| dist_batch_ndims = self.batch_shape.numel() |
| cat_batch_ndims = self.mixture_distribution.batch_shape.numel() |
| pad_ndims = 0 if cat_batch_ndims == 1 else dist_batch_ndims - cat_batch_ndims |
| xs = x.shape |
| x = x.reshape( |
| xs[:-1] |
| + torch.Size(pad_ndims * [1]) |
| + xs[-1:] |
| + torch.Size(self._event_ndims * [1]) |
| ) |
| return x |
| |
| def __repr__(self): |
| args_string = ( |
| f"\n {self.mixture_distribution},\n {self.component_distribution}" |
| ) |
| return "MixtureSameFamily" + "(" + args_string + ")" |
