def single_draw(elems):
values = dict(zip(trace_vars, elems))
_, st = evaluate_model_posterior_predictive(model,
values=values,
observed=observed)
return tuple([
(st.untransformed_values[k] if k in st.untransformed_values else
(st.deterministics[k]
if k in st.deterministics else st.transformed_values[k]))
for k in var_names
])
def sample_posterior_predictive(
model: ModelType,
trace: Dict[str, Any],
var_names: Optional[List[str]] = None,
observed: Optional[Dict[str, Any]] = None,
) -> Dict[str, np.ndarray]:
"""
Draw ``sample_shape`` values from the model for the desired ``var_names``.
Parameters
----------
model : types.GeneratorType, pymc4.Model
Model to draw samples from
trace: Dict[str, Any]
The samples drawn from the model's posterior distribution that should
be used for sampling from the posterior predictive
var_names: Optional[List[str]]
The list of variable names that will be included in the returned
samples. If ``None``, the samples drawn for all observed
distributions will be returned in the ``Samples`` dictionary.
observed : Optional[Dict[str, Any]]
A dictionary that can be used to override the distribution observed
values defined in the model.
Returns
-------
Samples: Dict[str, np.ndarray]
A dictionary of ``var_names`` keys and their corresponding drawn
samples.
Examples
--------
Lets define a simple model to sample from
>>> import pymc4 as pm
>>> @pm.model
... def model():
... sd = yield pm.HalfNormal("sd", 5.)
... norm = yield pm.Normal("n", 0, sd, observed=np.random.randn(100))
Now, we may want to draw samples from the model's posterior to then sample
from the posterior predictive.
>>> trace, stats = pm.inference.sampling.sample(model())
>>> ppc = pm.sample_posterior_predictive(model(), trace)
The samples are returned as a dictionary with the variable names as keys
>>> sorted(list(ppc))
['model/n']
The drawn values are the dictionary's values, and their shape will depend
on the supplied ``trace``
>>> ppc["model/n"].shape
(1000, 10, 100)
"""
# Get a copy of trace because we may manipulate the dictionary later in this
# function
trace = trace.copy()
if var_names is not None and len(var_names) == 0:
raise ValueError("Supplied an empty var_names list to sample from")
# We cannot assume that the model is vectorized, so we have batch the
# pm.evaluate_model_posterior_predictive calls across the trace entries
# This brings one big problem: we need to infer the batch dimensions from
# the trace. To do this, we will do
# 1) A regular single forward pass to determine the variable's shapes
# (we'll call these the core shapes)
# 2) Go through the supplied trace to get each variable's batch shapes
# (the shapes to the left of the core shapes)
# 3) Broadcast the encountered batch shapes between each other as a sanity
# check to get the global trace's batch_shape
# 4) Broadcast the values in the trace to the global batch_shape to get
# each variable's broadcasted value.
# 5) As tf.vectorized_map only iterates across the first dimension, we want
# to flatten the batch dimensions. To do this, we reshape the broadcasted
# values to (-1,) + core_shape. This way, tf.vectorized_map will be able
# to vectorize across the entire batch
# 6) Collect the samples from, reshape them to batch_shape + core_shape and
# return them
# Do a single forward pass to infer the distributions core shapes and
# default observeds
state = evaluate_model_posterior_predictive(model, observed=observed)[1]
if var_names is None:
var_names = list(state.posterior_predictives)
else:
defined_variables = set(state.all_values) | set(state.deterministics)
if not set(var_names) <= defined_variables:
raise KeyError(
"The supplied var_names = {} are not defined in the model.\n"
"Defined variables are = {}".format(
list(set(var_names) - defined_variables), list(defined_variables),
)
)
# Get the global batch_shape
batch_shape = tf.TensorShape([])
trace_names = list(trace)
for var_name in trace_names:
values = tf.convert_to_tensor(trace[var_name])
try:
core_shape = state.all_values[var_name].shape
except KeyError:
if var_name in state.deterministics:
# Remove the deterministics from the trace
del trace[var_name]
continue
else:
raise TypeError(
"Supplied the variable {} in the trace, yet this variable is "
"not defined in the model: {!r}".format(var_name, state)
)
assert_values_compatible_with_distribution_shape(var_name, values, core_shape)
batch_shape = tf.TensorShape(
tf.broadcast_static_shape(
values.shape[: len(values.shape) - len(core_shape)], batch_shape,
)
)
# Flatten the batch axis
flattened_trace = []
for k, v in trace.items():
core_shape = tf.TensorShape(state.all_values[k].shape)
batched_val = tf.broadcast_to(v, batch_shape + core_shape)
flattened_trace.append(tf.reshape(batched_val, shape=[-1] + core_shape.as_list()))
trace_vars = list(trace)
# Setup the function that makes a single draw
@tf.function(autograph=False)
def single_draw(elems):
values = dict(zip(trace_vars, elems))
_, st = evaluate_model_posterior_predictive(model, values=values, observed=observed)
return tuple(
[
(
st.untransformed_values[k]
if k in st.untransformed_values
else (
st.deterministics[k] if k in st.deterministics else st.transformed_values[k]
)
)
for k in var_names
]
)
# Make draws in parallel across the batch elements with tf.vectorized_map
samples = tf.vectorized_map(single_draw, flattened_trace)
# Convert the samples to ndarrays and make a dictionary with the correct
# batch_shape + core_shape
output = dict()
for name, sample in zip(var_names, samples):
sample = sample.numpy()
output[name] = np.reshape(sample, batch_shape + sample.shape[1:])
return output
def sample_posterior_predictive(
model: ModelType,
trace: InferenceData,
var_names: Optional[Union[str, List[str]]] = None,
observed: Optional[Dict[str, Any]] = None,
use_auto_batching: bool = True,
inplace: bool = True,
) -> InferenceData:
"""
Draw ``sample_shape`` values from the model for the desired ``var_names``.
Parameters
----------
model : types.GeneratorType, pymc4.Model
Model to draw samples from
trace: ArviZ's InferenceData object
The samples drawn from the model's posterior distribution that should be used for sampling
from the posterior predictive
var_names: Optional[Union[str, List[str]]]
The list of variable names that will be included in the returned samples. Strings can be
used to specify a single variable. If ``None``, the samples drawn for all observed
distributions will be returned in the ``Samples`` dictionary.
observed : Optional[Dict[str, Any]]
A dictionary that can be used to override the distribution observed values defined in the
model.
use_auto_batching: bool
A bool value that indicates whether ``sample_posterior_predictive`` should automatically
batch the draws or not. If you are sure you have manually tuned your model to be fully
vectorized, then you can set this to ``False``, and your sampling should be faster than the
auto batched counterpart. If you are not sure if your model is vectorized, then auto
batching will safely sample from it but with some additional overhead.
inplace: If True (default) it will add a posterior_predictive group to the provided ``trace``,
instead of returning a new InferenceData object. If a posterior_predictive group is already
present in ``trace`` it will be overwritten.
Returns
-------
Samples: InferenceDataType
An ArviZ's InferenceData object with a posterior_predictive group
Examples
--------
Lets define a simple model to sample from
>>> import pymc4 as pm
>>> @pm.model
... def model():
... sd = yield pm.HalfNormal("sd", 5.)
... norm = yield pm.Normal("n", 0, sd, observed=np.random.randn(100))
Now, we may want to draw samples from the model's posterior to then sample from the posterior
predictive.
>>> trace = pm.inference.sampling.sample(model())
>>> ppc = pm.sample_posterior_predictive(model(), trace).posterior_predictive
The samples are returned as a dictionary with the variable names as keys
>>> sorted(list(ppc))
['model/n']
The drawn values are the dictionary's values, and their shape will depend
on the supplied ``trace``
>>> ppc["model/n"].shape
(10, 1000, 100)
"""
if var_names is not None and len(var_names) == 0:
raise ValueError("Supplied an empty var_names list to sample from")
if isinstance(var_names, str):
var_names = [var_names]
# If we don't have to deal with auto-batching we can simply evaluate_model
# passing the trace as values
if not use_auto_batching:
values = {
var_name: tf.convert_to_tensor(value)
for var_name, value in trace.posterior.items()
}
# We need to pass the number of chains and draws as sample_shape for
# observed conditionally independent variables
sample_shape = (trace.posterior.sizes["chain"],
trace.posterior.sizes["draw"])
_, state = evaluate_model_posterior_predictive(
model, values=values, observed=observed, sample_shape=sample_shape)
all_values = collections.ChainMap(state.all_values,
state.deterministics_values)
if var_names is None:
var_names = list(state.posterior_predictives)
output = {k: all_values[k] for k in var_names}
return trace_to_arviz(trace=trace,
posterior_predictive=output,
inplace=inplace)
# We cannot assume that the model is vectorized, so we have batch the
# pm.evaluate_model_posterior_predictive calls across the trace entries
# This brings one big problem: we need to infer the batch dimensions from
# the trace. To do this, we will do
# 1) A single forward pass with the meta executor to determine the
# variable's shapes (we'll call these the core shapes)
# 2) Go through the supplied trace to get each variable's batch shapes
# (the shapes to the left of the core shapes)
# 3) Broadcast the encountered batch shapes between each other as a sanity
# check to get the global trace's batch_shape
# 4) Broadcast the values in the trace to the global batch_shape to get
# each variable's broadcasted value.
# 5) As tf.vectorized_map only iterates across the first dimension, we want
# to flatten the batch dimensions. To do this, we reshape the broadcasted
# values to (-1,) + core_shape. This way, tf.vectorized_map will be able
# to vectorize across the entire batch
# 6) Collect the samples from, reshape them to batch_shape + core_shape and
# return them
# Do a single forward pass to infer the distributions core shapes and
# default observeds
_, state = evaluate_meta_posterior_predictive_model(model,
observed=observed)
if var_names is None:
var_names = list(state.posterior_predictives)
else:
defined_variables = set(state.all_values) | set(
state.deterministics_values)
if not set(var_names) <= defined_variables:
raise KeyError(
"The supplied var_names = {} are not defined in the model.\n"
"Defined variables are = {}".format(
list(set(var_names) - defined_variables),
list(defined_variables)))
# Get the global batch_shape
batch_shape = tf.TensorShape([])
# Get a copy of trace because we may manipulate the dictionary later in this
# function
posterior = trace.posterior.copy() # type: ignore
posterior_names = list(posterior)
for var_name in posterior_names:
values = tf.convert_to_tensor(posterior[var_name].values)
try:
core_shape = state.all_values[var_name].shape
except KeyError:
if var_name in state.deterministics_values:
# Remove the deterministics from the trace
del posterior[var_name]
continue
else:
raise TypeError(
"Supplied the variable {} in the trace, yet this variable is "
"not defined in the model: {!r}".format(var_name, state))
assert_values_compatible_with_distribution_shape(
var_name,
values,
batch_shape=tf.TensorShape([]),
event_shape=core_shape)
batch_shape = tf.TensorShape(
tf.broadcast_static_shape(
values.shape[:len(values.shape) -
len(core_shape)], # type: ignore
batch_shape,
))
# Flatten the batch axis
flattened_posterior = []
for k, v in posterior.items():
core_shape = tf.TensorShape(state.all_values[k].shape)
batched_val = tf.broadcast_to(v.values, batch_shape + core_shape)
flattened_posterior.append(
tf.reshape(batched_val, shape=[-1] + core_shape.as_list()))
posterior_vars = list(posterior)
# Setup the function that makes a single draw
@tf.function(autograph=False)
def single_draw(elems):
values = dict(zip(posterior_vars, elems))
_, st = evaluate_model_posterior_predictive(model,
values=values,
observed=observed)
return tuple([
(st.untransformed_values[k] if k in st.untransformed_values else
(st.deterministics_values[k]
if k in st.deterministics_values else st.transformed_values[k]))
for k in var_names
])
# Make draws in parallel across the batch elements with tf.vectorized_map
samples = tf.vectorized_map(single_draw, flattened_posterior)
# Convert the samples to ndarrays and make a dictionary with the correct
# batch_shape + core_shape
output = dict()
for name, sample in zip(var_names, samples):
sample = sample.numpy()
output[name] = np.reshape(sample, batch_shape + sample.shape[1:])
return trace_to_arviz(trace=trace,
posterior_predictive=output,
inplace=inplace)