cd-dynamax 0.3.2__tar.gz → 0.3.3__tar.gz

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cd_dynamax
-Version: 0.3.2
+Version: 0.3.3
 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>
@@ -32,6 +32,7 @@ Requires-Dist: dm-tree>=0.1.8
 Requires-Dist: fastprogress>=1.0.0
 Requires-Dist: graphviz
 Requires-Dist: ipykernel
+Requires-Dist: orbax-checkpoint<0.11.3; sys_platform == "win32"
 Provides-Extra: dev
 Requires-Dist: pytest>=8.0; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
@@ -69,7 +70,7 @@ $$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) \\}$.
+We denote the collection of all parameters as $\theta = \\{f,\\  L,\\  \mu_0,\\  \Sigma_0,\\  Q,\\  h,\\  \textrm{Law}(\eta) \\}$.
 
 Note:
 
@@ -90,7 +91,7 @@ For a given set of observations $Y_K = [y(t_1),\\ \dots ,\\ y(t_K)]$, we wish to
 
 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 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)}$.
@@ -113,18 +114,19 @@ The `cd-dynamax` codebase extends the `dynamax` library to support continuous-di
 
 - 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
+.
+├── cd_dynamax/                  # Source code for cd-dynamax library
+│   ├── src/                     # Core source code
+│   │   ├── 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)
@@ -155,7 +157,7 @@ make test
 ```
 
 - For linting, we use `ruff`:
-```bashbash
+```bash
 make lint
 ```
 
@@ -171,40 +173,57 @@ make build_docs
 
 # Installation
 
-We support installation via **Conda** (recommended) or via a standard Python virtual environment.
+Install from **PyPI** (recommended), from source in editable mode, or with a Conda-managed environment.
 
 ---
 
-### Option 1: Conda (recommended)
+### Option 1: Install from PyPI (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
+# Create and activate a virtual environment
+python -m venv .venv        # or `uv venv`
+source .venv/bin/activate   # on macOS/Linux
+.venv\Scripts\activate      # on Windows
 
-# Install your package in editable mode (so local changes are picked up)
-pip install -e .[dev]
+# Upgrade pip
+pip install --upgrade pip
+
+# Install latest release from PyPI
+pip install cd-dynamax
 ```
 
-This installs the core dependencies listed in `pyproject.toml`, along with optional developer tools (`pytest`, etc.) if you use `[dev]`.
+`cd-dynamax` is currently **not available on Conda Forge**.
 
 ---
 
-### Option 2: Python venv + pip
+### Option 2: Install from source (editable)
 
 ```bash
 # Create and activate a virtual environment
-python -m venv .venv
+python -m venv .venv        # or `uv venv`
 source .venv/bin/activate   # on macOS/Linux
 .venv\Scripts\activate      # on Windows
 
 # Upgrade pip
 pip install --upgrade pip
 
-# Install in editable mode
+# Install in editable mode for local development
 pip install -e .[dev]
 ```
 
+---
+
+### Option 3: Conda environment + pip install
+
+```bash
+# Create and activate a Conda environment with Python 3.11
+conda create -n cd_dynamax python=3.11
+conda activate cd_dynamax
+
+# Install latest release from PyPI
+pip install cd-dynamax
+```
+
 #### GPU support
 If you want GPU acceleration with JAX, you must install a CUDA-enabled `jaxlib` wheel.  
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/README.md

@@ -25,7 +25,7 @@ $$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) \\}$.
+We denote the collection of all parameters as $\theta = \\{f,\\  L,\\  \mu_0,\\  \Sigma_0,\\  Q,\\  h,\\  \textrm{Law}(\eta) \\}$.
 
 Note:
 
@@ -46,7 +46,7 @@ For a given set of observations $Y_K = [y(t_1),\\ \dots ,\\ y(t_K)]$, we wish to
 
 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 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)}$.
@@ -69,18 +69,19 @@ The `cd-dynamax` codebase extends the `dynamax` library to support continuous-di
 
 - 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
+.
+├── cd_dynamax/                  # Source code for cd-dynamax library
+│   ├── src/                     # Core source code
+│   │   ├── 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)
@@ -111,7 +112,7 @@ make test
 ```
 
 - For linting, we use `ruff`:
