Dynestyx: A probabilistic programming library for dynamical systems

Joint/ongoing work with Matt Levine, Matthieu Darcy, Youssef Marzouk, and others.

Basis Research Institute
MIT

MIT CSE Community Seminar

September 18, 2026

Who am I

  • Postdoctoral Fellow @ Basis Research Institute
  • Research Affiliate @ LIDS, working in the MIT UQ Group
  • PhD in Electrical Engineering (statistical signal processing) @ Stony Brook University in 2025
  • Broadly interested in Bayesian stats for dynamical systems
dynestyx
  • dynestyx is a probabilistic programming library for dynamical systems (DW et al., 2026)
  • It provides a large suite of state-of-the-art methods for Bayesian inference in dynamical systems
  • It does this in an interpretable, principled way that enables “Bayesian workflow”

Dynamical systems provide complex systems & structure

  • Dynamical systems are a natural language for physical and biological systems
  • Even with simple rules, the results are often complex
  • The rich structure that yields emergent behavior also yields pathologies

State space models as a formalism

  • State space models (SSMs) will be our formalism
  • SSMs consist of a state dynamics process
    • Let's write \(x_t \in \mathbb{R}^{d_x}\), where \(t\) can be discrete-time (difference equation) or continuous-time (differential equation)
    • In either case, we specify \(x_0 \sim p(x_0)\)
  • We then observe \(y_t \sim p(y_t \,|\, x_t)\)
    • For us, \(y_t \in \mathbb{R}^{d_y}\), always in discrete time

\[ x_t \sim \overbrace{p(x_t \mid x_{t-1})}^{\text{transition density}} \quad\text{or}\quad \mathrm{d}x_t = \overbrace{f(x_t, t)\,\mathrm{d}t}^{\text{drift}} + \overbrace{g(x_t, t)\,\mathrm{d}W_t}^{\text{diffusion}} \]

\[ y_t \sim p(y_t \mid x_t) \]

