Factor graph models

A CouplingGraph declares a rooted tree of Gaussian-coupled variables. It holds the shape a chain cannot: a node with three or more neighbours, inferred through the ones around it. Coupling edges carry the within-slice drive child = W·parent + noise, and the factor types below supply each node's dynamics and its observation. Run one through a CouplingGraphBackend.

CouplingGraph

CouplingGraph(
    root: int,
    dims: Sequence[int],
    couplings: Sequence[Coupling],
    observations: Mapping[int, ObservationFactor],
)

A rooted tree of Gaussian-coupled variables.

The N nodes are indexed 0..N-1 with dimensions dims. couplings are the tree edges, directed away from root, and observations maps a node index to the observation factor attached to it (fixed GaussianObservation or state-dependent CallableGaussianObservation). Construction validates the wiring and raises if it is malformed.

Parameters:

Name	Type	Description	Default
root	int	index of the node the tree is rooted at.	required
dims	Sequence[int]	dims[i] is the dimension of node i; its length is the node count.	required
couplings	Sequence[Coupling]	the tree edges — one per non-root node, each that node's only parent.	required
observations	Mapping[int, ObservationFactor]	maps a node index to its ObservationFactor (fixed or R(x)).	required

Raises:

Type	Description
ValueError	if dims is empty or non-positive; if root, an edge, or an observation references an out-of-range node; if an edge's factor W is not (dim[child], dim[parent]); or if the edges do not form a tree rooted at root — the root has a parent, a node has two parents or none, or the edges contain a cycle.

Source code in src/cpomdp/ffg/graph.py

def __init__(
    self,
    root: int,
    dims: Sequence[int],
    couplings: Sequence[Coupling],
    observations: Mapping[int, ObservationFactor],
) -> None:
    self.root = int(root)
    self.dims = tuple(int(d) for d in dims)
    self.couplings = tuple(couplings)
    self.observations = dict(observations)
    self._validate()

infer

infer(
    prior: Belief, readings: Mapping[int, ArrayLike]
) -> Belief

Compute the marginal belief at the root from a prior and per-node readings.

Each reading becomes a message about its node; those messages are passed up the tree through the couplings and combined at the root with the prior, giving the root's posterior over every reading. Only the root is converted to and from moment form — once to lift the prior in, once to read the result out — while every message in between stays in canonical form.

Parameters:

Name	Type	Description	Default
prior	Belief	the belief on the root node, taken as its prior.	required
readings	Mapping[int, ArrayLike]	maps a node index to that node's observation; each such node must carry a fixed GaussianObservation (static inference has no predicted mean to linearize a state-dependent R(x) at).	required

Returns:

Type	Description
Belief	The marginal belief at the root.

Source code in src/cpomdp/ffg/graph.py

def infer(self, prior: Belief, readings: Mapping[int, ArrayLike]) -> Belief:
    """Compute the marginal belief at the root from a prior and per-node readings.

    Each reading becomes a message about its node; those messages are passed up the
    tree through the couplings and combined at the root with the prior, giving the
    root's posterior over every reading. Only the root is converted to and from
    moment form — once to lift the prior in, once to read the result out — while
    every message in between stays in canonical form.

    Args:
        prior: the belief on the root node, taken as its prior.
        readings: maps a node index to that node's observation; each such node must
            carry a fixed [`GaussianObservation`][cpomdp.GaussianObservation]
            (static inference has no predicted mean to linearize a state-dependent
            ``R(x)`` at).

    Returns:
        The marginal belief at the root.
    """

    def combine(acc, key, msg):
        """Add ``msg`` to ``acc[key]``, or start the slot with it if absent."""
        return acc[key] + msg if key in acc else msg

    def depth(node: int) -> int:
        """The number of edges from ``node`` up to the root."""
        hops = 0
        while node != self.root:
            node = parent_edge[node].parent
            hops += 1
        return hops

    # Lift the moment-form prior into a canonical message on the root.
    prior_precision = jnp.linalg.inv(prior.cov)  # Λ₀ = Σ⁻¹
    prior_msg = CanonicalGaussian._unchecked(
        prior_precision, prior_precision @ prior.mean
    )  # h₀ = Λ₀μ; invariant-preserving lift of a validated Belief — no re-validate

    # Seed each observed node with its reading's message, then fold in the prior at
    # the root (which may already hold the root's own observation).
    acc = {
        node: self.observations[node].message(reading)
        for node, reading in readings.items()
    }
    acc[self.root] = combine(acc, self.root, prior_msg)

    # Pass messages up to the root, deepest nodes first so every child folds into a
    # node before that node is itself sent up to its parent.
    parent_edge = {edge.child: edge for edge in self.couplings}
    order = sorted(parent_edge, key=depth, reverse=True)
    for node in order:
        if node not in acc:  # an unobserved leaf has nothing to send up
            continue
        edge = parent_edge[node]
        # Summarise everything known at `node` onto its parent, eliminating `node`.
        up = edge.factor.message_to_parent(acc[node])
        acc[edge.parent] = combine(acc, edge.parent, up)

    # Read the accumulated root message back into moment form.
    mean, cov = acc[self.root].to_moment()  # Σ = Λ⁻¹, μ = Λ⁻¹h
    return Belief(mean=mean, cov=cov)