-```bashbash
+```bash
 make lint
 ```
 
@@ -127,40 +128,57 @@ make build_docs
 
 # Installation
 
-We support installation via **Conda** (recommended) or via a standard Python virtual environment.
+Install from **PyPI** (recommended), from source in editable mode, or with a Conda-managed environment.
 
 ---
 
-### Option 1: Conda (recommended)
+### Option 1: Install from PyPI (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
+# Create and activate a virtual environment
+python -m venv .venv        # or `uv venv`
+source .venv/bin/activate   # on macOS/Linux
+.venv\Scripts\activate      # on Windows
 
-# Install your package in editable mode (so local changes are picked up)
-pip install -e .[dev]
+# Upgrade pip
+pip install --upgrade pip
+
+# Install latest release from PyPI
+pip install cd-dynamax
 ```
 
-This installs the core dependencies listed in `pyproject.toml`, along with optional developer tools (`pytest`, etc.) if you use `[dev]`.
+`cd-dynamax` is currently **not available on Conda Forge**.
 
 ---
 
-### Option 2: Python venv + pip
+### Option 2: Install from source (editable)
 
 ```bash
 # Create and activate a virtual environment
-python -m venv .venv
+python -m venv .venv        # or `uv venv`
 source .venv/bin/activate   # on macOS/Linux
 .venv\Scripts\activate      # on Windows
 
 # Upgrade pip
 pip install --upgrade pip
 
-# Install in editable mode
+# Install in editable mode for local development
 pip install -e .[dev]
 ```
 
+---
+
+### Option 3: Conda environment + pip install
+
+```bash
+# Create and activate a Conda environment with Python 3.11
+conda create -n cd_dynamax python=3.11
+conda activate cd_dynamax
+
+# Install latest release from PyPI
+pip install cd-dynamax
+```
+
 #### GPU support
 If you want GPU acceleration with JAX, you must install a CUDA-enabled `jaxlib` wheel.  
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_ekf.py

@@ -119,9 +119,9 @@ def _predict(
         f = params.dynamics.drift.f
 
         # Get time-varying parameters
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t)
+        Qc_t = params.dynamics.diffusion_cov.f(m, u, t)
         L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t)
+            params.dynamics.diffusion_coefficient.f(m, u, t)
             * filter_hyperparams.cov_rescaling
         )
 
@@ -185,9 +185,9 @@ def _predict(
             dt = filter_hyperparams.dt_average
 
         # Covariance parameters at time t0
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t0)
+        Qc_t = params.dynamics.diffusion_cov.f(m, u, t0)
         L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t0)
+            params.dynamics.diffusion_coefficient.f(m, u, t0)
             * filter_hyperparams.cov_rescaling
         )
         # Covariance update
@@ -527,12 +527,11 @@ def _smooth(
         m_smooth, P_smooth = y
         m_filter, P_filter = args
 
-        # TODO: possibly time- and parameter-dependent functions
         f = params.dynamics.drift.f
         # Get time-varying parameters
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t)
+        Qc_t = params.dynamics.diffusion_cov.f(m_filter, u, t)
         L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t)
+            params.dynamics.diffusion_coefficient.f(m_filter, u, t)
             * filter_hyperparams.cov_rescaling
         )
 
