Handlers

dynestyx is built using effectful, which operates using a primitive called a handler. The details of this can be abstracted away from the typical user experience, but impacts the implementation of the sample primitive. The long story short is that the basic implementation of sample is empty, and it is actually "interpreted" by context. For hierarchical models with multiple trajectories, use plate together with NumPyro sampling inside the plate context. For example,

with Filter(EKFConfig()):
    dsx.sample("f", dynamical_model, obs_times=obs_times, obs_values=obs_values)

will implement the dsx.sample primitive using an extended Kalman filter. For more details, see the corresponding developer API page.

Contains the sample primitive and effectful utilities for dynestyx.

plate

Bases: ObjectInterpretation

Hierarchical plate for batched trajectories. Allows for multiple levels of hierarchy.

Wraps numpyro.plate for parameter sampling semantics and intercepts dsx.sample to inject plate_shapes into the handler chain via fwd(). Nested plate handlers run from inner to outer under effectful dispatch, so each handler appends its size to preserve NumPyro's effective batch order (inner plate is the leftmost data batch dim).

Examples:

>>> with dsx.plate("trajectories", M):
...     theta = numpyro.sample("theta", dist.Normal(0, 1))  # shape (M,)
...     dynamics = DynamicalModel(...)  # built from theta
...     dsx.sample("f", dynamics, obs_times=t, obs_values=y)

>>> with dsx.plate("groups", G):
...     beta = numpyro.sample("beta", dist.Normal(0, 1))  # shape (G,)
...     with dsx.plate("trajectories", M):
...         alpha = numpyro.sample("alpha", dist.Normal(beta, 1))  # shape (M, G)
...         dynamics = DynamicalModel(...)  # built from alpha
...         dsx.sample("f", dynamics, obs_times=t, obs_values=y)

Note

The dim argument is not currently supported for dynestyx plates.

__enter__()

Enter both numpyro.plate context and dynestyx plate interpretation.

__exit__(exc_type, exc, tb)

Exit both numpyro.plate context and dynestyx plate interpretation.

__init__(name: str, size: int, dim: int | None = None)

Initialize the plate handler.

Parameters:

Name	Type	Description	Default
name	str	Name of the plate.	required
size	int	Size of the plate.	required
dim	int | None	Dimension of the plate.	None

sample(name: str, dynamics: DynamicalModel, *, obs_times: jax.Array | None = None, obs_values: jax.Array | None = None, ctrl_times: jax.Array | None = None, ctrl_values: jax.Array | None = None, predict_times: jax.Array | None = None, **kwargs) -> FunctionOfTime

Samples from a dynamical model. This is the main primitive of dynestyx.

The sample primitive is meant to mimic the numpyro.sample primitive in usage, but using a DynamicalModel instead of a Distribution.

The sample method calls _sample_intp, which is defined as a defop in effectful. This is where any real "work" is done, after input validation.

Parameters:

Name	Type	Description	Default
name	str	Name of the sample site.	required
dynamics	DynamicalModel	Dynamical model to sample from.	required
obs_times	Array | None	Times at which to sample the observations.	None
obs_values	Array | None	Values of the observations at the given times.	None
ctrl_times	Array | None	Times at which to sample the controls.	None
ctrl_values	Array | None	Values of the controls at the given times.	None
predict_times	Array | None	Times at which to predict the observations.	None
**kwargs		Additional keyword arguments.	{}

Returns:

Name	Type	Description
FunctionOfTime	FunctionOfTime	A function of time that samples from the dynamical model.