infer_all

infer_all(
    seeds: Mapping[int, CanonicalGaussian],
) -> dict[int, CanonicalGaussian]

Every node's exact marginal by two-pass belief propagation over the tree.

Where infer collects to the root and returns only that one marginal, this adds a downward distribute pass so every node's marginal comes back — the cheap, structure-exploiting alternative to a dense joint solve. Each seed is a node's already-formed canonical message (its local prior + evidence, combined by the caller); marginalisation stays in canonical form, so no node is inverted on this path.

Parameters:

Name	Type	Description	Default
seeds	Mapping[int, CanonicalGaussian]	a canonical message per node — the node's local information. Unlike infer's raw readings, these are ready-made CanonicalGaussian messages, not observations still needing a factor.	required

Returns:

Type	Description
dict[int, CanonicalGaussian]	A CanonicalGaussian marginal per node index; call .to_moment() on
dict[int, CanonicalGaussian]	any one for its (mean, cov).

Source code in src/cpomdp/ffg/graph.py

def infer_all(
    self, seeds: Mapping[int, CanonicalGaussian]
) -> dict[int, CanonicalGaussian]:
    """Every node's exact marginal by two-pass belief propagation over the tree.

    Where ``infer`` collects to the root and returns only that one marginal, this
    adds a downward *distribute* pass so **every** node's marginal comes back — the
    cheap, structure-exploiting alternative to a dense joint solve. Each seed is a
    node's already-formed canonical message (its local prior + evidence, combined by
    the caller); marginalisation stays in canonical form, so no node is inverted on
    this path.

    Args:
        seeds: a canonical message per node — the node's local information. Unlike
            ``infer``'s raw ``readings``, these are ready-made ``CanonicalGaussian``
            messages, not observations still needing a factor.

    Returns:
        A ``CanonicalGaussian`` marginal per node index; call ``.to_moment()`` on
        any one for its ``(mean, cov)``.
    """

    def combine(acc, key, msg):
        """Add ``msg`` to ``acc[key]``, or start the slot with it if absent."""
        return acc[key] + msg if key in acc else msg

    def depth(node: int) -> int:
        """The number of edges from ``node`` up to the root."""
        hops = 0
        while node != self.root:
            node = parent_edge[node].parent
            hops += 1
        return hops

    parent_edge = {edge.child: edge for edge in self.couplings}
    order = sorted(parent_edge, key=depth, reverse=True)  # deepest-first

    # Collect: fold each subtree up into its parent (as ``infer``), but cache every
    # edge's upward message for the distribute pass. ``below[node]`` ends holding
    # the node's own seed plus everything its children sent up.
    below = dict(seeds)
    up_msg = {}
    for node in order:
        edge = parent_edge[node]
        up = edge.factor.message_to_parent(below[node])
        up_msg[node] = up
        below[edge.parent] = combine(below, edge.parent, up)

    # Distribute: push each parent's finished marginal back down, shallowest-first
    # so a parent is done before its children. Divide out the child's own upward
    # message first (``- up_msg[node]``) so its evidence does not double-count.
    marginals = {self.root: below[self.root]}
    for node in reversed(order):
        edge = parent_edge[node]
        down = edge.factor.message_to_child(marginals[edge.parent] - up_msg[node])
        marginals[node] = combine(below, node, down)
    return marginals

Coupling dataclass

Coupling(
    parent: int,
    child: int,
    factor: GaussianCoupling,
    tau: float,
    efe_relevant: bool = False,
)

A directed edge from a parent node to a child node: child = W·parent + noise.

parent and child are node indices, oriented so the parent is the endpoint nearer the tree's root. factor is the GaussianCoupling holding this edge's W (shape (dim[child], dim[parent])) and its noise covariance. tau is a time-constant carried alongside the edge; it is metadata and does not affect the factor.

efe_relevant is a modeller's declaration that this edge carries information the instrumental epistemic depends on — a physics call, not a structural one (in chemotaxis the gradient rides receptor->CheA->CheY, while CheA->CheB methylation is observed and coupled yet gradient-blind). A carry partition (ADR-016) that severs a flagged edge drops the cross-temporal covariance it holds, breaking the integration of that information about a slow latent; the EFE selector refuses such a partition (ADR-018). Default False; it does not affect the factor or the filter.

GaussianCoupling dataclass

GaussianCoupling(
    coupling: ArrayLike, coupling_noise: ArrayLike
)

Tier-1 structural coupling factor N(child; W·parent, Q) — a graph edge.

Where GaussianTransition couples a state to its successor in time, this couples two variables joined by an edge of the factor graph (e.g. the shared CheA node to a branch latent). The maths is identical — a linear-Gaussian coupling — but a coupling carries no time semantics and W need not be square.

