Backends

The inference engine is swappable behind the InferenceBackend protocol. KalmanBackend is the default fast path; RxInferBackend — imported from cpomdp.backends.rxinfer and gated behind the optional rxinfer extra — re-derives the same answers through Julia and exists as an independent correctness oracle [^bagaev2023rxinfer].

InferenceBackend

Bases: Protocol

A swappable inference engine for a LinearGaussianModel.

A backend is built from a model: any expensive, data-independent work (front-loading — see DECISIONS.md ADR-002) happens at construction, so the per-step infer_states stays cheap. Each call advances the belief one recursive filter step: the current belief goes in as the prior and the updated belief comes back as the posterior.

The Protocol is structural: any class with a matching infer_states (and the model it was built from) is a backend, with no shared base class. This is the abstraction wall — the native Kalman fast path and the RxInfer oracle are interchangeable behind it, and neither's implementation (JAX, juliacall, …) leaks into this signature.

infer_states

infer_states(
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief

Advance the belief by one filter step: prior in, posterior out.

Given the current belief (prior) and a new observation (plus the action just taken, if the model has a control matrix), return the updated belief.

Source code in src/cpomdp/backends/base.py

def infer_states(
    self,
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief:
    """Advance the belief by one filter step: ``prior`` in, posterior out.

    Given the current belief (``prior``) and a new ``observation`` (plus the
    ``action`` just taken, if the model has a control matrix), return the
    updated belief.
    """
    ...

KalmanBackend

KalmanBackend(
    model: LinearGaussianModel,
    *,
    steady_state: bool = False,
    tol: float = 1e-12,
    max_iter: int = 1000,
)

Exact Kalman-filter inference for a LinearGaussianModel.

Implements the InferenceBackend protocol: constructed from a model, then advances a belief one step at a time (prior in, posterior out) via the standard predict/update recursion.

Two modes:

• Per-step (default): recomputes the Kalman gain and covariance every step from the incoming belief. Correct for any linear-Gaussian model, including transient (pre-convergence) behaviour. This is the analytic oracle the rest of the toolbox is validated against.
• Steady-state (steady_state=True): solves the covariance recursion once at construction to a fixed point, then reuses the frozen gain and covariance every step. Cheap (no per-step covariance maths), but only valid for time-invariant models with regular complete observations. Raises RuntimeError if the recursion does not converge within max_iter (i.e. the model is not stabilisable/detectable).

Parameters:

Name	Type	Description	Default
model	LinearGaussianModel	The linear-Gaussian generative model to filter under.	required
steady_state	bool	If True, precompute and freeze the steady-state gain.	False
tol	float	Convergence tolerance for the steady-state fixed point (absolute, on successive covariances).	1e-12
max_iter	int	Cap on steady-state iterations before giving up.	1000

Source code in src/cpomdp/backends/kalman.py

def __init__(
    self,
    model: LinearGaussianModel,
    *,
    steady_state: bool = False,
    tol: float = 1e-12,
    max_iter: int = 1000,
) -> None:
    self.model = model
    self.steady_state = steady_state
    if steady_state:
        sensor_fixed = (
            model.observation_model is None or model.observation_model.is_fixed
        )
        process_fixed = (
            model.dynamics_noise_model is None
            or model.dynamics_noise_model.is_fixed
        )
        if not (sensor_fixed and process_fixed):
            raise ValueError(
                "steady_state=True needs fixed sensor and process noise; a "
                "state-dependent R(x) or Q(x) has no constant fixed point — "
                "use steady_state=False."
            )
        self._steady_gain, self._steady_cov = self._converge_to_steady_state(
            tol, max_iter
        )

infer_states

infer_states(
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief

Advance the belief by one filter step.

Runs one predict/update cycle: step the prior through the dynamics (applying action if the model has a control matrix), then correct the prediction toward observation using the Kalman gain. In steady-state mode the gain and covariance are the frozen fixed-point values; otherwise they are recomputed from prior.cov on this step.

The numeric work is delegated to the jit-compiled module kernels (_gain_and_posterior_cov, _posterior_mean); this method stays the eager orchestrator that validates inputs and wraps the result in a Belief.

Parameters:

Name	Type	Description	Default
observation	ArrayLike	The latest sensor reading, shape (m,).	required
prior	Belief	The current belief, treated as this step's previous posterior. Never mutated.	required
action	ArrayLike | None	The action just taken, shape (p,). Required iff the model has a control matrix; ignored (pass None) for pure filtering.	None

Returns:

Type	Description
Belief	The posterior belief — a new Belief; the prior is left untouched.

Raises:

Type	Description
ValueError	If observation is not shape (m,), prior is not a belief over the model's n-D state, the model has a control matrix but action is None, or action is not shape (p,). (All enforced in _validate_inputs.)

Source code in src/cpomdp/backends/kalman.py

def infer_states(
    self,
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief:
    """Advance the belief by one filter step.

    Runs one predict/update cycle: step the prior through the dynamics
    (applying ``action`` if the model has a control matrix), then correct the
    prediction toward ``observation`` using the Kalman gain. In steady-state
    mode the gain and covariance are the frozen fixed-point values; otherwise
    they are recomputed from ``prior.cov`` on this step.

    The numeric work is delegated to the ``jit``-compiled module kernels
    (``_gain_and_posterior_cov``, ``_posterior_mean``); this method stays the
    eager orchestrator that validates inputs and wraps the result in a
    ``Belief``.

    Args:
        observation: The latest sensor reading, shape ``(m,)``.
        prior: The current belief, treated as this step's previous posterior.
            Never mutated.
        action: The action just taken, shape ``(p,)``. Required iff the model
            has a control matrix; ignored (pass ``None``) for pure filtering.

    Returns:
        The posterior belief — a new ``Belief``; the prior is left untouched.

    Raises:
        ValueError: If ``observation`` is not shape ``(m,)``, ``prior`` is not
            a belief over the model's ``n``-D state, the model has a control
            matrix but ``action`` is ``None``, or ``action`` is not shape
            ``(p,)``. (All enforced in ``_validate_inputs``.)
    """
    model = self.model
    observation, action = validate_step_inputs(model, observation, prior, action)
    control_matrix = model.control_matrix
    if control_matrix is None:
        control_term = jnp.zeros(model.n_states)
    else:
        # validate_step_inputs guarantees a non-None action when control exists
        assert action is not None
        control_term = control_matrix @ action

    sensor_is_fixed = (
        model.observation_model is None or model.observation_model.is_fixed
    )
    process_is_fixed = (
        model.dynamics_noise_model is None or model.dynamics_noise_model.is_fixed
    )

    # μ⁻ is needed only to linearize a state-dependent sensor and/or process
    # noise; the fully-fixed hot path computes no extra matvec.
    mean_pred = (
        model.dynamics_matrix @ prior.mean + control_term
        if not (sensor_is_fixed and process_is_fixed)
        else prior.mean  # placeholder, unused on the fixed path
    )

    if sensor_is_fixed:
        # fixed sensor: direct reads, byte-identical hot path (no linearize).
        observation_matrix, observation_noise = (
            model.observation_matrix,
            model.observation_noise,
        )
    else:
        # state-dependent R(x), linearized at μ⁻ (the EFE kernel's point).
        observation_matrix, observation_noise = model.observation_model.linearize(
            mean_pred
        )

    if process_is_fixed:
        dynamics_noise = model.dynamics_noise
    else:
        # state-dependent Q(x), evaluated at μ⁻ — the dual of the R(x) gate.
        dynamics_noise = model.dynamics_noise_model.noise_at(mean_pred)

    if self.steady_state:
        gain, cov_post = self._steady_gain, self._steady_cov  # frozen
    else:
        gain, cov_post = _gain_and_posterior_cov(
            model.dynamics_matrix,
            observation_matrix,
            dynamics_noise,
            observation_noise,
            prior.cov,
        )

    mean_post = _posterior_mean(
        model.dynamics_matrix,
        observation_matrix,
        prior.mean,
        control_term,
        gain,
        observation,
    )

    return Belief(mean=mean_post, cov=cov_post)

RxInferBackend

RxInferBackend(model: LinearGaussianModel)

Linear-Gaussian filtering via RxInfer.jl — the oracle backend.

Satisfies the InferenceBackend protocol: built from a model, advances a belief one step at a time. No steady-state mode — that belongs to the native fast path; this backend exists for correctness, not speed. The first instance built in a process loads the Julia runtime; later ones reuse it.

Source code in src/cpomdp/backends/rxinfer.py

def __init__(self, model: LinearGaussianModel) -> None:
    self.model = model
    self._jl = _julia()

infer_states

infer_states(
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief

Advance the belief one filter step: prior in, posterior out.

Parameters:

Name	Type	Description	Default
observation	ArrayLike	Latest sensor reading, shape (m,).	required
prior	Belief	Current belief; never mutated.	required
action	ArrayLike | None	Action just taken, shape (p,). Required iff the model has a control matrix; pass None for pure filtering.	None

Raises:

Type	Description
ValueError	On a shape/None mismatch (see validate_step_inputs).

Source code in src/cpomdp/backends/rxinfer.py

def infer_states(
    self,
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief:
    """Advance the belief one filter step: prior in, posterior out.

    Args:
        observation: Latest sensor reading, shape ``(m,)``.
        prior: Current belief; never mutated.
        action: Action just taken, shape ``(p,)``. Required iff the model has
            a control matrix; pass ``None`` for pure filtering.

    Raises:
        ValueError: On a shape/None mismatch (see ``validate_step_inputs``).
    """
    model = self.model
    observation, action = validate_step_inputs(model, observation, prior, action)
    control_matrix = model.control_matrix
    if control_matrix is None:
        control_term = jnp.zeros(model.n_states)
    else:
        # validate_step_inputs guarantees a non-None action when control exists
        assert action is not None
        control_term = control_matrix @ action

    # juliacall speaks numpy, not jax.Array, so coerce every array as it
    # crosses into Julia and coerce the posteriors back on the way out.
    mean_post, cov_post = self._jl.cpomdp_run_step(
        np.asarray(observation),
        np.asarray(model.dynamics_matrix),
        np.asarray(model.observation_matrix),
        np.asarray(model.dynamics_noise),
        np.asarray(model.observation_noise),
        np.asarray(prior.mean),
        np.asarray(prior.cov),
        np.asarray(control_term),
    )

    return Belief(mean=np.asarray(mean_post), cov=np.asarray(cov_post))

CouplingGraphBackend is the branching peer: message passing on a CouplingGraph rather than a chain. A state-dependent R(x) on a coupled node cannot be flattened to a fixed linear-Gaussian model, and asking it to raises IncompatibleLinearizationError.

CouplingGraphBackend

CouplingGraphBackend(
    graph: CouplingGraph,
    transitions: Sequence[GaussianTransition],
    *,
    control_matrix: ArrayLike | None = None,
    readout_node: int | None = None,
    partition: Sequence[Sequence[int]] | None = None,
)

FFG message-passing inference on a branching linear-Gaussian tree.

Implements the InferenceBackend protocol for a CouplingGraph whose nodes each carry their own temporal dynamics. Constructed once from the graph and the per-node transitions, then advances a joint belief over every node one step at a time (prior in, posterior out). Read a single node back with marginal / readout.

Parameters:

Name	Type	Description	Default
graph	CouplingGraph	the rooted tree — the structural couplings (the within-slice drive child = W·parent + noise) and the per-node observations.	required
transitions	Sequence[GaussianTransition]	one GaussianTransition (A_i, Q_i) per node, indexed by node, so transitions[i] is node i's own dynamics. Each Q_i must be positive-definite (the information form inverts it), the same divergence from moment-form Kalman that ChainBackend carries.	required
control_matrix	ArrayLike | None	B, shape (n_total, p), mapping an action into the joint None for a pure filtering model. n_total = sum(graph.dims).	None
readout_node	int | None	the node readout returns; defaults to graph.root. The latent of interest need not be the root (issue #25).	None

An observation passed to infer_states is the readings of the observed nodes stacked in ascending node-index order, each node contributing its sensor's rows.

Validate the wiring and front-load the data-independent work (ADR-002).

Everything here depends only on the graph and the transitions, never on a per-step observation / prior / action, so it is built once and reused across every infer_states call.

Source code in src/cpomdp/backends/coupling.py

def __init__(
    self,
    graph: CouplingGraph,
    transitions: Sequence[GaussianTransition],
    *,
    control_matrix: ArrayLike | None = None,
    readout_node: int | None = None,
    partition: Sequence[Sequence[int]] | None = None,
) -> None:
    """Validate the wiring and front-load the data-independent work (ADR-002).

    Everything here depends only on the graph and the transitions, never on a
    per-step observation / prior / action, so it is built once and reused across
    every ``infer_states`` call.
    """
    self._validate_transitions(graph, transitions)

    self.graph = graph
    self.dims = graph.dims
    self.transitions = tuple(transitions)
    self._offsets = tuple(int(o) for o in np.cumsum([0, *graph.dims]))
    self.n_total = self._offsets[-1]
    self.readout_node = self._resolve_readout_node(readout_node)
    self._control = self._coerce_control(control_matrix)
    self._partition = self._resolve_partition(partition)
    # Fully-factored (every cluster a singleton) → the carry keeps the belief
    # block-diagonal, so the within-slice solve is a tree and cheap two-pass BP
    # (``infer_all``) replaces the dense joint solve (ADR-016 energy lever).
    self._is_factored = all(len(cluster) == 1 for cluster in self._partition)

    # Front-loaded factor tier — built once, reused every step (ADR-002).
    self._transition = self._build_transition()
    self._structural_precision = self._assemble_structural_precision()
    self._obs_layout, self.n_observations = self._build_observation_layout()
    self._flat_model = self._build_validation_model()
    self._partition_mask = self._build_partition_mask()
    self._sensor_is_fixed = all(
        o.is_fixed for o in self.graph.observations.values()
    )
    # node -> its partition cluster index (ADR-018 admissibility diagnostic).
    self._cluster_of = {
        node: cid for cid, cluster in enumerate(self._partition) for node in cluster
    }

model property

model: LinearGaussianModel

The flat LinearGaussianModel of the real interface (the backend model).

Block-diagonal dynamics (F, Q), the real sensors (C, R), and the control — the joint state as one linear-Gaussian model, carrying the metadata an Agent reads (n_observations, control_matrix, n_controls). The Agent derives its own model from this, so passing the backend is enough (issue #26). Distinct from to_flat_model(), which augments this with the structural couplings as pseudo-observations for the oracle cross-check.

observation_model property

observation_model: tuple[Array, Array]

The real sensor (C, R) embedded in the joint state — what the EFE reads.

C (shape (n_observations, n_total)) reads the observed nodes out of the joint state, R their noise. The real sensors, not the structural pseudo- observations to_flat_model adds; the EFE pragmatic and epistemic terms both read them.

observation_noise_at

observation_noise_at(
    mean: ArrayLike,
) -> Float64[jax.Array, "m m"]

The stacked observation noise R at a predicted mean.

Each observed node's R is linearized at mean[node_block] (fixed sensors ignore the mean and return their constant R); the block-diagonal stack is the action-dependent R(μ⁺) the EFE selector folds per candidate (issue #27).

Source code in src/cpomdp/backends/coupling.py

def observation_noise_at(self, mean: ArrayLike) -> Float64[jax.Array, "m m"]:
    """The stacked observation noise R at a predicted mean.

    Each observed node's R is linearized at ``mean[node_block]`` (fixed sensors
    ignore the mean and return their constant R); the block-diagonal stack is the
    action-dependent ``R(μ⁺)`` the EFE selector folds per candidate (issue #27).
    """
    mean = jnp.asarray(mean, dtype=float)
    blocks = [
        self.graph.observations[node].linearize(mean[self._block(node)])[1]
        for node, _lo, _hi in self._obs_layout
    ]
    if not blocks:
        return jnp.zeros((0, 0))
    return jax.scipy.linalg.block_diag(*blocks)

infer_states

infer_states(
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief

Advance the joint belief by one filter step over the tree.

Validate the inputs, form the control shift b = control @ action, then run the driven-relaxation pipeline: lift the joint prior into canonical form, predict through the block-diagonal per-node dynamics (the temporal edges), add the structural coupling precision and the per-node observation messages (the within-slice update), and to_moment the joint posterior back.

Parameters:

Name	Type	Description	Default
observation	ArrayLike	the observed nodes' readings, stacked in ascending node-index order, shape (n_observations,).	required
prior	Belief	the current joint belief over all nodes. Never mutated.	required
action	ArrayLike | None	the action just taken, shape (p,); required iff the model has a control matrix, None for pure filtering.	None

Returns:

Type	Description
Belief	The posterior joint belief over all nodes; slice one node out with
Belief	marginal / readout.

Source code in src/cpomdp/backends/coupling.py

def infer_states(
    self,
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> Belief:
    """Advance the joint belief by one filter step over the tree.

    Validate the inputs, form the control shift ``b = control @ action``, then run
    the driven-relaxation pipeline: lift the joint prior into canonical form,
    ``predict`` through the block-diagonal per-node dynamics (the temporal edges),
    add the structural coupling precision and the per-node observation messages (the
    within-slice update), and ``to_moment`` the joint posterior back.

    Args:
        observation: the observed nodes' readings, stacked in ascending node-index
            order, shape ``(n_observations,)``.
        prior: the current *joint* belief over all nodes. Never mutated.
        action: the action just taken, shape ``(p,)``; required iff the model has a
            control matrix, ``None`` for pure filtering.

    Returns:
        The posterior *joint* belief over all nodes; slice one node out with
        ``marginal`` / ``readout``.
    """
    observation, action = validate_step_inputs(
        self._flat_model, observation, prior, action
    )
    if self._is_factored:
        # Fully-factored carry: skip the dense joint solve, run tree BP instead.
        return self._infer_states_factored(observation, prior, action)
    precision, potential = self._assemble_unchecked(observation, prior, action)
    mean, cov = CanonicalGaussian._unchecked(
        precision, potential
    ).to_moment()  # exact joint: Σ = Λ⁻¹, μ = Λ⁻¹h
    factored_cov, _severed = self._carry(cov)
    return Belief(mean=mean, cov=factored_cov)

partition_error

partition_error(
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> float

The severed mass a step under this partition drops (ADR-016 diagnostic).

The norm of the between-cluster covariance blocks the carry zeros — how much cross-cluster correlation the partition drops at the time boundary, the approximation cost of the cut. A covariance magnitude, not bits and not a rate. 0.0 for the full-joint [[all]] partition (exact), growing as the cut severs more correlation. The per-node marginals this slice are unaffected — the carry only drops what is carried forward (ADR-017).

This is the eager convenience surface: it forces a host float and so is not itself jit-able. For a per-run profile with no host syncs, stack the _carry scalar inside a traced rollout instead.

Source code in src/cpomdp/backends/coupling.py

def partition_error(
    self,
    observation: ArrayLike,
    prior: Belief,
    action: ArrayLike | None = None,
) -> float:
    """The severed mass a step under this partition drops (ADR-016 diagnostic).

    The norm of the between-cluster *covariance* blocks the carry zeros — how much
    cross-cluster correlation the partition drops at the time boundary, the
    approximation cost of the cut. A covariance magnitude, not bits and not a rate.
    ``0.0`` for the full-joint ``[[all]]`` partition (exact), growing as the cut
    severs more correlation. The per-node marginals this slice are unaffected — the
    carry only drops what is *carried forward* (ADR-017).

    This is the eager convenience surface: it forces a host ``float`` and so is not
    itself jit-able. For a per-run profile with no host syncs, stack the
    ``_carry`` scalar inside a traced rollout instead.
    """
    precision, potential = self._assemble_posterior(observation, prior, action)
    _mean, cov = CanonicalGaussian._unchecked(precision, potential).to_moment()
    _factored, severed = self._carry(cov)
    return float(severed)

rollout

rollout(
    prior: Belief,
    observations: ArrayLike,
    actions: ArrayLike | None = None,
) -> tuple[Belief, jax.Array]

Filter a whole sequence, profiling the severed mass at each step (ADR-016).

One traced lax.scan pass with no per-step host syncs: the joint belief is the scan carry, and each step emits its posterior and the severed mass its carry drops. This is the per-run diagnostic a mutable per-step field could not give — the profile is produced inside the traced rollout, not read off self.

Parameters:

Name	Type	Description	Default
prior	Belief	the joint belief the run starts from.	required
observations	ArrayLike	the stacked readings per step, shape (T, n_observations).	required
actions	ArrayLike | None	the per-step actions, shape (T, p), required iff the model has a control matrix; None for pure filtering.	None

Returns:

Type	Description
Belief	(beliefs, severed_masses) — the time-stacked posteriors (a
Array	Belief with a leading time axis: mean
tuple[Belief, Array]	(T, n_total), cov (T, n_total, n_total)) and the length-T
tuple[Belief, Array]	severed-mass profile. The full-joint [[all]] partition profiles
tuple[Belief, Array]	all-zero.

Source code in src/cpomdp/backends/coupling.py

def rollout(
    self,
    prior: Belief,
    observations: ArrayLike,
    actions: ArrayLike | None = None,
) -> tuple[Belief, jax.Array]:
    """Filter a whole sequence, profiling the severed mass at each step (ADR-016).

    One traced ``lax.scan`` pass with no per-step host syncs: the joint belief is
    the scan carry, and each step emits its posterior and the severed mass its carry
    drops. This is the per-run diagnostic a mutable per-step field could not give —
    the profile is produced *inside* the traced rollout, not read off ``self``.

    Args:
        prior: the joint belief the run starts from.
        observations: the stacked readings per step, shape ``(T, n_observations)``.
        actions: the per-step actions, shape ``(T, p)``, required iff the model has
            a control matrix; ``None`` for pure filtering.

    Returns:
        ``(beliefs, severed_masses)`` — the time-stacked posteriors (a
        [`Belief`][cpomdp.Belief] with a leading time axis: ``mean``
        ``(T, n_total)``, ``cov`` ``(T, n_total, n_total)``) and the length-``T``
        severed-mass profile. The full-joint ``[[all]]`` partition profiles
        all-zero.
    """
    observations = jnp.asarray(observations, dtype=float)
    if observations.ndim != 2 or observations.shape[1] != self.n_observations:
        raise ValueError(
            f"observations must have shape (T, {self.n_observations}), "
            f"got {observations.shape}"
        )

    if self._control is None:
        if actions is not None:
            raise ValueError("actions given, but this model has no control matrix")

        def step(belief: Belief, observation: jax.Array):
            return self._instrumented_step(belief, observation, None)

        _final, outputs = jax.lax.scan(step, prior, observations)
        return outputs

    if actions is None:
        raise ValueError("this model has a control matrix; actions are required")
    actions = jnp.asarray(actions, dtype=float)
    if actions.shape[0] != observations.shape[0]:
        raise ValueError(
            f"actions has {actions.shape[0]} steps, but observations "
            f"has {observations.shape[0]}"
        )

    def controlled_step(belief: Belief, step_inputs: tuple[jax.Array, jax.Array]):
        observation, action = step_inputs
        return self._instrumented_step(belief, observation, action)

    _final, outputs = jax.lax.scan(controlled_step, prior, (observations, actions))
    return outputs

marginal

marginal(node: int, belief: Belief) -> Belief

The marginal belief at a single node — a pure slice of the joint belief.

The joint carried by infer_states makes every node exact, so a chosen node's belief is a slice, no re-inference (issue #25: the target latent need not be the root).

Source code in src/cpomdp/backends/coupling.py

def marginal(self, node: int, belief: Belief) -> Belief:
    """The marginal belief at a single ``node`` — a pure slice of the joint belief.

    The joint carried by ``infer_states`` makes every node exact, so a chosen node's
    belief is a slice, no re-inference (issue #25: the target latent need not be the
    root).
    """
    block = self._block(node)
    return Belief(mean=belief.mean[block], cov=belief.cov[block, block])

readout

readout(belief: Belief) -> Belief

The marginal at readout_node (the root by default).

Source code in src/cpomdp/backends/coupling.py

def readout(self, belief: Belief) -> Belief:
    """The marginal at ``readout_node`` (the root by default)."""
    return self.marginal(self.readout_node, belief)

block

block(node: int) -> range

The state indices node node occupies in the joint (a public _block).

Lets a caller (e.g. the EFE selector) turn a node index into the joint-state block the epistemic term targets — info_node node → block (issue #26).

Source code in src/cpomdp/backends/coupling.py

def block(self, node: int) -> range:
    """The state indices node ``node`` occupies in the joint (a public ``_block``).

    Lets a caller (e.g. the EFE selector) turn a node index into the joint-state
    block the epistemic term targets — ``info_node`` node → ``block`` (issue #26).
    """
    return range(self._offsets[node], self._offsets[node + 1])

severed_efe_edges

severed_efe_edges() -> tuple[Coupling, ...]

The EFE-relevant edges this carry partition cuts (ADR-018 diagnostic).

Each efe_relevant coupling whose parent and child fall in different clusters — so the carry drops the cross-temporal covariance it holds, breaking the integration of that information about the targeted latent. Non-empty means the EFE selector must refuse this partition; empty for the exact [[all]] carry and for any cut that only severs unflagged edges (e.g. the methylation cut). A structural read, not a filter step — no covariance is formed here.

Source code in src/cpomdp/backends/coupling.py

def severed_efe_edges(self) -> tuple[Coupling, ...]:
    """The EFE-relevant edges this carry partition cuts (ADR-018 diagnostic).

    Each ``efe_relevant`` coupling whose parent and child fall in different clusters
    — so the carry drops the cross-temporal covariance it holds, breaking the
    integration of that information about the targeted latent. Non-empty means the
    EFE selector must refuse this partition; empty for the exact ``[[all]]`` carry
    and for any cut that only severs unflagged edges (e.g. the methylation cut). A
    structural read, not a filter step — no covariance is formed here.
    """
    return tuple(
        edge
        for edge in self.graph.couplings
        if edge.efe_relevant
        and self._cluster_of[edge.parent] != self._cluster_of[edge.child]
    )

predicted_belief

predicted_belief(
    prior: Belief, action: ArrayLike | None = None
) -> Belief

The joint belief after dynamics + structural couplings, before observing.

The predict side of one step (Σ⁺, μ⁺): push the joint prior through the per-node dynamics under action and fold in the structural couplings, but add no observation. This is the covariance the EFE reads — a candidate action's epistemic value is the info gain from this Σ⁺ to the post-observation covariance (issue #26). Pure (no host-side validation), so it rides jit / vmap over a grid of candidate actions.

Source code in src/cpomdp/backends/coupling.py

def predicted_belief(
    self, prior: Belief, action: ArrayLike | None = None
) -> Belief:
    """The joint belief after dynamics + structural couplings, before observing.

    The predict side of one step (Σ⁺, μ⁺): push the joint ``prior`` through the
    per-node dynamics under ``action`` and fold in the structural couplings, but add
    no observation. This is the covariance the EFE reads — a candidate action's
    epistemic value is the info gain from this Σ⁺ to the post-observation covariance
    (issue #26). Pure (no host-side validation), so it rides ``jit`` / ``vmap`` over
    a grid of candidate actions.
    """
    coerced = None if action is None else jnp.asarray(action, dtype=float)
    precision, potential = self._predicted_precision(prior, coerced)
    mean, cov = CanonicalGaussian._unchecked(precision, potential).to_moment()
    return Belief(mean=mean, cov=cov)

to_flat_model

to_flat_model() -> LinearGaussianModel

The tree flattened into one dense LinearGaussianModel (the oracle route).

The temporal edges become the block-diagonal transition (F, Q); the real observations and the structural couplings stack into one sensor, each coupling an always-zero pseudo-observation child − W·parent ~ N(0, Q_struct). Running KalmanBackend or RxInferBackend on this reproduces this backend's filter exactly — the independent cross-check, and what the Phase-3 demo contrasts the native FFG against. Pad a step's readings with flat_observation first; the returned prior is an unused placeholder (pass the real prior per step).

With couplings present a state-dependent R(x) is a category error for this route (IncompatibleLinearizationError — the flat Kalman linearizes at μ⁻, the FFG at μ⁺). Coupling-free, μ⁻ = μ⁺, so the emitted model carries the graph's own state-dependent sensor and stays faithful; it is a state-dependent flat model, not a fixed one, and no fixed one exists.

The refusal reads the sensor's declared state-dependence. Whether R really varies over the means a policy can reach is a property of the noise function and the reachable set, which this cannot decide — cpomdp.diagnostics.probe_model samples it.

Source code in src/cpomdp/backends/coupling.py

def to_flat_model(self) -> LinearGaussianModel:
    """The tree flattened into one dense ``LinearGaussianModel`` (the oracle route).

    The temporal edges become the block-diagonal transition (F, Q); the real
    observations and the structural couplings stack into one sensor, each coupling
    an always-zero pseudo-observation ``child − W·parent ~ N(0, Q_struct)``. Running
    [`KalmanBackend`][cpomdp.KalmanBackend] or ``RxInferBackend`` on this
    reproduces this backend's filter exactly — the independent cross-check, and what
    the Phase-3 demo contrasts the native FFG against. Pad a step's readings with
    ``flat_observation`` first; the returned prior is an unused placeholder (pass
    the real prior per step).

    With couplings present a state-dependent ``R(x)`` is a *category error* for this
    route (``IncompatibleLinearizationError`` — the flat Kalman linearizes at μ⁻,
    the FFG at μ⁺). Coupling-free, μ⁻ = μ⁺, so the emitted model carries the graph's
    own state-dependent sensor and stays faithful; it is a state-dependent flat
    model, not a fixed one, and no fixed one exists.

    The refusal reads the sensor's *declared* state-dependence. Whether ``R`` really
    varies over the means a policy can reach is a property of the noise function and
    the reachable set, which this cannot decide — ``cpomdp.diagnostics.probe_model``
    samples it.
    """
    if not self._sensor_is_fixed and self.graph.couplings:
        # A mean-shifting coupling makes μ⁺ ≠ μ⁻, so no fixed flat model
        # reproduces R(μ⁺) -- inherent, not a missing feature.
        raise IncompatibleLinearizationError(
            "Cannot flatten a state-dependent R(x) model with couplings: a "
            "mean-shifting coupling makes the prediction mean μ⁺ differ from "
            "the prior mean μ⁻, so no fixed linear-Gaussian model reproduces "
            "R(μ⁺). This reads the sensor's declared state-dependence; a "
            "declared R(x) that ignores the state is constant in fact and does "
            "flatten — see cpomdp.diagnostics.probe_model."
        )
    rows, noise_blocks = self._real_observation_blocks()
    # Coupling-free R(x): μ⁻ = μ⁺, so handing the graph's own sensor to a flat
    # backend reproduces this filter. With couplings the branch above has already
    # refused, so this is None whenever a pseudo-observation is about to be added.
    observation_model = (
        None if self._sensor_is_fixed else self._flat_model.observation_model
    )
    for edge in self.graph.couplings:  # child − W·parent ~ N(0, Q_struct)
        coupling = edge.factor.coupling  # W
        child_dim = coupling.shape[0]
        row = jnp.zeros((child_dim, self.n_total))
        row = row.at[:, self._block(edge.child)].set(jnp.eye(child_dim))
        row = row.at[:, self._block(edge.parent)].set(-coupling)
        rows.append(row)
        noise_blocks.append(edge.factor.coupling_noise)
    return LinearGaussianModel(
        dynamics_matrix=self._transition.dynamics_matrix,
        observation_matrix=jnp.vstack(rows),
        dynamics_noise=self._transition.dynamics_noise,
        observation_noise=jax.scipy.linalg.block_diag(*noise_blocks),
        prior=Belief(jnp.zeros(self.n_total), jnp.eye(self.n_total)),
        control_matrix=self._control,
        observation_model=observation_model,
    )

flat_observation

flat_observation(observation: ArrayLike) -> jax.Array

Pad a step's readings with the structural pseudo-observations' zeros.

to_flat_model stacks the real observations then the structural couplings, so a flat backend consumes [readings, zeros(n_structural)].

Source code in src/cpomdp/backends/coupling.py

def flat_observation(self, observation: ArrayLike) -> jax.Array:
    """Pad a step's readings with the structural pseudo-observations' zeros.

    ``to_flat_model`` stacks the real observations then the structural couplings, so
    a flat backend consumes ``[readings, zeros(n_structural)]``.
    """
    observation = jnp.asarray(observation, dtype=float)
    n_structural = sum(
        edge.factor.coupling.shape[0] for edge in self.graph.couplings
    )
    return jnp.concatenate([observation, jnp.zeros(n_structural)])

IncompatibleLinearizationError

Bases: ValueError

The flattened-Kalman route can't reproduce this backend's linearization.

to_flat_model encodes couplings as pseudo-observations, so the flat Kalman linearizes R at the pre-coupling predicted mean μ⁻, while the FFG linearizes at the coupling-resolved μ⁺ (issue #27, ADR-019). With a state-dependent R(x) and a mean-shifting coupling, μ⁻ ≠ μ⁺ and no fixed flat model reproduces the filter. The oracle is the standalone NumPy R(μ⁺) filter (backend tests) plus the single-node CallableSensor cross-check.

[^bagaev2023rxinfer]: Dmitry Bagaev, Albert Podusenko, and Bert de Vries. RxInfer: a Julia package for reactive real-time Bayesian inference. Journal of Open Source Software, 8(84):5161, 2023. URL: https://doi.org/10.21105/joss.05161, doi:10.21105/joss.05161.