Support structured diffusion specs in continuous-time models and enforce cd-dynamax full-diffusion constraints

dynestyx/inference/integrations/cd_dynamax/utils.py

@@ -22,40 +22,74 @@
     LinearGaussianObservation,
     LinearGaussianStateEvolution,
 )
+from dynestyx.models.diffusions import diffusion_as_matrix, evaluate_diffusion
 
 type SSMType = ContDiscreteNonlinearGaussianSSM | ContDiscreteNonlinearSSM
 
 
 def _normalize_cd_dynamax_diffusion(
-    diffusion_coefficient,
+    state_evolution: ContinuousTimeStateEvolution,
     state_dim: int,
 ):
     """Return a diffusion coeff compatible with cd-dynamax's EnKF SDE solve.
 
     cd-dynamax's internal diffrax wrapper builds Brownian controls with shape
-    equal to `y0.shape` (state_dim). For non-square diffusion coefficients
-    (state_dim, bm_dim) with bm_dim != state_dim, pad/truncate columns so the
-    returned matrix is always (state_dim, state_dim).
+    equal to `y0.shape` (state_dim). For diffusion with `bm_dim < state_dim`,
+    pad trailing Brownian columns with zeros to match `(state_dim, state_dim)`.
+    Diffusion with `bm_dim > state_dim` is rejected.
     """
 
     def _wrapped(x, u, t):
-        L = diffusion_coefficient(x, u, t)
-        if L.ndim == 1:
-            L = jnp.diag(L)
-        if L.ndim != 2:
+        diffusion = evaluate_diffusion(
+            state_evolution.diffusion_coefficient,
+            diffusion_type=state_evolution.diffusion_type,
+            bm_dim=state_evolution.bm_dim,
+            x=x,
+            u=u,
+            t=t,
+            state_dim=state_dim,
+        )
+        L = diffusion_as_matrix(diffusion, state_dim=state_dim)
+        n_cols = L.shape[-1]
+        if n_cols > state_dim:
             raise ValueError(
-                "diffusion_coefficient must return a vector or matrix for cd-dynamax."
+                "cd-dynamax continuous diffusion requires bm_dim <= state_dim. "
+                f"Got state_dim={state_dim}, bm_dim={n_cols}."
             )
-        n_cols = L.shape[-1]
-        if n_cols == state_dim:
-            return L
         if n_cols < state_dim:
-            return jnp.pad(L, ((0, 0), (0, state_dim - n_cols)))
-        return L[:, :state_dim]
+            pad_width = ((0, 0),) * (L.ndim - 1) + ((0, state_dim - n_cols),)
+            L = jnp.pad(L, pad_width)
+        return L
 
     return _wrapped
 
 
+def _validate_cd_dynamax_continuous_diffusion(
+    state_evolution: ContinuousTimeStateEvolution,
+    state_dim: int,
+) -> None:
+    """Eagerly validate diffusion shape constraints for cd-dynamax continuous filters."""
+    if state_evolution.diffusion_coefficient is None:
+        return
+
+    probe_x = jnp.zeros(state_dim)
+    diffusion = evaluate_diffusion(
+        state_evolution.diffusion_coefficient,
+        diffusion_type=state_evolution.diffusion_type,
+        bm_dim=state_evolution.bm_dim,
+        x=probe_x,
+        u=None,
+        t=jnp.array(0.0),
+        state_dim=state_dim,
+    )
+    bm_dim = diffusion_as_matrix(diffusion, state_dim=state_dim).shape[-1]
+    if bm_dim > state_dim:
+        raise ValueError(
+            "Continuous cd-dynamax filters require bm_dim <= state_dim. "
+            f"Got state_dim={state_dim}, bm_dim={bm_dim}."
+        )
+
+
 class _ConstantFunction(eqx.Module):
     value: Any
 
@@ -198,7 +232,16 @@ def dsx_to_cdlgssm_params(dsx_model: DynamicalModel) -> ParamsCDLGSSM:
 
     # Extract constant L and use inferred Brownian dimension.
     x0 = jnp.zeros(dsx_model.state_dim)