• coupling — W, shape (c, p): maps the p-D parent's mean to the c-D child.
• coupling_noise — Q, shape (c, c), positive-definite (it is inverted).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def __init__(self, coupling: ArrayLike, coupling_noise: ArrayLike) -> None:
    object.__setattr__(self, "coupling", jnp.asarray(coupling, dtype=float))
    object.__setattr__(
        self, "coupling_noise", jnp.asarray(coupling_noise, dtype=float)
    )
    self._validate()

message_to_parent

message_to_parent(
    child_message: CanonicalGaussian,
) -> CanonicalGaussian

Summarise what a child's belief says about the parent: eliminate the child.

The coupling is the joint Gaussian over z = [parent, child]::

Λ_J = [[ WᵀQ⁻¹W, −WᵀQ⁻¹ ],     h_J = 0   (a pure coupling has no bias)
       [ −Q⁻¹W,    Q⁻¹   ]]

The upward message:

1. Folds child_message into the child block — its precision into the bottom-right c×c of Λ_J, its potential into the trailing c of h_J (a block add during construction, not __add__).
2. Marginalizes the child out, leaving the message on the p-D parent.

This is the mirror of GaussianTransition.predict (which folds into the parent block and eliminates the parent, emitting downward onto the child); here we fold into the child block and eliminate the child, emitting upward.

Parameters:

Name	Type	Description	Default
child_message	CanonicalGaussian	the incoming belief on the c-D child, as a CanonicalGaussian.	required

Returns:

Type	Description
CanonicalGaussian	A CanonicalGaussian over the p-D parent.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def message_to_parent(self, child_message: CanonicalGaussian) -> CanonicalGaussian:
    """Summarise what a child's belief says about the parent: eliminate the child.

    The coupling is the joint Gaussian over ``z = [parent, child]``::

        Λ_J = [[ WᵀQ⁻¹W, −WᵀQ⁻¹ ],     h_J = 0   (a pure coupling has no bias)
               [ −Q⁻¹W,    Q⁻¹   ]]

    The upward message:

    1. Folds ``child_message`` into the *child* block — its precision into the
       bottom-right ``c×c`` of ``Λ_J``, its potential into the trailing ``c`` of
       ``h_J`` (a block add during construction, *not* ``__add__``).
    2. Marginalizes the child out, leaving the message on the p-D parent.

    This is the mirror of ``GaussianTransition.predict`` (which folds into the
    parent block and eliminates the parent, emitting downward onto the child);
    here we fold into the child block and eliminate the child, emitting upward.

    Args:
        child_message: the incoming belief on the c-D child, as a
            ``CanonicalGaussian``.

    Returns:
        A ``CanonicalGaussian`` over the p-D parent.
    """
    coupling, coupling_noise = self.coupling, self.coupling_noise  # W, Q
    c, p = coupling.shape  # W is (child, parent)
    noise_precision = jnp.linalg.inv(coupling_noise)  # Q⁻¹
    noise_weighted_coupling = noise_precision @ coupling  # Q⁻¹W

    # The incoming message is on the CHILD, so it folds into the child block —
    # the mirror of predict, where the message folds into the parent (state) block.
    parent_block = coupling.T @ noise_weighted_coupling  # WᵀQ⁻¹W
    child_block = noise_precision + child_message.precision  # Q⁻¹ + message

    precision = jnp.block(
        [
            [parent_block, -noise_weighted_coupling.T],
            [-noise_weighted_coupling, child_block],
        ]
    )
    # No bias and no parent message → the parent potential is zero; the child
    # slot carries the incoming message's potential.
    potential = jnp.concatenate([jnp.zeros(p), child_message.potential])

    joint = CanonicalGaussian._unchecked(precision, potential)
    return joint.marginalize(over=range(p, p + c))  # eliminate child, keep parent

message_to_child

message_to_child(
    parent_message: CanonicalGaussian,
) -> CanonicalGaussian

Push a parent's belief down the edge onto the child: eliminate the parent.

The distribute-pass mirror of message_to_parent. Over the same joint on z = [parent, child], but here the parent is known, so its message folds into the parent block of Λ_J and the parent is marginalized out, leaving the message on the c-D child.

Structurally this is GaussianTransition.predict (fold the incoming belief into the source block, eliminate the source, emit onto the target) with a non-square W and no control shift — a pure coupling carries no bias. So on its own the downward message is a full child belief, landing in moment form on mean = W·μ_parent and cov = W·Σ_parent·Wᵀ + Q.

Parameters:

Name	Type	Description	Default
parent_message	CanonicalGaussian	the incoming belief on the p-D parent, as a CanonicalGaussian.	required

Returns:

Type	Description
CanonicalGaussian	A CanonicalGaussian over the c-D child.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def message_to_child(self, parent_message: CanonicalGaussian) -> CanonicalGaussian:
    """Push a parent's belief down the edge onto the child: eliminate the parent.

    The distribute-pass mirror of ``message_to_parent``. Over the same joint on
    ``z = [parent, child]``, but here the *parent* is known, so its message folds
    into the *parent* block of ``Λ_J`` and the parent is marginalized out, leaving
    the message on the c-D child.

    Structurally this is ``GaussianTransition.predict`` (fold the incoming belief
    into the source block, eliminate the source, emit onto the target) with a
    non-square ``W`` and no control shift — a pure coupling carries no bias. So on
    its own the downward message is a full child belief, landing in moment form on
    ``mean = W·μ_parent`` and ``cov = W·Σ_parent·Wᵀ + Q``.

    Args:
        parent_message: the incoming belief on the p-D parent, as a
            ``CanonicalGaussian``.

    Returns:
        A ``CanonicalGaussian`` over the c-D child.
    """
    coupling, coupling_noise = self.coupling, self.coupling_noise  # W, Q
    c, p = coupling.shape  # W is (child, parent)

    noise_precision = jnp.linalg.inv(coupling_noise)  # Q⁻¹
    noise_weighted_coupling = noise_precision @ coupling  # Q⁻¹W

    parent_block = (
        coupling.T @ noise_weighted_coupling + parent_message.precision
    )  # WᵀQ⁻¹W + message
    child_block = noise_precision

    precision = jnp.block(
        [
            [parent_block, -noise_weighted_coupling.T],
            [-noise_weighted_coupling, child_block],
        ]
    )

    potential = jnp.concatenate([parent_message.potential, jnp.zeros(c)])

    joint = CanonicalGaussian._unchecked(precision, potential)
    return joint.marginalize(over=range(p))

tree_flatten

tree_flatten() -> tuple[
    tuple[Float64[Array, "c p"], Float64[Array, "c c"]],
    None,
]

Leaves for JAX: (coupling, coupling_noise), no static aux data.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def tree_flatten(
    self,
) -> tuple[tuple[Float64[Array, "c p"], Float64[Array, "c c"]], None]:
    """Leaves for JAX: ``(coupling, coupling_noise)``, no static aux data."""
    return (self.coupling, self.coupling_noise), None

tree_unflatten classmethod

tree_unflatten(
    aux_data: None,
    children: tuple[
        Float64[Array, "c p"], Float64[Array, "c c"]
    ],
) -> GaussianCoupling

Rebuild from leaves without validating — the leaves may be tracers.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: None,
    children: tuple[Float64[Array, "c p"], Float64[Array, "c c"]],
) -> "GaussianCoupling":
    """Rebuild from leaves without validating — the leaves may be tracers."""
    coupling, coupling_noise = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "coupling", coupling)
    object.__setattr__(obj, "coupling_noise", coupling_noise)
    return obj

GaussianTransition dataclass

GaussianTransition(
    dynamics_matrix: ArrayLike, *, dynamics_noise: ArrayLike
)