@@ -831,7 +830,7 @@ def extended_kalman_posterior_sample(
 
         # Get parameters and inputs for time t0
         u = inputs[t0_idx]
-        Q = params.dynamics.diffusion_cov.f(None, u, t0)
+        Q = params.dynamics.diffusion_cov.f(filtered_mean, u, t0)
 
         # Condition on next state
         smoothed_mean, smoothed_cov = _condition_on(

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_enkf.py

@@ -113,9 +113,9 @@ def _predict(
         # First order EnKF diffusion
         def diffusion(t, y, args):
             # Get parameters at time t and input u
-            Qc_t = params.dynamics.diffusion_cov.f(None, u, t)
+            Qc_t = params.dynamics.diffusion_cov.f(y, u, t)
             L_t = (
-                params.dynamics.diffusion_coefficient.f(None, u, t)
+                params.dynamics.diffusion_coefficient.f(y, u, t)
                 * filter_hyperparams.cov_rescaling
             )
             Q_sqrt = jnp.linalg.cholesky(Qc_t)
@@ -155,26 +155,20 @@ def _predict(
             # but the same amount of noise is added after each measurement.
             dt = filter_hyperparams.dt_average
 
-        # Get diffusion parameters at time t0 and input u
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t0)
-        L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t0)
-            * filter_hyperparams.cov_rescaling
-        )
+        key_array = jr.split(key, x.shape[0])
 
-        # Compute state noise covariance
-        state_noise_cov = dt * L_t @ Qc_t @ L_t.T  # D_hid x D_hid
+        def _sample_noise(x_i, key_i):
+            Qc_t = params.dynamics.diffusion_cov.f(x_i, u, t0)
+            L_t = (
+                params.dynamics.diffusion_coefficient.f(x_i, u, t0)
+                * filter_hyperparams.cov_rescaling
+            )
+            state_noise_cov = dt * L_t @ Qc_t @ L_t.T  # D_hid x D_hid
+            return jr.multivariate_normal(
+                key=key_i, mean=jnp.zeros(x.shape[1]), cov=state_noise_cov
+            )
 
-        # Noise realizations from MVN(0, state_noise_cov), for each particle particles (N x D_hid)
-        # Split keys for each particle
-        key_array = jr.split(key, x.shape[0])
-        # vmap over particles
-        noise = vmap(
-            lambda key: jr.multivariate_normal(
-                key=key, mean=jnp.zeros(x.shape[1]), cov=state_noise_cov
-            ),
-            in_axes=0,
-        )(key_array)
+        noise = vmap(_sample_noise, in_axes=(0, 0))(x_pred, key_array)
 
         # Add noise to predicted particles
         x_pred += noise

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_ukf.py

@@ -214,9 +214,9 @@ def _predict(
             dt = filter_hyperparams.dt_average
 
         # Get diffusion parameters at time t0 and input u
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t0)
+        Qc_t = params.dynamics.diffusion_cov.f(m, u, t0)
         L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t0)
+            params.dynamics.diffusion_coefficient.f(m, u, t0)
             * filter_hyperparams.cov_rescaling
         )
         # Compute state noise covariance
