Dynestyx: A Probabilistic Programming Library for Dynamical Systems

Abstract: State-space models (SSMs) are the standard formalism for Bayesian treatment of dynamical systems, with natural applications in statistics, signal processing, and machine learning. Despite their importance in both theory and application, dynamical systems have proven difficult to incorporate in modern probabilistic programming languages (PPLs), making state-of-the-art methods less accessible to practitioners and introducing friction in following the "Bayesian workflow." We introduce dynestyx, a probabilistic programming library with first-class support for SSMs, including state-of-the-art methods in the estimation of both states and parameters. Through a single, unified interface, users may specify arbitrary priors for discrete-time or continuous-time dynamical systems, perform inference over mixed-effect data, and make state and parameter estimates with principled uncertainty quantification.

• The paper introduces Dynestyx to integrate flexible Bayesian inference for state-space models within a composable probabilistic programming framework.
• It employs a modular architecture that separates model specification, likelihood estimation, and filtering methods to enable rapid prototyping and systematic comparisons.
• The framework supports a wide array of inference techniques, including particle filters and gradient-informed MCMC, for robust dynamical system identification.

Dynestyx: A Probabilistic Programming Library for Dynamical Systems

Introduction and Motivation

State-space models (SSMs) remain foundational for Bayesian modeling of dynamical systems, with prominent applications across signal processing, statistics, neuroscience, and general machine learning. Despite a rich history in SSM algorithms and theoretical developments, there has been a significant gap in integrating flexible, structure-exploiting SSM inference into general-purpose probabilistic programming languages (PPLs). This disconnect has traditionally forced researchers and practitioners to choose between ease-of-use and advanced computational techniques, with state-of-the-art inference methods distributed across disparate bespoke codebases.

"Dynestyx" (2606.16985) addresses this gap by introducing a dedicated PPL library that directly targets the needs of practitioners studying or deploying SSMs. It extends NumPyro's infrastructure to provide native, composable inference for both discrete- and continuous-time dynamical systems, supporting a modular "separation of concerns" for model specification, marginal likelihood estimation, and Bayesian inference. This separation facilitates systematic method comparison, supports rapid prototyping, and significantly reduces friction in the Bayesian workflow for dynamical systems.

Mathematical and Computational Framework

Dynestyx models latent state dynamics following the standard SSM formalism, with latent (typically Markovian) state variables, flexible prior specification, and the ability to include both discrete-time and continuous-time (SDE-driven) dynamics. It provides support for general model formulations:

• Discrete-time state evolution: $x_{t+1} \sim p(x_{t+1} \mid x_t, u_t, t)$.
• Continuous-time SDEs: $dx_t = f(x_t, u_t, t)\,dt + L(x_t,u_t,t)\,dB_t$, where $B_t$ is a multivariate Brownian motion.

Observation processes are specified as $y_t \sim p(y_t \mid x_t, u_t, t)$, with parameters for dynamics and observation likelihoods. Posterior inference in this setting is nontrivial, given that filtering and smoothing distributions are generally intractable for nonlinear/non-Gaussian systems. System identification, i.e., joint inference of dynamics, observation, and possible control parameters, proceeds via intractable marginal likelihoods—here tackled using a wide variety of filtering and Monte Carlo approximations.

Dynestyx emphasizes a strict modular interface, where:

• Model specification is entirely separated from choice of filtering/smoothing algorithm and the parameter inference scheme.
• Inference algorithms (e.g., SMC/particle filtering, Kalman filtering, ensemble approaches, MCMC, variational methods) can be interchanged without altering model code.
• Effect handlers provide a flexible mechanism for interpreting stochastic simulation primitives, supporting integration with NumPyro's composable inference handlers.

As such, users can build SSM models in idiomatic JAX/NumPyro code, run inference via any combination of implemented algorithms, and obtain uncertainty-quantified forecasts or smoothed state/parameter estimates.

Implemented Methods and Novel Combinations

A major contribution of dynestyx lies in its expansive support for both classical and state-of-the-art SSM inference algorithms. Through integration with dynamax, cd-dynamax, and cuthbert, it supports:

• Kalman filter and Rauch–Tung–Striebel (RTS) smoothing (linear/Gaussian models)
• Extended and unscented Kalman filtering for nonlinear systems
• Ensemble Kalman filtering and smoothing, including differentiable ensemble methods
• Particle filters and smoothers (including differentiable particle filtering)
• Support for both discrete- and continuous-discrete SSMs via modular discretization

For parameter inference, dynestyx models are compatible with:

• Standard NumPyro HMC/NUTS, variational inference (SVI, SVI with intractable likelihoods), and SG-MCMC techniques
• Gradient-informed MCMC algorithms such as MALA and Particle HMC

The library explicitly maps possible combinations of marginal likelihood estimators and parameter inference algorithms, documenting which are implemented natively and identifying several that (to the best of current knowledge) have not previously appeared in the literature. Notably, combinations such as continuous-discrete Ensemble Kalman filtering with gradient-informed MCMC are listed as novel.

Additionally, dynestyx supports mixed-effect dynamical systems via a "plate" primitive—enabling the modeling of individual heterogeneity in population-level SSMs and supporting applications such as hierarchical Bayesian modeling in neuroscience or systems biology.

Implications and Potential Impact

By exposing a unified, composable interface for SSM modeling and inference, dynestyx simplifies the practical workflow for Bayesian dynamical system identification. This lowers barriers for adopting principled Bayesian techniques in scientific and engineering domains where bespoke code or restricted inference methods have previously dominated. It also enables extensive empirical comparison and benchmarking of inference methodologies in complex SSMs, fostering reproducibility and transparency.

The modular orthogonalization of model, likelihood estimation, and parameter inference methods provides a natural testbed for future research, supporting new developments in:

• Differentiable filtering and smoothing for use with gradient-based Bayesian methods
• Hierarchical and mixed-effects SSMs in applied domains
• Systematic evaluation of computational and statistical trade-offs across algorithms (e.g., bias-variance, computation cost)
• Automated methodology selection or meta-learning for SSM inference

The design leverages composable algebraic effects, offering extensibility for future algorithmic techniques—e.g., learning-based proposal mechanisms, neural filters, automated discretization, or meta-inference strategies.

Conclusion

Dynestyx (2606.16985) establishes a rigorous foundation for Bayesian inference in dynamical systems within a probabilistic programming paradigm. Its comprehensive support for model specification, state-of-the-art inference algorithms, and orthogonal composability addresses longstanding challenges in the deployment of modern SSM methodologies. This work opens new avenues for practical use of Bayesian dynamical modeling, facilitates innovation in SSM inference, and provides a reproducible framework for empirical evaluation and development of novel methodology in the analysis of time-evolving systems.