def test_normal_context(self):
with _DrawValuesContext() as context0:
assert context0.parent is None
context0.drawn_vars['root_test'] = 1
with _DrawValuesContext() as context1:
assert id(context1.drawn_vars) == id(context0.drawn_vars)
assert context1.parent == context0
with _DrawValuesContext() as context2:
assert id(context2.drawn_vars) == id(context0.drawn_vars)
assert context2.parent == context1
context2.drawn_vars['leaf_test'] = 2
assert context1.drawn_vars['leaf_test'] == 2
context1.drawn_vars['root_test'] = 3
assert context0.drawn_vars['root_test'] == 3
assert context0.drawn_vars['leaf_test'] == 2
def test_mixed_contexts():
modelA = Model()
modelB = Model()
with raises((ValueError, TypeError)):
modelcontext(None)
with modelA:
with modelB:
assert Model.get_context() == modelB
assert modelcontext(None) == modelB
dvc = _DrawValuesContext()
with dvc:
assert Model.get_context() == modelB
assert modelcontext(None) == modelB
assert _DrawValuesContext.get_context() == dvc
dvcb = _DrawValuesContextBlocker()
with dvcb:
assert _DrawValuesContext.get_context() == dvcb
assert _DrawValuesContextBlocker.get_context() == dvcb
assert _DrawValuesContext.get_context() == dvc
assert _DrawValuesContextBlocker.get_context() is dvc
assert Model.get_context() == modelB
assert modelcontext(None) == modelB
assert _DrawValuesContext.get_context(error_if_none=False) is None
with raises(TypeError):
_DrawValuesContext.get_context()
assert Model.get_context() == modelB
assert modelcontext(None) == modelB
assert Model.get_context() == modelA
assert modelcontext(None) == modelA
assert Model.get_context(error_if_none=False) is None
with raises(TypeError):
Model.get_context(error_if_none=True)
with raises((ValueError, TypeError)):
modelcontext(None)
def random(self, point=None, size=None):
"""Sample from this distribution conditional on a given set of values.
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
"""
with _DrawValuesContext():
(states, ) = draw_values([self.states], point=point, size=size)
# This is a terrible thing to have to do here, but it's better than
# having to (know to) update `Distribution.shape` when/if dimensions
# change (e.g. when sampling new state sequences).
bcast_comps = np.broadcast(
states, *[dist.random(point=point) for dist in self.comp_dists])
self_shape = bcast_comps.shape
if size:
# `draw_values` will not honor the `size` parameter if its arguments
# don't contain random variables, so, when our `self.states` are
# constants, we have to broadcast `states` so that it matches `size +
# self.shape`.
expanded_states = np.broadcast_to(
states,
tuple(np.atleast_1d(size)) + self_shape)
else:
expanded_states = np.broadcast_to(states, self_shape)
samples = np.empty(expanded_states.shape)
for i, dist in enumerate(self.comp_dists):
# We want to sample from only the parts of our component
# distributions that are active given the states.
# This is only really relevant when the component distributions
# change over the state space (e.g. Poisson means that change
# over time).
# We could always sample such components over the entire space
# (e.g. time), but, for spaces with large dimension, that would
# be extremely costly and wasteful.
i_idx = np.where(expanded_states == i)
i_size = len(i_idx[0])
if i_size > 0:
subset_kwargs = distribution_subset_args(
dist, expanded_states.shape, i_idx)
state_dist = dist.dist(**subset_kwargs)
sample = state_dist.random(point=point)
samples[i_idx] = sample
return samples
def test_blocking_context(self):
with _DrawValuesContext() as context0:
assert context0.parent is None
context0.drawn_vars['root_test'] = 1
with _DrawValuesContext() as context1:
assert id(context1.drawn_vars) == id(context0.drawn_vars)
assert context1.parent == context0
with _DrawValuesContextBlocker() as blocker:
assert id(blocker.drawn_vars) != id(context0.drawn_vars)
assert blocker.parent is None
blocker.drawn_vars['root_test'] = 2
with _DrawValuesContext() as context2:
assert id(context2.drawn_vars) == id(blocker.drawn_vars)
assert context2.parent == blocker
context2.drawn_vars['root_test'] = 3
context2.drawn_vars['leaf_test'] = 4
assert blocker.drawn_vars['root_test'] == 3
assert 'leaf_test' not in context1.drawn_vars
assert context0.drawn_vars['root_test'] == 1
def random(self, point=None, size=None):
"""Sample from this distribution conditional on a given set of values.
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
"""
with _DrawValuesContext() as draw_context:
# TODO FIXME: Very, very lame...
term_smpl = draw_context.drawn_vars.get((self.states, 1), None)
if term_smpl is not None:
point[self.states.name] = term_smpl
# `draw_values` is inconsistent and will not use the `size`
# parameter if the variables aren't random variables.
if hasattr(self.states, "distribution"):
(states,) = draw_values([self.states], point=point, size=size)
else:
states = pm.Constant.dist(self.states).random(point=point, size=size)
# states = states.T
samples = np.empty(states.shape)
for i, dist in enumerate(self.comp_dists):
# We want to sample from only the parts of our component
# distributions that are active given the states.
# This is only really relevant when the component distributions
# change over the state space (e.g. Poisson means that change
# over time).
# We could always sample such components over the entire space
# (e.g. time), but, for spaces with large dimension, that would
# be extremely costly and wasteful.
i_idx = np.where(states == i)
i_size = len(i_idx[0])
if i_size > 0:
subset_args = distribution_subset_args(
dist, states.shape, i_idx, point=point
)
samples[i_idx] = dist.dist(*subset_args).random(point=point)
return samples
def random(self, point=None, size=None):
"""Sample from this distribution conditional on a given set of values.
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
"""
with _DrawValuesContext() as draw_context:
terms = [self.gamma_0, self.Gammas]
gamma_0, Gamma = draw_values(terms, point=point)
# Sample state 0 in each state sequence
state_n = pm.Categorical.dist(gamma_0, shape=self.shape[:-1]).random(
point=point, size=size
)
state_shape = state_n.shape
N = self.shape[-1]
states = np.empty(state_shape + (N,), dtype=self.dtype)
unif_samples = np.random.uniform(size=states.shape)
# Make sure we have a transition matrix for each element in a state
# sequence
Gamma = np.broadcast_to(Gamma, tuple(states.shape) + Gamma.shape[-2:])
# Slices across each independent/replication dimension
slices = [slice(None, d) for d in state_shape]
slices = tuple(np.ogrid[slices])
for n in range(0, N):
gamma_t = Gamma[..., n, :, :]
gamma_t = gamma_t[slices + (state_n,)]
state_n = vsearchsorted(gamma_t.cumsum(axis=-1), unif_samples[..., n])
states[..., n] = state_n
return states
def draw_values(self) -> list[np.ndarray]:
vars = self.vars
trace = self.trace
samples = self.samples
# size = self.size
params = dict(enumerate(vars))
with _DrawValuesContext() as context:
self.init()
self.make_graph()
drawn = context.drawn_vars
# Init givens and the stack of nodes to try to `_draw_value` from
givens = {
p.name: (p, v)
for (p, samples), v in drawn.items()
if getattr(p, "name", None) is not None
}
stack = list(
self.leaf_nodes.values()) # A queue would be more appropriate
while stack:
next_ = stack.pop(0)
if (next_, samples) in drawn:
# If the node already has a givens value, skip it
continue
elif isinstance(next_, (Constant, SharedVariable)):
# If the node is a aesara.tensor.TensorConstant or a
# aesara.tensor.sharedvar.SharedVariable, its value will be
# available automatically in _compile_aesara_function so
# we can skip it. Furthermore, if this node was treated as a
# TensorVariable that should be compiled by aesara in
# _compile_aesara_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 self.named_nodes_parents[next_]
if isinstance(node, (ObservedRV, MultiObservedRV)) and
(node, samples) 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 = self.named_nodes_children[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 = self.draw_value(next_,
trace=trace,
givens=temp_givens)
assert isinstance(value, np.ndarray)
givens[next_.name] = (next_, value)
drawn[(next_, samples)] = value
except aesara.graph.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 self.named_nodes_parents[next_]
if node is not None and (node,
samples) 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[int] = set()
missing_inputs: set[int] = {j for j, p in self.symbolic_params}
while to_eval or missing_inputs:
if to_eval == missing_inputs:
raise ValueError("Cannot resolve inputs for {}".format(
[get_var_name(trace.varnames[j]) for j in to_eval]))
to_eval = set(missing_inputs)
missing_inputs = set()
for param_idx in to_eval:
param = vars[param_idx]
drawn = context.drawn_vars
if (param, samples) in drawn:
self.evaluated[param_idx] = drawn[(param, samples)]
else:
try:
if param in self.named_nodes_children:
for node in self.named_nodes_children[param]:
if node.name not in givens and (
node, samples) in drawn:
givens[node.name] = (
node,
drawn[(node, samples)],
)
value = self.draw_value(param,
trace=self.trace,
givens=givens.values())
assert isinstance(value, np.ndarray)
self.evaluated[param_idx] = drawn[(
param, samples)] = value
givens[param.name] = (param, value)
except aesara.graph.fg.MissingInputError:
missing_inputs.add(param_idx)
return [self.evaluated[j] for j in params]
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
