def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
size_ = to_tuple(size)
outs, out_shape = get_broadcastable_dist_samples(
samples, size=size, return_out_shape=True)
assert out_shape == broadcast_shape
for i, o in zip(samples, outs):
ishape = i.shape
if ishape[:min([len(size_), len(ishape)])] == size_:
expected_shape = (size_ + (1, ) *
(len(broadcast_shape) - len(ishape)) +
ishape[len(size_):])
else:
expected_shape = ishape
assert o.shape == expected_shape
assert shapes_broadcasting(*[o.shape
for o in outs]) == broadcast_shape
else:
with pytest.raises(ValueError):
get_broadcastable_dist_samples(samples, size=size)
def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
size_ = to_tuple(size)
outs, out_shape = get_broadcastable_dist_samples(
samples, size=size, return_out_shape=True
)
assert out_shape == broadcast_shape
for i, o in zip(samples, outs):
ishape = i.shape
if ishape[: min([len(size_), len(ishape)])] == size_:
expected_shape = (
size_
+ (1,) * (len(broadcast_shape) - len(ishape))
+ ishape[len(size_) :]
)
else:
expected_shape = ishape
assert o.shape == expected_shape
assert shapes_broadcasting(*[o.shape for o in outs]) == broadcast_shape
else:
with pytest.raises(ValueError):
get_broadcastable_dist_samples(samples, size=size)
def generate_samples(generator, *args, **kwargs):
"""Generate samples from the distribution of a random variable.
Parameters
----------
generator: function
Function to generate the random samples. The function is
expected take parameters for generating samples and
a keyword argument ``size`` which determines the shape
of the samples.
The args and kwargs (stripped of the keywords below) will be
passed to the generator function.
keyword arguments
~~~~~~~~~~~~~~~~~
dist_shape: int or tuple of int
The shape of the random variable (i.e., the shape attribute).
size: int or tuple of int
The required shape of the samples.
broadcast_shape: tuple of int or None
The shape resulting from the broadcasting of the parameters.
If not specified it will be inferred from the shape of the
parameters. This may be required when the parameter shape
does not determine the shape of a single sample, for example,
the shape of the probabilities in the Categorical distribution.
not_broadcast_kwargs: dict or None
Key word argument dictionary to provide to the random generator, which
must not be broadcasted with the rest of the args and kwargs.
Any remaining args and kwargs are passed on to the generator function.
"""
dist_shape = kwargs.pop("dist_shape", ())
size = kwargs.pop("size", None)
broadcast_shape = kwargs.pop("broadcast_shape", None)
not_broadcast_kwargs = kwargs.pop("not_broadcast_kwargs", None)
if not_broadcast_kwargs is None:
not_broadcast_kwargs = dict()
# Parse out raw input parameters for the generator
args = tuple(p[0] if isinstance(p, tuple) else p for p in args)
for key in kwargs:
p = kwargs[key]
kwargs[key] = p[0] if isinstance(p, tuple) else p
# Convert size and dist_shape to tuples
size_tup = to_tuple(size)
dist_shape = to_tuple(dist_shape)
if dist_shape[: len(size_tup)] == size_tup:
# dist_shape is prepended with size_tup. This is not a consequence
# of the parameters being drawn size_tup times! By chance, the
# distribution's shape has its first elements equal to size_tup.
# This means that we must prepend the size_tup to dist_shape, and
# check if that broadcasts well with the parameters
_dist_shape = size_tup + dist_shape
else:
_dist_shape = dist_shape
if broadcast_shape is None:
# If broadcast_shape is not explicitly provided, it is inferred as the
# broadcasted shape of the input parameter and dist_shape, taking into
# account the potential size prefix
inputs = args + tuple(kwargs.values())
broadcast_shape = broadcast_dist_samples_shape(
[np.asarray(i).shape for i in inputs] + [_dist_shape], size=size_tup
)
# We do this instead of broadcast_distribution_samples to avoid
# creating a dummy array with dist_shape in memory
inputs = get_broadcastable_dist_samples(
inputs,
size=size_tup,
must_bcast_with=broadcast_shape,
)
# We modify the arguments with their broadcasted counterparts
args = tuple(inputs[: len(args)])
for offset, key in enumerate(kwargs):
kwargs[key] = inputs[len(args) + offset]
# Update kwargs with the keyword arguments that were not broadcasted
kwargs.update(not_broadcast_kwargs)
# We ensure that broadcast_shape is a tuple
broadcast_shape = to_tuple(broadcast_shape)
try:
dist_bcast_shape = broadcast_dist_samples_shape(
[_dist_shape, broadcast_shape],
size=size,
)
except (ValueError, TypeError):
raise TypeError(
"""Attempted to generate values with incompatible shapes:
size: {size}
size_tup: {size_tup}
broadcast_shape[:len(size_tup)] == size_tup: {size_prepended}
dist_shape: {dist_shape}
broadcast_shape: {broadcast_shape}
""".format(
size=size,
size_tup=size_tup,
dist_shape=dist_shape,
broadcast_shape=broadcast_shape,
size_prepended=broadcast_shape[: len(size_tup)] == size_tup,
)
)
if dist_bcast_shape[: len(size_tup)] == size_tup:
samples = generator(size=dist_bcast_shape, *args, **kwargs)
else:
samples = generator(size=size_tup + dist_bcast_shape, *args, **kwargs)
return np.asarray(samples)
def random(self, point=None, size=None):
"""
Draw random values from defined ``MixtureSameFamily`` 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
"""
sample_shape = to_tuple(size)
mixture_axis = self.mixture_axis
# First we draw values for the mixture component weights
(w,) = draw_values([self.w], point=point, size=size)
# We now draw random choices from those weights.
# However, we have to ensure that the number of choices has the
# sample_shape present.
w_shape = w.shape
batch_shape = self.comp_dists.shape[: mixture_axis + 1]
param_shape = np.broadcast(np.empty(w_shape), np.empty(batch_shape)).shape
event_shape = self.comp_dists.shape[mixture_axis + 1 :]
if np.asarray(self.shape).size != 0:
comp_dists_ndim = len(self.comp_dists.shape)
# If event_shape of both comp_dists and supplied shape matches,
# broadcast only batch_shape
# else broadcast the entire given shape with batch_shape.
if list(self.shape[mixture_axis - comp_dists_ndim + 1 :]) == list(event_shape):
dist_shape = np.broadcast(
np.empty(self.shape[:mixture_axis]), np.empty(param_shape[:mixture_axis])
).shape
else:
dist_shape = np.broadcast(
np.empty(self.shape), np.empty(param_shape[:mixture_axis])
).shape
else:
dist_shape = param_shape[:mixture_axis]
# 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[:mixture_axis]
psh = (1,) * (len(dist_shape) - len(param_shape) + 1) + param_shape[:mixture_axis]
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 sample_shape is not None and w_sample_size[: len(sample_shape)] != sample_shape:
w_sample_size = sample_shape + tuple(w_sample_size)
choices = random_choice(p=w, size=w_sample_size)
# We now draw samples from the mixture components random method
comp_samples = self.comp_dists.random(point=point, size=size)
if comp_samples.shape[: len(sample_shape)] != sample_shape:
comp_samples = np.broadcast_to(
comp_samples,
shape=sample_shape + comp_samples.shape,
)
# At this point the shapes of the arrays involved are:
# comp_samples.shape = (sample_shape, batch_shape, mixture_axis, event_shape)
# choices.shape = (sample_shape, batch_shape)
#
# To be able to take the choices along the mixture_axis of the
# comp_samples, we have to add in dimensions to the right of the
# choices array.
# We also need to make sure that the batch_shapes of both the comp_samples
# and choices broadcast with each other.
choices = np.reshape(choices, choices.shape + (1,) * (1 + len(event_shape)))
choices, comp_samples = get_broadcastable_dist_samples([choices, comp_samples], size=size)
# We now take the choices of the mixture components along the mixture_axis
# but we use the negative index representation to be able to handle the
# sample_shape
samples = np.take_along_axis(
comp_samples, choices, axis=mixture_axis - len(self.comp_dists.shape)
)
# The `samples` array still has the `mixture_axis`, so we must remove it:
output = samples[(..., 0) + (slice(None),) * len(event_shape)]
return output