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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cd_dynamax
-Version: 0.3.2
+Version: 0.3.4
 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.4}/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.4}/cd_dynamax/dynamax/linear_gaussian_ssm/inference.py

@@ -11,7 +11,7 @@ from tensorflow_probability.substrates.jax.distributions import (
 
 from jax.tree_util import tree_map
 from jaxtyping import Array, Float
-from typing import NamedTuple, Optional, Union, Tuple
+from typing import Collection, NamedTuple, Optional, Union, Tuple
 from ..utils.utils import psd_solve, symmetrize
 from ..parameters import ParameterProperties
 from ..types import PRNGKey, Scalar
@@ -112,9 +112,27 @@ class ParamsLGSSM(NamedTuple):
 class PosteriorGSSMFiltered(NamedTuple):
     r"""Marginals of the Gaussian filtering posterior.
 
+    Time indexing convention:
+    `filtered_*[t]` corresponds to $x_{t \mid t}$,
+    `predicted_*[t]` corresponds to $x_{t+1 \mid t}$,
+    and predictive observation quantities such as `y_pred_*[t]` and
+    `y_obs_pred_*[t]` correspond to the observation-time prior
+    $p(y_t \mid y_{1:t-1})$, i.e. they are computed from $x_{t \mid t-1}$.
+
     :param marginal_loglik: marginal log likelihood, $p(y_{1:T} \mid u_{1:T})$
-    :param filtered_means: array of filtered means $\mathbb{E}[z_t \mid y_{1:t}, u_{1:t}]$
-    :param filtered_covariances: array of filtered covariances $\mathrm{Cov}[z_t \mid y_{1:t}, u_{1:t}]$
+    :param filtered_means: array of filtered means $\mathbb{E}[x_t \mid y_{1:t}, u_{1:t}]$
+    :param filtered_covariances: array of filtered covariances $\mathrm{Cov}[x_t \mid y_{1:t}, u_{1:t}]$
+    :param predicted_means: optional array of one-step-ahead state means $\mathbb{E}[x_{t+1} \mid y_{1:t}, u_{1:t}]$
+    :param predicted_covariances: optional array of one-step-ahead state covariances $\mathrm{Cov}[x_{t+1} \mid y_{1:t}, u_{1:t}]$
+    :param filtered_ensembles: When filtering with EnKF, optional array of filtered state ensembles approximating $p(x_t \mid y_{1:t}, u_{1:t})$
+    :param predicted_ensembles: When filtering with EnKF, optional array of one-step-ahead predicted state ensembles approximating $p(x_{t+1} \mid y_{1:t}, u_{1:t})$
+    :param y_pred_mean: optional array of predictive emission means for $p(h(x_t) \mid y_{1:t-1}, u_{1:t})$,  i.e., before adding observation noise to h(x_t)
+    :param y_pred_cov: optional array of predictive emission covariances for $p(h(x_t) \mid y_{1:t-1}, u_{1:t})$, i.e., before adding observation noise to h(x_t)
+    :param y_obs_pred_mean: optional array of predictive observation means for $p(y_t \mid y_{1:t-1}, u_{1:t})$
+    :param y_obs_pred_cov: optional array of predictive observation covariances for $p(y_t \mid y_{1:t-1}, u_{1:t})$, including observation noise
+    :param y_ens_pred: When filtering with EnKF, optional array of predictive emission ensembles for observation time $t$, obtained from the predicted state ensemble for $x_{t \mid t-1}$
+    :param y_obs_ens_pred: When filtering with EnKF, optional array of predictive observation ensembles for observation time $t$, obtained by adding sampled observation noise to `y_ens_pred`
+    :param posterior_extras: When filtering with EnKF, optional dictionary of filter-specific per-step diagnostics that do not fit the standard filtered posterior schema
 
     """
     # Default attributes
@@ -123,9 +141,64 @@ class PosteriorGSSMFiltered(NamedTuple):
     filtered_covariances: Optional[Float[Array, "ntime state_dim state_dim"]] = None
     predicted_means: Optional[Float[Array, "ntime state_dim"]] = None
     predicted_covariances: Optional[Float[Array, "ntime state_dim state_dim"]] = None
+    filtered_ensembles: Optional[Float[Array, "ntime ensemble_size state_dim"]] = None
+    predicted_ensembles: Optional[Float[Array, "ntime ensemble_size state_dim"]] = None
+    y_pred_mean: Optional[Float[Array, "ntime emission_dim"]] = None
+    y_pred_cov: Optional[Float[Array, "ntime emission_dim emission_dim"]] = None
+    y_obs_pred_mean: Optional[Float[Array, "ntime emission_dim"]] = None
+    y_obs_pred_cov: Optional[Float[Array, "ntime emission_dim emission_dim"]] = None
+    y_ens_pred: Optional[Float[Array, "ntime ensemble_size emission_dim"]] = None
+    y_obs_ens_pred: Optional[Float[Array, "ntime ensemble_size emission_dim"]] = None
     # Additional extras
     posterior_extras: Optional[dict] = None
 
+
+FILTERED_POSTERIOR_FIELD_NAMES = (
+    "marginal_loglik",
+    "filtered_means",
+    "filtered_covariances",
+    "predicted_means",
+    "predicted_covariances",
+    "filtered_ensembles",
+    "predicted_ensembles",
+    "y_pred_mean",
+    "y_pred_cov",
+    "y_obs_pred_mean",
+    "y_obs_pred_cov",
+    "y_ens_pred",
+    "y_obs_ens_pred",
+    "posterior_extras",
+)
+
+
+def validate_filtered_posterior_output_fields(
+    filter_name: str,
+    output_fields: Optional[Collection[str]],
+    supported_fields: Collection[str],
+):
+    """Validate requested top-level filtered posterior fields."""
+
+    if output_fields is None:
+        return
+
+    valid_fields = set(FILTERED_POSTERIOR_FIELD_NAMES)
+    supported_fields = set(supported_fields)
+    requested_fields = list(output_fields)
+
+    unknown_fields = sorted(set(requested_fields) - valid_fields)
+    if unknown_fields:
+        raise ValueError(
+            f"Unknown output_fields for {filter_name}: {unknown_fields}. "
+            f"Valid top-level filtered posterior fields are: {sorted(valid_fields)}."
+        )
+
+    unsupported_fields = sorted(set(requested_fields) - supported_fields)
+    if unsupported_fields:
+        raise ValueError(
+            f"output_fields {unsupported_fields} are not available for {filter_name}. "
+            f"Available fields are: {sorted(supported_fields)}."
+        )
+
 class PosteriorGSSMSmoothed(NamedTuple):
     r"""Marginals of the Gaussian filtering and smoothing posterior.
 

{cd_dynamax-0.3.2 → cd_dynamax-0.3.4}/cd_dynamax/dynamax/nonlinear_gaussian_ssm/inference_ekf.py

@@ -101,8 +101,12 @@ def extended_kalman_filter(
         num_iter: number of linearizations around posterior for update step (default 1).
         inputs: optional array of inputs.
         output_fields: list of fields to return in posterior object.
-            These can take the values "filtered_means", "filtered_covariances",
-            "predicted_means", "predicted_covariances", and "marginal_loglik".
+            Options:
+            `"filtered_means"` (default)
+            `"filtered_covariances"` (default)
+            `"predicted_means"` (default)
+            `"predicted_covariances"` (default)
+            `"marginal_loglik"`
 
     Returns:
         post: posterior object.

{cd_dynamax-0.3.2 → cd_dynamax-0.3.4}/cd_dynamax/dynamax/nonlinear_gaussian_ssm/inference_ukf.py

@@ -151,6 +151,13 @@ def unscented_kalman_filter(
         emissions: array of observations.
         hyperparams: hyper-parameters.
         inputs: optional array of inputs.
+        output_fields: list of fields to return in posterior object.
+            Options:
+            `"filtered_means"` (default)
+            `"filtered_covariances"` (default)
+            `"predicted_means"` (default)
+            `"predicted_covariances"` (default)
+            `"marginal_loglik"`
 
     Returns:
         filtered_posterior: posterior object.

{cd_dynamax-0.3.2 → cd_dynamax-0.3.4}/cd_dynamax/src/continuous_discrete_linear_gaussian_ssm/inference.py

@@ -27,6 +27,7 @@ from cd_dynamax.dynamax.utils.utils import psd_solve
 from cd_dynamax.dynamax.linear_gaussian_ssm.inference import (
     PosteriorGSSMFiltered,
     PosteriorGSSMSmoothed,
+    validate_filtered_posterior_output_fields,
 )
 
 # Initial and emission parameter classes are equivalent
@@ -54,6 +55,19 @@ tfb = tfp.bijectors
 DEBUG = False
 
 
+CDLGSSM_FILTER_OUTPUT_FIELDS = (
+    "marginal_loglik",
+    "filtered_means",
+    "filtered_covariances",
+    "predicted_means",
+    "predicted_covariances",
+    "y_pred_mean",
+    "y_pred_cov",
+    "y_obs_pred_mean",
+    "y_obs_pred_cov",
+)
+
+
 #### Helper functions
 # Helper functions --- modified from dynamax
 def _get_params(x, dim, t):
@@ -137,20 +151,27 @@ def preprocess_args(f):
         filter_hyperparams = bound_args.arguments["filter_hyperparams"]
         inputs = bound_args.arguments["inputs"]
         warn = bound_args.arguments["warn"]
+        output_fields = bound_args.arguments.get("output_fields")
 
         num_timesteps = len(emissions)
         full_params, inputs = preprocess_params_and_inputs(
             params, num_timesteps, inputs
         )
 
-        return f(
-            full_params,
-            emissions,
-            t_emissions,
+        call_kwargs = dict(
+            params=full_params,
+            emissions=emissions,
+            t_emissions=t_emissions,
             filter_hyperparams=filter_hyperparams,
             inputs=inputs,
             warn=warn,
         )
+        if output_fields is not None:
+            call_kwargs["output_fields"] = output_fields
+
+        return f(
+            **call_kwargs,
+        )
 
     return wrapper
 
@@ -334,6 +355,14 @@ def _condition_on(m, P, H, D, d, R, u, y, warn: bool = True):
     return mu_cond, Sigma_cond
 
 
+def _emission_predicted_moments(m, P, H, D, d, u, warn: bool = True):
+    """Compute the predictive emission moments under the linear Gaussian model."""
+
+    y_pred_mean = H @ m + D @ u + d
+    y_pred_cov = psd(H @ P @ H.T, warn=warn)
+    return y_pred_mean, y_pred_cov
+
+
 # CD-LGSSM filtering implementation: Kalman filter
 @preprocess_args
 def cdlgssm_filter(
@@ -342,6 +371,12 @@ def cdlgssm_filter(
     t_emissions: Optional[Float[Array, "num_timesteps 1"]] = None,
     filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(),
     inputs: Optional[Float[Array, "num_timesteps input_dim"]] = None,
+    output_fields: Optional[List[str]] = [
+        "filtered_means",
+        "filtered_covariances",
+        "predicted_means",
+        "predicted_covariances",
+    ],
     warn: bool = True,
 ) -> PosteriorGSSMFiltered:
     r"""Run a Continuous Discrete Kalman filter
@@ -353,6 +388,17 @@ def cdlgssm_filter(
         t_emissions: continuous-time specific time instants of observations: if not None, it is an array
         filter_hyperparams: hyperparameters for the filter.
         inputs: optional array of inputs.
+        output_fields: list of top-level posterior fields to return.
+            Options:
+            `"filtered_means"` (default)
+            `"filtered_covariances"` (default)
+            `"predicted_means"` (default)
+            `"predicted_covariances"` (default)
+            `"marginal_loglik"`
+            `"y_pred_mean"`
+            `"y_pred_cov"`
+            `"y_obs_pred_mean"`
+            `"y_obs_pred_cov"`
         warn: whether to issue warnings during filtering (e.g., PSD issues).
 
     Returns:
@@ -362,6 +408,10 @@ def cdlgssm_filter(
     if filter_hyperparams is None:
         filter_hyperparams = KFHyperParams()
 
+    validate_filtered_posterior_output_fields(
+        "cdlgssm_filter", output_fields, CDLGSSM_FILTER_OUTPUT_FIELDS
+    )
+
     # Figure out timestamps, as vectors to scan over
     # t_emissions is of shape num_timesteps \times 1
     # t0 and t1 are num_timesteps \times 0
@@ -405,8 +455,14 @@ def cdlgssm_filter(
         u = inputs[t0_idx]
         y = emissions[t0_idx]
 
+        y_pred_mean, y_pred_cov = _emission_predicted_moments(
+            pred_mean, pred_cov, H, D, d, u, warn=warn
+        )
+        y_obs_pred_mean = y_pred_mean
+        y_obs_pred_cov = psd(y_pred_cov + R, warn=warn)
+
         # Update the log likelihood
-        ll += MVN(H @ pred_mean + D @ u + d, H @ pred_cov @ H.T + R).log_prob(y)
+        ll += MVN(y_obs_pred_mean, y_obs_pred_cov).log_prob(y)
 
         # Condition on this emission
         filtered_mean, filtered_cov = _condition_on(
@@ -427,30 +483,30 @@ def cdlgssm_filter(
             filtered_mean, filtered_cov, F, C, B, b, Q, u, warn=warn
         )
 
+        outputs = {
+            "filtered_means": filtered_mean,
+            "filtered_covariances": filtered_cov,
+            "predicted_means": pred_mean,
+            "predicted_covariances": pred_cov,
+            "y_pred_mean": y_pred_mean,
+            "y_pred_cov": y_pred_cov,
+            "y_obs_pred_mean": y_obs_pred_mean,
+            "y_obs_pred_cov": y_obs_pred_cov,
+        }
+        outputs = {key: val for key, val in outputs.items() if key in output_fields}
+
         # Return the carry and outputs
-        return (ll, pred_mean, pred_cov), (
-            filtered_mean,
-            filtered_cov,
-            pred_mean,
-            pred_cov,
-        )
+        return (ll, pred_mean, pred_cov), outputs
 
     # The Kalman filter
     # Initial carry
     carry = (0.0, params.initial.mean, params.initial.cov)
     # Scan over all time steps
-    (ll, _, _), (filtered_means, filtered_covs, pred_means, pred_covs) = lax.scan(
-        _step, carry, (t0, t1, t0_idx)
-    )
+    (ll, _, _), outputs = lax.scan(_step, carry, (t0, t1, t0_idx))
+    outputs = {"marginal_loglik": ll, **outputs}
 
     # Return the posterior object
-    return PosteriorGSSMFiltered(
-        marginal_loglik=ll,
-        filtered_means=filtered_means,
-        filtered_covariances=filtered_covs,
-        predicted_means=pred_means,
-        predicted_covariances=pred_covs,
-    )
+    return PosteriorGSSMFiltered(**outputs)
 
 
 # Kalmam Smoothing equations, in continuous-time
@@ -1110,12 +1166,10 @@ def cdlgssm_forecast(
         inputs: optional array of inputs, of shape (1 + num_timesteps) \times input_dim
             - The extra input is needed for the initial emission, i.e., it should be at time t_init
         output_fields: list of fields to return in posterior object.
-            These can take the values
-                If we forecast Gaussian distributions, based on filtering methods
-                    "forecasted_state_means",
-                    "forecasted_state_covariances",
-                If we forecast paths, based on solving the SDE
-                    "forecasted_state_path".
+            Options:
+            `"forecasted_state_means"` (default for Gaussian forecasts)
+            `"forecasted_state_covariances"` (default for Gaussian forecasts)
+            `"forecasted_state_path"` (for point-initialized path forecasts)
         key: random key (e.g., for Ensemble Kalman).
         diffeqsolve_settings: settings for the SDE solver
         warn: whether to issue warnings during filtering (e.g., PSD issues).

{cd_dynamax-0.3.2 → cd_dynamax-0.3.4}/cd_dynamax/src/continuous_discrete_linear_gaussian_ssm/models.py

@@ -4,7 +4,7 @@ import jax.random as jr
 from jaxtyping import Array, Float
 
 # Type annotations
-from typing import Any, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 from typing_extensions import Protocol
 
 # Distributions, compatible with JAX, from TensorFlow Probability
@@ -510,6 +510,12 @@ class ContDiscreteLinearGaussianSSM(SSM):
         t_emissions: Optional[Float[Array, "ntime 1"]] = None,
         filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(),
         inputs: Optional[Float[Array, "ntime input_dim"]] = None,
+        output_fields: Optional[List[str]] = [
+            "filtered_means",
+            "filtered_covariances",
+            "predicted_means",
+            "predicted_covariances",
+        ],
         warn: bool = True,
     ) -> PosteriorGSSMFiltered:
         r"""Run the CD-Kalman filter to compute the filtered posterior distribution over states.
@@ -519,6 +525,20 @@ class ContDiscreteLinearGaussianSSM(SSM):
             t_emissions: continuous-time specific time instants of observations: if not None, it is an array
             filter_hyperparams: hyperparameters for the Kalman filter.
             inputs: optional sequence of inputs.
+            output_fields: Which top-level posterior fields to return from the
+                filter. Options:
+                `"filtered_means"` (default)
+                `"filtered_covariances"` (default)
+                `"predicted_means"` (default)
+                `"predicted_covariances"` (default)
+                `"marginal_loglik"`
+                `"y_pred_mean"`
+                `"y_pred_cov"`
+                `"y_obs_pred_mean"`
+                `"y_obs_pred_cov"`
+                Predictive emission fields are available directly as
+                top-level filtered posterior outputs. Unsupported fields raise
+                a `ValueError`.
             warn: whether to warn about numerical issues.
         Returns:
             filtered posterior distribution over states.
@@ -526,7 +546,13 @@ class ContDiscreteLinearGaussianSSM(SSM):
 
         # Directly run the CD-Kalman filter
         return cdlgssm_filter(
-            params, emissions, t_emissions, filter_hyperparams, inputs, warn=warn
+            params,
+            emissions,
+            t_emissions,
+            filter_hyperparams,
+            inputs,
+            output_fields=output_fields,
+            warn=warn,
         )
 
     # High-level, user-friendly interface combining filtering and forecasting steps