🌀 nonlinear sciences

Dynestyx: A Probabilistic Programming Library for Dynamical Systems

The paper introduces Dynestyx, a probabilistic programming library that provides first-class support for state-space models, enabling users to easily specify priors, perform inference on mixed-effect data, and obtain principled uncertainty quantification for both discrete and continuous dynamical systems.

Original authors: Daniel Waxman, Dmitry Batenkov, John Feser, Andy Zane, Eli Bingham, Youssef Marzouk, Matthew E. Levine

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Daniel Waxman, Dmitry Batenkov, John Feser, Andy Zane, Eli Bingham, Youssef Marzouk, Matthew E. Levine

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand a complex machine, like a car engine, but you can't see inside it. You can only see the speedometer, the fuel gauge, and hear the noise it makes (these are your observations). You also know that the engine has hidden parts moving inside (the latent state) and that you can press the gas pedal to change how it runs (the control).

Dynestyx is a new software tool designed to help scientists and engineers figure out exactly what is happening inside that hidden engine, even when the data is messy or incomplete.

Here is a breakdown of what the paper claims, using simple analogies:

1. The Problem: "Siloed" Tools

Before this tool, figuring out how these hidden machines worked was like trying to build a house where the person who designs the blueprints, the person who lays the bricks, and the person who paints the walls all spoke different languages and used different tools.

The Issue: Scientists had to write custom, one-off code for every new problem. If they wanted to test a new way of guessing the engine's internal state, they often had to rebuild their entire software from scratch. This made it hard to compare different methods or follow a standard "recipe" for discovery.

2. The Solution: A Universal Adapter

Dynestyx acts like a universal adapter or a "Swiss Army Knife" for these problems. It is built on top of a popular programming language called NumPyro.

The "Separation of Concerns": The authors designed Dynestyx so that you can clearly separate what the machine is (the model) from how you try to figure it out (the math).
- Analogy: Imagine you are writing a recipe (the model). With Dynestyx, you can write the recipe once, and then choose to cook it in a microwave, an oven, or on a grill (different math algorithms) without rewriting the recipe.

3. How It Works: The "Black Box" and the "Guessers"

The tool handles two main types of time-based systems:

Discrete Time: Like a video game where the world updates frame-by-frame.
Continuous Time: Like a real-life river flowing smoothly, described by complex physics equations.

Dynestyx lets you define the rules of your system (the "drift" and "diffusion," which are fancy words for how the system naturally moves and how much it jitters randomly). Once you define the rules, Dynestyx uses effect handlers (a clever programming trick) to swap in different "guessing" algorithms on the fly.

4. The "Mix-and-Match" Menu

The paper highlights a massive benefit: Combinatorial Power.

The Menu: You can take a specific way of estimating the hidden state (like a "Kalman Filter," which is a standard way to track moving objects) and pair it with a specific way of guessing the parameters (like "MCMC," which is a way to explore all possible answers).
The Innovation: The authors created a table (Table 1 in the paper) showing all the possible combinations.
- Some combinations already existed in other software.
- The "Plus" Signs (+): The paper claims that Dynestyx is the first to make certain combinations work. For example, pairing a specific type of "Ensemble Kalman Filter" with a "Gradient-Informed MCMC" algorithm was previously impossible or required custom code. Dynestyx makes this a standard, plug-and-play feature.

5. Handling "Mixed Effects"

The tool also handles Mixed-Effect Models.

Analogy: Imagine you are studying how different cars (individuals) behave. Some cars are fast, some are slow, but they all follow the same basic engine rules. Dynestyx can learn the general rules for all cars while also figuring out the specific quirks of each individual car, all in one go.

Summary of Claims

The paper does not claim to solve a specific medical or engineering crisis. Instead, it claims to solve a workflow problem:

Unified Interface: You can describe a dynamical system once in Python.
Interchangeable Methods: You can swap out the math algorithms (filters, samplers, optimizers) without changing your model description.
New Combinations: It unlocks new ways of combining these algorithms that haven't been easily available before, allowing researchers to systematically compare which "guessing" method works best for their specific problem.

In short, Dynestyx is a library that turns the messy, custom-coded world of tracking hidden systems into a clean, modular, and experiment-friendly environment.

Technical Summary: Dynestyx – A Probabilistic Programming Library for Dynamical Systems

