Missing data: partial missingness, unroll_missing, diagonal obs models

dynestyx/models/__init__.py

@@ -12,6 +12,8 @@
 )
 from dynestyx.models.lti_dynamics import LTI_continuous, LTI_discrete
 from dynestyx.models.observations import (
+    DiagonalGaussianObservation,
+    DiagonalLinearGaussianObservation,
     DiracIdentityObservation,
     GaussianObservation,
     LinearGaussianObservation,
@@ -25,6 +27,8 @@
 __all__ = [
     "ContinuousTimeStateEvolution",
     "AffineDrift",
+    "DiagonalGaussianObservation",
+    "DiagonalLinearGaussianObservation",
     "DiracIdentityObservation",
     "DiscreteTimeStateEvolution",
     "DynamicalModel",

dynestyx/models/core.py

@@ -1,6 +1,7 @@
 """Core interfaces and base classes for dynamical models."""
 
 import dataclasses
+from abc import abstractmethod
 from collections.abc import Callable
 from typing import Any, Protocol
 
@@ -375,3 +376,42 @@ def sample(self, x, u, t, *args, **kwargs):
             seed = kwargs.pop("seed")
             kwargs["key"] = seed
         return dist.sample(*args, **kwargs)
+
+    @abstractmethod
+    def __call__(self, x, u, t) -> DistributionT: ...
+
+    def masked_log_prob(
+        self,
+        y: jax.Array,
+        obs_mask: jax.Array,
+        x: Any,
+        u: Any = None,
+        t: Any = None,
+    ) -> jax.Array:
+        """Log p(y_obs | x) scoring only observed dimensions.
+
+        Args:
+            y: Observation with NaN replaced by safe values. Shape (obs_dim,).
+            obs_mask: Boolean array, True = observed. Shape (obs_dim,).
+            x: Latent state.
+            u: Control or None.
+            t: Time or None.
+
+        Returns:
+            Scalar log-probability summed over observed dims only.
+        """
+        import numpyro.distributions as _dist_mod
+
+        d = self(x, u, t)
+        # Unwrap Independent(base, 1) to get per-element log_probs
+        base = d
+        if isinstance(d, _dist_mod.Independent) and d.reinterpreted_batch_ndims == 1:
+            base = d.base_dist
+        per_dim_lp = base.log_prob(y)  # (obs_dim,) if base is element-wise
+        if jnp.ndim(per_dim_lp) == 0:
+            raise NotImplementedError(
+                f"{type(self).__name__}.masked_log_prob: distribution "
+                f"{type(d).__name__} does not decompose per-dimension. "
+                "Override masked_log_prob in the subclass."
+            )
+        return jnp.sum(jnp.where(obs_mask, per_dim_lp, 0.0))

dynestyx/models/observations.py

@@ -95,6 +95,95 @@ def __call__(self, x, u, t):
         return dist.MultivariateNormal(loc=loc, covariance_matrix=self.R)
 
 
+class DiagonalLinearGaussianObservation(ObservationModel):
+    """
+    Linear-Gaussian observation model with diagonal noise covariance.
+
+    Observations are modeled as
+
+    $$
+    y_t \\sim \\mathcal{N}(H x_t + D u_t + b, \\mathrm{diag}(R_{\\mathrm{diag}})).
+    $$
+
+    Because the noise dims are independent, `masked_log_prob` is exact for any
+    obs mask — no sub-block extraction is required.
+    """
+
+    H: jax.Array
+    R_diag: jax.Array
+    D: jax.Array | None = None
+    bias: jax.Array | None = None
+
+    def __init__(
+        self,
+        H: jax.Array,
+        R_diag: jax.Array,
+        D: jax.Array | None = None,
+        bias: jax.Array | None = None,
+    ):
+        """
+        Args:
+            H (jax.Array): Observation matrix with shape $(d_y, d_x)$.
+            R_diag (jax.Array): Per-dimension noise variances with shape
+                $(d_y,)$.
+            D (jax.Array | None): Optional control matrix with shape
+                $(d_y, d_u)$.
+            bias (jax.Array | None): Optional additive bias with shape
+                $(d_y,)$.
+        """
+        self.H = H
+        self.R_diag = R_diag
+        self.D = D
+        self.bias = bias
+
+    def __call__(self, x, u, t):
+        loc = jnp.dot(self.H, x)
+        if self.D is not None and u is not None:
+            loc += jnp.dot(self.D, u)
+        if self.bias is not None:
+            loc += self.bias
+        return dist.Independent(dist.Normal(loc, jnp.sqrt(self.R_diag)), 1)
+
+
+class DiagonalGaussianObservation(ObservationModel):
+    """
+    Nonlinear Gaussian observation model with diagonal noise covariance.
+
+    Observations are modeled as
+
+    $$
+    y_t \\sim \\mathcal{N}(h(x_t, u_t, t), \\mathrm{diag}(R_{\\mathrm{diag}})),
+    $$
+
+    where $h$ is a user-provided measurement function.
+
+    Because the noise dims are independent, `masked_log_prob` is exact for any
+    obs mask — no sub-block extraction is required.
+    """
+
+    h: Callable[[State, Control, Time], Observation]
+    R_diag: jax.Array
+
+    def __init__(
+        self,
+        h: Callable[[State, Control, Time], jax.Array],
+        R_diag: jax.Array,
+    ):
+        """
+        Args:
+            h (Callable[[State, Control, Time], jax.Array]): Measurement
+                function mapping $(x, u, t)$ to the mean observation.
+            R_diag (jax.Array): Per-dimension noise variances with shape
+                $(d_y,)$.
+        """
+        self.h = h
+        self.R_diag = R_diag
+
+    def __call__(self, x, u, t):
+        loc = self.h(x, u, t)
+        return dist.Independent(dist.Normal(loc, jnp.sqrt(self.R_diag)), 1)
+
+
 class DiracIdentityObservation(ObservationModel):
     """
     Noise-free identity observation model.
@@ -110,3 +199,10 @@ class DiracIdentityObservation(ObservationModel):
 
     def __call__(self, x, u, t):
         return dist.Delta(x)
+
+    def masked_log_prob(self, y, obs_mask, x, u=None, t=None):
+        raise NotImplementedError(
+            "DiracIdentityObservation does not support partial missingness. "
+            "Use DiagonalLinearGaussianObservation or DiagonalGaussianObservation "
+            "with small sigma_obs instead."
+        )