def test_broadcast_distribution_samples(self, samples_to_broadcast):
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
outs = broadcast_distribution_samples(samples, size=size)
assert all(o.shape == broadcast_shape for o in outs)
else:
with pytest.raises(ValueError):
broadcast_distribution_samples(samples, size=size)
def test_broadcast_distribution_samples(self, samples_to_broadcast):
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
outs = broadcast_distribution_samples(samples, size=size)
assert all((o.shape == broadcast_shape for o in outs))
else:
with pytest.raises(ValueError):
broadcast_distribution_samples(samples, size=size)
def _comp_samples(self,
point=None,
size=None,
comp_dist_shapes=None,
broadcast_shape=None):
if self.comp_is_distribution:
samples = self._comp_dists.random(point=point, size=size)
else:
if comp_dist_shapes is None:
comp_dist_shapes = self._comp_dist_shapes
if broadcast_shape is None:
broadcast_shape = self._sample_shape
samples = []
for dist_shape, generator in zip(comp_dist_shapes,
self._generators):
sample = generate_samples(
generator=generator,
dist_shape=dist_shape,
broadcast_shape=broadcast_shape,
point=point,
size=size,
not_broadcast_kwargs={"raw_size_": size},
)
samples.append(sample)
samples = np.array(
broadcast_distribution_samples(samples, size=size))
# In the logp we assume the last axis holds the mixture components
# so we move the axis to the last dimension
samples = np.moveaxis(samples, 0, -1)
return samples.astype(self.dtype)
def random(self, point=None, size=None):
"""
Draw random values from defined Mixture distribution.
Parameters
----------
point: dict, optional
Dict of variable values on which random values are to be
conditioned (uses default point if not specified).
size: int, optional
Desired size of random sample (returns one sample if not
specified).
Returns
-------
array
"""
# Convert size to tuple
size = to_tuple(size)
# Draw mixture weights and infer the comp_dists shapes
with _DrawValuesContext() as draw_context:
# We first need to check w and comp_tmp shapes and re compute size
w = draw_values([self.w], point=point, size=size)[0]
comp_dist_shapes, broadcast_shape = self.infer_comp_dist_shapes(point=point)
# When size is not None, it's hard to tell the w parameter shape
if size is not None and w.shape[: len(size)] == size:
w_shape = w.shape[len(size) :]
else:
w_shape = w.shape
# Try to determine parameter shape and dist_shape
if self.comp_is_distribution:
param_shape = np.broadcast(np.empty(w_shape), np.empty(broadcast_shape)).shape
else:
param_shape = np.broadcast(np.empty(w_shape), np.empty(broadcast_shape + (1,))).shape
if np.asarray(self.shape).size != 0:
dist_shape = np.broadcast(np.empty(self.shape), np.empty(param_shape[:-1])).shape
else:
dist_shape = param_shape[:-1]
# Try to determine the size that must be used to get the mixture
# components (i.e. get random choices using w).
# 1. There must be size independent choices based on w.
# 2. There must also be independent draws for each non singleton axis
# of w.
# 3. There must also be independent draws for each dimension added by
# self.shape with respect to the w.ndim. These usually correspond to
# observed variables with batch shapes
wsh = (1,) * (len(dist_shape) - len(w_shape) + 1) + w_shape[:-1]
psh = (1,) * (len(dist_shape) - len(param_shape) + 1) + param_shape[:-1]
w_sample_size = []
# Loop through the dist_shape to get the conditions 2 and 3 first
for i in range(len(dist_shape)):
if dist_shape[i] != psh[i] and wsh[i] == 1:
# self.shape[i] is a non singleton dimension (usually caused by
# observed data)
sh = dist_shape[i]
else:
sh = wsh[i]
w_sample_size.append(sh)
if size is not None and w_sample_size[: len(size)] != size:
w_sample_size = size + tuple(w_sample_size)
# Broadcast w to the w_sample_size (add a singleton last axis for the
# mixture components)
w = broadcast_distribution_samples([w, np.empty(w_sample_size + (1,))], size=size)[0]
# Semiflatten the mixture weights. The last axis is the number of
# mixture mixture components, and the rest is all about size,
# dist_shape and broadcasting
w_ = np.reshape(w, (-1, w.shape[-1]))
w_samples = random_choice(p=w_, size=None) # w's shape already includes size
# Now we broadcast the chosen components to the dist_shape
w_samples = np.reshape(w_samples, w.shape[:-1])
if size is not None and dist_shape[: len(size)] != size:
w_samples = np.broadcast_to(w_samples, size + dist_shape)
else:
w_samples = np.broadcast_to(w_samples, dist_shape)
# When size is not None, maybe dist_shape partially overlaps with size
if size is not None:
if size == dist_shape:
size = None
elif size[-len(dist_shape) :] == dist_shape:
size = size[: len(size) - len(dist_shape)]
# We get an integer _size instead of a tuple size for drawing the
# mixture, then we just reshape the output
if size is None:
_size = None
else:
_size = int(np.prod(size))
# Compute the total size of the mixture's random call with size
if _size is not None:
output_size = int(_size * np.prod(dist_shape) * param_shape[-1])
else:
output_size = int(np.prod(dist_shape) * param_shape[-1])
# Get the size we need for the mixture's random call
if self.comp_is_distribution:
mixture_size = int(output_size // np.prod(broadcast_shape))
else:
mixture_size = int(output_size // (np.prod(broadcast_shape) * param_shape[-1]))
if mixture_size == 1 and _size is None:
mixture_size = None
# Sample from the mixture
with draw_context:
mixed_samples = self._comp_samples(
point=point,
size=mixture_size,
broadcast_shape=broadcast_shape,
comp_dist_shapes=comp_dist_shapes,
)
# Test that the mixture has the same number of "samples" as w
if w_samples.size != (mixed_samples.size // w.shape[-1]):
raise ValueError(
"Inconsistent number of samples from the "
"mixture and mixture weights. Drew {} mixture "
"weights elements, and {} samples from the "
"mixture components.".format(w_samples.size, mixed_samples.size // w.shape[-1])
)
# Semiflatten the mixture to be able to zip it with w_samples
w_samples = w_samples.flatten()
mixed_samples = np.reshape(mixed_samples, (-1, w.shape[-1]))
# Select the samples from the mixture
samples = np.array([mixed[choice] for choice, mixed in zip(w_samples, mixed_samples)])
# Reshape the samples to the correct output shape
if size is None:
samples = np.reshape(samples, dist_shape)
else:
samples = np.reshape(samples, size + dist_shape)
return samples