Tier-1 dynamics factor N(x'; Ax + b, Q) — emits the forward predict.

Holds the fixed transition and process noise; predict(message, b) pushes a belief on x through the dynamics to a belief on x'.

• dynamics_matrix — A, shape (n, n).
• dynamics_noise — Q, shape (n, n), positive-definite (it is inverted).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def __init__(
    self, dynamics_matrix: ArrayLike, *, dynamics_noise: ArrayLike
) -> None:
    object.__setattr__(
        self, "dynamics_matrix", jnp.asarray(dynamics_matrix, dtype=float)
    )
    object.__setattr__(
        self, "dynamics_noise", jnp.asarray(dynamics_noise, dtype=float)
    )
    self._validate()

from_ou classmethod

from_ou(
    tau: float, *, stationary_var: float, dt: float
) -> GaussianTransition

Build a 1-D transition from Ornstein–Uhlenbeck (OU) parameters.

An Ornstein–Uhlenbeck process is a scalar state that relaxes toward zero on a timescale tau while random noise keeps it wobbling with a stationary variance stationary_var (Σ_stat). Exactly discretising it over a step dt gives the linear-Gaussian transition x' = A·x + noise(Q) (ADR-017):

A = exp(−dt / tau)             — the fraction of the state surviving a step
Q = stationary_var · (1 − A²)  — the kick that holds the stationary variance

Scalar (1-D) only: the vector OU would need a matrix exponential and a Lyapunov solve, which no cpomdp node needs.

Parameters:

Name	Type	Description	Default
tau	float	the relaxation timescale τ (same time unit as dt).	required
stationary_var	float	the steady-state variance Σ_stat the node settles to; A (dynamics) and Q (dynamics_noise) are set so it holds this spread.	required
dt	float	the discretisation step.	required

Returns:

Type	Description
GaussianTransition	A GaussianTransition with 1×1 dynamics_matrix (A) and
GaussianTransition	dynamics_noise (Q).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

@classmethod
def from_ou(
    cls, tau: float, *, stationary_var: float, dt: float
) -> "GaussianTransition":
    """Build a 1-D transition from Ornstein–Uhlenbeck (OU) parameters.

    An Ornstein–Uhlenbeck process is a scalar state that relaxes toward zero on a
    timescale ``tau`` while random noise keeps it wobbling with a stationary
    variance ``stationary_var`` (Σ_stat). Exactly discretising it over a step ``dt``
    gives the linear-Gaussian transition ``x' = A·x + noise(Q)`` (ADR-017):

        A = exp(−dt / tau)             — the fraction of the state surviving a step
        Q = stationary_var · (1 − A²)  — the kick that holds the stationary variance

    Scalar (1-D) only: the vector OU would need a matrix exponential and a Lyapunov
    solve, which no cpomdp node needs.

    Args:
        tau: the relaxation timescale τ (same time unit as ``dt``).
        stationary_var: the steady-state variance Σ_stat the node settles to; A
            (dynamics) and Q (dynamics_noise) are set so it holds this spread.
        dt: the discretisation step.

    Returns:
        A ``GaussianTransition`` with 1×1 ``dynamics_matrix`` (A) and
        ``dynamics_noise`` (Q).
    """
    a = jnp.exp(-dt / tau)  # A = e^(−dt/τ)
    q = stationary_var * (1.0 - a * a)  # Q = Σ_stat (1 − A²)
    return cls(jnp.reshape(a, (1, 1)), dynamics_noise=jnp.reshape(q, (1, 1)))

predict

predict(
    message: CanonicalGaussian,
    control_term: ArrayLike | None = None,
) -> CanonicalGaussian

Push an incoming belief on x through the dynamics to a belief on x'.

The transition is the joint Gaussian over z = [x, x']::

Λ_J = [[ AᵀQ⁻¹A, −AᵀQ⁻¹ ],     h_J = [ −AᵀQ⁻¹b ,
       [ −Q⁻¹A,    Q⁻¹   ]]            Q⁻¹b ]

with b = control_term (the Bu shift; None → zero). The predict:

1. Folds the incoming message into the x block — its precision into the top-left n×n of Λ_J, its potential into the top n of h_J (a block add during construction, not __add__).
2. Marginalizes x out, leaving the predicted message on x'.

In moment form this lands exactly on cov_pred = AΣAᵀ + Q and mean_pred = Aμ + b.

Parameters:

Name	Type	Description	Default
message	CanonicalGaussian	the incoming belief on x, as a CanonicalGaussian (n-D).	required
control_term	ArrayLike | None	b = Bu, shape (n,); None for an uncontrolled step.	None

Returns:

Type	Description
CanonicalGaussian	A CanonicalGaussian over the n-D next state x'.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def predict(
    self,
    message: CanonicalGaussian,
    control_term: ArrayLike | None = None,
) -> CanonicalGaussian:
    """Push an incoming belief on x through the dynamics to a belief on x'.

    The transition is the joint Gaussian over ``z = [x, x']``::

        Λ_J = [[ AᵀQ⁻¹A, −AᵀQ⁻¹ ],     h_J = [ −AᵀQ⁻¹b ,
               [ −Q⁻¹A,    Q⁻¹   ]]            Q⁻¹b ]

    with ``b`` = ``control_term`` (the Bu shift; ``None`` → zero). The predict:

    1. Folds the incoming message into the x block — its precision into the
       top-left ``n×n`` of ``Λ_J``, its potential into the top ``n`` of ``h_J``
       (a block add during construction, *not* ``__add__``).
    2. Marginalizes x out, leaving the predicted message on x'.

    In moment form this lands exactly on ``cov_pred = AΣAᵀ + Q`` and
    ``mean_pred = Aμ + b``.

    Args:
        message: the incoming belief on x, as a ``CanonicalGaussian`` (n-D).
        control_term: b = Bu, shape ``(n,)``; ``None`` for an uncontrolled step.

    Returns:
        A ``CanonicalGaussian`` over the n-D next state x'.
    """
    dynamics_matrix, dynamics_noise = (
        self.dynamics_matrix,
        self.dynamics_noise,
    )  # A, Q
    n = dynamics_matrix.shape[0]
    # b = Bu, the control shift; None means no shift.
    if control_term is None:
        shift = jnp.zeros(n)
    else:
        shift = jnp.asarray(control_term, dtype=float)

    noise_precision = jnp.linalg.inv(dynamics_noise)  # Q⁻¹
    noise_weighted_dynamics = noise_precision @ dynamics_matrix  # Q⁻¹A
    # Joint precision over [x, x']: [[AᵀQ⁻¹A + Λ, −AᵀQ⁻¹], [−Q⁻¹A, Q⁻¹]], with
    # the incoming message's precision folded into the x (top-left) block.
    state_block = dynamics_matrix.T @ noise_weighted_dynamics + message.precision
    precision = jnp.block(
        [
            [state_block, -noise_weighted_dynamics.T],
            [-noise_weighted_dynamics, noise_precision],
        ]
    )
    # Joint potential [−AᵀQ⁻¹b + h, Q⁻¹b], message's potential folded into x.
    noise_weighted_shift = noise_precision @ shift  # Q⁻¹b
    state_potential = message.potential - dynamics_matrix.T @ noise_weighted_shift
    potential = jnp.concatenate([state_potential, noise_weighted_shift])

    joint = CanonicalGaussian._unchecked(precision, potential)
    return joint.marginalize(over=range(n))  # eliminate x, keep x'

tree_flatten

tree_flatten() -> tuple[
    tuple[Float64[Array, "n n"], Float64[Array, "n n"]],
    None,
]

Leaves for JAX: (dynamics, dynamics_noise), no static aux data.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def tree_flatten(
    self,
) -> tuple[tuple[Float64[Array, "n n"], Float64[Array, "n n"]], None]:
    """Leaves for JAX: ``(dynamics, dynamics_noise)``, no static aux data."""
    return (self.dynamics_matrix, self.dynamics_noise), None

tree_unflatten classmethod

tree_unflatten(
    aux_data: None,
    children: tuple[
        Float64[Array, "n n"], Float64[Array, "n n"]
    ],
) -> GaussianTransition

Rebuild from leaves without validating — the leaves may be tracers.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: None,
    children: tuple[Float64[Array, "n n"], Float64[Array, "n n"]],
) -> "GaussianTransition":
    """Rebuild from leaves without validating — the leaves may be tracers."""
    dynamics_matrix, dynamics_noise = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "dynamics_matrix", dynamics_matrix)
    object.__setattr__(obj, "dynamics_noise", dynamics_noise)
    return obj

