Spaces:
Running
Running
File size: 8,675 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
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 = len(self.batch_shape)
cat_batch_ndims = len(self.mixture_distribution.batch_shape)
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 + ")"
|