SSMs: blockdiag

BlockDiagLatentCond(A, noise, to_latent, to_observed)

Block-diagonal implementation of LatentCond operations.

apply_flat(x)

marginalise(rv)

merge(other: BlockDiagLatentCond) -> BlockDiagLatentCond

preconditioner_apply()

revert(rv, /, *, solve_triu)

BlockDiagNormal(mean_flat: Array, cholesky_flat: Array, tree_flatten: S)

Construct a block-diagonal normal distribution.

This assumes that the pytree is of the form [M_1, ..., M_{num_coeffs}], where each M_i is a pytree of the same structure.

Shapes: - mean: (d, num_coeffs) - cholesky: (d, num_coeffs, num_coeffs)

where d is the number of elements in each M.

from_dirac(mean, *, damp) classmethod

from_mean_and_std(mean, std) classmethod

identity_conditional() -> BlockDiagLatentCond

logpdf_flat(u_latent)

logpdf_scalar_flat(u)

logpdf_tree(u)

mean property

prototype_output_scale_calibrated()

register_pytree_node() -> None staticmethod

rescale_cholesky(factor)

residual_whitened_rms_flat(u_latent)

residual_whitened_rms_tree(u)

sample_flat(key)

sample_tree(key)

std property

to_derivative(i, std)

to_multivariate_normal()

BlockDiagOdeTs0(*, ode)

Block-diagonal ODE linearization via TS0 (zeroth-degree Taylor series: evaluate at the prior mean, no Jacobian).

init_linearization() -> None

linearize(rv, state: None, *, damp: float, t)

BlockDiagResidual(residual)

Block-diagonal residual linearization via TS1.

init_linearization()

linearize(rv, state, *, damp: float, t)

BlockDiagTreeFlatten

Flattening information for block-diagonal random variables.

flatten_tree(x)

from_example(x) classmethod

treedef: Any instance-attribute

unflatten_array(x)

unravel_leaf: Any instance-attribute

BlockDiagWienerIntegrated(init, output_scale, *, a, q_sqrtm, q0, tree_flatten, precon_fun)

a = a instance-attribute

precon_fun = precon_fun instance-attribute

q0 = q0 instance-attribute

q_sqrtm = q_sqrtm instance-attribute

register_pytree() staticmethod

transition(*, dt: float, output_scale: Array) -> BlockDiagLatentCond

tree_flatten = tree_flatten instance-attribute

C = TypeVar('C', bound=Sequence) module-attribute

A type-variable for Sequence types.

For example, this variable is used to type Taylor coefficients.

__all__ = ['state_space_model_blockdiag'] module-attribute

state_space_model_blockdiag

Implementation of block-diagonal SSM constructors.

constraint_ode_ts0(ode: problems.JetOde) -> BlockDiagOdeTs0

constraint_residual(residual: problems.JetResidual, *, taylor_point: taylor_points.TaylorPoint | None = None) -> BlockDiagResidual

prior_exponential(ode, tcoeffs_mean: C, /, *, is_exact: C | bool = True, inexact_eps: float = 1e-06, diffuse_derivatives: int = 0, diffuse_eps: float = 1.0, output_scale: Array | None = None)

prior_exponential_diffuse(ode, tcoeffs_mean: C, tcoeffs_std: C, /, *, diffuse_derivatives: int = 0, diffuse_eps: float = 1.0, output_scale: Array | None = None)

prior_wiener_integrated(tcoeffs_mean: C, /, *, is_exact: C | bool = True, inexact_eps: float = 1e-06, diffuse_derivatives: int = 0, diffuse_eps: float = 1.0, output_scale: Array | None = None) -> BlockDiagWienerIntegrated

prior_wiener_integrated_diffuse(tcoeffs_mean: C, tcoeffs_std: C, /, *, diffuse_derivatives: int = 0, diffuse_eps: float = 1.0, output_scale: Array | None = None) -> BlockDiagWienerIntegrated