GaussianObservation dataclass

GaussianObservation(
    observation_matrix: ArrayLike,
    *,
    observation_noise: ArrayLike,
)

Tier-1 likelihood factor N(y; Cx, R) — emits a message into the state.

Holds the fixed sensor map and noise; message(y) turns a reading into its canonical-form contribution to the belief on x.

• observation_matrix — C, shape (m, n).
• observation_noise — R, shape (m, m), positive-definite (it is inverted).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def __init__(
    self, observation_matrix: ArrayLike, *, observation_noise: ArrayLike
) -> None:
    object.__setattr__(
        self, "observation_matrix", jnp.asarray(observation_matrix, dtype=float)
    )
    object.__setattr__(
        self, "observation_noise", jnp.asarray(observation_noise, dtype=float)
    )
    self._validate()

message

message(
    observation: ArrayLike, state: ArrayLike | None = None
) -> CanonicalGaussian

The likelihood's message into x: Λ = CᵀR⁻¹C, h = CᵀR⁻¹y.

The information form of the reading — the evidence the observation injects about the state. The measurement update is then prior_message + this (CanonicalGaussian.__add__). A solve against R avoids forming R⁻¹; the result is valid by construction, so it builds via the no-validate seam.

Parameters:

Name	Type	Description	Default
observation	ArrayLike	the reading y, shape (m,).	required
state	ArrayLike | None	ignored — a fixed sensor's noise does not depend on the state. It is accepted so the fixed and state-dependent factors share one message interface (the backend can call either without a type branch).	None

Returns:

Type	Description
CanonicalGaussian	A CanonicalGaussian over the n-D state — precision (n, n),
CanonicalGaussian	potential (n,).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def message(
    self, observation: ArrayLike, state: ArrayLike | None = None
) -> CanonicalGaussian:
    """The likelihood's message into x: ``Λ = CᵀR⁻¹C``, ``h = CᵀR⁻¹y``.

    The information form of the reading — the evidence the observation injects
    about the state. The measurement update is then ``prior_message + this``
    (``CanonicalGaussian.__add__``). A solve against R avoids forming ``R⁻¹``;
    the result is valid by construction, so it builds via the no-validate seam.

    Args:
        observation: the reading y, shape ``(m,)``.
        state: ignored — a fixed sensor's noise does not depend on the state. It is
            accepted so the fixed and state-dependent factors share one ``message``
            interface (the backend can call either without a type branch).

    Returns:
        A ``CanonicalGaussian`` over the n-D state — precision ``(n, n)``,
        potential ``(n,)``.
    """
    observation_matrix, observation_noise = (
        self.observation_matrix,
        self.observation_noise,
    )  # C, R
    reading = jnp.asarray(observation, dtype=float)  # y
    # Λ = CᵀR⁻¹C, h = CᵀR⁻¹y — solved against R rather than forming R⁻¹.
    noise_weighted_model = jnp.linalg.solve(
        observation_noise, observation_matrix
    )  # R⁻¹C
    precision = observation_matrix.T @ noise_weighted_model  # CᵀR⁻¹C
    potential = observation_matrix.T @ jnp.linalg.solve(
        observation_noise, reading
    )  # CᵀR⁻¹y
    return CanonicalGaussian._unchecked(precision, potential)

linearize

linearize(
    state: ArrayLike | None = None,
) -> tuple[Float64[Array, "m n"], Float64[Array, "m m"]]

Local (C, R) — both constant; state is ignored (fixed sensor).

