def __getitem__(self, idx):
if isinstance(idx, slice):
return self._slice(idx)
try:
return self.point(int(idx))
except (ValueError, TypeError): # Passed variable or variable name.
pass
if isinstance(idx, tuple):
var, vslice = idx
burn, thin = vslice.start, vslice.step
if burn is None:
burn = 0
if thin is None:
thin = 1
else:
var = idx
burn, thin = 0, 1
var = get_var_name(var)
if var in self.varnames:
if var in self.stat_names:
warnings.warn(
"Attribute access on a trace object is ambigous. "
"Sampler statistic and model variable share a name. Use "
"trace.get_values or trace.get_sampler_stats.")
return self.get_values(var, burn=burn, thin=thin)
if var in self.stat_names:
return self.get_sampler_stats(var, burn=burn, thin=thin)
raise KeyError("Unknown variable %s" % var)
def get_values(self, varname, burn=0, thin=1, combine=True, chains=None, squeeze=True):
"""Get values from traces.
Parameters
----------
varname: str
burn: int
thin: int
combine: bool
If True, results from `chains` will be concatenated.
chains: int or list of ints
Chains to retrieve. If None, all chains are used. A single
chain value can also be given.
squeeze: bool
Return a single array element if the resulting list of
values only has one element. If False, the result will
always be a list of arrays, even if `combine` is True.
Returns
-------
A list of NumPy arrays or a single NumPy array (depending on
`squeeze`).
"""
if chains is None:
chains = self.chains
varname = get_var_name(varname)
try:
results = [self._straces[chain].get_values(varname, burn, thin) for chain in chains]
except TypeError: # Single chain passed.
results = [self._straces[chains].get_values(varname, burn, thin)]
return _squeeze_cat(results, combine, squeeze)
def setup_class(cls):
super().setup_class()
cls.model = cls.make_model()
with cls.model:
cls.step = cls.make_step()
cls.trace = pm.sample(cls.n_samples, tune=cls.tune, step=cls.step, cores=cls.chains)
cls.samples = {}
for var in cls.model.unobserved_RVs:
cls.samples[get_var_name(var)] = cls.trace.get_values(var, burn=cls.burn)
def __init__(self, vars, shared, blocked=True):
"""
Parameters
----------
vars: list of sampling variables
shared: dict of Aesara variable -> shared variable
blocked: Boolean (default True)
"""
self.vars = vars
self.shared = {
get_var_name(var): shared
for var, shared in shared.items()
}
self.blocked = blocked
def _filter_parents(self, var, parents) -> Set[VarName]:
"""Get direct parents of a var, as strings"""
keep = set() # type: Set[VarName]
for p in parents:
if p == var:
continue
elif p.name in self.var_names:
keep.add(p.name)
elif p in self.transform_map:
if self.transform_map[p] != var.name:
keep.add(self.transform_map[p])
else:
raise AssertionError("Do not know what to do with {}".format(
get_var_name(p)))
return keep
def __init__(self, vars, shared, blocked=True):
"""
Parameters
----------
vars: list of sampling variables
shared: dict of theano variable -> shared variable
blocked: Boolean (default True)
"""
self.vars = vars
self.ordering = ArrayOrdering(vars)
self.shared = {
get_var_name(var): shared
for var, shared in shared.items()
}
self.blocked = blocked
self.bij = None
def setup_class(cls):
super().setup_class()
cls.model = cls.make_model()
with cls.model:
cls.step = cls.make_step()
cls.trace = pm.sample(
cls.n_samples,
tune=cls.tune,
step=cls.step,
cores=cls.chains,
return_inferencedata=False,
compute_convergence_checks=False,
)
cls.samples = {}
for var in cls.model.unobserved_RVs:
cls.samples[get_var_name(var)] = cls.trace.get_values(
var, burn=cls.burn)
def __getattr__(self, name):
# Avoid infinite recursion when called before __init__
# variables are set up (e.g., when pickling).
if name in self._attrs:
raise AttributeError
name = get_var_name(name)
if name in self.varnames:
if name in self.stat_names:
warnings.warn(
"Attribute access on a trace object is ambigous. "
"Sampler statistic and model variable share a name. Use "
"trace.get_values or trace.get_sampler_stats.")
return self.get_values(name)
if name in self.stat_names:
return self.get_sampler_stats(name)
raise AttributeError("'{}' object has no attribute '{}'".format(
type(self).__name__, name))
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 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 flat_t(var):
x = trace[get_var_name(var)]
return x.reshape((x.shape[0], np.prod(x.shape[1:], dtype=int)))
def __init__(self, var, idx, dpoint):
self.var = get_var_name(var)
self.idx = idx
self.dpt = dpoint
