Gaussian relaxation

GaussianRelaxation transforms selected model components before inference or simulation. relax_dynamics applies the same transformation directly. state_evolution_cov can only be applied to discrete-time dynamics. See the public guide for examples, the model reference for covariance specifications, and handlers for nesting order.

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.

_condition(name: str, dynamics: DynamicalModel, **kwargs: Any) -> ConditionedResult | LatentStateResult | SimulatedResult

Forward the relaxed model to the next handler with the original arguments.

_stack_kind(**kwargs: Any) -> list[_DynestyxStackKind]

Register relaxation between discretization and inference.

_AddedCovariance

Bases: Module

Keep an existing covariance callable and added noise sliceable by plates.

__call__(*args: Real[Array, ...] | float | int | None) -> Float[Array, '*batch event_dim event_dim']

Evaluate the original model covariance and add the specified noise matrix.

_apply_covariance_setting(original: Covariance | Real[ArrayLike, '*batch event_dim event_dim'] | Callable[..., Covariance | Real[ArrayLike, '*batch event_dim event_dim']] | None, addition: Covariance, event_dim: int, mode: Literal['replace', 'add']) -> FullCovariance | _AddedCovariance

Apply a parsed covariance setting by replacing or adding to the original.

Expand the setting to event_dim. In add mode, keep existing covariance functions callable so the specified noise is added to each evaluated result.

_parse_covariance_setting(value: Covariance | Real[ArrayLike, ...] | None) -> Covariance | None

Interpret a relaxation setting as a covariance specification.

Scalar shorthand specifies a variance, a vector specifies diagonal variances, and a matrix specifies a full covariance. Batched settings use explicit covariance objects to distinguish leading plate axes from event axes. None selects no change; callable settings are rejected.

_relax_initial_condition(initial: dist.Distribution, covariance: Covariance, state_dim: int, mode: Literal['replace', 'add']) -> dist.Normal | dist.MultivariateNormal

Relax a supported scalar/vector initial distribution while retaining its mean.

_relax_observation(observation: ObservationModelLike, covariance: Covariance, observation_dim: int, mode: Literal['replace', 'add']) -> GaussianObservation | LinearGaussianObservation

Relax an observation while retaining its mean function and Gaussian class.

_relax_state_evolution(evolution: StateEvolutionLike, covariance: Covariance, state_dim: int, mode: Literal['replace', 'add']) -> GaussianStateEvolution | LinearGaussianStateEvolution

Relax a transition while retaining its mean function and Gaussian class.

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.