The shared seam with CallableGaussianObservation.linearize, so a caller can read (C, R) off either factor without a type branch.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def linearize(
    self, state: ArrayLike | None = None
) -> tuple[Float64[Array, "m n"], Float64[Array, "m m"]]:
    """Local ``(C, R)`` — both constant; ``state`` is ignored (fixed sensor).

    The shared seam with ``CallableGaussianObservation.linearize``, so a caller can
    read ``(C, R)`` off either factor without a type branch.
    """
    return self.observation_matrix, self.observation_noise

tree_flatten

tree_flatten() -> tuple[
    tuple[Float64[Array, "m n"], Float64[Array, "m m"]],
    None,
]

Leaves for JAX: (observation_matrix, observation_noise); no aux.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def tree_flatten(
    self,
) -> tuple[tuple[Float64[Array, "m n"], Float64[Array, "m m"]], None]:
    """Leaves for JAX: ``(observation_matrix, observation_noise)``; no aux."""
    return (self.observation_matrix, self.observation_noise), None

tree_unflatten classmethod

tree_unflatten(
    aux_data: None,
    children: tuple[
        Float64[Array, "m n"], Float64[Array, "m m"]
    ],
) -> GaussianObservation

Rebuild from leaves without validating — the leaves may be tracers.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: None,
    children: tuple[Float64[Array, "m n"], Float64[Array, "m m"]],
) -> "GaussianObservation":
    """Rebuild from leaves without validating — the leaves may be tracers."""
    observation_matrix, observation_noise = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "observation_matrix", observation_matrix)
    object.__setattr__(obj, "observation_noise", observation_noise)
    return obj

CallableGaussianObservation dataclass

CallableGaussianObservation(
    observation_matrix: ArrayLike,
    noise_fn: Callable[
        [Float64[Array, n], PyTree], Float64[Array, "m m"]
    ],
    noise_params: PyTree,
)