State space models as a formalism

  • State space models (SSMs) will be our formalism
  • SSMs consist of a state dynamics process
    • Let's write \(x_t \in \mathbb{R}^{d_x}\), where \(t\) can be discrete-time (difference equation) or continuous-time (differential equation)
    • In either case, we specify \(x_0 \sim p(x_0)\)
  • We then observe \(y_t \sim p(y_t \,|\, x_t)\)
    • For us, \(y_t \in \mathbb{R}^{d_y}\), always in discrete time
  • We may additionally parameterize the system by \(\textcolor{#2A9DF4}{\theta} \in \mathbb{R}^{d_{\textcolor{#2A9DF4}{\theta}}}\)

\[ x_t \sim \overbrace{p_{\textcolor{#2A9DF4}{\theta}}(x_t \mid x_{t-1})}^{\text{transition density}} \quad\text{or}\quad \mathrm{d}x_t = \overbrace{f_{\textcolor{#2A9DF4}{\theta}}(x_t, t)\,\mathrm{d}t}^{\text{drift}} + \overbrace{g_{\textcolor{#2A9DF4}{\theta}}(x_t, t)\,\mathrm{d}W_t}^{\text{diffusion}} \]

\[ y_t \sim p_{\textcolor{#2A9DF4}{\theta}}(y_t \mid x_t) \]

ODEs, HMMs, and other special cases

  • Our generic formulation encompasses other types of systems as special cases
    • Ordinary differential equations (ODEs) occur when the diffusion term is identically zero
    • Hidden Markov models (HMMs) occur when the state space \(x_t\) is discrete, i.e., \(x_t \in \mathbb{Z}_{H}\)
  • Some more general classes can also occasionally be written this way
    • We can discretize certain PDEs to an ODE or SDE in an appropriate basis

Different inference tasks

  • There are several different inference tasks we may consider
    • Filtering considers the state, given observations up to now: \(p(x_t \,|\, y_{1:t}, \textcolor{#2A9DF4}{\theta})\)
    • Smoothing considers the state, given observations past, present, and future: \(p(x_t \,|\, y_{1:T}, \textcolor{#2A9DF4}{\theta}),\ T \geq t\)
    • We can also forecast into the future, \(p(x_T \,|\, y_{1:t}, \textcolor{#2A9DF4}{\theta}),\ T \geq t\)
  • We can learn \(\textcolor{#2A9DF4}{\theta}\) (system identification): \(p(\textcolor{#2A9DF4}{\theta} \,|\, y_{1:T}) \propto p(\textcolor{#2A9DF4}{\theta})\, p(y_{1:T} \,|\, \textcolor{#2A9DF4}{\theta})\)
OU process timeline and prior over theta

The “probabilistic programming” part

  • dynestyx is a library that lets us perform each of these inference tasks
  • It is a probabilistic programming language (PPL)
    • This means it’s like a normal programming language, but the “variables” are distributions…
    • Or in our case, dynamical systems :)
  • We take a PPL “separation of concerns” seriously: model code describes a mathematical ideal, and all approximations come after
dynestyx

A dynestyx model from the ground up

  • Let’s write our first dynestyx model together; remember, we write down the idealized version
  • We’ll start by writing down the math, and piece-by-piece, see how this corresponds to code
def model(obs_times=None, obs_values=None, predict_times=None):
    dynamics = dsx.DynamicalModel(
\(x_0 \sim \mathcal{N}(0, 1)\)
        initial_condition=dist.MultivariateNormal(
            jnp.zeros(1), jnp.eye(1)),
\(\mathrm{d}x_t = \sin(x_t)\,\mathrm{d}t + \mathrm{d}W_t\)
        state_evolution=dsx.ContinuousTimeStateEvolution(
            drift=lambda x, u, t: jnp.sin(x),
            diffusion=dsx.ScalarDiffusion(1.0, bm_dim=1),
        ),
\(y_t \sim \mathcal{N}(x_t, 0.1^2)\)
        observation_model=dsx.LinearGaussianObservation(
            H=jnp.eye(1), R=0.1**2 * jnp.eye(1)),
    )
observe \(y_{1:T}\) at \(t_{1:T}\)
    dsx.sample("f", dynamics, obs_times=obs_times,
               obs_values=obs_values, predict_times=predict_times)

A dynestyx model from the ground up

def model(obs_times=None, obs_values=None, predict_times=None):
\(\textcolor{#2A9DF4}{\alpha} \sim \mathrm{HalfNormal}(0.5)\)
    alpha = numpyro.sample("alpha", dist.HalfNormal(0.5))
    dynamics = dsx.DynamicalModel(
\(x_0 \sim \mathcal{N}(0, 1)\)
        initial_condition=dist.MultivariateNormal(
            jnp.zeros(1), jnp.eye(1)),
\(\mathrm{d}x_t = \textcolor{#2A9DF4}{\alpha} \sin(x_t)\,\mathrm{d}t + \mathrm{d}W_t\)
        state_evolution=dsx.ContinuousTimeStateEvolution(
            drift=lambda x, u, t: alpha * jnp.sin(x),
            diffusion=dsx.ScalarDiffusion(1.0, bm_dim=1),
        ),
\(y_t \sim \mathcal{N}(x_t, 0.1^2)\)
        observation_model=dsx.LinearGaussianObservation(
            H=jnp.eye(1), R=0.1**2 * jnp.eye(1)),
    )
observe \(y_{1:T}\) at \(t_{1:T}\)
    dsx.sample("f", dynamics, obs_times=obs_times,
               obs_values=obs_values, predict_times=predict_times)
simulate \((\textcolor{#2A9DF4}{\alpha}, x, y) \sim p\)
with dsx.Simulator():
    sim = Predictive(model, num_samples=1)(key, predict_times=ts)

Forking paths of inference

θ x₀ x₁ x₂ xT y₁ y₂ yT ⋯

\[ p(\textcolor{#2A9DF4}{\theta} \mid y_{1:T}) \propto p(\textcolor{#2A9DF4}{\theta})\, p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) \]

← already hard! classical Bayesian inference

\[ p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) = \prod_{k=1}^{T} p(y_k \mid y_{1:k-1}, \textcolor{#2A9DF4}{\theta}) \]

← but where do these come from?

\[ p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) = \int \textcolor{#222222}{p(x_0)} \prod_{k=1}^{T} \textcolor{#00897B}{p_\theta(y_k \mid x_k)}\, \textcolor{#E64B35}{p_\theta(x_k \mid x_{k-1})}\, \mathrm{d}x_{0:T} \]

Forking paths of inference

Computation via “unrolling”

\[ p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) = \int \textcolor{#222222}{p(x_0)} \prod_{k=1}^{T} \textcolor{#00897B}{p_\theta(y_k \mid x_k)}\, \textcolor{#E64B35}{p_\theta(x_k \mid x_{k-1})}\, \mathrm{d}x_{0:T} \]

  • One option: we can sample from \(p(\theta, x_{0:T} \mid y_{1:T})\) jointly, then “drop” the \(x_t\) samples
  • This is a classic trick in Monte Carlo for marginalization!
    • e.g., it’s the one we use in slice sampling
  • This can work well for small problems
  • But as \(T\) gets large, the Monte Carlo problem is high-dimensional (hard!) and loses structure (less robust!)
  • In dynestyx, we call this a LatentPathBuilder
with LatentPathBuilder():
    with Discretizer(LocalLinearizationConfig()):
        mcmc = MCMC(NUTS(model), num_warmup=500,
                    num_samples=1_000)
        mcmc.run(key, obs_times, obs_values)
Posterior over alpha, latent path

Forking paths of inference

Computation via filtering

\[ p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) = \prod_{k=1}^{T} \int \textcolor{#00897B}{p_\theta(y_k \mid x_k)}\, p_\theta(x_k \mid y_{1:k-1})\, \mathrm{d}x_k \]

  • On the other hand, \(p_\theta(x_t \mid y_{1:t})\) is exactly what the filter computes
  • We can take expectations with respect to the one-step predictive density!
  • Thus, we use the filter for an approximate likelihood \(p(y_{1:T} \mid \theta)\)
  • We can pair this with MCMC like before
    • this can be exact marginalization, pseudo-marginal MCMC (Andrieu & Roberts, 2009; Andrieu et al., 2010), or just an approximation :)
with Filter(ContinuousTimeEKFConfig()):
    mcmc = MCMC(NUTS(model), num_warmup=500,
                num_samples=1_000)
    mcmc.run(key, obs_times, obs_values)
Posterior over alpha, latent path and filter

Forking paths of inference

Variational inference

\[ p(y_{1:T} \mid \textcolor{#2A9DF4}{\theta}) = \prod_{k=1}^{T} \int \textcolor{#00897B}{p_\theta(y_k \mid x_k)}\, p_\theta(x_k \mid y_{1:k-1})\, \mathrm{d}x_k \]

  • In either case, we can also use variational inference
  • Related to “variational Bayes with an intractable likelihood” (Tran et al., 2017)
with Filter(ContinuousTimeEKFConfig()):
    guide = AutoNormal(model)
    svi = SVI(model, guide, optax.adam(1e-2),
              loss=Trace_ELBO())
    result = svi.run(key, 2_000, obs_times, obs_values)
Posterior over alpha, latent path, filter, and variational

To approximate or not to approximate

An example with discretization

Continuous-time EKF

Euler–Maruyama

Local linearization

Ensemble Kalman filter

Variational inference

\(\textcolor{#2A9DF4}{\alpha} \sim \mathrm{HalfNormal}(0.5)\)

\(x_0 \sim \mathcal{N}(0, 1)\)

\(\mathrm{d}x_t = \textcolor{#2A9DF4}{\alpha} \sin(x_t)\,\mathrm{d}t + \mathrm{d}W_t\)

\(y_t \sim \mathcal{N}(x_t, 0.1^2)\)

\(x_{k+1} \mid x_k \sim \mathcal{N}\big(x_k + \textcolor{#2A9DF4}{\alpha} \sin(x_k)\,\Delta t,\; \Delta t\big)\)

\(J_k = \textcolor{#2A9DF4}{\alpha} \cos(x_k)\)

\(x_{k+1} \mid x_k \sim \mathcal{N}\Big(x_k + \tfrac{e^{J_k \Delta t} - 1}{J_k}\, \textcolor{#2A9DF4}{\alpha} \sin(x_k),\; \tfrac{e^{2 J_k \Delta t} - 1}{2 J_k}\Big)\)

with Filter(ContinuousTimeEKFConfig()):
    mcmc = MCMC(NUTS(model), num_warmup=500,
                num_samples=1_000)
    mcmc.run(key, obs_times, obs_values)
with Filter(EKFConfig()):
    with Discretizer(EulerMaruyamaConfig()):
        mcmc = MCMC(NUTS(model), num_warmup=500,
                    num_samples=1_000)
        mcmc.run(key, obs_times, obs_values)
with Filter(EKFConfig()):
    with Discretizer(LocalLinearizationConfig()):
        mcmc = MCMC(NUTS(model), num_warmup=500,
                    num_samples=1_000)
        mcmc.run(key, obs_times, obs_values)
with Filter(ContinuousTimeEnKFConfig()):
    mcmc = MCMC(NUTS(model), num_warmup=500,
                num_samples=1_000)
    mcmc.run(key, obs_times, obs_values)
with Filter(ContinuousTimeEnKFConfig()):
    guide = AutoNormal(model)
    svi = SVI(model, guide, optax.adam(1e-2),
              loss=Trace_ELBO())
    result = svi.run(key, 300,
                     obs_times, obs_values)
Posterior over alpha with a continuous-time EKF
Posterior over alpha, continuous-time EKF and Euler-Maruyama
Posterior over alpha, continuous-time EKF, Euler-Maruyama, and local linearization
Posterior over alpha with ensemble filters: continuous-time, Euler-Maruyama, and local linearization
Posterior over alpha from SVI: continuous-time, Euler-Maruyama, and local linearization

Can we use neural networks?

Yes.

\(\mathrm{d}x_t =\) θ \((x_t)\,\mathrm{d}t + \sigma\,\mathrm{d}W_t\)

\(y_t \mid x_t \sim \mathcal{N}\!\left(x_t,\,0.01 I_2\right)\)

layers = sample_layers()  # Gaussian priors
dynamics = dsx.ContinuousTimeStateEvolution(
    drift=lambda x, u, t: mlp(layers, x),
    diffusion=dsx.ScalarDiffusion(0.01, bm_dim=2),
)
True FHN drift and final neural-network mean drift near the training observations, shown with a common arrow scale

Can we do SINDy?

Yes, we can do (a Bayesian version of) SINDy (Brunton et al., 2016; Hirsh et al., 2022).

Drift valuesLibrary Sparse coefficients =· F(X)Φ(X)Θᵀ
theta = numpyro.sample("theta",
    dist.Laplace(0, .1).expand((2, 10)).to_event(2))
dynamics = dsx.ContinuousTimeStateEvolution(
    drift=lambda x, u, t: theta @ library(x),
    diffusion=dsx.ScalarDiffusion(.01, bm_dim=2))
guide = AutoMultivariateNormal(filtered_model, **guide_init)
Two horizontal interval plots with shared polynomial term labels: teal posterior means and 50, 80, and 95 percent intervals for all twenty coefficients, compared with black diamond true FHN coefficients
Held-out v and w filtering with truth, available observations and nested uncertainty bands, including two gaps with no observations

Can we write hybrid models?

Yes... even for PDEs :)

True system

\(\displaystyle \frac{\partial u}{\partial t}=\nu\frac{\partial^2u}{\partial x^2}+1.1u(1-u)\)

\(x\in(0,1),\quad u(t,0)=u(t,1)=0,\quad \nu=0.01\)

Hybrid model

\(\displaystyle \dot{\mathbf u}=\underbrace{\nu L_h\mathbf u}_{\text{known diffusion}}+\underbrace{\color{#00897B}{c_\theta(\mathbf u)}}_{\text{learned reaction}}\)

\(\mathbf y_k\mid\mathbf u(t_k)\sim\mathcal N\!\left(\mathbf u(t_k),\sigma_y^2 I\right)\)

  • We learn the missing term with a Gaussian process (GP)
    • Specifically, we use a Hilbert space GP (HSGP) approximation (Solin & Särkkä, 2020)

Can we write hybrid models?

Yes... even for PDEs :)

c = sample_hsgp()  # GP prior for the reaction
dynamics = dsx.ContinuousTimeStateEvolution(
    drift=lambda u, _, t: nu * Lh @ u + c(u))
guide = AutoMultivariateNormal(filtered_model, **guide_init)

Can we do hierarchical/mixed-effect inference?

Yes.

\(\boldsymbol\mu_i\mid\boldsymbol\mu_{\rm global},\boldsymbol\sigma\sim\mathcal N(\boldsymbol\mu_{\rm global},\operatorname{diag}(\boldsymbol\sigma^2))\)

\(d\mathbf x_{i,t}=A(\mathbf x_{i,t}-\boldsymbol\mu_i)\,dt+L\,d\mathbf W_{i,t}\)

\(\mathbf y_{i,k}\mid\mathbf x_{i,t_k}\sim\mathcal N(\mathbf x_{i,t_k},R)\)

with dsx.plate('individuals', N):
    mu = numpyro.sample('mu', group_prior)
    dynamics = make_ou(mu, mu0_global, sigma0)
    return dsx.sample('x', dynamics, **data)

Other things that can be done

Inference in hidden Markov models (with arbitrary missingness)

QR link to the hmm example

Evaluation via proper scoring rules

QR link to the scores example
Evaluation via proper scoring rules

Closed-loop & model-predictive control

QR link to the control example

More advanced PDEs (Coming soon...)

Ensemble filters with localization... learned via gradient descent

QR link to the localization example
Ensemble filters with localization... learned via gradient descent

There's lots left to do

  • dynestyx provides the state-of-the-art in inference for practitioners to use, helping to close a gap between theory and application
  • Our thesis: closing this gap will uncover a wide array of new methodological work
    • Applications will have their own pathologies, and we need methods that effectively detect and address them
  • We want to provide support to people who want to use dynestyx
    • Develop benchmarks based on challenges we see
    • Develop diagnostics that can detect pathologies
    • Ultimately, use our general framework + knowledge we gain to develop new methods

Questions?

danwaxman.github.io

dwaxman@mit.edu

dynestyx

Paper

GitHub