calculate#

class pangolin.calculate.Calculate(sample_flat, *, frozen=(), **options)[source]#

A Calculate object just remembers a set of options and then offers inference methods.

The idea is that the user provides a sample_flat function. This can be any function that takes a triple of (vars, given, values) where vars and given are sequences of RV and values is a sequence of constants (arrays / floats) with the same length and shapes as given. The function then returns a list of arrays of the same length as vars but each has one more dimension at the beginning corresponding to the samples.

Given this, Calculate provides three conveniences:

It allows for binding default and/or frozen values for the sample_flat function.
It “lifts” sample_flat into a function sample that can run on arbitrary pytrees rather than lists.
It provides convenience functions E, var, std, sample_arviz that automatically take expectations and so on without manual intervention from the user.

Parameters:

sample_flat (callable) – function that performs inference.
frozen (iterable of str, optional) – names of options that cannot be overridden at call time
**options – default option values for sample_flat

sample(vars, given_vars=None, given_vals=None, reduce_fn=None, **options)[source]#

Draw samples!

Parameters:

vars (PyTree[pangolin.ir.RV]) – A RV or list/tuple of RV or pytree of RV to sample.
given_vars (PyTree[pangolin.ir.RV]) – A RV or list/tuple of RV or pytree of RV to condition on. None indicates no conditioning variables.
given_vals (PyTree[ArrayLike]) – An ArrayLike or list/tuple of ArrayLike or pytree of ArrayLike representing observed values. Must match the structure and shape of given_vars.
reduce_fn (Callable | None) – Function to apply to each leaf node in samples before returning. This is used to create E, var, etc. (If None, does nothing.)
options – extra options to pass to sampler

Returns:

Pytree of JAX arrays matching structure and shape of vars but with one extra dimension at the start, containing the samples.

Examples

>>> from pangolin.blackjax import sample_flat, run_nuts
>>> zero    = ir.RV(ir.Constant(0))
>>> one     = ir.RV(ir.Constant(1))
>>> x       = ir.RV(ir.Normal(), zero, one)
>>> y       = ir.RV(ir.Normal(), x, one)
>>> calc    = Calculate(sample_flat, run_inf=run_nuts, num_samples=529)
>>> x_samps = calc.sample(x,y,2)
>>> x_samps.shape
(529,)
>>> np.mean(x_samps) # something close to 1.0
Array(...)

E(vars, given_vars=None, given_vals=None, **options)[source]#

Compute (conditional) expected values. This is just a thin wrapper that calls sample and then reduces by taking the mean.

Parameters:

vars (PyTree[pangolin.ir.RV]) – A RV or list/tuple of RV or pytree of RV to sample.
given_vars (PyTree[pangolin.ir.RV]) – A RV or list/tuple of RV or pytree of RV to condition on. None indicates no conditioning variables.
given_vals (PyTree[ArrayLike]) – An ArrayLike or list/tuple of ArrayLike or pytree of ArrayLike representing observed values. Must match the structure and shape of given_vars.
reduce_fn – Function to apply to each leaf node in samples before returning. This is used to create E, var, etc. (If None, does nothing.)
options – extra options to pass to sampler

Returns:

Pytree of JAX arrays matching structure and shape of vars, containing: the expectations.

Examples

>>> from pangolin.blackjax import sample_flat, run_nuts
>>> zero    = ir.RV(ir.Constant(0))
>>> one     = ir.RV(ir.Constant(1))
>>> x       = ir.RV(ir.Normal(), zero, one)
>>> y       = ir.RV(ir.Normal(), x, one)
>>> calc    = Calculate(sample_flat, run_inf=run_nuts, num_samples=529)
>>> calc.E(x,y,2) # something close to 1.0
Array(...)

var(vars, given_vars=None, given_vals=None, **options)[source]#

Parameters:

vars (PyTree[pangolin.ir.RV])
given_vars (PyTree[pangolin.ir.RV])
given_vals (PyTree[ArrayLike])

std(vars, given_vars=None, given_vals=None, **options)[source]#

Parameters:

vars (PyTree[pangolin.ir.RV])
given_vars (PyTree[pangolin.ir.RV])
given_vals (PyTree[ArrayLike])

sample_arviz(vars, given_vars=None, given_vals=None, **options)[source]#

This is an experimental function to draw samples in ArviZ format.

Note: ArviZ is not installed with pangolin by default: You must install it manually.

Parameters:

vars (dict[str, RV]) – dictionary mapping names to individual random variables given_vars: A RV or list/tuple of RV or pytree of RV to condition on. None indicates no conditioning variables.
given_vars (PyTree[pangolin.ir.RV]) – A RV or list/tuple of RV or pytree of RV to condition on. given_vals: An ArrayLike or list/tuple of ArrayLike or pytree of ArrayLike representing observed values. Must match the structure and shape of given_vars.
reduce_fn – Function to apply to each leaf node in samples before returning. This is used to create E, var, etc. (If None, does nothing.)
options – extra options to pass to sampler
given_vals (PyTree[ArrayLike])

calculate

Contents

calculate#