cd-dynamax 0.2.5__tar.gz

cd_dynamax-0.2.5/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Iñigo Urteaga and Matthew Levine
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cd_dynamax-0.2.5/PKG-INFO

@@ -0,0 +1,228 @@
+Metadata-Version: 2.4
+Name: cd_dynamax
+Version: 0.2.5
+Summary: Continuous-discrete dynamical systems with JAX and related libraries.
+Author: Matthew Levine, Iñigo Urteaga
+Maintainer-email: Matthew Levine <matt@basis.ai>, Iñigo Urteaga <iurteaga@bcamath.org>
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: scipy>=1.13
+Requires-Dist: pandas>=2.2
+Requires-Dist: scikit-learn>=1.5
+Requires-Dist: matplotlib>=3.9
+Requires-Dist: seaborn>=0.13
+Requires-Dist: pillow>=10
+Requires-Dist: tqdm>=4.66
+Requires-Dist: jax>=0.6.2
+Requires-Dist: jaxlib>=0.6.2
+Requires-Dist: optax>=0.2.2
+Requires-Dist: equinox>=0.13
+Requires-Dist: diffrax==0.6.2
+Requires-Dist: lineax>=0.0.8
+Requires-Dist: flax>=0.11
+Requires-Dist: blackjax>=1.2
+Requires-Dist: jaxopt>=0.8
+Requires-Dist: chex>=0.1.88
+Requires-Dist: distrax>=0.1.3
+Requires-Dist: orbax-checkpoint>=0.9
+Requires-Dist: etils>=1.5
+Requires-Dist: dm-tree>=0.1.8
+Requires-Dist: graphviz
+Requires-Dist: ipykernel
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mkdocs-material>=9.7.1; extra == "dev"
+Requires-Dist: mkdocs<2,>=1.6.1; extra == "dev"
+Requires-Dist: mypy>=1.19.1; extra == "dev"
+Requires-Dist: mkdocs-jupyter>=0.25.1; extra == "dev"
+Requires-Dist: mkdocstrings-python; extra == "dev"
+Dynamic: license-file
+
+# Overview of cd-dynamax
+
+The primary goal of this codebase is to extend [dynamax](https://github.com/probml/dynamax) to a continuous-discrete (CD) state-space-modeling setting, that is, to problems where 
+
+- the underlying dynamics are continuous in time,
+- and measurements can arise at arbitrary (i.e., non-regular) discrete times.
+
+To address these gaps, `cd-dynamax` modifies `dynamax` to accept irregularly sampled data and implements classical algorithms for continuous-discrete filtering and smoothing.
+
+## Mathematical Framework: continuous-discrete state-space models
+
+In this repository, we build an expanded toolkit for filtering, forecasting and learning dynamical systems that underpin real-world messy time-series data.
+
+We move towards this goal by working with the following flexible mathematical setting:
+
+- We assume there exists a (possibly unknown) stochastic dynamical system of form
+
+$$dx(t) = f(x(t),t)dt + L(x(t),t) dw(t)$$
+
+where $x \in \mathbb{R}^{d_x}$, $x(0) \sim \mathcal{N}(\mu_0, \Sigma_0)$, $f$ a possibly time-dependent drift function, $L$ a possibly state and/or time-dependent diffusion coefficient, and $dw$ is the derivative of a $d_x$-dimensional Brownian motion with a covariance $Q$.
+
+- We assume data are available at arbitrary times $\\{t_k\\}_{k=1}^K$ and observed via a measurement process dictated by
+
+$$y(t) = h(x(t)) + \eta(t)$$
+
+where $h: \mathbb{R}^{d_x} \mapsto \mathbb{R}^{d_y}$ creates a $d_y$-dimensional observation from the $d_x$-dimensional state of the dynamical system $x(t)$ (a realization of the above SDE), and $\eta(t)$ applies additive Gaussian noise to the observation.
+
+We denote the collection of all parameters as $\theta = \\{f,\\  L,\\  \mu_0,\\  \Sigma_0,\\  L,\\  Q,\\  h,\\  \textrm{Law}(\eta) \\}$.
+
+Note:
+
+- We assume $\eta(t)$ i.i.d. w.r.t. $t$:
+    - This assumption places us in the *continuous (dynamics) - discrete (observation)* setting.
+    - If $\eta(t)$ had temporal correlations, we would likely adopt a mathematical setting that defines the observation process continuously in time via its own SDE.
+
+- Other extensions of the above paradigm include categorical state-spaces and non-additive observation noise distributions
+    - These can fit into our code framework (indeed, some are covered in `dynamax`), but have not been our focus.
+
+## cd-dynamax goals and approach
+
+For a given set of observations $Y_K = [y(t_1),\\ \dots ,\\ y(t_K)]$, we wish to:
+- Filter: estimate $x(t_K) \\ | \\ Y_K, \\ \theta$
+- Smooth: estimate $\\{x(t)\\}_t \\ | \\ Y_K, \\ \theta$
+- Predict: estimate $x(t > t_K)\\ |\\ Y_K, \\ \theta$
+- Infer parameters: estimate $\theta \\ |\\ Y_K$
+
+All of these problems are deeply interconnected.
+
+- In cd-dynamax, we enable filtering, smoothing, and parameter inference for a single system under multiple trajectory observations ($[Y^{(1)}, \\ \dots \\, \\ Y^{(N)}]$.
+   
+    - In these cases, we assume that each trajectory represents an independent realization of the same dynamics-data model, which we may be interested in learning, filtering, smoothing, or predicting.
+        - In the future, we would like to have options to perform hierarchical inference, where we assume that each trajectory came from a different, yet similar set of system-defining parameters $\theta^{(n)}$.
+
+- We implement such filtering/smoothing algorithms in an efficient, autodifferentiable framework.
+    - We enable usage of modern general-purpose tools for parameter inference (e.g., stochastic gradient descent, Hamiltonian Monte Carlo).
+
+- In cd-dynamax, we take onto the parameter inference case by relying on marginalizing out unobserved states $\\{x(t)\\}_t$
+    - this is a design choice of ours, other alternatives are possible.
+    - This marginalization is performed (approximately, in cases of non-linear dynamics) via filtering/smoothing algorithms.
+
+## Codebase description and status
+
+The `cd-dynamax` codebase extends the `dynamax` library to support continuous-discrete state space models, where observations are made at specified discrete times rather than at regular intervals.
+
+- We leverage [dynamax](https://github.com/probml/dynamax) code
+    - Currently, based on a local directory with [Dynamax release 0.1.5](https://github.com/probml/dynamax/releases/tag/0.1.5)
+
+- We have implemented the [`cd-dynamax` codebase](./cd_dynamax/README.md) to deal with [continuous-discrete linear and non-linear models](./cd_dynamax/src/README.md), along with several filtering and smoothing algorithms.
+
+- The codebase is organized into several key directories:
+```
+cd_dynamax/
+├── src/                       # Source code for cd-dynamax library
+│   ├── continuous_discrete_linear_gaussian_ssm/  # CD-LGSSM models and algorithms
+│   ├── continuous_discrete_nonlinear_gaussian_ssm/ # CD-NLGSSM models and algorithms
+│   ├── ssm_temissions.py      # Modified SSM class for discrete emissions
+│   └── utils/               # Utility functions and example models
+├── dynamax/                     # Original dynamax library (as a submodule)
+demos/                       # Python demos showcasing cd-dynamax functionality
+├── python/scripts/          # Python scripts for running demos
+├── python/notebooks/        # Jupyter notebooks for interactive demos
+├── python/configs/          # Configuration files for demos
+tests/                       # Tests for cd-dynamax functionality
+```
+
+## [Demos](./demos/python)
+
+We provide a set of [demos](./demos/python/README.md) that showcase key functionality of `cd-dynamax`.
+
+These [scripts](./demos/python/scripts) and [notebooks](./demos/python/notebooks) illustrate how to learn components of continuous-discrete SDEs from data.
+
+For instance:
+
+- [Filtering-based likelihood tutorial](./demos/python/notebooks/lorenz63_filter_based_likelihood_tutorial.ipynb) to filtering-based likelihood computation for continuous-discrete SDEs.
+
+- [SGD-based model fitting tutorial](./demos/python/notebooks/lorenz63_sgd_fit_to_data_tutorial.ipynb) to SGD-based fitting of continuous-discrete SDE model to data.
+
+- [MCMC-based model fitting tutorial](./demos/python/notebooks/lorenz63_mcmc_fit_to_data_tutorial.ipynb) to MCMC-based fitting of continuous-discrete SDE model to data.
+
+## [Tests](./tests)
+
+- Several [tests](./tests/README.md) to establish cd-dynamax general functionality, as well as linear and non-linear filters/smoothers tests: e.g., checks that non-linear algorithms applied to linear problems return similar results as linear algorithms.
+
+## [Makefile](./Makefile)
+
+- We provide a [Makefile](./Makefile) to automate common tasks, such as running tests and demos.
+
+- To run all tests, simply execute:
+```bash
+make test
+```
+
+- For linting, we use `ruff`:
+```bashbash
+make lint
+```
+
+- We can also format files using `ruff`:
+```bash
+make clean
+```
+
+- The docs can be built using `mkdocs` as:
+```bash
+make build_docs
+```
+
+# Installation
+
+We support installation via **Conda** (recommended) or via a standard Python virtual environment.
+
+---
+
+### Option 1: Conda (recommended)
+
+```bash
+# Create and activate a new environment with Python 3.11
+conda create -n cd_dynamax_joss python=3.11
+conda activate cd_dynamax_joss
+
+# Install your package in editable mode (so local changes are picked up)
+pip install -e .[dev]
+```
+
+This installs the core dependencies listed in `pyproject.toml`, along with optional developer tools (`pytest`, etc.) if you use `[dev]`.
+
+---
+
+### Option 2: Python venv + pip
+
+```bash
+# Create and activate a virtual environment
+python -m venv .venv
+source .venv/bin/activate   # on macOS/Linux
+.venv\Scripts\activate      # on Windows
+
+# Upgrade pip
+pip install --upgrade pip
+
+# Install in editable mode
+pip install -e .[dev]
+```
+
+#### GPU support
+If you want GPU acceleration with JAX, you must install a CUDA-enabled `jaxlib` wheel.  
+
+Check the [JAX installation docs](https://jax.readthedocs.io/en/latest/installation.html#installation) for the exact commands for your system.
+
+---
+
+### Notes
+
+- `pip install -e .` puts the repo in *editable mode*, so changes to source code are immediately available without reinstalling.
+
+- If you plan to use plotting features that rely on `graphviz`, make sure the system binary is installed:
+  - **macOS:** `brew install graphviz`  
+  - **Ubuntu/Debian:** `sudo apt install graphviz`  
+  - **Windows (conda):** `conda install graphviz`
+  
+- The `[dev]` extra installs additional developer tools (like `pytest`).
+    - Once your environment is installed, you can run automated tests:
+    ```bash
+    pytest
+    ```
+

cd_dynamax-0.2.5/README.md

@@ -0,0 +1,185 @@
+# Overview of cd-dynamax
+
+The primary goal of this codebase is to extend [dynamax](https://github.com/probml/dynamax) to a continuous-discrete (CD) state-space-modeling setting, that is, to problems where 
+
+- the underlying dynamics are continuous in time,
+- and measurements can arise at arbitrary (i.e., non-regular) discrete times.
+
+To address these gaps, `cd-dynamax` modifies `dynamax` to accept irregularly sampled data and implements classical algorithms for continuous-discrete filtering and smoothing.
+
+## Mathematical Framework: continuous-discrete state-space models
+
+In this repository, we build an expanded toolkit for filtering, forecasting and learning dynamical systems that underpin real-world messy time-series data.
+
+We move towards this goal by working with the following flexible mathematical setting:
+
+- We assume there exists a (possibly unknown) stochastic dynamical system of form
+
+$$dx(t) = f(x(t),t)dt + L(x(t),t) dw(t)$$
+
+where $x \in \mathbb{R}^{d_x}$, $x(0) \sim \mathcal{N}(\mu_0, \Sigma_0)$, $f$ a possibly time-dependent drift function, $L$ a possibly state and/or time-dependent diffusion coefficient, and $dw$ is the derivative of a $d_x$-dimensional Brownian motion with a covariance $Q$.
+
+- We assume data are available at arbitrary times $\\{t_k\\}_{k=1}^K$ and observed via a measurement process dictated by
+
+$$y(t) = h(x(t)) + \eta(t)$$
+
+where $h: \mathbb{R}^{d_x} \mapsto \mathbb{R}^{d_y}$ creates a $d_y$-dimensional observation from the $d_x$-dimensional state of the dynamical system $x(t)$ (a realization of the above SDE), and $\eta(t)$ applies additive Gaussian noise to the observation.
+
+We denote the collection of all parameters as $\theta = \\{f,\\  L,\\  \mu_0,\\  \Sigma_0,\\  L,\\  Q,\\  h,\\  \textrm{Law}(\eta) \\}$.
+
+Note:
+
+- We assume $\eta(t)$ i.i.d. w.r.t. $t$:
+    - This assumption places us in the *continuous (dynamics) - discrete (observation)* setting.
+    - If $\eta(t)$ had temporal correlations, we would likely adopt a mathematical setting that defines the observation process continuously in time via its own SDE.
+
+- Other extensions of the above paradigm include categorical state-spaces and non-additive observation noise distributions
+    - These can fit into our code framework (indeed, some are covered in `dynamax`), but have not been our focus.
+
+## cd-dynamax goals and approach
+
+For a given set of observations $Y_K = [y(t_1),\\ \dots ,\\ y(t_K)]$, we wish to:
+- Filter: estimate $x(t_K) \\ | \\ Y_K, \\ \theta$
+- Smooth: estimate $\\{x(t)\\}_t \\ | \\ Y_K, \\ \theta$
+- Predict: estimate $x(t > t_K)\\ |\\ Y_K, \\ \theta$
+- Infer parameters: estimate $\theta \\ |\\ Y_K$
+
+All of these problems are deeply interconnected.
+
+- In cd-dynamax, we enable filtering, smoothing, and parameter inference for a single system under multiple trajectory observations ($[Y^{(1)}, \\ \dots \\, \\ Y^{(N)}]$.
+   
+    - In these cases, we assume that each trajectory represents an independent realization of the same dynamics-data model, which we may be interested in learning, filtering, smoothing, or predicting.
+        - In the future, we would like to have options to perform hierarchical inference, where we assume that each trajectory came from a different, yet similar set of system-defining parameters $\theta^{(n)}$.
+
+- We implement such filtering/smoothing algorithms in an efficient, autodifferentiable framework.
+    - We enable usage of modern general-purpose tools for parameter inference (e.g., stochastic gradient descent, Hamiltonian Monte Carlo).
+
+- In cd-dynamax, we take onto the parameter inference case by relying on marginalizing out unobserved states $\\{x(t)\\}_t$
+    - this is a design choice of ours, other alternatives are possible.
+    - This marginalization is performed (approximately, in cases of non-linear dynamics) via filtering/smoothing algorithms.
+
+## Codebase description and status
+
+The `cd-dynamax` codebase extends the `dynamax` library to support continuous-discrete state space models, where observations are made at specified discrete times rather than at regular intervals.
+
+- We leverage [dynamax](https://github.com/probml/dynamax) code
+    - Currently, based on a local directory with [Dynamax release 0.1.5](https://github.com/probml/dynamax/releases/tag/0.1.5)
+
+- We have implemented the [`cd-dynamax` codebase](./cd_dynamax/README.md) to deal with [continuous-discrete linear and non-linear models](./cd_dynamax/src/README.md), along with several filtering and smoothing algorithms.
+
+- The codebase is organized into several key directories:
+```
+cd_dynamax/
+├── src/                       # Source code for cd-dynamax library
+│   ├── continuous_discrete_linear_gaussian_ssm/  # CD-LGSSM models and algorithms
+│   ├── continuous_discrete_nonlinear_gaussian_ssm/ # CD-NLGSSM models and algorithms
+│   ├── ssm_temissions.py      # Modified SSM class for discrete emissions
+│   └── utils/               # Utility functions and example models
+├── dynamax/                     # Original dynamax library (as a submodule)
+demos/                       # Python demos showcasing cd-dynamax functionality
+├── python/scripts/          # Python scripts for running demos
+├── python/notebooks/        # Jupyter notebooks for interactive demos
+├── python/configs/          # Configuration files for demos
+tests/                       # Tests for cd-dynamax functionality
+```
+
+## [Demos](./demos/python)
+
+We provide a set of [demos](./demos/python/README.md) that showcase key functionality of `cd-dynamax`.
+
+These [scripts](./demos/python/scripts) and [notebooks](./demos/python/notebooks) illustrate how to learn components of continuous-discrete SDEs from data.
+
+For instance:
+
+- [Filtering-based likelihood tutorial](./demos/python/notebooks/lorenz63_filter_based_likelihood_tutorial.ipynb) to filtering-based likelihood computation for continuous-discrete SDEs.
+
+- [SGD-based model fitting tutorial](./demos/python/notebooks/lorenz63_sgd_fit_to_data_tutorial.ipynb) to SGD-based fitting of continuous-discrete SDE model to data.
+
+- [MCMC-based model fitting tutorial](./demos/python/notebooks/lorenz63_mcmc_fit_to_data_tutorial.ipynb) to MCMC-based fitting of continuous-discrete SDE model to data.
+
+## [Tests](./tests)
+
+- Several [tests](./tests/README.md) to establish cd-dynamax general functionality, as well as linear and non-linear filters/smoothers tests: e.g., checks that non-linear algorithms applied to linear problems return similar results as linear algorithms.
+
+## [Makefile](./Makefile)
+
+- We provide a [Makefile](./Makefile) to automate common tasks, such as running tests and demos.
+
+- To run all tests, simply execute:
+```bash
+make test
+```
+
+- For linting, we use `ruff`:
+```bashbash
+make lint
+```
+
+- We can also format files using `ruff`:
+```bash
+make clean
+```
+
+- The docs can be built using `mkdocs` as:
+```bash
+make build_docs
+```
+
+# Installation
+
+We support installation via **Conda** (recommended) or via a standard Python virtual environment.
+
+---
+
+### Option 1: Conda (recommended)
+
+```bash
+# Create and activate a new environment with Python 3.11
+conda create -n cd_dynamax_joss python=3.11
+conda activate cd_dynamax_joss
+
+# Install your package in editable mode (so local changes are picked up)
+pip install -e .[dev]
+```
+
+This installs the core dependencies listed in `pyproject.toml`, along with optional developer tools (`pytest`, etc.) if you use `[dev]`.
+
+---
+
+### Option 2: Python venv + pip
+
+```bash
+# Create and activate a virtual environment
+python -m venv .venv
+source .venv/bin/activate   # on macOS/Linux
+.venv\Scripts\activate      # on Windows
+
+# Upgrade pip
+pip install --upgrade pip
+
+# Install in editable mode
+pip install -e .[dev]
+```
+
+#### GPU support
+If you want GPU acceleration with JAX, you must install a CUDA-enabled `jaxlib` wheel.  
+
+Check the [JAX installation docs](https://jax.readthedocs.io/en/latest/installation.html#installation) for the exact commands for your system.
+
+---
+
+### Notes
+
+- `pip install -e .` puts the repo in *editable mode*, so changes to source code are immediately available without reinstalling.
+
+- If you plan to use plotting features that rely on `graphviz`, make sure the system binary is installed:
+  - **macOS:** `brew install graphviz`  
+  - **Ubuntu/Debian:** `sudo apt install graphviz`  
+  - **Windows (conda):** `conda install graphviz`
+  
+- The `[dev]` extra installs additional developer tools (like `pytest`).
+    - Once your environment is installed, you can run automated tests:
+    ```bash
+    pytest
+    ```
+

cd_dynamax-0.2.5/cd_dynamax/__init__.py

@@ -0,0 +1,82 @@
+# cd_dynamax/__init__.py
+
+# Nonlinear SSM
+from .src.continuous_discrete_nonlinear_gaussian_ssm import (
+    ContDiscreteNonlinearGaussianSSM,
+    ParamsCDNLGSSM,
+    cdnlgssm_filter,
+    cdnlgssm_smoother,
+    cdnlgssm_forecast,
+    cdnlgssm_emissions,
+    EKFHyperParams,
+    UKFHyperParams,
+    EnKFHyperParams,
+)
+
+from .src.continuous_discrete_nonlinear_ssm import (
+    ContDiscreteNonlinearSSM,
+    ParamsCDNLSSM,
+    DPFHyperParams,
+    cdnlssm_filter,
+)
+
+# Linear SSM
+from .src.continuous_discrete_linear_gaussian_ssm import (
+    ContDiscreteLinearGaussianSSM,
+    ParamsCDLGSSM,
+    cdlgssm_filter,
+    cdlgssm_smoother,
+    cdlgssm_forecast,
+    cdlgssm_emissions,
+    cdlgssm_posterior_sample,
+    cdlgssm_joint_sample,
+    KFHyperParams,
+)
+
+# Discrete-Discrete Linear SSM
+from .dynamax.linear_gaussian_ssm import LinearGaussianSSM
+
+# Shared pieces
+from .src.ssm_temissions import SSM, Prior
+
+# Utilities (the ones your demos use most)
+from .src.utils.diffrax_utils import adjust_rhs
+from .src.utils.optimize_utils import make_optimizer
+from .src.utils.simulation_utils import make_key_sequence
+
+__all__ = [
+    # Models
+    "ContDiscreteNonlinearGaussianSSM",
+    "ContDiscreteNonlinearSSM",
+    "ContDiscreteLinearGaussianSSM",
+    "LinearGaussianSSM",
+    # Params
+    "ParamsCDNLGSSM",
+    "ParamsCDNLSSM",
+    "ParamsCDLGSSM",
+    # Nonlinear algos
+    "cdnlgssm_filter",
+    "cdnlgssm_smoother",
+    "cdnlgssm_forecast",
+    "cdnlgssm_emissions",
+    "cdnlssm_filter",
+    "EKFHyperParams",
+    "UKFHyperParams",
+    "EnKFHyperParams",
+    "DPFHyperParams",
+    # Linear algos
+    "cdlgssm_filter",
+    "cdlgssm_smoother",
+    "cdlgssm_forecast",
+    "cdlgssm_emissions",
+    "cdlgssm_posterior_sample",
+    "cdlgssm_joint_sample",
+    "KFHyperParams",
+    # SSM/emissions
+    "SSM",
+    "Prior",
+    # Utils
+    "adjust_rhs",
+    "make_optimizer",
+    "make_key_sequence",
+]

cd_dynamax-0.2.5/cd_dynamax/dynamax/__init__.py

@@ -0,0 +1,9 @@
+from . import _version
+__version__ = _version.get_versions()['version']
+
+# Catch expected warnings from TFP
+from . import warnings
+
+# Default to float32 matrix multiplication on TPUs and GPUs
+import jax
+jax.config.update('jax_default_matmul_precision', 'float32')