Likelihood factor with state-dependent noise N(y; Cx, R(x)) (issue #27).

The state-dependent sibling of GaussianObservation: the observation map stays linear (constant C), but the noise covariance varies with the state through noise_fn(x, params) -> R(x). Evaluated at the predicted mean μ⁺ — which the action moves — R is no longer action-invariant, so the FFG epistemic term stops collapsing to LQR (ADR-003) and the chosen action can seek states where the sensor is sharper (the dual effect, ADR-014 finding #1). message(y, state) emits the same information-form message as the fixed factor, at the plugged-in R(state).

• observation_matrix — C, shape (m, n) (constant); a traced pytree leaf.
• noise_fn — (x, params) -> R(x), a positive-definite (m, m) covariance; static aux (a callable cannot be a traced leaf, and keeping it static lets jit cache on it). Pass a module-level function, not a closure.
• noise_params — the sensor's tunables; a traced leaf, so the EFE is grad-able w.r.t. them (sensor learning). Keep every tunable here, not in a closure over noise_fn, or jit caching breaks.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def __init__(
    self,
    observation_matrix: ArrayLike,
    noise_fn: Callable[[Float64[Array, "n"], PyTree], Float64[Array, "m m"]],
    noise_params: PyTree,
) -> None:
    object.__setattr__(
        self, "observation_matrix", jnp.asarray(observation_matrix, dtype=float)
    )
    object.__setattr__(self, "noise_fn", noise_fn)
    object.__setattr__(self, "noise_params", noise_params)
    self._validate()

message

message(
    observation: ArrayLike, state: ArrayLike | None = None
) -> CanonicalGaussian

The likelihood's message into x, with R evaluated at the plug-in state.

Identical to GaussianObservation.message (Λ = CᵀR⁻¹C, h = CᵀR⁻¹y) but for the one thing that makes the sensor state-dependent: R is taken at state — the predicted mean μ⁺ — rather than fixed. A solve against R(state) avoids forming its inverse; the result is valid by construction, so it builds via the no-validate seam. A constant noise_fn reproduces the fixed factor's message exactly (the reduction gate).

Parameters:

Name	Type	Description	Default
observation	ArrayLike	the reading y, shape (m,).	required
state	ArrayLike | None	the state R is evaluated at (the predicted mean μ⁺), shape (n,). Required here (the shared interface makes it optional): without a linearization point R(x) is undefined, so a static/factored inference context — which has no μ⁺ — is rejected.	None

Returns:

Type	Description
CanonicalGaussian	A CanonicalGaussian over the n-D state — precision (n, n),
CanonicalGaussian	potential (n,).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def message(
    self, observation: ArrayLike, state: ArrayLike | None = None
) -> CanonicalGaussian:
    """The likelihood's message into x, with ``R`` evaluated at the plug-in state.

    Identical to ``GaussianObservation.message`` (``Λ = CᵀR⁻¹C``, ``h = CᵀR⁻¹y``)
    but for the one thing that makes the sensor state-dependent: ``R`` is taken at
    ``state`` — the predicted mean ``μ⁺`` — rather than fixed. A solve against
    ``R(state)`` avoids forming its inverse; the result is valid by construction, so
    it builds via the no-validate seam. A constant ``noise_fn`` reproduces the fixed
    factor's message exactly (the reduction gate).

    Args:
        observation: the reading y, shape ``(m,)``.
        state: the state R is evaluated at (the predicted mean μ⁺), shape ``(n,)``.
            Required here (the shared interface makes it optional): without a
            linearization point ``R(x)`` is undefined, so a static/factored
            inference context — which has no ``μ⁺`` — is rejected.

    Returns:
        A ``CanonicalGaussian`` over the n-D state — precision ``(n, n)``,
        potential ``(n,)``.
    """
    if state is None:
        raise ValueError(
            "CallableGaussianObservation.message needs the plug-in state (the "
            "predicted mean μ⁺) to evaluate R(x); a static or factored inference "
            "context has no such linearization point."
        )
    observation_matrix = self.observation_matrix  # C
    reading = jnp.asarray(observation, dtype=float)  # y
    state = jnp.asarray(state, dtype=float)
    observation_noise = self.noise_fn(state, self.noise_params)  # R(state)
    # Λ = CᵀR⁻¹C, h = CᵀR⁻¹y — solved against R rather than forming R⁻¹.
    noise_weighted_model = jnp.linalg.solve(
        observation_noise, observation_matrix
    )  # R⁻¹C
    precision = observation_matrix.T @ noise_weighted_model  # CᵀR⁻¹C
    potential = observation_matrix.T @ jnp.linalg.solve(
        observation_noise, reading
    )  # CᵀR⁻¹y
    return CanonicalGaussian._unchecked(precision, potential)

linearize

linearize(
    state: ArrayLike,
) -> tuple[Float64[Array, "m n"], Float64[Array, "m m"]]

Local (C, R(state)) — constant C, noise evaluated at the plug-in state.

The seam the FFG backend and EFE selector read R(μ⁺) from, per candidate action, without reconstructing a message (mirrors CallableSensor.linearize).

Parameters:

Name	Type	Description	Default
state	ArrayLike	the state R is evaluated at (the predicted mean μ⁺), shape (n,).	required

Returns:

Type	Description
Float64[Array, 'm n']	(observation_matrix, R(state)) — C (m, n) and the noise covariance
Float64[Array, 'm m']	(m, m).

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def linearize(
    self, state: ArrayLike
) -> tuple[Float64[Array, "m n"], Float64[Array, "m m"]]:
    """Local ``(C, R(state))`` — constant C, noise evaluated at the plug-in state.

    The seam the FFG backend and EFE selector read ``R(μ⁺)`` from, per candidate
    action, without reconstructing a message (mirrors ``CallableSensor.linearize``).

    Args:
        state: the state R is evaluated at (the predicted mean μ⁺), shape ``(n,)``.

    Returns:
        ``(observation_matrix, R(state))`` — C ``(m, n)`` and the noise covariance
        ``(m, m)``.
    """
    state = jnp.asarray(state, dtype=float)
    return self.observation_matrix, self.noise_fn(state, self.noise_params)

tree_flatten

tree_flatten() -> tuple[
    tuple[Float64[Array, "m n"], PyTree], Callable
]

Leaves (traced): (observation_matrix, noise_params); aux noise_fn.

The callable cannot be a traced leaf, so it rides as aux (and staying static lets jit cache on it); the sensor map and the tunable params are leaves, the params grad-able for sensor learning.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

def tree_flatten(
    self,
) -> tuple[tuple[Float64[Array, "m n"], PyTree], Callable]:
    """Leaves (traced): ``(observation_matrix, noise_params)``; aux ``noise_fn``.

    The callable cannot be a traced leaf, so it rides as aux (and staying static
    lets ``jit`` cache on it); the sensor map and the tunable params are leaves, the
    params grad-able for sensor learning.
    """
    return (self.observation_matrix, self.noise_params), self.noise_fn

tree_unflatten classmethod

tree_unflatten(
    aux_data: Callable[
        [Float64[Array, n], PyTree], Float64[Array, "m m"]
    ],
    children: tuple[Float64[Array, "m n"], PyTree],
) -> CallableGaussianObservation

Rebuild from leaves without validating — the leaves may be tracers.

Under jit/grad/vmap the leaves arrive as tracers, so the construction-time PD probe (which needs a concrete R) is skipped here; it already ran once when the factor was first built.

Source code in src/cpomdp/ffg/factors/linear_gaussian.py

@classmethod
def tree_unflatten(
    cls,
    aux_data: Callable[[Float64[Array, "n"], PyTree], Float64[Array, "m m"]],
    children: tuple[Float64[Array, "m n"], PyTree],
) -> "CallableGaussianObservation":
    """Rebuild from leaves without validating — the leaves may be tracers.

    Under ``jit``/``grad``/``vmap`` the leaves arrive as tracers, so the
    construction-time PD probe (which needs a concrete ``R``) is skipped here; it
    already ran once when the factor was first built.
    """
    observation_matrix, noise_params = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "observation_matrix", observation_matrix)
    object.__setattr__(obj, "noise_params", noise_params)
    object.__setattr__(obj, "noise_fn", aux_data)
    return obj