@@ -237,9 +237,9 @@ def _predict(
             f = (
                 params.dynamics.drift.f
             )  # TODO: reconsider when we want time-varying dynamics functions
-            Qc_t = params.dynamics.diffusion_cov.f(None, u, t)
+            Qc_t = params.dynamics.diffusion_cov.f(m_t, u, t)
             L_t = (
-                params.dynamics.diffusion_coefficient.f(None, u, t)
+                params.dynamics.diffusion_coefficient.f(m_t, u, t)
                 * filter_hyperparams.cov_rescaling
             )
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/models.py

@@ -111,8 +111,8 @@ def compute_pushforward(
         f = params.dynamics.drift.f
 
         # Get time-varying parameters
-        Qc_t = params.dynamics.diffusion_cov.f(None, inputs, t)
-        L_t = params.dynamics.diffusion_coefficient.f(None, inputs, t)
+        Qc_t = params.dynamics.diffusion_cov.f(x, inputs, t)
+        L_t = params.dynamics.diffusion_coefficient.f(x, inputs, t)
 
         # Different SDE approximations to the dynamics
         # Zeroth-order (no gradient information),
@@ -964,8 +964,8 @@ def cdnlgssm_path_sample(
             return params.dynamics.drift.f(y, inpt, t)
 
         def diffusion(t, y, args):
-            Qc_t = params.dynamics.diffusion_cov.f(None, inpt, t)
-            L_t = params.dynamics.diffusion_coefficient.f(None, inpt, t)
+            Qc_t = params.dynamics.diffusion_cov.f(y, inpt, t)
+            L_t = params.dynamics.diffusion_coefficient.f(y, inpt, t)
             Q_sqrt = jnp.linalg.cholesky(Qc_t)
             combined_diffusion = L_t @ Q_sqrt
             return combined_diffusion
@@ -1261,9 +1261,9 @@ def cdnlgssm_forecast(
                 return params.dynamics.drift.f(y, inputs[t0_idx], t)
 
             def diffusion(t, y, args):
-                Qc_t = params.dynamics.diffusion_cov.f(None, inputs[t0_idx], t)
+                Qc_t = params.dynamics.diffusion_cov.f(y, inputs[t0_idx], t)
                 Q_sqrt = jnp.linalg.cholesky(Qc_t)
-                L_t = params.dynamics.diffusion_coefficient.f(None, inputs[t0_idx], t)
+                L_t = params.dynamics.diffusion_coefficient.f(y, inputs[t0_idx], t)
                 combined_diffusion = L_t @ Q_sqrt
                 return combined_diffusion
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_ssm/inference_dpf.py

@@ -112,9 +112,9 @@ def _predict(
     else:
 
         def diffusion(t, y, args):
-            Qc_t = params.dynamics.diffusion_cov.f(None, u, t)
+            Qc_t = params.dynamics.diffusion_cov.f(y, u, t)
             L_t = (
-                params.dynamics.diffusion_coefficient.f(None, u, t)
+                params.dynamics.diffusion_coefficient.f(y, u, t)
                 * filter_hyperparams.cov_rescaling
             )
             Q_sqrt = jnp.linalg.cholesky(Qc_t)
@@ -143,20 +143,20 @@ def _predict(
         else:
             dt = filter_hyperparams.dt_average
 
-        Qc_t = params.dynamics.diffusion_cov.f(None, u, t0)
-        L_t = (
-            params.dynamics.diffusion_coefficient.f(None, u, t0)
-            * filter_hyperparams.cov_rescaling
-        )
-        state_noise_cov = dt * L_t @ Qc_t @ L_t.T  # D_hid x D_hid
-
         key_array = jr.split(key, x.shape[0])
-        noise = vmap(
-            lambda key: jr.multivariate_normal(
-                key=key, mean=jnp.zeros(x.shape[1]), cov=state_noise_cov
-            ),
-            in_axes=0,
-        )(key_array)
+
+        def _sample_noise(x_i, key_i):
+            Qc_t = params.dynamics.diffusion_cov.f(x_i, u, t0)
+            L_t = (
+                params.dynamics.diffusion_coefficient.f(x_i, u, t0)
+                * filter_hyperparams.cov_rescaling
+            )
+            state_noise_cov = dt * L_t @ Qc_t @ L_t.T
+            return jr.multivariate_normal(
+                key=key_i, mean=jnp.zeros(x.shape[1]), cov=state_noise_cov
+            )
+
+        noise = vmap(_sample_noise, in_axes=(0, 0))(x_pred, key_array)
         x_pred += noise
     return x_pred
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax/src/continuous_discrete_nonlinear_ssm/models.py

@@ -387,8 +387,8 @@ class ContDiscreteNonlinearSSM(SSM):
                 return params.dynamics.drift.f(y, u_prev_t, t)
 
             def diffusion(t, y, _):
-                Qc_t = params.dynamics.diffusion_cov.f(None, u_prev_t, t)
-                L_t = params.dynamics.diffusion_coefficient.f(None, u_prev_t, t)
+                Qc_t = params.dynamics.diffusion_cov.f(y, u_prev_t, t)
+                L_t = params.dynamics.diffusion_coefficient.f(y, u_prev_t, t)
                 Q_sqrt = jnp.linalg.cholesky(Qc_t)
                 return L_t @ Q_sqrt
 
@@ -818,10 +818,10 @@ def cdnlssm_forecast(
                 return params.dynamics.drift.f(y, forecast_inputs[t0_idx], t)
 
             def diffusion(t, y, args):
-                Qc_t = params.dynamics.diffusion_cov.f(None, forecast_inputs[t0_idx], t)
+                Qc_t = params.dynamics.diffusion_cov.f(y, forecast_inputs[t0_idx], t)
                 Q_sqrt = jnp.linalg.cholesky(Qc_t)
                 L_t = params.dynamics.diffusion_coefficient.f(
-                    None, forecast_inputs[t0_idx], t
+                    y, forecast_inputs[t0_idx], t
                 )
                 combined_diffusion = L_t @ Q_sqrt
                 return combined_diffusion

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cd_dynamax
-Version: 0.3.2
+Version: 0.3.3
 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>
@@ -32,6 +32,7 @@ Requires-Dist: dm-tree>=0.1.8
 Requires-Dist: fastprogress>=1.0.0
 Requires-Dist: graphviz
 Requires-Dist: ipykernel
+Requires-Dist: orbax-checkpoint<0.11.3; sys_platform == "win32"
 Provides-Extra: dev
 Requires-Dist: pytest>=8.0; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
@@ -69,7 +70,7 @@ $$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) \\}$.
+We denote the collection of all parameters as $\theta = \\{f,\\  L,\\  \mu_0,\\  \Sigma_0,\\  Q,\\  h,\\  \textrm{Law}(\eta) \\}$.
 
 Note:
 
@@ -90,7 +91,7 @@ For a given set of observations $Y_K = [y(t_1),\\ \dots ,\\ y(t_K)]$, we wish to
 
 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 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)}$.
@@ -113,18 +114,19 @@ The `cd-dynamax` codebase extends the `dynamax` library to support continuous-di
 
 - 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
+.
+├── cd_dynamax/                  # Source code for cd-dynamax library
+│   ├── src/                     # Core source code
+│   │   ├── 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)
@@ -155,7 +157,7 @@ make test
 ```
 
 - For linting, we use `ruff`:
-```bashbash
+```bash
 make lint
 ```
 
@@ -171,40 +173,57 @@ make build_docs
 
 # Installation
 
-We support installation via **Conda** (recommended) or via a standard Python virtual environment.
+Install from **PyPI** (recommended), from source in editable mode, or with a Conda-managed environment.
 
 ---
 
-### Option 1: Conda (recommended)
+### Option 1: Install from PyPI (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
+# Create and activate a virtual environment
+python -m venv .venv        # or `uv venv`
+source .venv/bin/activate   # on macOS/Linux
+.venv\Scripts\activate      # on Windows
 
-# Install your package in editable mode (so local changes are picked up)
-pip install -e .[dev]
+# Upgrade pip
+pip install --upgrade pip
+
+# Install latest release from PyPI
+pip install cd-dynamax
 ```
 
-This installs the core dependencies listed in `pyproject.toml`, along with optional developer tools (`pytest`, etc.) if you use `[dev]`.
+`cd-dynamax` is currently **not available on Conda Forge**.
 
 ---
 
-### Option 2: Python venv + pip
+### Option 2: Install from source (editable)
 
 ```bash
 # Create and activate a virtual environment
-python -m venv .venv
+python -m venv .venv        # or `uv venv`
 source .venv/bin/activate   # on macOS/Linux
 .venv\Scripts\activate      # on Windows
 
 # Upgrade pip
 pip install --upgrade pip
 
-# Install in editable mode
+# Install in editable mode for local development
 pip install -e .[dev]
 ```
 
+---
+
+### Option 3: Conda environment + pip install
+
+```bash
+# Create and activate a Conda environment with Python 3.11
+conda create -n cd_dynamax python=3.11
+conda activate cd_dynamax
+
+# Install latest release from PyPI
+pip install cd-dynamax
+```
+
 #### GPU support
 If you want GPU acceleration with JAX, you must install a CUDA-enabled `jaxlib` wheel.  
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax.egg-info/SOURCES.txt

@@ -107,4 +107,5 @@ tests/test_cdnonlinear_cdlinear_match.py
 tests/test_filter_forecast_emissions.py
 tests/test_imports.py
 tests/test_models.py
+tests/test_state_dependent_diffusion.py
 tests/test_utils_imports.py

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/cd_dynamax.egg-info/requires.txt

@@ -24,6 +24,9 @@ fastprogress>=1.0.0
 graphviz
 ipykernel
 
+[:sys_platform == "win32"]
+orbax-checkpoint<0.11.3
+
 [dev]
 pytest>=8.0
 ruff

{cd_dynamax-0.3.2 → cd_dynamax-0.3.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cd_dynamax"
-version = "0.3.2"
+version = "0.3.3"
 requires-python = ">=3.11"
 description = "Continuous-discrete dynamical systems with JAX and related libraries."
 readme = "README.md"
@@ -42,7 +42,8 @@ dependencies = [
     "dm-tree>=0.1.8",
     "fastprogress>=1.0.0",
     "graphviz",
-    "ipykernel"
+    "ipykernel",
+    "orbax-checkpoint<0.11.3; sys_platform == 'win32'"
 ]
 
 [tool.setuptools.packages.find]

cd_dynamax-0.3.3/tests/test_state_dependent_diffusion.py

@@ -0,0 +1,93 @@
+import jax.numpy as jnp
+import jax.random as jr
+from typing import NamedTuple
+
+from cd_dynamax.src.continuous_discrete_nonlinear_gaussian_ssm import (
+    ContDiscreteNonlinearGaussianSSM,
+)
+from cd_dynamax.src.continuous_discrete_nonlinear_gaussian_ssm.inference_ekf import (
+    EKFHyperParams,
+    extended_kalman_filter,
+)
+from cd_dynamax.src.continuous_discrete_nonlinear_ssm import ContDiscreteNonlinearSSM
+from cd_dynamax.src.continuous_discrete_nonlinear_ssm.inference_dpf import (
+    DPFHyperParams,
+    diff_particle_filter,
+)
+
+
+class StateDependentDiagonalDiffusion(NamedTuple):
+    scale: float = 0.1
+
+    def f(self, x, u=None, t=None):
+        if x is None:
+            raise ValueError("state-dependent diffusion requires a state argument")
+        x = jnp.atleast_1d(x)
+        return jnp.diag(1.0 + self.scale * jnp.square(x))
+
+
+def test_cdnlgssm_sample_path_supports_state_dependent_diffusion():
+    model = ContDiscreteNonlinearGaussianSSM(state_dim=1, emission_dim=1)
+    params, _ = model.initialize()
+    params = params._replace(
+        dynamics=params.dynamics._replace(
+            diffusion_coefficient=StateDependentDiagonalDiffusion()
+        )
+    )
+
+    states, emissions = model.sample_path(params, key=jr.PRNGKey(0), num_timesteps=3)
+
+    assert states.shape == (3, 1)
+    assert emissions.shape == (3, 1)
+    assert jnp.all(jnp.isfinite(states))
+    assert jnp.all(jnp.isfinite(emissions))
+
+
+def test_extended_kalman_filter_supports_state_dependent_diffusion():
+    model = ContDiscreteNonlinearGaussianSSM(state_dim=1, emission_dim=1)
+    params, _ = model.initialize()
+    params = params._replace(
+        dynamics=params.dynamics._replace(
+            diffusion_coefficient=StateDependentDiagonalDiffusion()
+        )
+    )
+
+    emissions = jnp.zeros((2, 1))
+    t_emissions = jnp.array([[0.0], [0.1]])
+
+    posterior = extended_kalman_filter(
+        params,
+        emissions,
+        t_emissions=t_emissions,
+        filter_hyperparams=EKFHyperParams(state_order="first"),
+    )
+
+    assert posterior.filtered_means.shape == (2, 1)
+    assert posterior.predicted_covariances.shape == (2, 1, 1)
+    assert jnp.all(jnp.isfinite(posterior.filtered_means))
+    assert jnp.all(jnp.isfinite(posterior.predicted_covariances))
+
+
+def test_diff_particle_filter_supports_state_dependent_diffusion():
+    model = ContDiscreteNonlinearSSM(state_dim=1, emission_dim=1)
+    params, _ = model.initialize()
+    params = params._replace(
+        dynamics=params.dynamics._replace(
+            diffusion_coefficient=StateDependentDiagonalDiffusion()
+        )
+    )
+
+    emissions = jnp.zeros((2, 1))
+    ts = jnp.array([0.0, 0.1])
+
+    particles, log_weights, log_evidence = diff_particle_filter(
+        jr.PRNGKey(0),
+        params,
+        emissions,
+        ts=ts,
+        hyperparams=DPFHyperParams(N_particles=8),
+    )
+
+    assert particles.shape == (2, 8, 1)
+    assert log_weights.shape == (2, 8)
+    assert jnp.isfinite(log_evidence)