-    L = state_evo.diffusion_coefficient(x0, None, jnp.array(0.0))
+    diffusion = evaluate_diffusion(
+        state_evo.diffusion_coefficient,
+        diffusion_type=state_evo.diffusion_type,
+        bm_dim=state_evo.bm_dim,
+        x=x0,
+        u=None,
+        t=jnp.array(0.0),
+        state_dim=dsx_model.state_dim,
+    )
+    L = diffusion_as_matrix(diffusion, state_dim=dsx_model.state_dim)
     if state_evo.bm_dim is None:
         raise ValueError(
             "state_evolution.bm_dim is not set on ContinuousTimeStateEvolution."
@@ -273,8 +316,12 @@ def dsx_to_cd_dynamax(
                 raise ValueError(
                     "state_evolution.bm_dim is not set on ContinuousTimeStateEvolution."
                 )
+            _validate_cd_dynamax_continuous_diffusion(
+                state_evo,
+                dsx_model.state_dim,
+            )
             diffusion_coeff = _normalize_cd_dynamax_diffusion(
-                state_evo.diffusion_coefficient,
+                state_evo,
                 dsx_model.state_dim,
             )
             shared_params.update(

dynestyx/models/__init__.py

@@ -10,6 +10,7 @@
     DynamicalModel,
     ObservationModel,
 )
+from dynestyx.models.diffusions import DiffusionType
 from dynestyx.models.lti_dynamics import LTI_continuous, LTI_discrete
 from dynestyx.models.observations import (
     DiracIdentityObservation,
@@ -27,6 +28,7 @@
     "AffineDrift",
     "DiracIdentityObservation",
     "DiscreteTimeStateEvolution",
+    "DiffusionType",
     "DynamicalModel",
     "Drift",
     "GaussianObservation",

dynestyx/models/checkers.py

@@ -8,6 +8,10 @@
 import numpyro.distributions as dist
 import numpyro.primitives
 
+from dynestyx.models.diffusions import (
+    evaluate_diffusion_value,
+    resolve_diffusion_metadata,
+)
 from dynestyx.types import Control, State, Time
 
 
@@ -100,27 +104,16 @@ def _infer_bm_dim(
         return None
 
     diffusion_shape = jax.eval_shape(
-        lambda: state_evolution.diffusion_coefficient(x0, u0, t0)
-    ).shape
-    if len(diffusion_shape) < 2:
-        raise ValueError(
-            "diffusion_coefficient must return shape (..., state_dim, bm_dim). "
-            f"Got shape {diffusion_shape}."
-        )
-    if int(diffusion_shape[-2]) != state_dim:
-        raise ValueError(
-            "diffusion_coefficient penultimate dimension must match state_dim. "
-            f"Got diffusion shape {diffusion_shape}, state_dim={state_dim}."
-        )
-    inferred_bm_dim = int(diffusion_shape[-1])
-    if (
-        state_evolution.bm_dim is not None
-        and int(state_evolution.bm_dim) != inferred_bm_dim
-    ):
-        raise ValueError(
-            "bm_dim does not match inferred diffusion_coefficient output shape. "
-            f"Got bm_dim={state_evolution.bm_dim}, inferred={inferred_bm_dim}."
+        lambda: evaluate_diffusion_value(
+            state_evolution.diffusion_coefficient, x0, u0, t0
         )
+    ).shape
+    _, inferred_bm_dim = resolve_diffusion_metadata(
+        diffusion_shape,
+        state_dim=state_dim,
+        diffusion_type=state_evolution.diffusion_type,
+        bm_dim=state_evolution.bm_dim,
+    )
     return inferred_bm_dim
 
 

dynestyx/models/core.py

@@ -21,7 +21,15 @@
     _validate_state_dim,
     _validate_state_evolution_output_shape,
 )
-from dynestyx.types import Control, State, Time, TimeLike, as_scalar_time_array, dState
+from dynestyx.models.diffusions import DiffusionSpec, DiffusionType
+from dynestyx.types import (
+    Control,
+    State,
+    Time,
+    TimeLike,
+    as_scalar_time_array,
+    dState,
+)
 
 
 class DynamicalModel(eqx.Module):
@@ -298,18 +306,25 @@ class ContinuousTimeStateEvolution(eqx.Module):
             At least one of `drift` or `potential` must be non-None.
         use_negative_gradient (bool): If True, use $-\\nabla_x V$ (e.g., gradient descent on potential);
             otherwise use $+\\nabla_x V$. Default is False.
-        diffusion_coefficient (Drift | None): Diffusion coefficient $L(x, u, t)$ mapping to a matrix;
-            multiplies the Brownian increment $dW_t$.
-            Defaults to zero if None (i.e., deterministic ODE).
+        diffusion_coefficient (DiffusionSpec | None): Diffusion coefficient specification.
+            This may be a callable `L(x, u, t)`, a constant scalar/vector/matrix, or `None`
+            for deterministic dynamics.
+        diffusion_type ("full" | "diag" | "scalar" | None): Optional explicit diffusion semantics.
+            Use `"full"` for matrix-valued diffusion, `"diag"` for diagonal shorthand with
+            trailing shape `(..., state_dim)`, and `"scalar"` for scalar shorthand with shape
+            `()` or `(..., 1)`. If omitted, legacy behavior infers `"full"` from trailing
+            shape `(..., state_dim, bm_dim)` and otherwise infers scalar/diagonal shorthand
+            from the trailing dimension.
         bm_dim (int | None): Dimension of the Brownian motion $W_t$.
-            Inferred automatically from the output shape of `diffusion_coefficient`;
-            if passed by the user, it must match diffusion_coefficient(...).shape[1].
+            Inferred automatically only for full matrix diffusion. Scalar and diagonal
+            diffusion require explicit `bm_dim`, which must be either `1` or `state_dim`.
     """
 
     drift: Drift | None = None
     potential: Potential | None = None
     use_negative_gradient: bool = eqx.field(static=True, default=False)
-    diffusion_coefficient: Drift | None = None
+    diffusion_coefficient: DiffusionSpec | None = None
+    diffusion_type: DiffusionType | None = eqx.field(static=True, default=None)
     bm_dim: int | None = eqx.field(static=True, default=None)
 
     def total_drift(self, x: State, u: Control | None, t: Time) -> dState: