def comp_dists(self, comp_dists):
self._comp_dists = comp_dists
if isinstance(comp_dists, Distribution):
self._comp_dist_shapes = to_tuple(comp_dists.shape)
self._broadcast_shape = self._comp_dist_shapes
self.comp_is_distribution = True
else:
# Now we check the comp_dists distribution shape, see what
# the broadcast shape would be. This shape will be the dist_shape
# used by generate samples (the shape of a single random sample)
# from the mixture
self._comp_dist_shapes = [to_tuple(d.shape) for d in comp_dists]
# All component distributions must broadcast with each other
try:
self._broadcast_shape = np.broadcast(
*[np.empty(shape)
for shape in self._comp_dist_shapes]).shape
except Exception:
raise TypeError("Supplied comp_dists shapes do not broadcast "
"with each other. comp_dists shapes are: "
"{}".format(self._comp_dist_shapes))
# We wrap the _comp_dist.random by adding the kwarg raw_size_,
# which will be the size attribute passed to _comp_samples.
# _comp_samples then calls generate_samples, which may change the
# size value to make it compatible with scipy.stats.*.rvs
self._generators = []
for comp_dist in comp_dists:
generator = Mixture._comp_dist_random_wrapper(comp_dist.random)
self._generators.append(generator)
self.comp_is_distribution = False
def random(self, point=None, size=None, **kwargs):
if self.rand is not None:
not_broadcast_kwargs = dict(point=point)
not_broadcast_kwargs.update(**kwargs)
if self.wrap_random_with_dist_shape:
size = to_tuple(size)
with _DrawValuesContextBlocker():
test_draw = generate_samples(
self.rand,
size=None,
not_broadcast_kwargs=not_broadcast_kwargs,
)
test_shape = test_draw.shape
if self.shape[:len(size)] == size:
dist_shape = size + self.shape
else:
dist_shape = self.shape
broadcast_shape = broadcast_dist_samples_shape(
[dist_shape, test_shape], size=size)
broadcast_shape = broadcast_shape[:len(broadcast_shape) -
len(test_shape)]
samples = generate_samples(
self.rand,
broadcast_shape=broadcast_shape,
size=size,
not_broadcast_kwargs=not_broadcast_kwargs,
)
else:
samples = self.rand(point=point, size=size, **kwargs)
if self.check_shape_in_random:
expected_shape = self.shape if size is None else to_tuple(
size) + self.shape
if not expected_shape == samples.shape:
raise RuntimeError(
"DensityDist encountered a shape inconsistency "
"while drawing samples using the supplied random "
"function. Was expecting to get samples of shape "
"{expected} but got {got} instead.\n"
"Whenever possible wrap_random_with_dist_shape = True "
"is recommended.\n"
"Be aware that the random callable provided as the "
"DensityDist random method cannot "
"adapt to shape changes in the distribution's "
"shape, which sometimes are necessary for sampling "
"when the model uses pymc3.Data or aesara shared "
"tensors, or when the DensityDist has observed "
"values.\n"
"This check can be disabled by passing "
"check_shape_in_random=False when the DensityDist "
"is initialized.".format(
expected=expected_shape,
got=samples.shape,
))
return samples
else:
raise ValueError(
"Distribution was not passed any random method. "
"Define a custom random method and pass it as kwarg random")
def random(self, point=None, size=None):
"""Draw random values from GaussianRandomWalk.
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
"""
sigma, mu = distribution.draw_values([self.sigma, self.mu],
point=point,
size=size)
return distribution.generate_samples(
self._random,
sigma=sigma,
mu=mu,
size=size,
dist_shape=self.shape,
not_broadcast_kwargs={"sample_shape": to_tuple(size)},
)
def test_sample_generate_values(fixture_model, fixture_sizes):
model, RVs = fixture_model
size = to_tuple(fixture_sizes)
with model:
prior = pm.sample_prior_predictive(samples=fixture_sizes)
for rv in RVs:
assert prior[rv.name].shape == size + tuple(rv.distribution.shape)
def generate_poisson_mixture_data(w, mu, size=1000):
component = np.random.choice(w.size, size=size, p=w)
mu = np.atleast_1d(mu)
out_size = to_tuple(size) + mu.shape[:-1]
mu_ = np.array([mu[..., comp] for comp in component.ravel()])
mu_ = np.reshape(mu_, out_size)
x = np.random.poisson(mu_, size=out_size)
return x
def samples_to_broadcast_to(request, samples_to_broadcast):
to_shape = request.param
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
try:
broadcast_shape = broadcast_dist_samples_shape(
[broadcast_shape, to_tuple(to_shape)], size=size)
except ValueError:
broadcast_shape = None
return to_shape, size, samples, broadcast_shape
def test_sample_generate_values(fixture_model, fixture_sizes):
model, RVs = fixture_model
size = to_tuple(fixture_sizes)
if size == (1,):
# Single draws are interpreted as scalars for backwards compatibility
size = tuple()
with model:
prior = pm.sample_prior_predictive(samples=fixture_sizes)
for rv in RVs:
assert prior[rv.name].shape == size + tuple(rv.distribution.shape)
def test_sample_generate_values(fixture_model, fixture_sizes):
model, RVs = fixture_model
size = to_tuple(fixture_sizes)
if size == (1,):
# Single draws are interpreted as scalars for backwards compatibility
size = tuple()
with model:
prior = pm.sample_prior_predictive(samples=fixture_sizes)
for rv in RVs:
assert prior[rv.name].shape == size + tuple(rv.distribution.shape)
def generate_normal_mixture_data(w, mu, sd, size=1000):
component = np.random.choice(w.size, size=size, p=w)
mu, sd = np.broadcast_arrays(mu, sd)
out_size = to_tuple(size) + mu.shape[:-1]
mu_ = np.array([mu[..., comp] for comp in component.ravel()])
sd_ = np.array([sd[..., comp] for comp in component.ravel()])
mu_ = np.reshape(mu_, out_size)
sd_ = np.reshape(sd_, out_size)
x = np.random.normal(mu_, sd_, size=out_size)
return x
def samples_to_broadcast_to(request, samples_to_broadcast):
to_shape = request.param
size, samples, broadcast_shape = samples_to_broadcast
if broadcast_shape is not None:
try:
broadcast_shape = broadcast_dist_samples_shape(
[broadcast_shape, to_tuple(to_shape)], size=size
)
except ValueError:
broadcast_shape = None
return to_shape, size, samples, broadcast_shape
def __init__(self, w, comp_dists, mixture_axis=-1, *args, **kwargs):
self.w = at.as_tensor_variable(w)
if not isinstance(comp_dists, Distribution):
raise TypeError(
"The MixtureSameFamily distribution only accepts Distribution "
f"instances as its components. Got {type(comp_dists)} instead."
)
self.comp_dists = comp_dists
if mixture_axis < 0:
mixture_axis = len(comp_dists.shape) + mixture_axis
if mixture_axis < 0:
raise ValueError(
"`mixture_axis` is supposed to be in shape of components' distribution. "
f"Got {mixture_axis + len(comp_dists.shape)} axis instead out of the bounds."
)
comp_shape = to_tuple(comp_dists.shape)
self.shape = comp_shape[:mixture_axis] + comp_shape[mixture_axis + 1:]
self.mixture_axis = mixture_axis
kwargs.setdefault("dtype", self.comp_dists.dtype)
# Compute the mode so we don't always have to pass a initval
defaults = kwargs.pop("defaults", [])
event_shape = self.comp_dists.shape[mixture_axis + 1:]
_w = at.shape_padleft(
at.shape_padright(w, len(event_shape)),
len(self.comp_dists.shape) - w.ndim - len(event_shape),
)
mode = take_along_axis(
self.comp_dists.mode,
at.argmax(_w, keepdims=True),
axis=mixture_axis,
)
self.mode = mode[(..., 0) + (slice(None), ) * len(event_shape)]
if not all_discrete(comp_dists):
mean = at.as_tensor_variable(self.comp_dists.mean)
self.mean = (_w * mean).sum(axis=mixture_axis)
if "mean" not in defaults:
defaults.append("mean")
defaults.append("mode")
super().__init__(defaults=defaults, *args, **kwargs)
def test_broadcast_dist_samples_shape(self, fixture_sizes, fixture_shapes):
size = fixture_sizes
shapes = fixture_shapes
size_ = to_tuple(size)
shapes_ = [
s if s[:min([len(size_), len(s)])] != size_ else s[len(size_):]
for s in shapes
]
try:
expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape
except ValueError:
expected_out = None
if expected_out is not None and any(
s[:min([len(size_), len(s)])] == size_ for s in shapes):
expected_out = size_ + expected_out
if expected_out is None:
with pytest.raises(ValueError):
broadcast_dist_samples_shape(shapes, size=size)
else:
out = broadcast_dist_samples_shape(shapes, size=size)
assert out == expected_out
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_broadcast_dist_samples_shape(self, fixture_sizes, fixture_shapes):
size = fixture_sizes
shapes = fixture_shapes
size_ = to_tuple(size)
shapes_ = [
s if s[: min([len(size_), len(s)])] != size_ else s[len(size_) :]
for s in shapes
]
try:
expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape
except ValueError:
expected_out = None
if expected_out is not None and any(
(s[: min([len(size_), len(s)])] == size_ for s in shapes)
):
expected_out = size_ + expected_out
if expected_out is None:
with pytest.raises(ValueError):
broadcast_dist_samples_shape(shapes, size=size)
else:
out = broadcast_dist_samples_shape(shapes, size=size)
assert out == expected_out
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 random_choice(*args, **kwargs):
"""Return draws from a categorial probability functions
Args:
p: array
Probability of each class. If p.ndim > 1, the last axis is
interpreted as the probability of each class, and numpy.random.choice
is iterated for every other axis element.
size: int or tuple
Shape of the desired output array. If p is multidimensional, size
should broadcast with p.shape[:-1].
Returns:
random sample: array
"""
p = kwargs.pop("p")
size = kwargs.pop("size")
k = p.shape[-1]
if p.ndim > 1:
# If p is an nd-array, the last axis is interpreted as the class
# probability. We must iterate over the elements of all the other
# dimensions.
# We first ensure that p is broadcasted to the output's shape
size = to_tuple(size) + (1, )
p = np.broadcast_arrays(p, np.empty(size))[0]
out_shape = p.shape[:-1]
# np.random.choice accepts 1D p arrays, so we semiflatten p to
# iterate calls using the last axis as the category probabilities
p = np.reshape(p, (-1, p.shape[-1]))
samples = np.array([np.random.choice(k, p=p_) for p_ in p])
# We reshape to the desired output shape
samples = np.reshape(samples, out_shape)
else:
samples = np.random.choice(k, p=p, size=size)
return samples
def random(self, point=None, size=None):
"""
Draw random values from MvGaussianRandomWalk.
Parameters
----------
point: dict, optional
Dict of variable values on which random values are to be
conditioned (uses default point if not specified).
size: int or tuple of ints, optional
Desired size of random sample (returns one sample if not
specified).
Returns
-------
array
Examples
-------
.. code-block:: python
with pm.Model():
mu = np.array([1.0, 0.0])
cov = np.array([[1.0, 0.0],
[0.0, 2.0]])
# draw one sample from a 2-dimensional Gaussian random walk with 10 timesteps
sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random()
# draw three samples from a 2-dimensional Gaussian random walk with 10 timesteps
sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=3)
# draw four samples from a 2-dimensional Gaussian random walk with 10 timesteps,
# indexed with a (2, 2) array
sample = MvGaussianRandomWalk(mu, cov, shape=(10, 2)).random(size=(2, 2))
"""
# for each draw specified by the size input, we need to draw time_steps many
# samples from MvNormal.
size = to_tuple(size)
multivariate_samples = self.innov.random(point=point, size=size)
# this has shape (size, self.shape)
if len(self.shape) == 2:
# have time dimension in first slot of shape. Therefore the time
# component can be accessed with the index equal to the length of size.
time_axis = len(size)
multivariate_samples = multivariate_samples.cumsum(axis=time_axis)
if time_axis != 0:
# this for loop covers the case where size is a tuple
for idx in np.ndindex(size):
multivariate_samples[idx] = (multivariate_samples[idx] -
multivariate_samples[idx][0])
else:
# size was passed as None
multivariate_samples = multivariate_samples - multivariate_samples[
0]
# if the above statement fails, then only a spatial dimension was passed in for self.shape.
# Therefore don't subtract off the initial value since otherwise you get all zeros
# as your output.
return multivariate_samples
def test_normal_mixture_nd(self, nd, ncomp):
nd = to_tuple(nd)
ncomp = int(ncomp)
comp_shape = nd + (ncomp, )
test_mus = np.random.randn(*comp_shape)
test_taus = np.random.gamma(1, 1, size=comp_shape)
observed = generate_normal_mixture_data(w=np.ones(ncomp) / ncomp,
mu=test_mus,
sd=1 / np.sqrt(test_taus),
size=10)
with Model() as model0:
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp, ))
mixture0 = NormalMixture("m",
w=ws,
mu=mus,
tau=taus,
shape=nd,
comp_shape=comp_shape)
obs0 = NormalMixture("obs",
w=ws,
mu=mus,
tau=taus,
shape=nd,
comp_shape=comp_shape,
observed=observed)
with Model() as model1:
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp, ))
comp_dist = [
Normal.dist(mu=mus[..., i], tau=taus[..., i], shape=nd)
for i in range(ncomp)
]
mixture1 = Mixture("m", w=ws, comp_dists=comp_dist, shape=nd)
obs1 = Mixture("obs",
w=ws,
comp_dists=comp_dist,
shape=nd,
observed=observed)
with Model() as model2:
# Expected to fail if comp_shape is not provided,
# nd is multidim and it does not broadcast with ncomp. If by chance
# it does broadcast, an error is raised if the mixture is given
# observed data.
# Furthermore, the Mixture will also raise errors when the observed
# data is multidimensional but it does not broadcast well with
# comp_dists.
mus = Normal("mus", shape=comp_shape)
taus = Gamma("taus", alpha=1, beta=1, shape=comp_shape)
ws = Dirichlet("ws", np.ones(ncomp), shape=(ncomp, ))
if len(nd) > 1:
if nd[-1] != ncomp:
with pytest.raises(ValueError):
NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd)
mixture2 = None
else:
mixture2 = NormalMixture("m",
w=ws,
mu=mus,
tau=taus,
shape=nd)
else:
mixture2 = NormalMixture("m", w=ws, mu=mus, tau=taus, shape=nd)
observed_fails = False
if len(nd) >= 1 and nd != (1, ):
try:
np.broadcast(np.empty(comp_shape), observed)
except Exception:
observed_fails = True
if observed_fails:
with pytest.raises(ValueError):
NormalMixture("obs",
w=ws,
mu=mus,
tau=taus,
shape=nd,
observed=observed)
obs2 = None
else:
obs2 = NormalMixture("obs",
w=ws,
mu=mus,
tau=taus,
shape=nd,
observed=observed)
testpoint = model0.test_point
testpoint["mus"] = test_mus
testpoint["taus"] = test_taus
assert_allclose(model0.logp(testpoint), model1.logp(testpoint))
assert_allclose(mixture0.logp(testpoint), mixture1.logp(testpoint))
assert_allclose(obs0.logp(testpoint), obs1.logp(testpoint))
if mixture2 is not None and obs2 is not None:
assert_allclose(model0.logp(testpoint), model2.logp(testpoint))
if mixture2 is not None:
assert_allclose(mixture0.logp(testpoint), mixture2.logp(testpoint))
if obs2 is not None:
assert_allclose(obs0.logp(testpoint), obs2.logp(testpoint))
def draw_values(params, point=None, size=None):
"""
Draw (fix) parameter values. Handles a number of cases:
1) The parameter is a scalar
2) The parameter is an RV
a) parameter can be fixed to the value in the point
b) parameter can be fixed by sampling from the RV
c) parameter can be fixed using tag.test_value (last resort)
3) The parameter is a tensor variable/constant. Can be evaluated using
theano.function, but a variable may contain nodes which
a) are named parameters in the point
b) are RVs with a random method
"""
# The following check intercepts and redirects calls to
# draw_values in the context of sample_posterior_predictive
size = to_tuple(size)
ppc_sampler = vectorized_ppc.get(None)
if ppc_sampler is not None:
# this is being done inside new, vectorized sample_posterior_predictive
return ppc_sampler(params, trace=point, samples=size)
if point is None:
point = {}
# Get fast drawable values (i.e. things in point or numbers, arrays,
# constants or shares, or things that were already drawn in related
# contexts)
with _DrawValuesContext() as context:
params = dict(enumerate(params))
drawn = context.drawn_vars
evaluated = {}
symbolic_params = []
for i, p in params.items():
# If the param is fast drawable, then draw the value immediately
if is_fast_drawable(p):
v = _draw_value(p, point=point, size=size)
evaluated[i] = v
continue
name = getattr(p, "name", None)
if (p, size) in drawn:
# param was drawn in related contexts
v = drawn[(p, size)]
evaluated[i] = v
# We filter out Deterministics by checking for `model` attribute
elif name is not None and hasattr(p, "model") and name in point:
# param.name is in point
v = point[name]
evaluated[i] = drawn[(p, size)] = v
else:
# param still needs to be drawn
symbolic_params.append((i, p))
if not symbolic_params:
# We only need to enforce the correct order if there are symbolic
# params that could be drawn in variable order
return [evaluated[i] for i in params]
# Distribution parameters may be nodes which have named node-inputs
# specified in the point. Need to find the node-inputs, their
# parents and children to replace them.
leaf_nodes, named_nodes_descendents, named_nodes_ancestors = build_named_node_tree(
(param for _, param in symbolic_params if hasattr(param, "name"))
)
# Init givens and the stack of nodes to try to `_draw_value` from
givens = {
p.name: (p, v) for (p, size), v in drawn.items() if getattr(p, "name", None) is not None
}
stack = list(leaf_nodes.values())
while stack:
next_ = stack.pop(0)
if (next_, size) in drawn:
# If the node already has a givens value, skip it
continue
elif isinstance(next_, (theano_constant, tt.sharedvar.SharedVariable)):
# If the node is a theano.tensor.TensorConstant or a
# theano.tensor.sharedvar.SharedVariable, its value will be
# available automatically in _compile_theano_function so
# we can skip it. Furthermore, if this node was treated as a
# TensorVariable that should be compiled by theano in
# _compile_theano_function, it would raise a `TypeError:
# ('Constants not allowed in param list', ...)` for
# TensorConstant, and a `TypeError: Cannot use a shared
# variable (...) as explicit input` for SharedVariable.
# ObservedRV and MultiObservedRV instances are ViewOPs
# of TensorConstants or SharedVariables, we must add them
# to the stack or risk evaluating deterministics with the
# wrong values (issue #3354)
stack.extend(
[
node
for node in named_nodes_descendents[next_]
if isinstance(node, (ObservedRV, MultiObservedRV))
and (node, size) not in drawn
]
)
continue
else:
# If the node does not have a givens value, try to draw it.
# The named node's children givens values must also be taken
# into account.
children = named_nodes_ancestors[next_]
temp_givens = [givens[k] for k in givens if k in children]
try:
# This may fail for autotransformed RVs, which don't
# have the random method
value = _draw_value(next_, point=point, givens=temp_givens, size=size)
givens[next_.name] = (next_, value)
drawn[(next_, size)] = value
except theano.gof.fg.MissingInputError:
# The node failed, so we must add the node's parents to
# the stack of nodes to try to draw from. We exclude the
# nodes in the `params` list.
stack.extend(
[
node
for node in named_nodes_descendents[next_]
if node is not None and (node, size) not in drawn
]
)
# the below makes sure the graph is evaluated in order
# test_distributions_random::TestDrawValues::test_draw_order fails without it
# The remaining params that must be drawn are all hashable
to_eval = set()
missing_inputs = {j for j, p in symbolic_params}
while to_eval or missing_inputs:
if to_eval == missing_inputs:
raise ValueError(
"Cannot resolve inputs for {}".format(
[get_var_name(params[j]) for j in to_eval]
)
)
to_eval = set(missing_inputs)
missing_inputs = set()
for param_idx in to_eval:
param = params[param_idx]
if (param, size) in drawn:
evaluated[param_idx] = drawn[(param, size)]
else:
try: # might evaluate in a bad order,
# Sometimes _draw_value recurrently calls draw_values.
# This may set values for certain nodes in the drawn
# dictionary, but they don't get added to the givens
# dictionary. Here, we try to fix that.
if param in named_nodes_ancestors:
for node in named_nodes_ancestors[param]:
if node.name not in givens and (node, size) in drawn:
givens[node.name] = (node, drawn[(node, size)])
value = _draw_value(param, point=point, givens=givens.values(), size=size)
evaluated[param_idx] = drawn[(param, size)] = value
givens[param.name] = (param, value)
except theano.gof.fg.MissingInputError:
missing_inputs.add(param_idx)
return [evaluated[j] for j in params] # set the order back
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 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
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
