Gaussian relaxation

GaussianRelaxation adds Gaussian noise to a deterministic or Gaussian component of a DynamicalModel. In particular, Delta/Normal/MultivariateNormal initial conditions can be relaxed, as can a DeterministicStateEvolution/GaussianStateEvolution and DeterministicObservation/DiracIdentityObservation/GaussianObservation.

Linear-Gaussian components retain their structure and filter eligibility. state_evolution_cov can only be applied to discrete-time dynamics; discretize continuous-time dynamics first to relax transition covariance.

There are two ways to apply the same transformation:

• Effect-handler form: use GaussianRelaxation around a model call so inference or simulation receives the relaxed model.
• Direct form: call relax_dynamics to obtain a model for direct simulation, scoring, or use in your own inference code.

Effect-handler form

Use GaussianRelaxation with the handler order:

import dynestyx as dsx
import jax
import jax.numpy as jnp
import numpyro.distributions as dist
from dynestyx.inference.configs.filter import EnKFConfig

dynamics = dsx.DynamicalModel(
    initial_condition=dist.Delta(jnp.array([1.0, 2.0]), event_dim=1),
    state_evolution=dsx.DeterministicStateEvolution(
        lambda x, u, t_now, t_next: x + (t_next - t_now) * jnp.sin(x)
    ),
    observation_model=dsx.DeterministicObservation(lambda x, u, t: x**2),
)

with dsx.Filter(EnKFConfig()):
    with dsx.GaussianRelaxation(
        initial_condition_cov=0.05,
        state_evolution_cov=dsx.DiagonalCovariance(variance=jnp.array([0.01, 0.02])),
        observation_model_cov=dsx.FullCovariance(0.1 * jnp.eye(2)),
    ):
        result = dsx.condition(
            "trajectory", dynamics,
            obs_times=jnp.arange(3.0),
            obs_values=jnp.array([[1.0, 4.0], [3.4, 8.5], [7.9, 9.9]]),
        )

Direct form

Using the model above, call relax_dynamics before direct simulation or scoring:

relaxed = dsx.relax_dynamics(dynamics, observation_model_cov=0.1)
draws = dsx.simulate(relaxed, rng_key=jax.random.key(0), predict_times=jnp.arange(3.0))

Both entry points preserve everything other than the components you ask to change.

Selecting components and noise

Each setting defaults to None, leaving that component unchanged. See covariance specifications for accepted representations, units, and batched settings.

mode="replace" replaces the selected covariances. mode="add" adds the specified covariance to each selected component's covariance, treating deterministic components as having zero covariance. One mode applies to all selected components.

For example, if a GaussianStateEvolution already defines a covariance function cov(x, u, t_now, t_next), adding covariance C gives cov(x, u, t_now, t_next) + C at each transition. Time-varying linear-model functions cov(t_now, t_next) and R(t) work the same way. The added noise settings accept constants or learned JAX/NumPyro arrays, not callables.

API reference

Gaussian relaxation of selected model components.

GaussianRelaxation

Bases: ObjectInterpretation, HandlesSelf

Relax selected components before inference or simulation.

Nest this inside Filter/Smoother/Simulator and outside plates and any Discretizer. Arguments and covariance units match :func:relax_dynamics.

Example

with Filter(config): ... with GaussianRelaxation(observation_model_cov=0.1): ... result = condition("trajectory", dynamics, obs_times=t, obs_values=y)

__init__(initial_condition_cov: Covariance | Real[ArrayLike, ...] | None = None, state_evolution_cov: Covariance | Real[ArrayLike, ...] | None = None, observation_model_cov: Covariance | Real[ArrayLike, ...] | None = None, mode: Literal['replace', 'add'] = 'replace') -> None

Configure the components to relax and the replacement/addition mode.

relax_dynamics(dynamics: DynamicalModel, *, initial_condition_cov: Covariance | Real[ArrayLike, ...] | None = None, state_evolution_cov: Covariance | Real[ArrayLike, ...] | None = None, observation_model_cov: Covariance | Real[ArrayLike, ...] | None = None, mode: Literal['replace', 'add'] = 'replace') -> DynamicalModel

Return a Gaussian relaxation of selected model components.

state_evolution_cov requires discrete-time dynamics; discretize continuous-time dynamics first to relax transition covariance.

None leaves a component unchanged. A scalar setting is a variance; a vector contains diagonal variances; a matrix is a covariance. Use explicit :class:ScalarCovariance, :class:DiagonalCovariance, or :class:FullCovariance objects for batched settings. Settings may be learned JAX arrays, but may not be callables.

mode='replace' sets the selected covariance. mode='add' adds it to the existing covariance (zero for deterministic components), preserving the state/control/time dependence of existing model covariance functions. Means and all model metadata are preserved; the input model is never mutated. Covariance objects check shape, finite values, and symmetry; callers supply the positive-definite effective covariance needed by ordinary Gaussian densities. Filter jitter settings retain their existing behavior.

See the developer API for implementation details.