Problem Statement
State-space models (SSMs) serve as the standard formalism for the Bayesian treatment of dynamical systems, with critical applications in neuroscience, signal processing, and machine learning. Despite their theoretical importance and practical utility, dynamical systems have proven difficult to integrate into modern probabilistic programming languages (PPLs). Consequently, state-of-the-art inference methods often remain confined to bespoke codebases, creating friction for practitioners attempting to follow a rigorous "Bayesian workflow." This separation hinders the systematic comparison of inference algorithms and limits the accessibility of advanced methods for estimating both latent states and model parameters.

Methodology
The authors introduce dynestyx, a probabilistic programming library built as an extension of NumPyro. Dynestyx is designed to provide first-class support for SSMs while adhering to the PPL principle of "separation of concerns."

Unified Interface: The library allows users to specify arbitrary priors for both discrete-time and continuous-time dynamical systems through a single, unified interface.
- Discrete-time: Defined by $x_{t+1} \sim p(x_{t+1} | x_t, u_t, t)$ .
- Continuous-time: Defined by stochastic differential equations (SDEs) $dx_t = f(x_t, u_t, t) dt + L(x_t, u_t, t) d\beta_t$ .
- Observations: Coupled to the state via $y_t \sim p(y_t | x_t, u_t, t)$ .
Orthogonalization of Components: Dynestyx decouples three distinct components:
1. Model Specification: A Python function defining the joint prior $p(\theta, \phi)$ and the dynamical structure (using DynamicalModel).
2. Marginal Likelihood (ML) Estimation: The user selects an inference algorithm (e.g., Kalman Filter, Particle Filter, Ensemble Kalman Filter) to approximate the intractable marginal likelihood $p(y_{1:T} | \theta, \phi)$ .
3. Parameter Inference: The resulting log-posterior is passed to standard NumPyro inference routines (e.g., HMC, NUTS, Variational Inference, SG-MCMC).
Implementation via Effect Handlers: The library utilizes effect handlers (inspired by Plotkin and Pretnar, 2009) to interpret the dsx.sample statement. This allows the model definition to remain agnostic to the specific numerical discretizer or filtering algorithm used, deferring all approximations to the inference phase.
Mixed-Effect Support: The library implements mixed-effect dynamical systems via a plate primitive, enabling inference over hierarchical data structures compatible with any combination of ML estimators and parameter inference methods.

Key Contributions

First-Class SSM Support in PPL: Dynestyx bridges the gap between complex dynamical system theory and modern probabilistic programming, making state-of-the-art methods accessible within a standard NumPyro workflow.
Combinatorial Flexibility: By orthogonalizing model specification, ML estimation, and parameter inference, dynestyx enables a vast design space of algorithmic combinations. Users can swap filtering algorithms (e.g., EKF, UKF, Particle Filters) and parameter inference methods (e.g., HMC, Variational Bayes) without rewriting the model code.
Novel Method Combinations: The library implements combinations of estimators and inference methods that, to the authors' knowledge, have not appeared in the literature. These include continuous-discrete SSMs paired with gradient-informed MCMC (e.g., CD EnKF with MALA) and various other hybrid configurations, as detailed in the paper's mapping table.
Integration with Existing Ecosystems: Dynestyx integrates with existing software libraries such as dynamax, cd-dynamax, and cuthbert to provide a wide range of state estimation algorithms, including Kalman filtering, RTS smoothing, and differentiable particle filtering.

Results and Capabilities
The paper demonstrates that dynestyx successfully unifies the specification of idealized data-generating processes with flexible inference strategies. Key capabilities include:

State and Parameter Estimation: The ability to perform inference over mixed-effect data and generate estimates for both latent states and parameters with principled uncertainty quantification.
Algorithmic Comparison: The framework facilitates the direct comparison of different inference method compositions, allowing practitioners to identify "rules of thumb" and pathologies in complex, flexible, or potentially mis-specified models.
Compatibility: The resulting models are fully compatible with generic JAX routines, including Variational Inference, HMC/NUTS, and stochastic gradient MCMC.

Significance
The authors position dynestyx as a tool that removes the friction currently associated with applying Bayesian workflows to dynamical systems. By providing a unified representation for SSMs, the library not only makes advanced methods more accessible to practitioners but also creates a natural testbed for methodological innovation. The ability to systematically compose and compare inference algorithms allows for the discovery of new insights regarding the behavior of complex models, potentially driving further advancements in the field of Bayesian dynamical systems.