LGM API

pyrox_lgm holds latent Gaussian models (LGMs) in precision form: GMRF latent components, penalised-complexity priors and the inla() driver, built on gaussx's sparse precision operators and Laplace modes and kernellib's graph and GMRF structure matrices. GPs in covariance form stay in pyrox-gp; pyrox_lgm imports nothing from it.

This page is a stub: the package is a scaffold (P6 of the inla() epic), and the API lands with its phases.

Phase	Adds
P7	Components (RW1, RW2, AR1, ICAR, BYM2, SPDE, generic, combinators) and PC priors
P8	LGM, inla() and INLAResult
P9	f(...) formula sugar, diagnostics, simplified Laplace, MCMC-INLA hybrid

Penalised-complexity priors

Hyperpriors calibrated by a tail statement (Simpson et al., 2017). Each is a NumPyro distribution, usable under NUTS as well as in inla().

PCPrecision

Bases: Distribution

PC prior on the precision \(\tau\) of a Gaussian effect.

The distance to the base model \(\sigma = 0\) is proportional to the standard deviation \(\sigma = \tau^{-1/2}\), so \(\sigma \sim \mathrm{Exp}(\lambda)\) with \(\lambda = -\log\alpha / U\), i.e. \(P(\sigma > U) = \alpha\). On \(\tau\):

\[ \pi(\tau) = \tfrac{\lambda}{2}\,\tau^{-3/2}\,e^{-\lambda\tau^{-1/2}}. \]

This is R-INLA's pc.prec.

Parameters:

Name	Type	Description	Default
U	ArrayLike	Upper reference value for \(\sigma\), U > 0.	1.0
alpha	ArrayLike	Tail probability \(P(\sigma > U)\), in (0, 1).	0.01

Examples:

>>> import jax
>>> import pyrox_lgm as lgm
>>> prior = lgm.PCPrecision(U=1.0, alpha=0.01)  # P(sigma > 1) = 0.01
>>> tau = prior.sample(jax.random.key(0), (20_000,))
>>> round(float((tau < 1.0).mean()), 2)  # sigma > 1  <=>  tau < 1
0.01

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

class PCPrecision(dist.Distribution):
    r"""PC prior on the precision $\tau$ of a Gaussian effect.

    The distance to the base model $\sigma = 0$ is proportional to the
    standard deviation $\sigma = \tau^{-1/2}$, so $\sigma \sim
    \mathrm{Exp}(\lambda)$ with $\lambda = -\log\alpha / U$, i.e.
    $P(\sigma > U) = \alpha$. On $\tau$:

    $$
    \pi(\tau) = \tfrac{\lambda}{2}\,\tau^{-3/2}\,e^{-\lambda\tau^{-1/2}}.
    $$

    This is R-INLA's ``pc.prec``.

    Args:
        U: Upper reference value for $\sigma$, ``U > 0``.
        alpha: Tail probability $P(\sigma > U)$, in ``(0, 1)``.

    Examples:
        >>> import jax
        >>> import pyrox_lgm as lgm
        >>> prior = lgm.PCPrecision(U=1.0, alpha=0.01)  # P(sigma > 1) = 0.01
        >>> tau = prior.sample(jax.random.key(0), (20_000,))
        >>> round(float((tau < 1.0).mean()), 2)  # sigma > 1  <=>  tau < 1
        0.01
    """

    arg_constraints = {  # noqa: RUF012 (NumPyro's class-level contract)
        "U": constraints.positive,
        "alpha": constraints.open_interval(0.0, 1.0),
    }
    support = constraints.positive
    reparametrized_params = ["U", "alpha"]  # noqa: RUF012

    def __init__(
        self,
        U: ArrayLike = 1.0,
        alpha: ArrayLike = 0.01,
        *,
        validate_args: bool | None = None,
    ) -> None:
        self.U, self.alpha = promote_shapes(U, alpha)
        batch_shape = jax.lax.broadcast_shapes(jnp.shape(U), jnp.shape(alpha))
        super().__init__(batch_shape=batch_shape, validate_args=validate_args)

    @property
    def rate(self) -> Array:
        """The rate $\\lambda = -\\log\\alpha / U$ of the exponential on $\\sigma$."""
        return -jnp.log(self.alpha) / self.U

    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        sigma = jax.random.exponential(key, shape) / self.rate
        return sigma**-2.0

    @validate_sample
    def log_prob(self, value):
        lam = self.rate
        return jnp.log(lam / 2.0) - 1.5 * jnp.log(value) - lam * jax.lax.rsqrt(value)

    def cdf(self, value):
        # P(tau <= t) = P(sigma >= t^{-1/2}) = exp(-lam t^{-1/2}).
        return jnp.exp(-self.rate * jax.lax.rsqrt(value))

    def icdf(self, q):
        return (-jnp.log(q) / self.rate) ** -2.0

rate: Array property

The rate \(\lambda = -\log\alpha / U\) of the exponential on \(\sigma\).

PCAR1Rho

Bases: Distribution

PC prior on the lag-one correlation \(\rho\) of an AR(1) process.

Base model \(\rho = 0\) (Sørbye & Rue, 2017): the distance is \(d(\rho) = \sqrt{-\log(1-\rho^2)}\), and the prior is symmetric,

\[ \pi(\rho) = \frac{\lambda}{2}\,e^{-\lambda d(\rho)}\, \frac{\lvert\rho\rvert}{(1-\rho^2)\,d(\rho)},\qquad \lambda = -\log\alpha / d(U), \]

so that \(P(\lvert\rho\rvert > U) = \alpha\). This is R-INLA's pc.cor0 (which states its calibration on \(\lvert\rho\rvert\)).

Parameters:

Name	Type	Description	Default
U	ArrayLike	Reference correlation, in (0, 1).	0.5
alpha	ArrayLike	Tail probability \(P(\lvert\rho\rvert > U)\), in (0, 1).	0.5

Examples:

>>> import jax, jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> prior = lgm.PCAR1Rho(U=0.5, alpha=0.1)  # P(|rho| > 0.5) = 0.1
>>> rho = prior.sample(jax.random.key(0), (20_000,))
>>> round(float((jnp.abs(rho) > 0.5).mean()), 2)
0.1

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

class PCAR1Rho(dist.Distribution):
    r"""PC prior on the lag-one correlation $\rho$ of an AR(1) process.

    Base model $\rho = 0$ (Sørbye & Rue, 2017): the distance is
    $d(\rho) = \sqrt{-\log(1-\rho^2)}$, and the prior is symmetric,

    $$
    \pi(\rho) = \frac{\lambda}{2}\,e^{-\lambda d(\rho)}\,
        \frac{\lvert\rho\rvert}{(1-\rho^2)\,d(\rho)},\qquad
    \lambda = -\log\alpha / d(U),
    $$

    so that $P(\lvert\rho\rvert > U) = \alpha$. This is R-INLA's ``pc.cor0``
    (which states its calibration on $\lvert\rho\rvert$).

    Args:
        U: Reference correlation, in ``(0, 1)``.
        alpha: Tail probability $P(\lvert\rho\rvert > U)$, in ``(0, 1)``.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> prior = lgm.PCAR1Rho(U=0.5, alpha=0.1)  # P(|rho| > 0.5) = 0.1
        >>> rho = prior.sample(jax.random.key(0), (20_000,))
        >>> round(float((jnp.abs(rho) > 0.5).mean()), 2)
        0.1
    """

    arg_constraints = {  # noqa: RUF012 (NumPyro's class-level contract)
        "U": constraints.open_interval(0.0, 1.0),
        "alpha": constraints.open_interval(0.0, 1.0),
    }
    support = constraints.open_interval(-1.0, 1.0)
    reparametrized_params = ["U", "alpha"]  # noqa: RUF012

    def __init__(
        self,
        U: ArrayLike = 0.5,
        alpha: ArrayLike = 0.5,
        *,
        validate_args: bool | None = None,
    ) -> None:
        self.U, self.alpha = promote_shapes(U, alpha)
        batch_shape = jax.lax.broadcast_shapes(jnp.shape(U), jnp.shape(alpha))
        super().__init__(batch_shape=batch_shape, validate_args=validate_args)

    @property
    def rate(self) -> Array:
        """The rate $\\lambda = -\\log\\alpha / d(U)$."""
        return -jnp.log(self.alpha) / jnp.sqrt(-jnp.log1p(-(self.U**2)))

    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        k_d, k_sign = jax.random.split(key)
        d = jax.random.exponential(k_d, shape) / self.rate
        magnitude = jnp.sqrt(-jnp.expm1(-(d**2)))
        sign = jnp.where(jax.random.bernoulli(k_sign, 0.5, shape), 1.0, -1.0)
        return sign * magnitude

    @validate_sample
    def log_prob(self, value):
        lam = self.rate
        ratio = _safe_ratio_rho_over_d(value)
        # d = |rho| / ratio: unlike sqrt(-log1p(-rho^2)) it has a finite
        # gradient at rho = 0.
        d = jnp.abs(value) / ratio
        return jnp.log(lam / 2.0) - lam * d + jnp.log(ratio) - jnp.log1p(-(value**2))

rate: Array property

The rate \(\lambda = -\log\alpha / d(U)\).

PCBYM2Phi

Bases: Distribution

PC prior on the BYM2 mixing parameter \(\phi \in (0, 1)\).

BYM2 (Riebler et al., 2016) mixes unstructured and scaled-structured effects, \(b = \sigma(\sqrt{1-\phi}\,v + \sqrt{\phi}\,u_\ast)\). The base model is \(\phi = 0\) (no spatial structure); with \(\gamma_i\) the eigenvalues of \(R_\ast^{+}\),

\[ 2\,\mathrm{KLD}(\phi) = n\phi\big(\tfrac1n\operatorname{tr}R_\ast^{+} - 1\big) - \sum_i \log\big(1 - \phi + \phi\gamma_i\big), \]

and \(d(\phi) = \sqrt{2\,\mathrm{KLD}(\phi)}\) gets an exponential prior with \(\lambda = -\log(1-\alpha)/d(U)\), so \(P(\phi < U) = \alpha\). This is R-INLA's pc prior on BYM2's phi. The null-space eigenvalues (\(\gamma = 0\)) send \(d(\phi) \to \infty\) as \(\phi \to 1\), which makes the prior proper on \((0, 1)\), so the spectrum must contain at least one (a BYM2 structure always has a null space).

Parameters:

Name	Type	Description	Default
U	ArrayLike	Reference value of \(\phi\), in (0, 1).	0.5
alpha	ArrayLike	\(P(\phi < U)\), in (0, 1).	2.0 / 3.0
structure_spectrum	StructureSpectrum	The spectrum of \(R_\ast^{+}\), from structure_spectrum (computed once per graph).	required

Examples:

>>> import jax
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> g = kl.grid_graph((6, 6))
>>> R = kl.structure_matrix(g, scaled=True)
>>> spec = lgm.structure_spectrum(R, kl.graph_null_space(g))
>>> prior = lgm.PCBYM2Phi(U=0.5, alpha=2 / 3, structure_spectrum=spec)
>>> phi = prior.sample(jax.random.key(0), (20_000,))
>>> round(float((phi < 0.5).mean()), 2)  # P(phi < 0.5) = 2/3
0.67

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

class PCBYM2Phi(dist.Distribution):
    r"""PC prior on the BYM2 mixing parameter $\phi \in (0, 1)$.

    BYM2 (Riebler et al., 2016) mixes unstructured and scaled-structured
    effects, $b = \sigma(\sqrt{1-\phi}\,v + \sqrt{\phi}\,u_\ast)$. The base
    model is $\phi = 0$ (no spatial structure); with $\gamma_i$ the
    eigenvalues of $R_\ast^{+}$,

    $$
    2\,\mathrm{KLD}(\phi) = n\phi\big(\tfrac1n\operatorname{tr}R_\ast^{+} - 1\big)
        - \sum_i \log\big(1 - \phi + \phi\gamma_i\big),
    $$

    and $d(\phi) = \sqrt{2\,\mathrm{KLD}(\phi)}$ gets an exponential prior
    with $\lambda = -\log(1-\alpha)/d(U)$, so $P(\phi < U) = \alpha$. This is
    R-INLA's ``pc`` prior on BYM2's ``phi``. The null-space eigenvalues
    ($\gamma = 0$) send $d(\phi) \to \infty$ as $\phi \to 1$, which makes the
    prior proper on $(0, 1)$, so the spectrum must contain at least one (a
    BYM2 structure always has a null space).

    Args:
        U: Reference value of $\phi$, in ``(0, 1)``.
        alpha: $P(\phi < U)$, in ``(0, 1)``.
        structure_spectrum: The spectrum of $R_\ast^{+}$, from
            `structure_spectrum` (computed once per graph).

    Examples:
        >>> import jax
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> g = kl.grid_graph((6, 6))
        >>> R = kl.structure_matrix(g, scaled=True)
        >>> spec = lgm.structure_spectrum(R, kl.graph_null_space(g))
        >>> prior = lgm.PCBYM2Phi(U=0.5, alpha=2 / 3, structure_spectrum=spec)
        >>> phi = prior.sample(jax.random.key(0), (20_000,))
        >>> round(float((phi < 0.5).mean()), 2)  # P(phi < 0.5) = 2/3
        0.67
    """

    arg_constraints = {  # noqa: RUF012 (NumPyro's class-level contract)
        "U": constraints.open_interval(0.0, 1.0),
        "alpha": constraints.open_interval(0.0, 1.0),
    }
    support = constraints.open_interval(0.0, 1.0)
    pytree_data_fields = ("U", "alpha", "gamma", "weights")

    def __init__(
        self,
        U: ArrayLike = 0.5,
        alpha: ArrayLike = 2.0 / 3.0,
        *,
        structure_spectrum: StructureSpectrum,
        validate_args: bool | None = None,
    ) -> None:
        self.U = jnp.asarray(U)
        self.alpha = jnp.asarray(alpha)
        if self.U.ndim or self.alpha.ndim:
            raise ValueError("PCBYM2Phi takes scalar U and alpha")
        self.gamma = jnp.asarray(structure_spectrum.eigenvalues)
        self.weights = jnp.asarray(structure_spectrum.weights)
        try:
            has_null = bool(jnp.any((self.gamma == 0.0) & (self.weights > 0.0)))
        except jax.errors.ConcretizationTypeError:  # traced: checked eagerly
            has_null = True
        if not has_null:
            raise ValueError(
                "structure_spectrum has no null space (no zero eigenvalue): "
                "BYM2's structure is intrinsic, and the PC prior on phi is "
                "proper only through its null space"
            )
        super().__init__(batch_shape=(), validate_args=validate_args)

    def distance(self, phi: ArrayLike) -> Array:
        r"""$d(\phi) = \sqrt{2\,\mathrm{KLD}(\phi)}$."""
        phi = jnp.asarray(phi)
        return phi * jnp.sqrt(self._g(phi))

    def _g(self, phi):
        # 2 KLD = sum_i w_i [x_i - log1p(x_i)] with x_i = phi (gamma_i - 1),
        # = phi^2 sum_i w_i (gamma_i - 1)^2 h(x_i) / x_i^2: cancellation-free.
        a = self.gamma - 1.0
        x = jnp.expand_dims(phi, -1) * a
        return jnp.sum(self.weights * a**2 * _h_over_x2(x), axis=-1)

    def _ddistance(self, phi):
        a = self.gamma - 1.0
        x = jnp.expand_dims(phi, -1) * a
        num = jnp.sum(self.weights * a**2 / (1.0 + x), axis=-1)
        return num / (2.0 * jnp.sqrt(self._g(phi)))

    @property
    def rate(self) -> Array:
        """The rate $\\lambda = -\\log(1-\\alpha) / d(U)$."""
        return -jnp.log1p(-self.alpha) / self.distance(self.U)

    @validate_sample
    def log_prob(self, value):
        lam = self.rate
        return (
            jnp.log(lam) - lam * self.distance(value) + jnp.log(self._ddistance(value))
        )

    def cdf(self, value):
        return -jnp.expm1(-self.rate * self.distance(value))

    def icdf(self, q):
        return self._invert(-jnp.log1p(-q) / self.rate)

    def sample(self, key, sample_shape=()):
        q = jax.random.uniform(key, sample_shape, minval=1e-12, maxval=1.0)
        return self.icdf(q)

    def _invert(self, d: Array) -> Array:
        """Solve ``distance(phi) = d`` for phi by bisection (d is increasing)."""

        def body(_, bounds):
            lo, hi = bounds
            mid = 0.5 * (lo + hi)
            below = self.distance(mid) < d
            return jnp.where(below, mid, lo), jnp.where(below, hi, mid)

        lo = jnp.zeros_like(d)
        hi = jnp.ones_like(d)
        lo, hi = jax.lax.fori_loop(0, 60, body, (lo, hi))
        # Mass within rounding of 1 is real (d grows like sqrt(-log(1-phi)));
        # keep draws inside the open support.
        eps = jnp.finfo(lo.dtype).eps
        return jnp.clip(0.5 * (lo + hi), eps, 1.0 - eps)

rate: Array property

The rate \(\lambda = -\log(1-\alpha) / d(U)\).

distance(phi: ArrayLike) -> Array

\(d(\phi) = \sqrt{2\,\mathrm{KLD}(\phi)}\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

def distance(self, phi: ArrayLike) -> Array:
    r"""$d(\phi) = \sqrt{2\,\mathrm{KLD}(\phi)}$."""
    phi = jnp.asarray(phi)
    return phi * jnp.sqrt(self._g(phi))

PCMatern

Bases: Distribution

Joint PC prior on the range \(\rho\) and marginal sd \(\sigma\) of a Matérn field.

Fuglstad et al. (2019): in dimension \(d\), the base model is infinite range and \(\sigma = 0\), and

\[ \pi(\rho, \sigma) = \tfrac{d}{2}\,\lambda_\rho\,\rho^{-d/2-1} e^{-\lambda_\rho\rho^{-d/2}}\cdot\lambda_\sigma e^{-\lambda_\sigma\sigma}, \]

with \(\lambda_\rho = -\log(\alpha_\rho)\,\rho_0^{d/2}\) and \(\lambda_\sigma = -\log(\alpha_\sigma)/\sigma_0\), so \(P(\rho < \rho_0) = \alpha_\rho\) and \(P(\sigma > \sigma_0) = \alpha_\sigma\). The value is the pair [range, sigma] (event shape (2,)). This is R-INLA's inla.spde2.pcmatern prior.

Parameters:

Name	Type	Description	Default
range0	ArrayLike	Reference range \(\rho_0 > 0\).	1.0
alpha_range	ArrayLike	\(P(\rho < \rho_0)\), in (0, 1).	0.05
sigma0	ArrayLike	Reference standard deviation \(\sigma_0 > 0\).	1.0
alpha_sigma	ArrayLike	\(P(\sigma > \sigma_0)\), in (0, 1).	0.05
d	int	Spatial dimension of the field.	2

Examples:

>>> import jax
>>> import pyrox_lgm as lgm
>>> prior = lgm.PCMatern(
...     range0=10.0, alpha_range=0.05, sigma0=2.0, alpha_sigma=0.05, d=2
... )
>>> x = prior.sample(jax.random.key(0), (20_000,))
>>> x.shape
(20000, 2)
>>> round(float((x[:, 0] < 10.0).mean()), 2)  # P(range < 10)
0.05
>>> round(float((x[:, 1] > 2.0).mean()), 2)  # P(sigma > 2)
0.05

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

class PCMatern(dist.Distribution):
    r"""Joint PC prior on the range $\rho$ and marginal sd $\sigma$ of a Matérn field.

    Fuglstad et al. (2019): in dimension $d$, the base model is infinite
    range and $\sigma = 0$, and

    $$
    \pi(\rho, \sigma) = \tfrac{d}{2}\,\lambda_\rho\,\rho^{-d/2-1}
        e^{-\lambda_\rho\rho^{-d/2}}\cdot\lambda_\sigma e^{-\lambda_\sigma\sigma},
    $$

    with $\lambda_\rho = -\log(\alpha_\rho)\,\rho_0^{d/2}$ and
    $\lambda_\sigma = -\log(\alpha_\sigma)/\sigma_0$, so
    $P(\rho < \rho_0) = \alpha_\rho$ and $P(\sigma > \sigma_0) = \alpha_\sigma$.
    The value is the pair ``[range, sigma]`` (event shape ``(2,)``). This is
    R-INLA's ``inla.spde2.pcmatern`` prior.

    Args:
        range0: Reference range $\rho_0 > 0$.
        alpha_range: $P(\rho < \rho_0)$, in ``(0, 1)``.
        sigma0: Reference standard deviation $\sigma_0 > 0$.
        alpha_sigma: $P(\sigma > \sigma_0)$, in ``(0, 1)``.
        d: Spatial dimension of the field.

    Examples:
        >>> import jax
        >>> import pyrox_lgm as lgm
        >>> prior = lgm.PCMatern(
        ...     range0=10.0, alpha_range=0.05, sigma0=2.0, alpha_sigma=0.05, d=2
        ... )
        >>> x = prior.sample(jax.random.key(0), (20_000,))
        >>> x.shape
        (20000, 2)
        >>> round(float((x[:, 0] < 10.0).mean()), 2)  # P(range < 10)
        0.05
        >>> round(float((x[:, 1] > 2.0).mean()), 2)  # P(sigma > 2)
        0.05
    """

    arg_constraints = {  # noqa: RUF012 (NumPyro's class-level contract)
        "range0": constraints.positive,
        "alpha_range": constraints.open_interval(0.0, 1.0),
        "sigma0": constraints.positive,
        "alpha_sigma": constraints.open_interval(0.0, 1.0),
    }
    support = constraints.independent(constraints.positive, 1)
    pytree_aux_fields = ("d",)

    def __init__(
        self,
        range0: ArrayLike = 1.0,
        alpha_range: ArrayLike = 0.05,
        sigma0: ArrayLike = 1.0,
        alpha_sigma: ArrayLike = 0.05,
        d: int = 2,
        *,
        validate_args: bool | None = None,
    ) -> None:
        self.range0 = jnp.asarray(range0)
        self.alpha_range = jnp.asarray(alpha_range)
        self.sigma0 = jnp.asarray(sigma0)
        self.alpha_sigma = jnp.asarray(alpha_sigma)
        if any(jnp.ndim(a) for a in (range0, alpha_range, sigma0, alpha_sigma)):
            raise ValueError("PCMatern takes scalar parameters")
        self.d = int(d)
        super().__init__(batch_shape=(), event_shape=(2,), validate_args=validate_args)

    @property
    def rate_range(self) -> Array:
        """$\\lambda_\\rho = -\\log(\\alpha_\\rho)\\,\\rho_0^{d/2}$."""
        return -jnp.log(self.alpha_range) * self.range0 ** (self.d / 2.0)

    @property
    def rate_sigma(self) -> Array:
        """$\\lambda_\\sigma = -\\log(\\alpha_\\sigma)/\\sigma_0$."""
        return -jnp.log(self.alpha_sigma) / self.sigma0

    def sample(self, key, sample_shape=()):
        k_r, k_s = jax.random.split(key)
        e = jax.random.exponential(k_r, sample_shape) / self.rate_range
        range_ = e ** (-2.0 / self.d)
        sigma = jax.random.exponential(k_s, sample_shape) / self.rate_sigma
        return jnp.stack([range_, sigma], axis=-1)

    @validate_sample
    def log_prob(self, value):
        range_, sigma = value[..., 0], value[..., 1]
        h = self.d / 2.0
        lr, ls = self.rate_range, self.rate_sigma
        log_range = jnp.log(h * lr) - (h + 1.0) * jnp.log(range_) - lr * range_**-h
        return log_range + jnp.log(ls) - ls * sigma

rate_range: Array property

\(\lambda_\rho = -\log(\alpha_\rho)\,\rho_0^{d/2}\).

rate_sigma: Array property

\(\lambda_\sigma = -\log(\alpha_\sigma)/\sigma_0\).

Structure spectra

PCBYM2Phi needs the spectrum of the scaled structure's generalised inverse, computed once per graph: exactly for a grid (gaussx.KroneckerSum) or n ≤ 5000, else by deflated stochastic Lanczos quadrature.

structure_spectrum(structure: lx.AbstractLinearOperator, null_space: ArrayLike | None = None, *, method: Literal['auto', 'dense', 'lanczos'] = 'auto', num_probes: int = 64, order: int = 50, deflate: int = 32, key: jax.Array | None = None) -> StructureSpectrum

The spectrum of \(R^{+}\) that PCBYM2Phi needs, computed once.

• "dense": exact eigenvalues. A gaussx.KroneckerSum (the structure of a kernellib.GridGraph) is diagonalised factor by factor, so a grid costs only its 1-D eigenproblems; anything else is a dense eigh.
• "lanczos": a deflated stochastic Lanczos quadrature (SLQ) estimate. One Lanczos run of max(order, 8 * deflate) steps converges up to deflate of the largest eigenvalues of \(R_\ast^{+}\) (the smallest of \(R_\ast\)), which are kept exactly: they dominate \(\operatorname{tr}R_\ast^{+}\), and leaving them to the probes gives a large variance. The rest of the measure comes from num_probes Rademacher probes, projected off the null space and the deflated eigenvectors, with order Lanczos steps each (full reorthogonalisation); the null space contributes exact zeros. Cost: about num_probes * order + 8 * deflate matvecs. Runs eagerly (it picks the converged Ritz pairs), so call it once, outside jit.
• "auto": "dense" for a gaussx.KroneckerSum or n <= 5000, else "lanczos".

Parameters:

Name	Type	Description	Default
structure	AbstractLinearOperator	The (scaled) structure matrix \(R_\ast\), symmetric PSD.	required
null_space	ArrayLike | None	Orthonormal basis of its null space, (n, c) or (n,); None if it has none. Its columns are the eigenvalues reported as zero.	None
method	Literal['auto', 'dense', 'lanczos']	See above.	'auto'
num_probes	int	SLQ probes ("lanczos" only).	64
order	int	Lanczos steps per probe ("lanczos" only).	50
deflate	int	Most eigenvalues to deflate exactly ("lanczos" only).	32
key	Array | None	PRNG key for the probes; None means jax.random.key(0).	None

Returns:

Type	Description
StructureSpectrum	The spectral measure of \(R_\ast^{+}\).

Examples:

>>> import jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> g = kl.grid_graph((4,))
>>> R = kl.structure_matrix(g, scaled=True)
>>> spec = lgm.structure_spectrum(R, kl.graph_null_space(g))
>>> float(spec.size)
4.0
>>> int(jnp.sum(spec.eigenvalues == 0.0))  # one connected component
1

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

def structure_spectrum(
    structure: lx.AbstractLinearOperator,
    null_space: ArrayLike | None = None,
    *,
    method: Literal["auto", "dense", "lanczos"] = "auto",
    num_probes: int = 64,
    order: int = 50,
    deflate: int = 32,
    key: jax.Array | None = None,
) -> StructureSpectrum:
    r"""The spectrum of $R^{+}$ that `PCBYM2Phi` needs, computed once.

    - ``"dense"``: exact eigenvalues. A `gaussx.KroneckerSum` (the structure
      of a `kernellib.GridGraph`) is diagonalised factor by factor, so a grid
      costs only its 1-D eigenproblems; anything else is a dense ``eigh``.
    - ``"lanczos"``: a deflated stochastic Lanczos quadrature (SLQ) estimate.
      One Lanczos run of ``max(order, 8 * deflate)`` steps converges up to
      ``deflate`` of the largest eigenvalues of $R_\ast^{+}$ (the smallest
      of $R_\ast$), which are kept exactly: they dominate
      $\operatorname{tr}R_\ast^{+}$, and leaving them to the probes gives a
      large variance. The rest of the measure comes from ``num_probes``
      Rademacher probes, projected off the null space and the deflated
      eigenvectors, with ``order`` Lanczos steps each (full
      reorthogonalisation); the null space contributes exact zeros. Cost:
      about ``num_probes * order + 8 * deflate`` matvecs. Runs eagerly (it picks the
      converged Ritz pairs), so call it once, outside ``jit``.
    - ``"auto"``: ``"dense"`` for a `gaussx.KroneckerSum` or ``n <= 5000``,
      else ``"lanczos"``.

    Args:
        structure: The (scaled) structure matrix $R_\ast$, symmetric PSD.
        null_space: Orthonormal basis of its null space, ``(n, c)`` or
            ``(n,)``; ``None`` if it has none. Its columns are the
            eigenvalues reported as zero.
        method: See above.
        num_probes: SLQ probes (``"lanczos"`` only).
        order: Lanczos steps per probe (``"lanczos"`` only).
        deflate: Most eigenvalues to deflate exactly (``"lanczos"`` only).
        key: PRNG key for the probes; ``None`` means ``jax.random.key(0)``.

    Returns:
        The spectral measure of $R_\ast^{+}$.

    Examples:
        >>> import jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> g = kl.grid_graph((4,))
        >>> R = kl.structure_matrix(g, scaled=True)
        >>> spec = lgm.structure_spectrum(R, kl.graph_null_space(g))
        >>> float(spec.size)
        4.0
        >>> int(jnp.sum(spec.eigenvalues == 0.0))  # one connected component
        1
    """
    n = structure.in_size()
    V = None if null_space is None else _as_columns(null_space)
    c = 0 if V is None else V.shape[1]
    if method == "auto":
        dense = isinstance(structure, gx.KroneckerSum) or n <= 5000
        method = "dense" if dense else "lanczos"
    if method == "dense":
        lam = jnp.sort(_dense_eigvals(structure))
        is_null = jnp.arange(n) < c  # the c smallest are the null space
        safe = jnp.where(is_null, 1.0, lam)
        gamma = jnp.where(is_null, 0.0, 1.0 / safe)
        return StructureSpectrum(gamma, jnp.ones(n, dtype=gamma.dtype))
    if method != "lanczos":
        raise ValueError(f"method must be 'auto', 'dense' or 'lanczos', got {method!r}")
    return _lanczos_spectrum(structure, V, n, c, num_probes, order, deflate, key)

StructureSpectrum

Bases: NamedTuple

Spectral measure of \(R_\ast^{+}\), the generalised inverse of a structure.

eigenvalues are the eigenvalues \(\gamma_i\) of \(R_\ast^{+}\) (zero on the null space of \(R_\ast\)) and weights their multiplicities, which sum to \(n\). A dense spectrum has unit weights; a Lanczos (SLQ) spectrum has quadrature nodes and weights. Anything linear in the spectrum (\(\operatorname{tr}R_\ast^{+}\), \(\log|(1-\phi)I + \phi R_\ast^{+}|\)) is a weighted sum over it.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_priors/_pc.py

class StructureSpectrum(NamedTuple):
    r"""Spectral measure of $R_\ast^{+}$, the generalised inverse of a structure.

    ``eigenvalues`` are the eigenvalues $\gamma_i$ of $R_\ast^{+}$ (zero on
    the null space of $R_\ast$) and ``weights`` their multiplicities, which
    sum to $n$. A dense spectrum has unit weights; a Lanczos (SLQ) spectrum
    has quadrature nodes and weights. Anything linear in the spectrum
    ($\operatorname{tr}R_\ast^{+}$, $\log|(1-\phi)I + \phi R_\ast^{+}|$) is a
    weighted sum over it.
    """

    eigenvalues: Float[Array, " m"]
    weights: Float[Array, " m"]

    @property
    def size(self) -> Array:
        """The dimension $n$, the total weight."""
        return jnp.sum(self.weights)

size: Array property

The dimension \(n\), the total weight.

Components

Each component is a GMRF over its own nodes with hyperpriors and a projector to observations. prior(theta) gives the gaussx distribution that inla() uses; sample(index) is the NumPyro face, for NUTS.

AbstractComponent

Bases: Module

A latent field with hyperpriors and a projector to observations.

Subclasses define n_nodes, theta_spec and prior. The default projector selects nodes by index; the default sample is the NumPyro face (see pyrox_lgm._numpyro.sample_component).

Attributes:

Name	Type	Description
name	AbstractVar[str]	Site-name prefix, unique within a model: the field is the NumPyro site name and hyperparameter k the site f"{name}_{k}".

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

class AbstractComponent(eqx.Module):
    """A latent field with hyperpriors and a projector to observations.

    Subclasses define `n_nodes`, `theta_spec` and `prior`. The default
    `projector` selects nodes by index; the default `sample` is the NumPyro
    face (see `pyrox_lgm._numpyro.sample_component`).

    Attributes:
        name: Site-name prefix, unique within a model: the field is the
            NumPyro site ``name`` and hyperparameter ``k`` the site
            ``f"{name}_{k}"``.
    """

    name: eqx.AbstractVar[str]

    @property
    @abc.abstractmethod
    def n_nodes(self) -> int:
        """Size of the latent field (including any padding node)."""

    @property
    def n_index(self) -> int:
        """Number of addressable nodes: what an observation index refers to.

        Equal to `n_nodes` except where the field carries padding nodes
        (an odd-length `RW2`), which observations never touch.
        """
        return self.n_nodes

    @abc.abstractmethod
    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        """Hyperparameters: name to (prior, bijection from the reals).

        The transform maps an unconstrained real to the parameter's support,
        which is how ``inla()`` explores $\\theta$.
        """

    @abc.abstractmethod
    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF | gx.IntrinsicGMRF:
        r"""The GMRF of the field given its hyperparameters.

        Args:
            theta: Hyperparameter values keyed as in `theta_spec`.
            constraint: For intrinsic fields, how the null space is handled
                (``"hard"`` for ``inla()``, ``"soft"`` for NUTS); ignored by
                proper fields.
            soft_constraint_scale: Standard deviation $s$ of the soft
                constraint $V^\top x \sim \mathcal N(0, s^2 I)$.
        """

    def assembly_precision(
        self, theta: dict[str, Array], gmrf: gx.GaussianMRF | gx.IntrinsicGMRF
    ) -> lx.AbstractLinearOperator:
        """The precision an `LGM` assembles into its sparse block-diagonal.

        ``gmrf`` is ``prior(theta)``; its precision by default (``scale *
        structure`` for an intrinsic field). A component whose prior uses an
        operator that cannot be assembled entry-wise (a grid SPDE's spectral
        function) returns an equivalent sparse one here.
        """
        if isinstance(gmrf, gx.IntrinsicGMRF):
            return scale_operator(gmrf.structure, gmrf.precision_scale)
        return gmrf.precision

    def constraint_basis(
        self, gmrf: gx.GaussianMRF | gx.IntrinsicGMRF
    ) -> Float[Array, "n c"] | None:
        """Orthonormal basis of the hard constraints ``inla()`` imposes.

        The intrinsic field's null space by default (sum-to-zero per
        connected component, say). A component may constrain fewer
        directions and leave the rest to the data, as R-INLA's ``rw2``
        constrains only the sum and keeps the linear trend; the density's
        rank is unchanged.
        """
        return gmrf.null_space if isinstance(gmrf, gx.IntrinsicGMRF) else None

    def projector(self, index: Int[ArrayLike, " n_obs"]) -> gx.SparseOperator:
        """``(n_obs, n_nodes)`` selector with a one where observation i sits.

        Args:
            index: Node of each observation, integers in ``[0, n_index)``.
        """
        index = np.asarray(index)
        if index.ndim != 1:
            raise ValueError(f"index must be 1-D, got shape {index.shape}")
        if index.size and (index.min() < 0 or index.max() >= self.n_index):
            raise ValueError(
                f"index out of range for {self.n_index} nodes: "
                f"[{index.min()}, {index.max()}]"
            )
        m = index.shape[0]
        return gx.SparseOperator.from_coo(
            np.arange(m), index, jnp.ones(m), (m, self.n_nodes)
        )

    def sample(
        self,
        index: Int[ArrayLike, " n_obs"] | None = None,
        *,
        soft_constraint_scale: float = 1e-2,
    ) -> Array:
        """NumPyro face: sample $\\theta$ and the field, return it at ``index``.

        Draws each hyperparameter from its prior (site ``f"{name}_{k}"``),
        then the field (site ``name``) with a *soft* constraint for intrinsic
        fields, so it works under NUTS. Returns the field at ``index``, or all
        `n_index` addressable nodes when ``index`` is ``None``. See
        `pyrox_lgm._numpyro.sample_component` for ``soft_constraint_scale``.
        """
        from pyrox_lgm._numpyro import sample_component

        return sample_component(
            self, index, soft_constraint_scale=soft_constraint_scale
        )

n_nodes: int abstractmethod property

Size of the latent field (including any padding node).

n_index: int property

Number of addressable nodes: what an observation index refers to.

Equal to n_nodes except where the field carries padding nodes (an odd-length RW2), which observations never touch.

theta_spec() -> dict[str, tuple[dist.Distribution, Transform]] abstractmethod

Hyperparameters: name to (prior, bijection from the reals).

The transform maps an unconstrained real to the parameter's support, which is how inla() explores \(\theta\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

@abc.abstractmethod
def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
    """Hyperparameters: name to (prior, bijection from the reals).

    The transform maps an unconstrained real to the parameter's support,
    which is how ``inla()`` explores $\\theta$.
    """

prior(theta: dict[str, Array], *, constraint: Constraint = 'hard', soft_constraint_scale: float = 0.001) -> gx.GaussianMRF | gx.IntrinsicGMRF abstractmethod

The GMRF of the field given its hyperparameters.

Parameters:

Name	Type	Description	Default
theta	dict[str, Array]	Hyperparameter values keyed as in theta_spec.	required
constraint	Constraint	For intrinsic fields, how the null space is handled ("hard" for inla(), "soft" for NUTS); ignored by proper fields.	'hard'
soft_constraint_scale	float	Standard deviation \(s\) of the soft constraint \(V^\top x \sim \mathcal N(0, s^2 I)\).	0.001

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

@abc.abstractmethod
def prior(
    self,
    theta: dict[str, Array],
    *,
    constraint: Constraint = "hard",
    soft_constraint_scale: float = 1e-3,
) -> gx.GaussianMRF | gx.IntrinsicGMRF:
    r"""The GMRF of the field given its hyperparameters.

    Args:
        theta: Hyperparameter values keyed as in `theta_spec`.
        constraint: For intrinsic fields, how the null space is handled
            (``"hard"`` for ``inla()``, ``"soft"`` for NUTS); ignored by
            proper fields.
        soft_constraint_scale: Standard deviation $s$ of the soft
            constraint $V^\top x \sim \mathcal N(0, s^2 I)$.
    """

assembly_precision(theta: dict[str, Array], gmrf: gx.GaussianMRF | gx.IntrinsicGMRF) -> lx.AbstractLinearOperator

The precision an LGM assembles into its sparse block-diagonal.

gmrf is prior(theta); its precision by default (scale * structure for an intrinsic field). A component whose prior uses an operator that cannot be assembled entry-wise (a grid SPDE's spectral function) returns an equivalent sparse one here.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

def assembly_precision(
    self, theta: dict[str, Array], gmrf: gx.GaussianMRF | gx.IntrinsicGMRF
) -> lx.AbstractLinearOperator:
    """The precision an `LGM` assembles into its sparse block-diagonal.

    ``gmrf`` is ``prior(theta)``; its precision by default (``scale *
    structure`` for an intrinsic field). A component whose prior uses an
    operator that cannot be assembled entry-wise (a grid SPDE's spectral
    function) returns an equivalent sparse one here.
    """
    if isinstance(gmrf, gx.IntrinsicGMRF):
        return scale_operator(gmrf.structure, gmrf.precision_scale)
    return gmrf.precision

constraint_basis(gmrf: gx.GaussianMRF | gx.IntrinsicGMRF) -> Float[Array, 'n c'] | None

Orthonormal basis of the hard constraints inla() imposes.

The intrinsic field's null space by default (sum-to-zero per connected component, say). A component may constrain fewer directions and leave the rest to the data, as R-INLA's rw2 constrains only the sum and keeps the linear trend; the density's rank is unchanged.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

def constraint_basis(
    self, gmrf: gx.GaussianMRF | gx.IntrinsicGMRF
) -> Float[Array, "n c"] | None:
    """Orthonormal basis of the hard constraints ``inla()`` imposes.

    The intrinsic field's null space by default (sum-to-zero per
    connected component, say). A component may constrain fewer
    directions and leave the rest to the data, as R-INLA's ``rw2``
    constrains only the sum and keeps the linear trend; the density's
    rank is unchanged.
    """
    return gmrf.null_space if isinstance(gmrf, gx.IntrinsicGMRF) else None

projector(index: Int[ArrayLike, ' n_obs']) -> gx.SparseOperator

(n_obs, n_nodes) selector with a one where observation i sits.

Parameters:

Name	Type	Description	Default
index	Int[ArrayLike, ' n_obs']	Node of each observation, integers in [0, n_index).	required

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

def projector(self, index: Int[ArrayLike, " n_obs"]) -> gx.SparseOperator:
    """``(n_obs, n_nodes)`` selector with a one where observation i sits.

    Args:
        index: Node of each observation, integers in ``[0, n_index)``.
    """
    index = np.asarray(index)
    if index.ndim != 1:
        raise ValueError(f"index must be 1-D, got shape {index.shape}")
    if index.size and (index.min() < 0 or index.max() >= self.n_index):
        raise ValueError(
            f"index out of range for {self.n_index} nodes: "
            f"[{index.min()}, {index.max()}]"
        )
    m = index.shape[0]
    return gx.SparseOperator.from_coo(
        np.arange(m), index, jnp.ones(m), (m, self.n_nodes)
    )

sample(index: Int[ArrayLike, ' n_obs'] | None = None, *, soft_constraint_scale: float = 0.01) -> Array

NumPyro face: sample \(\theta\) and the field, return it at index.

Draws each hyperparameter from its prior (site f"{name}_{k}"), then the field (site name) with a soft constraint for intrinsic fields, so it works under NUTS. Returns the field at index, or all n_index addressable nodes when index is None. See pyrox_lgm._numpyro.sample_component for soft_constraint_scale.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_base.py

def sample(
    self,
    index: Int[ArrayLike, " n_obs"] | None = None,
    *,
    soft_constraint_scale: float = 1e-2,
) -> Array:
    """NumPyro face: sample $\\theta$ and the field, return it at ``index``.

    Draws each hyperparameter from its prior (site ``f"{name}_{k}"``),
    then the field (site ``name``) with a *soft* constraint for intrinsic
    fields, so it works under NUTS. Returns the field at ``index``, or all
    `n_index` addressable nodes when ``index`` is ``None``. See
    `pyrox_lgm._numpyro.sample_component` for ``soft_constraint_scale``.
    """
    from pyrox_lgm._numpyro import sample_component

    return sample_component(
        self, index, soft_constraint_scale=soft_constraint_scale
    )

Temporal and unstructured

IID

Bases: AbstractComponent

Independent effects \(x_i \sim \mathcal N(0, 1/\tau)\), \(i < n\).

Parameters:

Name	Type	Description	Default
n	int	Number of nodes.	required
name	str	Site-name prefix.	'iid'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None

Examples:

>>> import jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> iid = lgm.IID(4, name="u")
>>> q = iid.prior({"tau": jnp.asarray(2.0)}).precision.as_matrix()
>>> [float(v) for v in jnp.diag(q)]
[2.0, 2.0, 2.0, 2.0]

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_temporal.py

class IID(AbstractComponent):
    r"""Independent effects $x_i \sim \mathcal N(0, 1/\tau)$, $i < n$.

    Args:
        n: Number of nodes.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> iid = lgm.IID(4, name="u")
        >>> q = iid.prior({"tau": jnp.asarray(2.0)}).precision.as_matrix()
        >>> [float(v) for v in jnp.diag(q)]
        [2.0, 2.0, 2.0, 2.0]
    """

    n: int = eqx.field(static=True)
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution

    def __init__(
        self, n: int, *, name: str = "iid", tau_prior: dist.Distribution | None = None
    ) -> None:
        self.n = int(n)
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)

    @property
    def n_nodes(self) -> int:
        return self.n

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {"tau": (self.tau_prior, default_transform(self.tau_prior))}

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF:
        del constraint, soft_constraint_scale  # proper
        return gx.GaussianMRF(jnp.zeros(self.n), gx.iid_precision(self.n, theta["tau"]))

RW1

Bases: _IntrinsicWalk

First-order random walk, \(x_{t+1} - x_t \sim \mathcal N(0, 1/\tau)\).

Intrinsic of order 1: the constant is its null space, removed by a sum-to-zero constraint. A cyclic walk also links \(x_{n-1}\) and \(x_0\).

Parameters:

Name	Type	Description	Default
n	int	Number of time points.	required
name	str	Site-name prefix.	'rw1'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
scale_model	bool	Rescale the structure so the geometric mean of the constrained marginal variances is 1 (R-INLA's scale.model).	True
cyclic	bool	Join the ends.	False

Examples:

>>> import jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> rw = lgm.RW1(20)
>>> var = rw.prior({"tau": jnp.asarray(1.0)}).marginal_variances()
>>> round(float(jnp.exp(jnp.mean(jnp.log(var)))), 4)  # scale_model
1.0

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_temporal.py

class RW1(_IntrinsicWalk):
    r"""First-order random walk, $x_{t+1} - x_t \sim \mathcal N(0, 1/\tau)$.

    Intrinsic of order 1: the constant is its null space, removed by a
    sum-to-zero constraint. A cyclic walk also links $x_{n-1}$ and $x_0$.

    Args:
        n: Number of time points.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        scale_model: Rescale the structure so the geometric mean of the
            constrained marginal variances is 1 (R-INLA's ``scale.model``).
        cyclic: Join the ends.

    Examples:
        >>> import jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> rw = lgm.RW1(20)
        >>> var = rw.prior({"tau": jnp.asarray(1.0)}).marginal_variances()
        >>> round(float(jnp.exp(jnp.mean(jnp.log(var)))), 4)  # scale_model
        1.0
    """

    n: int = eqx.field(static=True)
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    structure: lx.AbstractLinearOperator
    null_space: Float[Array, "N c"]
    scale: Float[Array, ""]

    def __init__(
        self,
        n: int,
        *,
        name: str = "rw1",
        tau_prior: dist.Distribution | None = None,
        scale_model: bool = True,
        cyclic: bool = False,
    ) -> None:
        self.n = int(n)
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        self.structure = gx.rw1_structure(self.n, cyclic=cyclic)
        self.null_space = jnp.ones((self.n, 1)) / jnp.sqrt(self.n)
        self.scale = (
            gx.generalized_variance_scale(self.structure, self.null_space)
            if scale_model
            else jnp.asarray(1.0)
        )

RW2

Bases: _IntrinsicWalk

Second-order random walk, \(x_{t+1} - 2x_t + x_{t-1} \sim \mathcal N(0, 1/\tau)\).

Intrinsic of order 2: constants and linear trends form its null space. Inside inla() only the sum-to-zero constraint is imposed, as R-INLA's rw2 does, and the linear trend is left to the data (the density's rank stays n - 2); the NumPyro face soft-constrains both directions, so add a linear fixed effect there if the data have a trend. For odd n, gaussx's banded structure carries one decoupled padding node (n_nodes = n + 1); it never enters a projector, and its own scale is taken out of scale_model.

Parameters:

Name	Type	Description	Default
n	int	Number of time points (at least 3).	required
name	str	Site-name prefix.	'rw2'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
scale_model	bool	Rescale as for RW1.	True

Examples:

>>> import jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> rw = lgm.RW2(7)  # odd: one padding node
>>> rw.n_nodes, rw.n_index
(8, 7)
>>> var = rw.prior({"tau": jnp.asarray(1.0)}).marginal_variances()[:7]
>>> round(float(jnp.exp(jnp.mean(jnp.log(var)))), 4)
1.0

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_temporal.py

class RW2(_IntrinsicWalk):
    r"""Second-order random walk, $x_{t+1} - 2x_t + x_{t-1} \sim \mathcal N(0, 1/\tau)$.

    Intrinsic of order 2: constants and linear trends form its null space.
    Inside ``inla()`` only the sum-to-zero constraint is imposed, as R-INLA's
    ``rw2`` does, and the linear trend is left to the data (the density's
    rank stays ``n - 2``); the NumPyro face soft-constrains both
    directions, so add a linear fixed effect there if the data have a trend.
    For odd ``n``,
    gaussx's banded structure carries one decoupled padding node
    (``n_nodes = n + 1``); it never enters a projector, and its own scale
    is taken out of ``scale_model``.

    Args:
        n: Number of time points (at least 3).
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        scale_model: Rescale as for `RW1`.

    Examples:
        >>> import jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> rw = lgm.RW2(7)  # odd: one padding node
        >>> rw.n_nodes, rw.n_index
        (8, 7)
        >>> var = rw.prior({"tau": jnp.asarray(1.0)}).marginal_variances()[:7]
        >>> round(float(jnp.exp(jnp.mean(jnp.log(var)))), 4)
        1.0
    """

    n: int = eqx.field(static=True)
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    structure: lx.AbstractLinearOperator
    null_space: Float[Array, "N c"]
    scale: Float[Array, ""]

    def __init__(
        self,
        n: int,
        *,
        name: str = "rw2",
        tau_prior: dist.Distribution | None = None,
        scale_model: bool = True,
    ) -> None:
        if n < 3:
            raise ValueError(f"RW2 needs at least 3 time points, got {n}")
        self.n = int(n)
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        self.structure = gx.rw2_structure(self.n)
        N = self.structure.in_size()
        t = jnp.arange(self.n, dtype=float)
        basis = jnp.zeros((N, 2)).at[: self.n, 0].set(1.0)
        basis = basis.at[: self.n, 1].set(t - t.mean())
        self.null_space, _ = jnp.linalg.qr(basis)  # zero on the padding row
        if scale_model:
            s_all = gx.generalized_variance_scale(self.structure, self.null_space)
            # The padding node has unit variance under R, so it pulls the
            # geometric mean over all N nodes towards 1: undo it to get the
            # scale over the n real nodes, s_n = s_N^(N/n).
            self.scale = s_all ** (N / self.n)
        else:
            self.scale = jnp.asarray(1.0)

    def constraint_basis(self, gmrf):
        # R-INLA's rw2 constrains the sum only: the linear trend stays in the
        # model (flat, identified by the data) instead of being removed.
        return self.null_space[:, :1]

AR1

Bases: AbstractComponent

Stationary AR(1), \(x_t = \rho x_{t-1} + \varepsilon_t\).

The marginal precision is \(\tau\): every marginal variance is \(1/\tau\) (R-INLA's ar1 convention, so \(\tau\) and \(\rho\) are separately interpretable).

Parameters:

Name	Type	Description	Default
n	int	Number of time points (at least 2).	required
name	str	Site-name prefix.	'ar1'
tau_prior	Distribution | None	Prior on the marginal precision; default PCPrecision(1, 0.01).	None
rho_prior	Distribution | None	Prior on \(\rho\); default PCAR1Rho(0.5, 0.5).	None

Examples:

>>> import jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> ar = lgm.AR1(30)
>>> gmrf = ar.prior({"tau": jnp.asarray(4.0), "rho": jnp.asarray(0.7)})
>>> [round(float(v), 4) for v in gmrf.marginal_variances()[:3]]
[0.25, 0.25, 0.25]

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_temporal.py

class AR1(AbstractComponent):
    r"""Stationary AR(1), $x_t = \rho x_{t-1} + \varepsilon_t$.

    The marginal precision is $\tau$: every marginal variance is $1/\tau$
    (R-INLA's ``ar1`` convention, so $\tau$ and $\rho$ are separately
    interpretable).

    Args:
        n: Number of time points (at least 2).
        name: Site-name prefix.
        tau_prior: Prior on the marginal precision; default
            ``PCPrecision(1, 0.01)``.
        rho_prior: Prior on $\rho$; default ``PCAR1Rho(0.5, 0.5)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> ar = lgm.AR1(30)
        >>> gmrf = ar.prior({"tau": jnp.asarray(4.0), "rho": jnp.asarray(0.7)})
        >>> [round(float(v), 4) for v in gmrf.marginal_variances()[:3]]
        [0.25, 0.25, 0.25]
    """

    n: int = eqx.field(static=True)
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    rho_prior: dist.Distribution

    def __init__(
        self,
        n: int,
        *,
        name: str = "ar1",
        tau_prior: dist.Distribution | None = None,
        rho_prior: dist.Distribution | None = None,
    ) -> None:
        if n < 2:
            raise ValueError(f"AR1 needs at least 2 time points, got {n}")
        self.n = int(n)
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        self.rho_prior = PCAR1Rho(0.5, 0.5) if rho_prior is None else rho_prior

    @property
    def n_nodes(self) -> int:
        return self.n

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {
            "tau": (self.tau_prior, default_transform(self.tau_prior)),
            "rho": (self.rho_prior, default_transform(self.rho_prior)),
        }

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF:
        del constraint, soft_constraint_scale  # proper
        Q = gx.ar1_precision(self.n, theta["rho"], theta["tau"])
        return gx.GaussianMRF(jnp.zeros(self.n), Q)

Generic

Generic

Bases: AbstractComponent

A field with precision \(\tau R\) for a user-supplied structure \(R\).

With null_space=None, \(R\) must be positive definite and the field is the proper GMRF \(\mathcal N(0, (\tau R)^{-1})\). Otherwise \(R\) is positive semidefinite with the given null space and the field is intrinsic, constrained to be orthogonal to it (R-INLA's generic0 with constr).

Pass \(R\) as a gaussx / lineax operator: a gaussx.SparseOperator keeps the sparse-Cholesky path.

Parameters:

Name	Type	Description	Default
structure	AbstractLinearOperator	\(R\), symmetric, size n.	required
null_space	ArrayLike | None	Basis of \(\ker R\), (n, c) or (n,); None for a proper field.	None
name	str	Site-name prefix.	'generic'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import pyrox_lgm as lgm
>>> R = lx.MatrixLinearOperator(
...     jnp.array([[2.0, -1.0], [-1.0, 2.0]]), lx.positive_semidefinite_tag
... )
>>> comp = lgm.Generic(R, name="g")
>>> comp.prior({"tau": jnp.asarray(3.0)}).precision.as_matrix().tolist()
[[6.0, -3.0], [-3.0, 6.0]]

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_generic.py

class Generic(AbstractComponent):
    r"""A field with precision $\tau R$ for a user-supplied structure $R$.

    With ``null_space=None``, $R$ must be positive definite and the field is
    the proper GMRF $\mathcal N(0, (\tau R)^{-1})$. Otherwise $R$ is
    positive semidefinite with the given null space and the field is
    intrinsic, constrained to be orthogonal to it (R-INLA's ``generic0``
    with ``constr``).

    Pass $R$ as a gaussx / lineax operator: a `gaussx.SparseOperator` keeps
    the sparse-Cholesky path.

    Args:
        structure: $R$, symmetric, size ``n``.
        null_space: Basis of $\ker R$, ``(n, c)`` or ``(n,)``; ``None`` for a
            proper field.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import pyrox_lgm as lgm
        >>> R = lx.MatrixLinearOperator(
        ...     jnp.array([[2.0, -1.0], [-1.0, 2.0]]), lx.positive_semidefinite_tag
        ... )
        >>> comp = lgm.Generic(R, name="g")
        >>> comp.prior({"tau": jnp.asarray(3.0)}).precision.as_matrix().tolist()
        [[6.0, -3.0], [-3.0, 6.0]]
    """

    structure: lx.AbstractLinearOperator
    null_space: Float[Array, "n c"] | None
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution

    def __init__(
        self,
        structure: lx.AbstractLinearOperator,
        null_space: ArrayLike | None = None,
        *,
        name: str = "generic",
        tau_prior: dist.Distribution | None = None,
    ) -> None:
        if structure.in_size() != structure.out_size():
            raise ValueError("structure must be square")
        self.structure = structure
        if null_space is None:
            self.null_space = None
        else:
            V = jnp.asarray(null_space)
            V = V[:, None] if V.ndim == 1 else V
            if V.shape[0] != structure.in_size():
                raise ValueError(
                    f"null_space has {V.shape[0]} rows for a structure of "
                    f"size {structure.in_size()}"
                )
            self.null_space = V
        self.name = name
        self.tau_prior = PCPrecision(1.0, 0.01) if tau_prior is None else tau_prior

    @property
    def n_nodes(self) -> int:
        return self.structure.in_size()

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {"tau": (self.tau_prior, default_transform(self.tau_prior))}

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF | gx.IntrinsicGMRF:
        loc = jnp.zeros(self.n_nodes)
        if self.null_space is None:
            return gx.GaussianMRF(loc, scale_operator(self.structure, theta["tau"]))
        return gx.IntrinsicGMRF(
            loc,
            theta["tau"],
            self.structure,
            self.null_space,
            constraint=constraint,
            soft_constraint_scale=soft_constraint_scale,
        )

Areal

On a kernellib graph (from edges, polygons' contiguity, or a grid).

Besag

Bases: AbstractComponent

Intrinsic CAR (ICAR / Besag): precision \(\tau R\), \(R = D - W\).

One sum-to-zero constraint per connected component. With scale_model=True (default) \(R\) is BYM2-scaled per component, so \(\tau\) is comparable across graphs (R-INLA's scale.model).

Parameters:

Name	Type	Description	Default
graph	AbstractGraph	A kernellib graph.	required
name	str	Site-name prefix.	'besag'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
scale_model	bool	Scale the structure per connected component.	True

Examples:

>>> import jax, jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> g = kl.graph_from_edges([0, 1, 3], [1, 2, 4], 5)  # two components
>>> icar = lgm.Besag(g)
>>> x = icar.prior({"tau": jnp.asarray(1.0)}).sample(jax.random.key(0))
>>> [round(float(s), 6) + 0.0 for s in (x[:3].sum(), x[3:].sum())]
[0.0, 0.0]

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_areal.py

class Besag(AbstractComponent):
    r"""Intrinsic CAR (ICAR / Besag): precision $\tau R$, $R = D - W$.

    One sum-to-zero constraint per connected component. With
    ``scale_model=True`` (default) $R$ is BYM2-scaled per component, so
    $\tau$ is comparable across graphs (R-INLA's ``scale.model``).

    Args:
        graph: A kernellib graph.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        scale_model: Scale the structure per connected component.

    Examples:
        >>> import jax, jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> g = kl.graph_from_edges([0, 1, 3], [1, 2, 4], 5)  # two components
        >>> icar = lgm.Besag(g)
        >>> x = icar.prior({"tau": jnp.asarray(1.0)}).sample(jax.random.key(0))
        >>> [round(float(s), 6) + 0.0 for s in (x[:3].sum(), x[3:].sum())]
        [0.0, 0.0]
    """

    graph: kl.AbstractGraph
    structure: lx.AbstractLinearOperator
    null_space: Float[Array, "n c"]
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution

    def __init__(
        self,
        graph: kl.AbstractGraph,
        *,
        name: str = "besag",
        tau_prior: dist.Distribution | None = None,
        scale_model: bool = True,
    ) -> None:
        self.graph = graph
        self.structure = kl.structure_matrix(graph, scaled=scale_model)
        self.null_space = kl.graph_null_space(graph)
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)

    @property
    def n_nodes(self) -> int:
        return self.structure.in_size()

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {"tau": (self.tau_prior, default_transform(self.tau_prior))}

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.IntrinsicGMRF:
        return gx.IntrinsicGMRF(
            jnp.zeros(self.n_nodes),
            theta["tau"],
            self.structure,
            self.null_space,
            constraint=constraint,
            soft_constraint_scale=soft_constraint_scale,
        )

BYM2

Bases: AbstractComponent

BYM2 (Riebler et al., 2016), a mixture of unstructured and ICAR effects.

\(b = \tau^{-1/2}(\sqrt{1-\phi}\,v + \sqrt\phi\,u_\ast)\), with \(v\) unstructured and \(u_\ast\) a scaled ICAR, so \(\tau\) is the total precision of \(b\) and \(\phi\) the fraction of its variance that is spatially structured. The field is the pair \((b, u_\ast)\) (n_nodes = 2n), distributed as gaussx.BYM2GMRF with its exact constrained density; observations see \(b\) only (n_index = n). The PC prior on \(\phi\) needs the spectrum of \(R_\ast^{+}\), computed once here (structure_spectrum), as is \(\log|R_\ast|_+\) for the normaliser.

Parameters:

Name	Type	Description	Default
graph	AbstractGraph	A kernellib graph.	required
name	str	Site-name prefix.	'bym2'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
phi_prior	Distribution | None	Prior on \(\phi\); default PCBYM2Phi(0.5, 2/3) on this graph's spectrum.	None
spectrum_method	Literal['auto', 'dense', 'lanczos']	How to compute that spectrum (see structure_spectrum).	'auto'

Examples:

>>> import jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> bym2 = lgm.BYM2(kl.grid_graph((4, 4)))
>>> bym2.n_nodes, bym2.n_index
(32, 16)
>>> gmrf = bym2.prior({"tau": jnp.asarray(2.0), "phi": jnp.asarray(0.5)})
>>> var = gmrf.marginal_variances()
>>> var_b, var_u = var[:16], var[16:]  # var(b) = (1 - phi + phi var(u*)) / tau
>>> bool(jnp.allclose(var_b, (0.5 + 0.5 * var_u) / 2.0))
True

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_areal.py

class BYM2(AbstractComponent):
    r"""BYM2 (Riebler et al., 2016), a mixture of unstructured and ICAR effects.

    $b = \tau^{-1/2}(\sqrt{1-\phi}\,v + \sqrt\phi\,u_\ast)$, with $v$
    unstructured and $u_\ast$ a scaled ICAR, so $\tau$ is the total
    precision of $b$ and $\phi$ the fraction of its variance that is spatially
    structured. The field is the pair $(b, u_\ast)$ (``n_nodes = 2n``),
    distributed as `gaussx.BYM2GMRF` with its exact constrained density;
    observations see $b$ only (``n_index = n``). The PC prior on $\phi$
    needs the spectrum of $R_\ast^{+}$, computed once here
    (`structure_spectrum`), as is $\log|R_\ast|_+$ for the normaliser.

    Args:
        graph: A kernellib graph.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        phi_prior: Prior on $\phi$; default ``PCBYM2Phi(0.5, 2/3)`` on this
            graph's spectrum.
        spectrum_method: How to compute that spectrum (see
            `structure_spectrum`).

    Examples:
        >>> import jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> bym2 = lgm.BYM2(kl.grid_graph((4, 4)))
        >>> bym2.n_nodes, bym2.n_index
        (32, 16)
        >>> gmrf = bym2.prior({"tau": jnp.asarray(2.0), "phi": jnp.asarray(0.5)})
        >>> var = gmrf.marginal_variances()
        >>> var_b, var_u = var[:16], var[16:]  # var(b) = (1 - phi + phi var(u*)) / tau
        >>> bool(jnp.allclose(var_b, (0.5 + 0.5 * var_u) / 2.0))
        True
    """

    graph: kl.AbstractGraph
    structure: lx.AbstractLinearOperator
    null_space: Float[Array, "n c"]
    log_pdet: Float[Array, ""]
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    phi_prior: dist.Distribution

    def __init__(
        self,
        graph: kl.AbstractGraph,
        *,
        name: str = "bym2",
        tau_prior: dist.Distribution | None = None,
        phi_prior: dist.Distribution | None = None,
        spectrum_method: Literal["auto", "dense", "lanczos"] = "auto",
    ) -> None:
        self.graph = graph
        scaled = kl.structure_matrix(graph, scaled=True)
        # gaussx's BYM2 precision is assembled sparse, so a grid's Kronecker
        # sum is used only for its (exact, cheap) spectrum.
        self.structure = (
            kl.structure_matrix(kl.Graph(graph.topology, graph.weights), scaled=True)
            if isinstance(graph, kl.GridGraph)
            else scaled
        )
        self.null_space = kl.graph_null_space(graph)
        self.log_pdet = gx.pseudo_logdet(self.structure, structure="laplacian")
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        if phi_prior is None:
            spectrum = structure_spectrum(
                scaled, self.null_space, method=spectrum_method
            )
            phi_prior = PCBYM2Phi(0.5, 2.0 / 3.0, structure_spectrum=spectrum)
        self.phi_prior = phi_prior

    @property
    def n_nodes(self) -> int:
        return 2 * self.structure.in_size()

    @property
    def n_index(self) -> int:
        return self.structure.in_size()

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {
            "tau": (self.tau_prior, default_transform(self.tau_prior)),
            "phi": (self.phi_prior, default_transform(self.phi_prior)),
        }

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.BYM2GMRF:
        return gx.BYM2GMRF(
            self.structure,
            theta["tau"],
            theta["phi"],
            self.null_space,
            constraint=constraint,
            include_normalizer=True,
            log_pdet=self.log_pdet,
        )

CAR

Bases: _ProperAreal

Proper CAR: precision \(\tau(D - \rho W)\), \(\rho \in [0, 1)\).

\(D\) is the diagonal of degrees and \(W\) the weighted adjacency; \(\rho \to 1\) approaches the ICAR. Every node needs a neighbour. The log-determinant comes from the spectrum \(\mu_i\) of \(D^{-1/2} W D^{-1/2}\) (one dense eigh at construction, for n ≤ 10⁴): \(\log|Q| = n\log\tau + \sum_i \log d_i + \sum_i \log(1 - \rho\mu_i)\); larger graphs use gaussx's sparse Cholesky at each \(\theta\).

Parameters:

Name	Type	Description	Default
graph	AbstractGraph	A kernellib graph without isolated nodes.	required
name	str	Site-name prefix.	'car'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
rho_prior	Distribution | None	Prior on \(\rho\); default Uniform(0, 1).	None

Examples:

>>> import jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> car = lgm.CAR(kl.graph_from_edges([0, 1], [1, 2], 3))
>>> Q = car.prior({"tau": jnp.asarray(1.0), "rho": jnp.asarray(0.5)})
>>> Q.precision.as_matrix().tolist()
[[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0]]

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_areal.py

class CAR(_ProperAreal):
    r"""Proper CAR: precision $\tau(D - \rho W)$, $\rho \in [0, 1)$.

    $D$ is the diagonal of degrees and $W$ the weighted adjacency; $\rho \to
    1$ approaches the ICAR. Every node needs a neighbour. The
    log-determinant comes from the spectrum $\mu_i$ of
    $D^{-1/2} W D^{-1/2}$ (one dense ``eigh`` at construction, for
    ``n ≤ 10⁴``): $\log|Q| = n\log\tau + \sum_i \log d_i + \sum_i \log(1 - \rho\mu_i)$;
    larger graphs use gaussx's sparse Cholesky at each $\theta$.

    Args:
        graph: A kernellib graph without isolated nodes.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        rho_prior: Prior on $\rho$; default ``Uniform(0, 1)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> car = lgm.CAR(kl.graph_from_edges([0, 1], [1, 2], 3))
        >>> Q = car.prior({"tau": jnp.asarray(1.0), "rho": jnp.asarray(0.5)})
        >>> Q.precision.as_matrix().tolist()
        [[1.0, -0.5, 0.0], [-0.5, 2.0, -0.5], [0.0, -0.5, 1.0]]
    """

    graph: kl.AbstractGraph
    laplacian: gx.SparseOperator
    log_degree: Float[Array, ""]
    spectrum: Float[Array, " n"] | None
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    rho_prior: dist.Distribution

    def __init__(
        self,
        graph: kl.AbstractGraph,
        *,
        name: str = "car",
        tau_prior: dist.Distribution | None = None,
        rho_prior: dist.Distribution | None = None,
    ) -> None:
        _check_no_isolated(graph, "CAR")
        self.graph = graph
        self.laplacian = _sparse_laplacian(graph)
        d = jnp.asarray(graph.degree())
        self.log_degree = jnp.sum(jnp.log(d))
        n = d.shape[0]
        if n <= _DENSE_SPECTRUM_MAX:
            L = self.laplacian.as_matrix()
            s = 1.0 / jnp.sqrt(d)
            # D^-1/2 W D^-1/2 = I - D^-1/2 L D^-1/2.
            self.spectrum = 1.0 - jnp.linalg.eigvalsh(s[:, None] * L * s[None, :])
        else:
            self.spectrum = None
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        self.rho_prior = _rho_prior(rho_prior)

    @property
    def n_nodes(self) -> int:
        return self.laplacian.in_size()

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF:
        del constraint, soft_constraint_scale  # proper
        tau, rho = theta["tau"], theta["rho"]
        # tau (D - rho W) = tau ((1 - rho) D + rho L): degrees on the
        # diagonal, rho times the Laplacian's off-diagonal.
        Q = eqx.tree_at(
            lambda o: o.values,
            self.laplacian,
            tau
            * jnp.where(
                _diag_mask(self.laplacian),
                self.laplacian.values,
                rho * self.laplacian.values,
            ),
        )
        log_det = None
        if self.spectrum is not None:
            n = self.n_nodes
            log_det = (
                n * jnp.log(tau)
                + self.log_degree
                + jnp.sum(jnp.log1p(-rho * self.spectrum))
            )
        return gx.GaussianMRF(jnp.zeros(self.n_nodes), Q, log_det_precision=log_det)

Leroux

Bases: _ProperAreal

Leroux CAR: precision \(\tau(\rho R + (1 - \rho) I)\), \(\rho \in [0, 1]\).

Interpolates between independent effects (\(\rho = 0\)) and the ICAR (\(\rho \to 1\)). On a face-connected kernellib.GridGraph the precision stays a gaussx.KroneckerSum (exact eigen-structured solves and log-determinant); elsewhere it is a gaussx.SparseOperator, with \(\log|Q| = n\log\tau + \sum_i \log(\rho\lambda_i + 1 - \rho)\) from the Laplacian spectrum \(\lambda_i\) cached at construction (n ≤ 10⁴).

Parameters:

Name	Type	Description	Default
graph	AbstractGraph	A kernellib graph.	required
name	str	Site-name prefix.	'leroux'
tau_prior	Distribution | None	Prior on \(\tau\); default PCPrecision(1, 0.01).	None
rho_prior	Distribution | None	Prior on \(\rho\); default Uniform(0, 1).	None
scale_model	bool	Use the BYM2-scaled \(R\) (per connected component).	False

Examples:

>>> import jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> lr = lgm.Leroux(kl.grid_graph((3, 3)))
>>> gmrf = lr.prior({"tau": jnp.asarray(1.0), "rho": jnp.asarray(0.0)})
>>> bool(jnp.allclose(gmrf.precision.as_matrix(), jnp.eye(9)))  # IID
True

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_areal.py

class Leroux(_ProperAreal):
    r"""Leroux CAR: precision $\tau(\rho R + (1 - \rho) I)$, $\rho \in [0, 1]$.

    Interpolates between independent effects ($\rho = 0$) and the ICAR
    ($\rho \to 1$). On a face-connected `kernellib.GridGraph` the precision
    stays a `gaussx.KroneckerSum` (exact eigen-structured solves and
    log-determinant); elsewhere it is a `gaussx.SparseOperator`, with
    $\log|Q| = n\log\tau + \sum_i \log(\rho\lambda_i + 1 - \rho)$ from the
    Laplacian spectrum $\lambda_i$ cached at construction (``n ≤ 10⁴``).

    Args:
        graph: A kernellib graph.
        name: Site-name prefix.
        tau_prior: Prior on $\tau$; default ``PCPrecision(1, 0.01)``.
        rho_prior: Prior on $\rho$; default ``Uniform(0, 1)``.
        scale_model: Use the BYM2-scaled $R$ (per connected component).

    Examples:
        >>> import jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> lr = lgm.Leroux(kl.grid_graph((3, 3)))
        >>> gmrf = lr.prior({"tau": jnp.asarray(1.0), "rho": jnp.asarray(0.0)})
        >>> bool(jnp.allclose(gmrf.precision.as_matrix(), jnp.eye(9)))  # IID
        True
    """

    graph: kl.AbstractGraph
    structure: gx.KroneckerSum | gx.SparseOperator
    spectrum: Float[Array, " n"] | None
    name: str = eqx.field(static=True)
    tau_prior: dist.Distribution
    rho_prior: dist.Distribution

    def __init__(
        self,
        graph: kl.AbstractGraph,
        *,
        name: str = "leroux",
        tau_prior: dist.Distribution | None = None,
        rho_prior: dist.Distribution | None = None,
        scale_model: bool = False,
    ) -> None:
        self.graph = graph
        R = kl.structure_matrix(graph, scaled=scale_model)
        if not isinstance(R, gx.KroneckerSum | gx.SparseOperator):
            raise TypeError(f"unsupported structure type {type(R)}")
        self.structure = R
        n = R.in_size()
        if isinstance(R, gx.KroneckerSum) or n <= _DENSE_SPECTRUM_MAX:
            self.spectrum = _dense_eigvals(R)
        else:
            self.spectrum = None
        self.name = name
        self.tau_prior = _tau_prior(tau_prior)
        self.rho_prior = _rho_prior(rho_prior)

    @property
    def n_nodes(self) -> int:
        return self.structure.in_size()

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF:
        del constraint, soft_constraint_scale  # proper
        tau, rho = theta["tau"], theta["rho"]
        if isinstance(self.structure, gx.KroneckerSum):
            Q = _affine_kronecker_sum(self.structure, tau * rho, tau * (1.0 - rho))
        else:
            Q = _affine_sparse(self.structure, tau * rho, tau * (1.0 - rho))
        log_det = None
        if self.spectrum is not None:
            log_det = self.n_nodes * jnp.log(tau) + jnp.sum(
                jnp.log(rho * self.spectrum + 1.0 - rho)
            )
        return gx.GaussianMRF(jnp.zeros(self.n_nodes), Q, log_det_precision=log_det)

SPDE (Matérn)

SPDE

Bases: AbstractComponent

Matérn field via the SPDE, on a triangle mesh or a regular grid.

Exactly one of mesh and grid is given:

• mesh=(vertices, triangles): P1 finite elements on a triangulated surface (vertices of shape (V, 2) or (V, 3) on a manifold), with lumped mass C̃ and stiffness G from gaussx.fem_matrices; the precision is gaussx.spde_precision (sparse). Nodes are the vertices; project_points builds the barycentric projector for arbitrary locations. The field is 2-D (d = 2).
• grid=shape: the finite-difference SPDE on a regular lattice with spacing (gaussx.spde_precision_grid, diagonalised exactly by its Kronecker structure). Nodes are the grid cells in C order; d is len(shape).

The hyperparameter "range_sigma" is the pair [range, sigma] under the joint PCMatern prior.

Parameters:

Name	Type	Description	Default
mesh	tuple[ArrayLike, ArrayLike] | None	(vertices, triangles).	None
grid	tuple[int, ...] | None	Lattice shape.	None
spacing	float	Grid spacing (grid only).	1.0
alpha	int	SPDE order \(\alpha = \nu + d/2\), an integer (2 is the default Matérn-1 field in 2-D).	2
name	str	Site-name prefix.	'spde'
prior	Distribution | None	Joint prior on [range, sigma]; default PCMatern(range0, 0.5, 1, 0.01, d) with range0 a fifth of the domain's diameter.	None

Examples:

>>> import jax.numpy as jnp
>>> import pyrox_lgm as lgm
>>> spde = lgm.SPDE(grid=(30, 30), alpha=2)
>>> g = spde.prior({"range_sigma": jnp.array([8.0, 2.0])})
>>> var = g.marginal_variances().reshape(30, 30)
>>> # ~ sigma^2 = 4 away from the border, up to the lattice discretisation
>>> bool(abs(float(var[15, 15]) - 4.0) < 0.4)
True

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_spde.py

class SPDE(AbstractComponent):
    r"""Matérn field via the SPDE, on a triangle mesh or a regular grid.

    Exactly one of ``mesh`` and ``grid`` is given:

    - ``mesh=(vertices, triangles)``: P1 finite elements on a triangulated
      surface (``vertices`` of shape ``(V, 2)`` or ``(V, 3)`` on a manifold),
      with lumped mass ``C̃`` and stiffness ``G`` from
      `gaussx.fem_matrices`; the precision is `gaussx.spde_precision`
      (sparse). Nodes are the vertices; `project_points` builds the
      barycentric projector for arbitrary locations. The field is 2-D
      (``d = 2``).
    - ``grid=shape``: the finite-difference SPDE on a regular lattice with
      ``spacing`` (`gaussx.spde_precision_grid`, diagonalised exactly by its
      Kronecker structure). Nodes are the grid cells in C order; ``d`` is
      ``len(shape)``.

    The hyperparameter ``"range_sigma"`` is the pair ``[range, sigma]``
    under the joint `PCMatern` prior.

    Args:
        mesh: ``(vertices, triangles)``.
        grid: Lattice shape.
        spacing: Grid spacing (``grid`` only).
        alpha: SPDE order $\alpha = \nu + d/2$, an integer (2 is the
            default Matérn-1 field in 2-D).
        name: Site-name prefix.
        prior: Joint prior on ``[range, sigma]``; default
            ``PCMatern(range0, 0.5, 1, 0.01, d)`` with ``range0`` a fifth of
            the domain's diameter.

    Examples:
        >>> import jax.numpy as jnp
        >>> import pyrox_lgm as lgm
        >>> spde = lgm.SPDE(grid=(30, 30), alpha=2)
        >>> g = spde.prior({"range_sigma": jnp.array([8.0, 2.0])})
        >>> var = g.marginal_variances().reshape(30, 30)
        >>> # ~ sigma^2 = 4 away from the border, up to the lattice discretisation
        >>> bool(abs(float(var[15, 15]) - 4.0) < 0.4)
        True
    """

    alpha: int = eqx.field(static=True)
    d: int = eqx.field(static=True)
    shape: tuple[int, ...] | None = eqx.field(static=True)
    spacing: float = eqx.field(static=True)
    vertices: Float[Array, "V D"] | None
    triangles: Int[Array, "T 3"] | None
    mass: lx.DiagonalLinearOperator | None
    stiffness: gx.SparseOperator | None
    name: str = eqx.field(static=True)
    prior_range_sigma: dist.Distribution

    def __init__(
        self,
        *,
        mesh: tuple[ArrayLike, ArrayLike] | None = None,
        grid: tuple[int, ...] | None = None,
        spacing: float = 1.0,
        alpha: int = 2,
        name: str = "spde",
        prior: dist.Distribution | None = None,
    ) -> None:
        if (mesh is None) == (grid is None):
            raise ValueError(
                "give exactly one of mesh=(vertices, triangles) or grid=shape"
            )
        self.alpha = int(alpha)
        self.spacing = float(spacing)
        self.name = name
        if mesh is not None:
            V = jnp.asarray(mesh[0])
            T = jnp.asarray(mesh[1])
            self.vertices, self.triangles = V, T
            self.mass, self.stiffness = gx.fem_matrices(V, T)
            self.shape = None
            self.d = 2
            extent = np.asarray(V.max(axis=0) - V.min(axis=0))
        else:
            assert grid is not None
            self.shape = tuple(int(s) for s in grid)
            self.vertices = self.triangles = None
            self.mass = self.stiffness = None
            self.d = len(self.shape)
            extent = (np.asarray(self.shape) - 1) * self.spacing
        if self.alpha - self.d / 2 <= 0:
            raise ValueError(
                f"alpha={self.alpha} gives smoothness nu = alpha - d/2 <= 0 "
                f"in d={self.d}"
            )
        if prior is None:
            range0 = float(np.linalg.norm(extent)) / 5.0
            prior = PCMatern(range0, 0.5, 1.0, 0.01, d=self.d)
        self.prior_range_sigma = prior

    @property
    def nu(self) -> float:
        """Smoothness $\\nu = \\alpha - d/2$."""
        return self.alpha - self.d / 2

    @property
    def n_nodes(self) -> int:
        if self.shape is not None:
            return int(np.prod(self.shape))
        assert self.vertices is not None
        return int(self.vertices.shape[0])

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        p = self.prior_range_sigma
        return {"range_sigma": (p, default_transform(p))}

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF:
        del constraint, soft_constraint_scale  # proper
        range_, sigma = theta["range_sigma"][0], theta["range_sigma"][1]
        kappa, tau, alpha = gx.matern_spde_params(range_, sigma, self.nu, self.d)
        if self.shape is not None:
            Q = gx.spde_precision_grid(
                self.shape, kappa, tau, alpha, spacing=self.spacing
            )
        else:
            assert self.mass is not None and self.stiffness is not None
            Q = gx.spde_precision(self.mass, self.stiffness, kappa, tau, alpha)
        return gx.GaussianMRF(jnp.zeros(self.n_nodes), Q)

    def assembly_precision(self, theta, gmrf):
        """The grid precision as a `gaussx.SparseOperator` (mesh: unchanged).

        ``Q = tau^2 h^d (kappa^2 I + h^-2 L)^alpha``, ``L`` the Kronecker sum
        of the axes' path-graph Laplacians, exactly `gaussx.spde_precision_grid`
        written entry-wise: ``alpha`` sparse products on a theta-free pattern.
        """
        if self.shape is None:
            return gmrf.precision
        range_, sigma = theta["range_sigma"][0], theta["range_sigma"][1]
        kappa, tau, alpha = gx.matern_spde_params(range_, sigma, self.nu, self.d)
        h = self.spacing
        laplacian = _path_laplacians(self.shape)
        r, c, v = full_coo(laplacian)
        n = self.n_nodes
        diag = np.arange(n)
        base = (
            np.concatenate([r, diag]),
            np.concatenate([c, diag]),
            jnp.concatenate([v / h**2, jnp.full(n, kappa**2)]),
        )
        rows, cols, vals = coo_power(base, alpha, n)
        return gx.SparseOperator.from_coo(rows, cols, tau**2 * h**self.d * vals, (n, n))

    def project_points(self, points: Float[ArrayLike, "n D"]) -> gx.SparseOperator:
        """Barycentric ``(n, V)`` projector from mesh vertices to ``points``."""
        if self.vertices is None or self.triangles is None:
            raise ValueError("project_points needs a mesh; index grid cells directly")
        return gx.fem_projector(self.vertices, self.triangles, points)

nu: float property

Smoothness \(\nu = \alpha - d/2\).

assembly_precision(theta, gmrf)

The grid precision as a gaussx.SparseOperator (mesh: unchanged).

Q = tau^2 h^d (kappa^2 I + h^-2 L)^alpha, L the Kronecker sum of the axes' path-graph Laplacians, exactly gaussx.spde_precision_grid written entry-wise: alpha sparse products on a theta-free pattern.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_spde.py

def assembly_precision(self, theta, gmrf):
    """The grid precision as a `gaussx.SparseOperator` (mesh: unchanged).

    ``Q = tau^2 h^d (kappa^2 I + h^-2 L)^alpha``, ``L`` the Kronecker sum
    of the axes' path-graph Laplacians, exactly `gaussx.spde_precision_grid`
    written entry-wise: ``alpha`` sparse products on a theta-free pattern.
    """
    if self.shape is None:
        return gmrf.precision
    range_, sigma = theta["range_sigma"][0], theta["range_sigma"][1]
    kappa, tau, alpha = gx.matern_spde_params(range_, sigma, self.nu, self.d)
    h = self.spacing
    laplacian = _path_laplacians(self.shape)
    r, c, v = full_coo(laplacian)
    n = self.n_nodes
    diag = np.arange(n)
    base = (
        np.concatenate([r, diag]),
        np.concatenate([c, diag]),
        jnp.concatenate([v / h**2, jnp.full(n, kappa**2)]),
    )
    rows, cols, vals = coo_power(base, alpha, n)
    return gx.SparseOperator.from_coo(rows, cols, tau**2 * h**self.d * vals, (n, n))

project_points(points: Float[ArrayLike, 'n D']) -> gx.SparseOperator

Barycentric (n, V) projector from mesh vertices to points.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_spde.py

def project_points(self, points: Float[ArrayLike, "n D"]) -> gx.SparseOperator:
    """Barycentric ``(n, V)`` projector from mesh vertices to ``points``."""
    if self.vertices is None or self.triangles is None:
        raise ValueError("project_points needs a mesh; index grid cells directly")
    return gx.fem_projector(self.vertices, self.triangles, points)

Combinators

Kronecker

Bases: AbstractComponent

Separable product \(Q = Q_\text{main} \otimes Q_\text{group}\) (R-INLA group).

The precisions of the two factors multiply, so only the product of their \(\tau\)'s is identified; as R-INLA does, the group factor's tau is fixed at 1 and the group contributes only its correlation structure (an AR1 group keeps rho). Hyperparameter k of a factor named f is f"{f}_{k}".

Parameters:

Name	Type	Description	Default
main	AbstractComponent	The main component (e.g. a spatial field).	required
group	AbstractComponent	The grouping component (e.g. an AR1 over time).	required
name	str	Site-name prefix.	'kronecker'

Examples:

>>> import jax.numpy as jnp
>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> st = lgm.Kronecker(
...     lgm.Leroux(kl.grid_graph((3, 3)), name="space"),
...     lgm.AR1(4, name="time"),
... )
>>> sorted(st.theta_spec())
['space_rho', 'space_tau', 'time_rho']
>>> st.n_nodes
36

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_combinators.py

class Kronecker(AbstractComponent):
    r"""Separable product $Q = Q_\text{main} \otimes Q_\text{group}$ (R-INLA ``group``).

    The precisions of the two factors multiply, so only the product of their
    $\tau$'s is identified; as R-INLA does, the group factor's ``tau`` is
    fixed at 1 and the group contributes only its correlation structure (an
    `AR1` group keeps ``rho``). Hyperparameter ``k`` of a factor named ``f``
    is ``f"{f}_{k}"``.

    Args:
        main: The main component (e.g. a spatial field).
        group: The grouping component (e.g. an `AR1` over time).
        name: Site-name prefix.

    Examples:
        >>> import jax.numpy as jnp
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> st = lgm.Kronecker(
        ...     lgm.Leroux(kl.grid_graph((3, 3)), name="space"),
        ...     lgm.AR1(4, name="time"),
        ... )
        >>> sorted(st.theta_spec())
        ['space_rho', 'space_tau', 'time_rho']
        >>> st.n_nodes
        36
    """

    main: AbstractComponent
    group: AbstractComponent
    name: str = eqx.field(static=True)

    def __init__(
        self,
        main: AbstractComponent,
        group: AbstractComponent,
        *,
        name: str = "kronecker",
    ) -> None:
        _check_addressable(main, "main")
        _check_addressable(group, "group")
        if "tau" not in group.theta_spec():
            raise ValueError("the group component must have a 'tau' to fix at 1")
        self.main, self.group, self.name = main, group, name

    @property
    def n_nodes(self) -> int:
        return self.main.n_nodes * self.group.n_nodes

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        spec = {f"{self.main.name}_{k}": v for k, v in self.main.theta_spec().items()}
        spec |= {
            f"{self.group.name}_{k}": v
            for k, v in self.group.theta_spec().items()
            if k != "tau"
        }
        return spec

    def _split(self, theta: dict[str, Array]):
        main = {k: theta[f"{self.main.name}_{k}"] for k in self.main.theta_spec()}
        group = {
            k: (jnp.asarray(1.0) if k == "tau" else theta[f"{self.group.name}_{k}"])
            for k in self.group.theta_spec()
        }
        return main, group

    def constraint_basis(self, gmrf):
        # The intrinsic factor's own constraints, lifted (an RW2 main keeps
        # its sum-only rule), not the product's whole null space.
        main, group = _probe_basis(self.main), _probe_basis(self.group)
        if main is not None:
            return jnp.kron(main, jnp.eye(self.group.n_nodes))
        if group is not None:
            return jnp.kron(jnp.eye(self.main.n_nodes), group)
        return None

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF | gx.IntrinsicGMRF:
        main, group = self._split(theta)
        return _product(
            self.main.prior(main, constraint=constraint),
            self.group.prior(group, constraint=constraint),
            constraint,
            soft_constraint_scale,
        )

Replicate

Bases: AbstractComponent

n_rep independent copies of a component, sharing its hyperparameters.

\(Q = I_{n_\text{rep}} \otimes Q_\text{component}\) (R-INLA's replicate); an intrinsic component keeps its constraints in every copy.

Parameters:

Name	Type	Description	Default
component	AbstractComponent	The replicated component.	required
n_rep	int	Number of copies.	required
name	str | None	Site-name prefix; defaults to the component's.	None

Examples:

>>> import pyrox_lgm as lgm
>>> rep = lgm.Replicate(lgm.RW1(10, name="trend"), 3)
>>> rep.n_nodes, list(rep.theta_spec())
(30, ['tau'])

Source code in packages/pyrox-lgm/src/pyrox_lgm/_components/_combinators.py

class Replicate(AbstractComponent):
    r"""``n_rep`` independent copies of a component, sharing its hyperparameters.

    $Q = I_{n_\text{rep}} \otimes Q_\text{component}$ (R-INLA's
    ``replicate``); an intrinsic component keeps its constraints in every
    copy.

    Args:
        component: The replicated component.
        n_rep: Number of copies.
        name: Site-name prefix; defaults to the component's.

    Examples:
        >>> import pyrox_lgm as lgm
        >>> rep = lgm.Replicate(lgm.RW1(10, name="trend"), 3)
        >>> rep.n_nodes, list(rep.theta_spec())
        (30, ['tau'])
    """

    component: AbstractComponent
    n_rep: int = eqx.field(static=True)
    name: str = eqx.field(static=True)

    def __init__(
        self, component: AbstractComponent, n_rep: int, *, name: str | None = None
    ) -> None:
        _check_addressable(component, "replicated")
        self.component = component
        self.n_rep = int(n_rep)
        self.name = component.name if name is None else name

    @property
    def n_nodes(self) -> int:
        return self.n_rep * self.component.n_nodes

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return self.component.theta_spec()

    def constraint_basis(self, gmrf):
        inner = _probe_basis(self.component)
        return None if inner is None else jnp.kron(jnp.eye(self.n_rep), inner)

    def prior(
        self,
        theta: dict[str, Array],
        *,
        constraint: Constraint = "hard",
        soft_constraint_scale: float = 1e-3,
    ) -> gx.GaussianMRF | gx.IntrinsicGMRF:
        identity = gx.GaussianMRF(
            jnp.zeros(self.n_rep), gx.iid_precision(self.n_rep, 1.0)
        )
        return _product(
            identity,
            self.component.prior(theta, constraint=constraint),
            constraint,
            soft_constraint_scale,
        )

Models and inference

An LGM assembles components, fixed effects and an observation model; inla() integrates over the hyperparameters (Rue, Martino & Chopin, 2009).

LGM

Bases: Module

A latent Gaussian model: components, fixed effects and an observation model.

\(\theta\) collects every component's and the likelihood's hyperparameters under f"{name}.{k}". inla() works on the unconstrained vector \(u\), \(\theta = T(u)\) with each prior's bijection, whose log-prior includes the Jacobian; log_posterior_theta is \(\log\tilde\pi(u\mid y)\) in those coordinates.

data holds "y", an optional "offset", the node index of each observation for every component (data[component.name]; an SPDE on a mesh also takes (n, 2) points) and the fixed-effect covariates.

Parameters:

Name	Type	Description	Default
components	tuple[AbstractComponent, ...]	Latent components, with distinct names.	required
fixed	FixedEffects | None	Fixed effects, or None.	None
likelihood	AbstractObservation | None	Observation model.	None

Examples:

>>> import jax.numpy as jnp
>>> import numpy as np
>>> import pyrox_lgm as lgm
>>> model = lgm.LGM(
...     components=(lgm.RW1(5, name="t"),),
...     fixed=lgm.FixedEffects(("intercept",)),
...     likelihood=lgm.Poisson(),
... )
>>> list(model.theta_spec())
['t.tau']
>>> data = {"y": np.array([1, 0, 2, 3, 1]), "t": np.arange(5)}
>>> model.projector(data).as_matrix().shape  # 5 nodes + intercept
(5, 6)

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

class LGM(eqx.Module):
    r"""A latent Gaussian model: components, fixed effects and an observation model.

    $\theta$ collects every component's and the likelihood's hyperparameters
    under ``f"{name}.{k}"``. `inla()` works on the *unconstrained* vector
    $u$, $\theta = T(u)$ with each prior's bijection, whose log-prior
    includes the Jacobian; `log_posterior_theta` is $\log\tilde\pi(u\mid y)$
    in those coordinates.

    ``data`` holds ``"y"``, an optional ``"offset"``, the node index of each
    observation for every component (``data[component.name]``; an
    `SPDE` on a mesh also takes ``(n, 2)`` points) and the fixed-effect
    covariates.

    Args:
        components: Latent components, with distinct names.
        fixed: Fixed effects, or ``None``.
        likelihood: Observation model.

    Examples:
        >>> import jax.numpy as jnp
        >>> import numpy as np
        >>> import pyrox_lgm as lgm
        >>> model = lgm.LGM(
        ...     components=(lgm.RW1(5, name="t"),),
        ...     fixed=lgm.FixedEffects(("intercept",)),
        ...     likelihood=lgm.Poisson(),
        ... )
        >>> list(model.theta_spec())
        ['t.tau']
        >>> data = {"y": np.array([1, 0, 2, 3, 1]), "t": np.arange(5)}
        >>> model.projector(data).as_matrix().shape  # 5 nodes + intercept
        (5, 6)
    """

    components: tuple[AbstractComponent, ...]
    fixed: FixedEffects | None
    likelihood: AbstractObservation
    consts: tuple[Float[Array, ""], ...]

    def __init__(
        self,
        components: tuple[AbstractComponent, ...],
        fixed: FixedEffects | None = None,
        likelihood: AbstractObservation | None = None,
    ) -> None:
        from pyrox_lgm._likelihood import Gaussian

        names = [c.name for c in components]
        if fixed is not None:
            names += list(fixed.names)
        if len(set(names)) != len(names):
            raise ValueError(
                f"component and fixed-effect names must be distinct, got {names}"
            )
        # Each effect reads its data column by name, so it cannot share one
        # with the response or the observation model's inputs.
        reserved = sorted(set(names) & _RESERVED)
        if reserved:
            raise ValueError(
                f"{reserved} name observation data; rename the effect "
                f"(reserved: {sorted(_RESERVED)})"
            )
        self.components = tuple(components)
        self.fixed = fixed
        self.likelihood = Gaussian() if likelihood is None else likelihood
        # Hyperparameters are keyed f"{name}.{k}": the observation model's
        # namespace must not be a component's.
        if self.likelihood.name in {c.name for c in self.components}:
            raise ValueError(
                f"the observation model's name {self.likelihood.name!r} is also "
                "a component's; rename one"
            )
        self.consts = tuple(self._normaliser(c) for c in self.components)

    # -- hyperparameters --------------------------------------------------

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        """Every hyperparameter: ``f"{name}.{k}"`` to (prior, bijection)."""
        spec = {}
        for c in self.components:
            spec |= {f"{c.name}.{k}": v for k, v in c.theta_spec().items()}
        lik = self.likelihood
        spec |= {f"{lik.name}.{k}": v for k, v in lik.theta_spec().items()}
        return spec

    def _sizes(self) -> list[tuple[str, int, tuple[int, ...]]]:
        out = []
        for key, (prior, _) in self.theta_spec().items():
            shape = prior_shape(prior)
            out.append((key, int(np.prod(shape)) if shape else 1, shape))
        return out

    @property
    def n_theta(self) -> int:
        """Length $m$ of the unconstrained vector $u$."""
        return sum(size for _, size, _ in self._sizes())

    def unflatten(self, u: Float[Array, " m"]) -> dict[str, Array]:
        """Constrained $\\theta$ by key, from the unconstrained vector."""
        spec = self.theta_spec()
        theta, i = {}, 0
        for key, size, shape in self._sizes():
            block = u[i : i + size].reshape(shape) if shape else u[i]
            theta[key] = spec[key][1](block)
            i += size
        return theta

    def log_prior(self, u: Float[Array, " m"]) -> Array:
        """$\\log\\pi(u) = \\sum_k \\log\\pi_k(T_k(u_k)) + \\log|T_k'(u_k)|$."""
        spec = self.theta_spec()
        total, i = jnp.zeros(()), 0
        for key, size, shape in self._sizes():
            block = u[i : i + size].reshape(shape) if shape else u[i]
            prior, transform = spec[key]
            value = transform(block)
            total = total + jnp.sum(prior.log_prob(value))
            total = total + jnp.sum(transform.log_abs_det_jacobian(block, value))
            i += size
        return total

    def _component_theta(self, theta: Mapping[str, Array], comp) -> dict[str, Array]:
        return {k: theta[f"{comp.name}.{k}"] for k in comp.theta_spec()}

    def _likelihood_theta(self, theta: Mapping[str, Array]) -> dict[str, Array]:
        lik = self.likelihood
        return {k: theta[f"{lik.name}.{k}"] for k in lik.theta_spec()}

    # -- latent field -----------------------------------------------------

    @property
    def n_fixed(self) -> int:
        return 0 if self.fixed is None else len(self.fixed.names)

    def slices(self) -> dict[str, tuple[int, int]]:
        """``name -> (start, stop)`` in the latent vector; fixed effects last."""
        out, start = {}, 0
        for c in self.components:
            out[c.name] = (start, start + c.n_nodes)
            start += c.n_nodes
        if self.fixed is not None:
            for j, name in enumerate(self.fixed.names):
                out[name] = (start + j, start + j + 1)
        return out

    @property
    def n_latent(self) -> int:
        return sum(c.n_nodes for c in self.components) + self.n_fixed

    def _normaliser(self, comp: AbstractComponent) -> Float[Array, ""]:
        """The theta-free constant an intrinsic component's log_prob omits.

        Computed once, so it must not depend on theta: probed at two points,
        a structure whose pseudo-determinant moves is rejected.
        """
        spec = comp.theta_spec()
        first, second = (
            self._intrinsic_constant(comp.prior(probe_theta(spec, o)))
            for o in (0.0, 0.5)
        )
        if first is None or second is None:
            return jnp.asarray(0.0)
        if not np.isclose(float(first), float(second), rtol=1e-8, atol=1e-8):
            raise ValueError(
                f"component {comp.name!r} returns an IntrinsicGMRF whose structure "
                "depends on its hyperparameters beyond precision_scale; return it "
                "with include_normalizer=True so its normaliser follows theta"
            )
        return first

    @staticmethod
    def _intrinsic_constant(gmrf) -> Float[Array, ""] | None:
        if (
            not isinstance(gmrf, gx.IntrinsicGMRF)
            or gmrf.include_normalizer
            or getattr(gmrf, "normalized", False)
        ):
            return None
        V = gmrf.null_space
        rank = gmrf.structure.in_size() - V.shape[1]
        log_pdet = gx.pseudo_logdet(gmrf.structure, null_space=V)
        return 0.5 * log_pdet - 0.5 * rank * _LOG_2PI

    def latent_prior(
        self, theta: Mapping[str, Array]
    ) -> LatentGMRF | LatentIntrinsicGMRF:
        """The assembled prior of $(x_1, \\dots, x_k, \\beta)$ at ``theta``."""
        parts, blocks, nulls, slices, start = [], [], [], [], 0
        for c in self.components:
            theta_c = self._component_theta(theta, c)
            gmrf = c.prior(theta_c, constraint="hard")
            op, V = c.assembly_precision(theta_c, gmrf), c.constraint_basis(gmrf)
            parts.append(gmrf)
            blocks.append(op)
            nulls.append((start, V))
            slices.append((start, start + c.n_nodes))
            start += c.n_nodes
        lam = 0.0
        if self.fixed is not None:
            lam = self.fixed.prior_precision
            blocks.append(lx.DiagonalLinearOperator(jnp.full(self.n_fixed, lam)))
        Q = block_diagonal(blocks)
        n = Q.in_size()
        constrained = [(s, V) for s, V in nulls if V is not None]
        if not constrained:
            return LatentGMRF(parts, slices, self.consts, self.n_fixed, lam, Q)
        c_total = sum(V.shape[1] for _, V in constrained)
        null = jnp.zeros((n, c_total))
        col = 0
        for s, V in constrained:
            null = null.at[s : s + V.shape[0], col : col + V.shape[1]].set(V)
            col += V.shape[1]
        return LatentIntrinsicGMRF(
            parts, slices, self.consts, self.n_fixed, lam, Q, null
        )

    # -- observations ------------------------------------------------------

    def projector(self, data: Mapping[str, ArrayLike]) -> gx.SparseOperator:
        """``[A_1 | ... | A_k | X]`` for ``data``, built on the host."""
        y = np.asarray(data["y"])
        m = y.shape[0]
        blocks: list[lx.AbstractLinearOperator] = []
        for c in self.components:
            if c.name not in data:
                raise KeyError(f"data has no index for component {c.name!r}")
            idx = np.asarray(data[c.name])
            if idx.ndim == 2:  # locations (any dtype), not node indices
                project = getattr(c, "project_points", None)
                if project is None:
                    raise ValueError(f"{c.name!r} takes node indices, not points")
                blocks.append(project(idx.astype(float)))
            else:
                blocks.append(c.projector(idx))
        if self.fixed is not None:
            X = self.fixed.design(data, m)
            r, cidx = np.divmod(np.arange(X.size), X.shape[1])
            blocks.append(
                gx.SparseOperator.from_coo(
                    r, cidx, jnp.asarray(X.ravel()), (m, X.shape[1])
                )
            )
        return hstack(blocks, m)

    def offset(self, data: Mapping[str, ArrayLike]) -> Array | None:
        return None if "offset" not in data else jnp.asarray(data["offset"])

    def log_posterior_theta(
        self,
        u: Float[Array, " m"],
        data: Mapping[str, ArrayLike],
        *,
        projector: gx.SparseOperator | None = None,
        max_newton: int = 50,
    ) -> Array:
        """The unnormalised log-posterior of the unconstrained hyperparameters.

        $\\log\\tilde\\pi(u \\mid y) = \\log\\tilde\\pi(y\\mid\\theta)
        + \\log\\pi(u)$. Exact gradients flow through gaussx's
        implicit-differentiated Laplace mode. Pass a prebuilt ``projector`` to
        keep host work out of a traced loop.
        """
        fit = self.laplace(u, data, projector=projector, max_newton=max_newton)
        return fit.log_marginal + self.log_prior(u)

    def laplace(
        self,
        u: Float[Array, " m"],
        data: Mapping[str, ArrayLike],
        *,
        projector: gx.SparseOperator | None = None,
        max_newton: int = 50,
    ) -> gx.LaplaceResult:
        """The Laplace approximation of $x \\mid y, \\theta = T(u)$."""
        theta = self.unflatten(u)
        A = self.projector(data) if projector is None else projector
        lik = self.likelihood.build(
            jnp.asarray(data["y"]), self._likelihood_theta(theta), data
        )
        return gx.laplace_mode(
            self.latent_prior(theta),
            lik,
            projector=A,
            offset=self.offset(data),
            max_iter=max_newton,
        )

n_theta: int property

Length \(m\) of the unconstrained vector \(u\).

theta_spec() -> dict[str, tuple[dist.Distribution, Transform]]

Every hyperparameter: f"{name}.{k}" to (prior, bijection).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
    """Every hyperparameter: ``f"{name}.{k}"`` to (prior, bijection)."""
    spec = {}
    for c in self.components:
        spec |= {f"{c.name}.{k}": v for k, v in c.theta_spec().items()}
    lik = self.likelihood
    spec |= {f"{lik.name}.{k}": v for k, v in lik.theta_spec().items()}
    return spec

unflatten(u: Float[Array, ' m']) -> dict[str, Array]

Constrained \(\theta\) by key, from the unconstrained vector.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def unflatten(self, u: Float[Array, " m"]) -> dict[str, Array]:
    """Constrained $\\theta$ by key, from the unconstrained vector."""
    spec = self.theta_spec()
    theta, i = {}, 0
    for key, size, shape in self._sizes():
        block = u[i : i + size].reshape(shape) if shape else u[i]
        theta[key] = spec[key][1](block)
        i += size
    return theta

log_prior(u: Float[Array, ' m']) -> Array

\(\log\pi(u) = \sum_k \log\pi_k(T_k(u_k)) + \log|T_k'(u_k)|\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def log_prior(self, u: Float[Array, " m"]) -> Array:
    """$\\log\\pi(u) = \\sum_k \\log\\pi_k(T_k(u_k)) + \\log|T_k'(u_k)|$."""
    spec = self.theta_spec()
    total, i = jnp.zeros(()), 0
    for key, size, shape in self._sizes():
        block = u[i : i + size].reshape(shape) if shape else u[i]
        prior, transform = spec[key]
        value = transform(block)
        total = total + jnp.sum(prior.log_prob(value))
        total = total + jnp.sum(transform.log_abs_det_jacobian(block, value))
        i += size
    return total

slices() -> dict[str, tuple[int, int]]

name -> (start, stop) in the latent vector; fixed effects last.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def slices(self) -> dict[str, tuple[int, int]]:
    """``name -> (start, stop)`` in the latent vector; fixed effects last."""
    out, start = {}, 0
    for c in self.components:
        out[c.name] = (start, start + c.n_nodes)
        start += c.n_nodes
    if self.fixed is not None:
        for j, name in enumerate(self.fixed.names):
            out[name] = (start + j, start + j + 1)
    return out

latent_prior(theta: Mapping[str, Array]) -> LatentGMRF | LatentIntrinsicGMRF

The assembled prior of \((x_1, \dots, x_k, \beta)\) at theta.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def latent_prior(
    self, theta: Mapping[str, Array]
) -> LatentGMRF | LatentIntrinsicGMRF:
    """The assembled prior of $(x_1, \\dots, x_k, \\beta)$ at ``theta``."""
    parts, blocks, nulls, slices, start = [], [], [], [], 0
    for c in self.components:
        theta_c = self._component_theta(theta, c)
        gmrf = c.prior(theta_c, constraint="hard")
        op, V = c.assembly_precision(theta_c, gmrf), c.constraint_basis(gmrf)
        parts.append(gmrf)
        blocks.append(op)
        nulls.append((start, V))
        slices.append((start, start + c.n_nodes))
        start += c.n_nodes
    lam = 0.0
    if self.fixed is not None:
        lam = self.fixed.prior_precision
        blocks.append(lx.DiagonalLinearOperator(jnp.full(self.n_fixed, lam)))
    Q = block_diagonal(blocks)
    n = Q.in_size()
    constrained = [(s, V) for s, V in nulls if V is not None]
    if not constrained:
        return LatentGMRF(parts, slices, self.consts, self.n_fixed, lam, Q)
    c_total = sum(V.shape[1] for _, V in constrained)
    null = jnp.zeros((n, c_total))
    col = 0
    for s, V in constrained:
        null = null.at[s : s + V.shape[0], col : col + V.shape[1]].set(V)
        col += V.shape[1]
    return LatentIntrinsicGMRF(
        parts, slices, self.consts, self.n_fixed, lam, Q, null
    )

projector(data: Mapping[str, ArrayLike]) -> gx.SparseOperator

[A_1 | ... | A_k | X] for data, built on the host.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def projector(self, data: Mapping[str, ArrayLike]) -> gx.SparseOperator:
    """``[A_1 | ... | A_k | X]`` for ``data``, built on the host."""
    y = np.asarray(data["y"])
    m = y.shape[0]
    blocks: list[lx.AbstractLinearOperator] = []
    for c in self.components:
        if c.name not in data:
            raise KeyError(f"data has no index for component {c.name!r}")
        idx = np.asarray(data[c.name])
        if idx.ndim == 2:  # locations (any dtype), not node indices
            project = getattr(c, "project_points", None)
            if project is None:
                raise ValueError(f"{c.name!r} takes node indices, not points")
            blocks.append(project(idx.astype(float)))
        else:
            blocks.append(c.projector(idx))
    if self.fixed is not None:
        X = self.fixed.design(data, m)
        r, cidx = np.divmod(np.arange(X.size), X.shape[1])
        blocks.append(
            gx.SparseOperator.from_coo(
                r, cidx, jnp.asarray(X.ravel()), (m, X.shape[1])
            )
        )
    return hstack(blocks, m)

log_posterior_theta(u: Float[Array, ' m'], data: Mapping[str, ArrayLike], *, projector: gx.SparseOperator | None = None, max_newton: int = 50) -> Array

The unnormalised log-posterior of the unconstrained hyperparameters.

\(\log\tilde\pi(u \mid y) = \log\tilde\pi(y\mid\theta) + \log\pi(u)\). Exact gradients flow through gaussx's implicit-differentiated Laplace mode. Pass a prebuilt projector to keep host work out of a traced loop.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def log_posterior_theta(
    self,
    u: Float[Array, " m"],
    data: Mapping[str, ArrayLike],
    *,
    projector: gx.SparseOperator | None = None,
    max_newton: int = 50,
) -> Array:
    """The unnormalised log-posterior of the unconstrained hyperparameters.

    $\\log\\tilde\\pi(u \\mid y) = \\log\\tilde\\pi(y\\mid\\theta)
    + \\log\\pi(u)$. Exact gradients flow through gaussx's
    implicit-differentiated Laplace mode. Pass a prebuilt ``projector`` to
    keep host work out of a traced loop.
    """
    fit = self.laplace(u, data, projector=projector, max_newton=max_newton)
    return fit.log_marginal + self.log_prior(u)

laplace(u: Float[Array, ' m'], data: Mapping[str, ArrayLike], *, projector: gx.SparseOperator | None = None, max_newton: int = 50) -> gx.LaplaceResult

The Laplace approximation of \(x \mid y, \theta = T(u)\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def laplace(
    self,
    u: Float[Array, " m"],
    data: Mapping[str, ArrayLike],
    *,
    projector: gx.SparseOperator | None = None,
    max_newton: int = 50,
) -> gx.LaplaceResult:
    """The Laplace approximation of $x \\mid y, \\theta = T(u)$."""
    theta = self.unflatten(u)
    A = self.projector(data) if projector is None else projector
    lik = self.likelihood.build(
        jnp.asarray(data["y"]), self._likelihood_theta(theta), data
    )
    return gx.laplace_mode(
        self.latent_prior(theta),
        lik,
        projector=A,
        offset=self.offset(data),
        max_iter=max_newton,
    )

FixedEffects

Bases: Module

Fixed effects \(\beta \sim \mathcal N(0, \lambda^{-1} I)\), columns of \(X\).

Column "intercept" is all ones unless data has its own; every other name is read from data[name].

Attributes:

Name	Type	Description
names	tuple[str, ...]	Fixed-effect names, in the order they enter the latent vector.
prior_precision	float	\(\lambda\) (R-INLA's default is 0.001).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

class FixedEffects(eqx.Module):
    r"""Fixed effects $\beta \sim \mathcal N(0, \lambda^{-1} I)$, columns of $X$.

    Column ``"intercept"`` is all ones unless ``data`` has its own; every
    other name is read from ``data[name]``.

    Attributes:
        names: Fixed-effect names, in the order they enter the latent vector.
        prior_precision: $\lambda$ (R-INLA's default is 0.001).
    """

    names: tuple[str, ...] = eqx.field(static=True)
    prior_precision: float = eqx.field(static=True, default=1e-3)

    def design(self, data: Mapping[str, ArrayLike], n_obs: int) -> Array:
        """The ``(n_obs, p)`` design matrix.

        Covariates may be traced JAX values (a covariate that depends on a
        parameter sampled outside the LGM, as in an MCMC-INLA hybrid): the
        design is dense, so its pattern does not depend on them.
        """
        cols = []
        for name in self.names:
            if name in data:
                col = jnp.asarray(data[name], dtype=float)
            elif name == "intercept":
                col = jnp.ones(n_obs)
            else:
                raise KeyError(f"fixed effect {name!r} is not in data")
            if col.shape != (n_obs,):
                raise ValueError(f"{name!r} has shape {col.shape}, expected ({n_obs},)")
            cols.append(col)
        return jnp.stack(cols, axis=1)

design(data: Mapping[str, ArrayLike], n_obs: int) -> Array

The (n_obs, p) design matrix.

Covariates may be traced JAX values (a covariate that depends on a parameter sampled outside the LGM, as in an MCMC-INLA hybrid): the design is dense, so its pattern does not depend on them.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_model.py

def design(self, data: Mapping[str, ArrayLike], n_obs: int) -> Array:
    """The ``(n_obs, p)`` design matrix.

    Covariates may be traced JAX values (a covariate that depends on a
    parameter sampled outside the LGM, as in an MCMC-INLA hybrid): the
    design is dense, so its pattern does not depend on them.
    """
    cols = []
    for name in self.names:
        if name in data:
            col = jnp.asarray(data[name], dtype=float)
        elif name == "intercept":
            col = jnp.ones(n_obs)
        else:
            raise KeyError(f"fixed effect {name!r} is not in data")
        if col.shape != (n_obs,):
            raise ValueError(f"{name!r} has shape {col.shape}, expected ({n_obs},)")
        cols.append(col)
    return jnp.stack(cols, axis=1)

inla(model: LGM, data: Mapping[str, ArrayLike], *, strategy: Literal['vb', 'gaussian', 'sla'] = 'vb', integration: Literal['auto', 'eb', 'grid', 'ccd'] = 'auto', key: jax.Array | None = None, max_newton: int = 50, theta_init: ArrayLike | None = None, max_theta_iter: int = 200, theta_tol: float = 1e-05, grid_step: float = 1.0, grid_threshold: float = 2.5, verbose: bool = False) -> INLAResult

Integrated nested Laplace approximation of a latent Gaussian model.

Parameters:

Name	Type	Description	Default
model	LGM	The LGM.	required
data	Mapping[str, ArrayLike]	"y", optional "offset" (and "n_trials" for a binomial model), the node index of each observation per component, and the fixed-effect covariates.	required
strategy	Literal['vb', 'gaussian', 'sla']	"vb" (default, as R-INLA since 22.11) corrects the Gaussian approximation's mean by gaussx.vb_mean_correction in the span of R-INLA's nodes (_vb_subspace: the fixed effects, every node of a component with at most 30, 30 evenly spaced nodes of a larger one); "gaussian" keeps the Laplace mode; "sla" (simplified Laplace) shifts every latent marginal's mean and gives it the skewness of the third-order Laplace expansion, as a skew-normal (_simplified_laplace; it costs one solve per observation and leaves a Gaussian likelihood's marginals unchanged).	'vb'
integration	Literal['auto', 'eb', 'grid', 'ccd']	The design over \(\theta\). "auto" (default) is R-INLA's (int.strategy = "auto", _rinla_design): a fixed 11-point grid for one hyperparameter, a 45-point one for two, a CCD above, each stretched along the Hessian's eigen-axes by the skewness corrections, with R-INLA's weights, early stop and pruning; the latent and predictor summaries are then of each mixture on its mean \(\pm 5\) sds, as R-INLA reports them. This reproduces R-INLA 26.8.7 to about 1e-3 of the latent means (tests/test_rinla_fixtures.py), including its departures from the exact integral over \(\theta\): the weights leave out the stretch's Jacobian, its two-hyperparameter grid scaled by \(f_0\sqrt2\) integrates a Gaussian's variance as 0.87, and the range limit drops the tails that wide design points add. "grid" and "ccd" are gaussx.theta_design's unstretched designs, summarised as plain mixtures ("grid" refined by grid_step and grid_threshold converges to the exact integral); "eb" is the mode alone. Under every design but "eb" the hyperparameter marginals are R-INLA's (_theta_marginals), and the log marginal likelihood is the Gaussian approximation corrected by the design (_log_marginal_likelihood), not R-INLA's integrated estimate, whose unnormalised weights bias it by a constant per dimension.	'auto'
key	Array | None	Unused by the deterministic design; reserved for stochastic strategies. Accepted so call sites stay stable.	None
max_newton	int	Newton iterations per inner Laplace fit; a design point that does not converge is re-run with four times as many, then dropped with a warning.	50
theta_init	ArrayLike | None	Starting unconstrained \(u\) (default zeros: every positive hyperparameter at 1, every interval one at its centre).	None
max_theta_iter	int	L-BFGS iterations for the mode; not reaching theta_tol within them raises.	200
theta_tol	float	Stop when \(\max|\nabla_u| <\) theta_tol.	1e-05
grid_step	float	Step in standardised \(z\) of integration="grid".	1.0
grid_threshold	float	Log-density drop from the mode at which integration="grid" stops.	2.5
verbose	bool	Print the mode search.	False

Returns:

Type	Description
INLAResult	An INLAResult.

Raises:

Type	Description
RuntimeError	If the theta-mode search misses theta_tol within max_theta_iter, or the inner Newton fits along it or at the mode do not converge even at four times max_newton.

Examples:

>>> import jax.numpy as jnp
>>> import numpy as np
>>> import pyrox_lgm as lgm
>>> t = np.arange(30)
>>> y = np.sin(t / 5.0) + 0.1 * np.cos(7.0 * t)  # smooth + wiggle
>>> model = lgm.LGM(
...     components=(lgm.RW2(30, name="trend"),),
...     fixed=lgm.FixedEffects(("intercept",)),
...     likelihood=lgm.Gaussian(),
... )
>>> res = lgm.inla(model, {"y": y, "trend": t})
>>> sorted(res.hyperpar)
['lik.prec', 'trend.tau']
>>> res.random["trend"].mean.shape
(30,)

Source code in packages/pyrox-lgm/src/pyrox_lgm/_inla.py

def inla(
    model: LGM,
    data: Mapping[str, ArrayLike],
    *,
    strategy: Literal["vb", "gaussian", "sla"] = "vb",
    integration: Literal["auto", "eb", "grid", "ccd"] = "auto",
    key: jax.Array | None = None,
    max_newton: int = 50,
    theta_init: ArrayLike | None = None,
    max_theta_iter: int = 200,
    theta_tol: float = 1e-5,
    grid_step: float = 1.0,
    grid_threshold: float = 2.5,
    verbose: bool = False,
) -> INLAResult:
    r"""Integrated nested Laplace approximation of a latent Gaussian model.

    Args:
        model: The `LGM`.
        data: ``"y"``, optional ``"offset"`` (and ``"n_trials"`` for a
            binomial model), the node index of each observation per
            component, and the fixed-effect covariates.
        strategy: ``"vb"`` (default, as R-INLA since 22.11) corrects the
            Gaussian approximation's mean by `gaussx.vb_mean_correction` in
            the span of R-INLA's nodes (`_vb_subspace`: the fixed effects,
            every node of a component with at most 30, 30 evenly spaced
            nodes of a larger one);
            ``"gaussian"`` keeps the Laplace mode; ``"sla"`` (simplified
            Laplace) shifts every latent marginal's mean and gives it the
            skewness of the third-order Laplace expansion, as a skew-normal
            (`_simplified_laplace`; it costs one solve per observation and
            leaves a Gaussian likelihood's marginals unchanged).
        integration: The design over $\theta$. ``"auto"`` (default) is
            R-INLA's (``int.strategy = "auto"``, `_rinla_design`): a fixed
            11-point grid for one hyperparameter, a 45-point one for two, a
            CCD above, each stretched along the Hessian's eigen-axes by the
            skewness corrections, with R-INLA's weights, early stop and
            pruning; the latent and predictor summaries are then of each
            mixture on its mean $\pm 5$ sds, as R-INLA reports them. This
            reproduces R-INLA 26.8.7 to about 1e-3 of the latent means
            (``tests/test_rinla_fixtures.py``), including its departures
            from the exact integral over $\theta$: the weights leave out the
            stretch's Jacobian, its two-hyperparameter grid scaled by
            $f_0\sqrt2$ integrates a Gaussian's variance as 0.87, and the
            range limit drops the tails that wide design points add.
            ``"grid"`` and ``"ccd"`` are `gaussx.theta_design`'s unstretched
            designs, summarised as plain mixtures (``"grid"`` refined by
            ``grid_step`` and ``grid_threshold`` converges to the exact
            integral); ``"eb"`` is the mode alone. Under every design but
            ``"eb"`` the hyperparameter marginals are R-INLA's
            (`_theta_marginals`), and the log marginal likelihood is the
            Gaussian approximation corrected by the design
            (`_log_marginal_likelihood`), not R-INLA's integrated estimate,
            whose unnormalised weights bias it by a constant per
            dimension.
        key: Unused by the deterministic design; reserved for stochastic
            strategies. Accepted so call sites stay stable.
        max_newton: Newton iterations per inner Laplace fit; a design point
            that does not converge is re-run with four times as many, then
            dropped with a warning.
        theta_init: Starting unconstrained $u$ (default zeros: every
            positive hyperparameter at 1, every interval one at its centre).
        max_theta_iter: L-BFGS iterations for the mode; not reaching
            ``theta_tol`` within them raises.
        theta_tol: Stop when $\max|\nabla_u| <$ ``theta_tol``.
        grid_step: Step in standardised $z$ of ``integration="grid"``.
        grid_threshold: Log-density drop from the mode at which
            ``integration="grid"`` stops.
        verbose: Print the mode search.

    Returns:
        An `INLAResult`.

    Raises:
        RuntimeError: If the theta-mode search misses ``theta_tol`` within
            ``max_theta_iter``, or the inner Newton fits along it or at the
            mode do not converge even at four times ``max_newton``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import numpy as np
        >>> import pyrox_lgm as lgm
        >>> t = np.arange(30)
        >>> y = np.sin(t / 5.0) + 0.1 * np.cos(7.0 * t)  # smooth + wiggle
        >>> model = lgm.LGM(
        ...     components=(lgm.RW2(30, name="trend"),),
        ...     fixed=lgm.FixedEffects(("intercept",)),
        ...     likelihood=lgm.Gaussian(),
        ... )
        >>> res = lgm.inla(model, {"y": y, "trend": t})
        >>> sorted(res.hyperpar)
        ['lik.prec', 'trend.tau']
        >>> res.random["trend"].mean.shape
        (30,)
    """

    def run(budget: int) -> tuple[INLAResult | None, bool]:
        return _inla_once(
            model,
            data,
            strategy=strategy,
            integration=integration,
            key=key,
            max_newton=budget,
            theta_init=theta_init,
            max_theta_iter=max_theta_iter,
            theta_tol=theta_tol,
            grid_step=grid_step,
            grid_threshold=grid_threshold,
            verbose=verbose,
        )

    result, mode_ok = run(max_newton)
    if mode_ok and result is not None:
        return result
    budget = 4 * max_newton
    warnings.warn(
        f"the inner Newton fit at a skewness probe, at the theta-mode or along "
        f"its search did not converge in {max_newton} iterations; re-running "
        f"the whole fit with "
        f"max_newton={budget}",
        stacklevel=2,
    )
    result, mode_ok = run(budget)
    if not mode_ok or result is None:
        raise RuntimeError(
            f"the inner Newton fit at a skewness probe, at the theta-mode or along "
            f"its search did not converge in {budget} iterations; raise "
            f"max_newton or check the model"
        )
    return result

INLAResult

Bases: Module

What inla() returns.

Attributes:

Name	Type	Description
fixed	dict[str, Summary]	Fixed effect name to its Summary (scalars).
random	dict[str, Summary]	Component name to the Summary of its whole field (arrays over n_nodes; a BYM2 field is (b, u*)).
hyperpar	dict[str, Summary]	Hyperparameter key (f"{name}.{k}") to its Summary, on the user (constrained) scale.
theta_mode	Float[Array, ' m']	Unconstrained mode \(u^\ast\).
theta_points	Float[Array, 'K m']	Design points (K, m), unconstrained.
theta_weights	Float[Array, ' K']	Their normalised weights (K,).
latent_means	Float[Array, 'K N']	Per-point latent means (K, N) (VB-corrected under strategy="vb").
latent_variances	Float[Array, 'K N']	Per-point marginal variances (K, N).
latent_skewness	Float[Array, 'K N']	Per-point marginal skewness (K, N): zero but under strategy="sla", whose summaries are skew-normal mixtures (sample_latent stays Gaussian, at the shifted means).
predictor_means	Float[Array, 'K M']	Per-point means of the linear predictor \(\eta = Ax + o\) at the observations (K, n_obs).
predictor_variances	Float[Array, 'K M']	Their variances (K, n_obs), from the Takahashi selected inverse (no solves).
linear_predictor	Summary	Summary of the mixed predictor marginals.
log_marginal_likelihood	Float[Array, '']	\(\log\tilde\pi(y)\), Gaussian approximation over \(\theta\) (R-INLA's "Gaussian" mlik).
n_dropped	int	Design points dropped because their inner Newton did not converge.
newton_iters	tuple[int, ...]	The Newton budget each kept point converged with (max_newton, or four times it after a retry).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

class INLAResult(eqx.Module):
    """What `inla()` returns.

    Attributes:
        fixed: Fixed effect name to its `Summary` (scalars).
        random: Component name to the `Summary` of its whole field (arrays
            over ``n_nodes``; a `BYM2` field is ``(b, u*)``).
        hyperpar: Hyperparameter key (``f"{name}.{k}"``) to its `Summary`,
            on the user (constrained) scale.
        theta_mode: Unconstrained mode $u^\\ast$.
        theta_points: Design points ``(K, m)``, unconstrained.
        theta_weights: Their normalised weights ``(K,)``.
        latent_means: Per-point latent means ``(K, N)`` (VB-corrected under
            ``strategy="vb"``).
        latent_variances: Per-point marginal variances ``(K, N)``.
        latent_skewness: Per-point marginal skewness ``(K, N)``: zero but
            under ``strategy="sla"``, whose summaries are skew-normal
            mixtures (`sample_latent` stays Gaussian, at the shifted means).
        predictor_means: Per-point means of the linear predictor
            $\\eta = Ax + o$ at the observations ``(K, n_obs)``.
        predictor_variances: Their variances ``(K, n_obs)``, from the
            Takahashi selected inverse (no solves).
        linear_predictor: `Summary` of the mixed predictor marginals.
        log_marginal_likelihood: $\\log\\tilde\\pi(y)$, Gaussian approximation
            over $\\theta$ (R-INLA's "Gaussian" mlik).
        n_dropped: Design points dropped because their inner Newton did not
            converge.
        newton_iters: The Newton budget each kept point converged with
            (``max_newton``, or four times it after a retry).
    """

    fixed: dict[str, Summary]
    random: dict[str, Summary]
    hyperpar: dict[str, Summary]
    theta_mode: Float[Array, " m"]
    theta_points: Float[Array, "K m"]
    theta_weights: Float[Array, " K"]
    latent_means: Float[Array, "K N"]
    latent_variances: Float[Array, "K N"]
    latent_skewness: Float[Array, "K N"]
    predictor_means: Float[Array, "K M"]
    predictor_variances: Float[Array, "K M"]
    linear_predictor: Summary
    log_marginal_likelihood: Float[Array, ""]
    n_dropped: int = eqx.field(static=True)
    model: Any
    data: Mapping[str, Any]
    newton_iters: tuple[int, ...] = eqx.field(static=True)

    def sample_latent(self, key: jax.Array, n: int) -> Float[Array, "n N"]:
        """Draws of the latent vector from the mixture over design points.

        Each draw picks a design point by its weight, then samples that
        point's Gaussian approximation $\\mathcal N(\\bar x_k, H_k^{-1})$
        (kriged onto the hard constraints) through the sparse Cholesky
        factor of $H_k$.
        """
        k_pick, k_draw = jax.random.split(key)
        picks = np.asarray(
            jax.random.choice(
                k_pick, self.theta_weights.shape[0], (n,), p=self.theta_weights
            )
        )
        out = np.empty((n, self.latent_means.shape[1]))
        keys = jax.random.split(k_draw, len(self.theta_weights))
        A = self.model.projector(self.data)
        for k in np.unique(picks):
            rows = np.flatnonzero(picks == k)
            fit = self.model.laplace(
                self.theta_points[k],
                self.data,
                projector=A,
                max_newton=self.newton_iters[k],
            )
            if not bool(fit.converged):
                raise RuntimeError(f"refit of design point {k} did not converge")
            prior = self.model.latent_prior(self.model.unflatten(self.theta_points[k]))
            z = jax.random.normal(keys[k], (rows.size, fit.mode.shape[0]))
            draws = jax.vmap(fit.factor.solve_lower_transpose)(z)
            if isinstance(prior, gx.IntrinsicGMRF):
                V = prior.null_space
                HV = jax.vmap(fit.factor.solve, in_axes=1, out_axes=1)(V)
                S = V.T @ HV
                draws = draws - (HV @ jnp.linalg.solve(S, V.T @ draws.T)).T
            out[rows] = np.asarray(self.latent_means[k] + draws)
        return jnp.asarray(out)

    def diagnostics(self, order: int = 80):
        """DIC, WAIC, CPO and PIT, without refits (see `pyrox_lgm.diagnostics`)."""
        from pyrox_lgm._diagnostics import diagnostics

        return diagnostics(self, order=order)

    def predict(
        self, new_data: Mapping[str, ArrayLike], key: jax.Array, n_samples: int = 1000
    ) -> Summary:
        """Monte Carlo summary of the linear predictor $\\eta = Ax + o$ at ``new_data``.

        ``new_data`` has the component indices and covariates of the new
        locations (and ``"y"``, any placeholder of the right length, which
        sizes the projector); ``"offset"`` is optional.
        """
        A = self.model.projector(new_data)
        x = self.sample_latent(key, n_samples)
        eta = jax.vmap(A.mv)(x)
        if "offset" in new_data:
            eta = eta + jnp.asarray(new_data["offset"])
        q = jnp.quantile(eta, jnp.asarray(_PROBS), axis=0)
        return Summary(eta.mean(0), eta.std(0), q[0], q[1], q[2])

    def to_xarray(self):
        """The summaries as an ``xarray.Dataset`` (needs the ``xarray`` extra)."""
        try:
            import xarray as xr
        except ImportError as err:  # pragma: no cover - optional extra
            raise ImportError(
                "to_xarray needs `pip install pyrox-lgm[xarray]`"
            ) from err
        variables = {}
        for group in ("fixed", "random", "hyperpar"):
            for name, s in getattr(self, group).items():
                for field, value in s._asdict().items():
                    arr = np.asarray(value)
                    dims = (
                        ()
                        if arr.ndim == 0
                        else tuple(f"{name}_{i}" for i in range(arr.ndim))
                    )
                    variables[f"{group}.{name}.{field}"] = (dims, arr)
        for field, value in self.linear_predictor._asdict().items():
            variables[f"linear_predictor.{field}"] = ("obs", np.asarray(value))
        ds = xr.Dataset(variables)
        ds.attrs["log_marginal_likelihood"] = float(self.log_marginal_likelihood)
        ds.attrs["n_dropped"] = self.n_dropped
        return ds

sample_latent(key: jax.Array, n: int) -> Float[Array, 'n N']

Draws of the latent vector from the mixture over design points.

Each draw picks a design point by its weight, then samples that point's Gaussian approximation \(\mathcal N(\bar x_k, H_k^{-1})\) (kriged onto the hard constraints) through the sparse Cholesky factor of \(H_k\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

def sample_latent(self, key: jax.Array, n: int) -> Float[Array, "n N"]:
    """Draws of the latent vector from the mixture over design points.

    Each draw picks a design point by its weight, then samples that
    point's Gaussian approximation $\\mathcal N(\\bar x_k, H_k^{-1})$
    (kriged onto the hard constraints) through the sparse Cholesky
    factor of $H_k$.
    """
    k_pick, k_draw = jax.random.split(key)
    picks = np.asarray(
        jax.random.choice(
            k_pick, self.theta_weights.shape[0], (n,), p=self.theta_weights
        )
    )
    out = np.empty((n, self.latent_means.shape[1]))
    keys = jax.random.split(k_draw, len(self.theta_weights))
    A = self.model.projector(self.data)
    for k in np.unique(picks):
        rows = np.flatnonzero(picks == k)
        fit = self.model.laplace(
            self.theta_points[k],
            self.data,
            projector=A,
            max_newton=self.newton_iters[k],
        )
        if not bool(fit.converged):
            raise RuntimeError(f"refit of design point {k} did not converge")
        prior = self.model.latent_prior(self.model.unflatten(self.theta_points[k]))
        z = jax.random.normal(keys[k], (rows.size, fit.mode.shape[0]))
        draws = jax.vmap(fit.factor.solve_lower_transpose)(z)
        if isinstance(prior, gx.IntrinsicGMRF):
            V = prior.null_space
            HV = jax.vmap(fit.factor.solve, in_axes=1, out_axes=1)(V)
            S = V.T @ HV
            draws = draws - (HV @ jnp.linalg.solve(S, V.T @ draws.T)).T
        out[rows] = np.asarray(self.latent_means[k] + draws)
    return jnp.asarray(out)

diagnostics(order: int = 80)

DIC, WAIC, CPO and PIT, without refits (see pyrox_lgm.diagnostics).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

def diagnostics(self, order: int = 80):
    """DIC, WAIC, CPO and PIT, without refits (see `pyrox_lgm.diagnostics`)."""
    from pyrox_lgm._diagnostics import diagnostics

    return diagnostics(self, order=order)

predict(new_data: Mapping[str, ArrayLike], key: jax.Array, n_samples: int = 1000) -> Summary

Monte Carlo summary of the linear predictor \(\eta = Ax + o\) at new_data.

new_data has the component indices and covariates of the new locations (and "y", any placeholder of the right length, which sizes the projector); "offset" is optional.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

def predict(
    self, new_data: Mapping[str, ArrayLike], key: jax.Array, n_samples: int = 1000
) -> Summary:
    """Monte Carlo summary of the linear predictor $\\eta = Ax + o$ at ``new_data``.

    ``new_data`` has the component indices and covariates of the new
    locations (and ``"y"``, any placeholder of the right length, which
    sizes the projector); ``"offset"`` is optional.
    """
    A = self.model.projector(new_data)
    x = self.sample_latent(key, n_samples)
    eta = jax.vmap(A.mv)(x)
    if "offset" in new_data:
        eta = eta + jnp.asarray(new_data["offset"])
    q = jnp.quantile(eta, jnp.asarray(_PROBS), axis=0)
    return Summary(eta.mean(0), eta.std(0), q[0], q[1], q[2])

to_xarray()

The summaries as an xarray.Dataset (needs the xarray extra).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

def to_xarray(self):
    """The summaries as an ``xarray.Dataset`` (needs the ``xarray`` extra)."""
    try:
        import xarray as xr
    except ImportError as err:  # pragma: no cover - optional extra
        raise ImportError(
            "to_xarray needs `pip install pyrox-lgm[xarray]`"
        ) from err
    variables = {}
    for group in ("fixed", "random", "hyperpar"):
        for name, s in getattr(self, group).items():
            for field, value in s._asdict().items():
                arr = np.asarray(value)
                dims = (
                    ()
                    if arr.ndim == 0
                    else tuple(f"{name}_{i}" for i in range(arr.ndim))
                )
                variables[f"{group}.{name}.{field}"] = (dims, arr)
    for field, value in self.linear_predictor._asdict().items():
        variables[f"linear_predictor.{field}"] = ("obs", np.asarray(value))
    ds = xr.Dataset(variables)
    ds.attrs["log_marginal_likelihood"] = float(self.log_marginal_likelihood)
    ds.attrs["n_dropped"] = self.n_dropped
    return ds

Summary

Bases: NamedTuple

Posterior summary: mean, sd and the 2.5 / 50 / 97.5 % quantiles.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_result.py

class Summary(NamedTuple):
    """Posterior summary: mean, sd and the 2.5 / 50 / 97.5 % quantiles."""

    mean: Array
    sd: Array
    q025: Array
    q50: Array
    q975: Array

Formula sugar

f(column, model, ...) builds a component named after the data column that indexes it, as R-INLA's f() does.

f(column: str, model: str, *, n: int | None = None, graph: kl.AbstractGraph | None = None, hyper: Mapping[str, dist.Distribution] | None = None, **kwargs: Any) -> AbstractComponent

Build a latent component named after the data column it indexes.

Parameters:

Name	Type	Description	Default
column	str	The data column with each observation's node index; becomes the component's name.	required
model	str	"iid", "rw1", "rw2", "ar1" (sized by n), "besag", "bym2", "car", "leroux" (on graph), "spde" (mesh= / grid= in kwargs) or "generic" (the structure as the first of kwargs' structure=).	required
n	int | None	Number of nodes, for the temporal and unstructured models.	None
graph	AbstractGraph | None	A kernellib graph, for the areal models.	None
hyper	Mapping[str, Distribution] | None	Hyperparameter priors by name ({"tau": ..., "phi": ...}; {"range_sigma": ...} for an SPDE).	None
**kwargs	Any	Passed through to the constructor (scale_model=, cyclic=, alpha=, ...).	{}

Returns:

Type	Description
AbstractComponent	The component.

Raises:

Type	Description
ValueError	For an unknown model, a missing size or graph, or an unknown hyperparameter.

Examples:

>>> import kernellib as kl
>>> import pyrox_lgm as lgm
>>> g = kl.grid_graph((4, 4))
>>> tau = lgm.PCPrecision(1.0, 0.01)
>>> region = lgm.f("region", "bym2", graph=g, hyper={"tau": tau})
>>> region.name, type(region).__name__
('region', 'BYM2')
>>> lgm.f("week", "rw2", n=52, scale_model=True).n_index
52

Source code in packages/pyrox-lgm/src/pyrox_lgm/_formula.py

def f(
    column: str,
    model: str,
    *,
    n: int | None = None,
    graph: kl.AbstractGraph | None = None,
    hyper: Mapping[str, dist.Distribution] | None = None,
    **kwargs: Any,
) -> AbstractComponent:
    """Build a latent component named after the data column it indexes.

    Args:
        column: The data column with each observation's node index; becomes
            the component's ``name``.
        model: ``"iid"``, ``"rw1"``, ``"rw2"``, ``"ar1"`` (sized by ``n``),
            ``"besag"``, ``"bym2"``, ``"car"``, ``"leroux"`` (on ``graph``),
            ``"spde"`` (``mesh=`` / ``grid=`` in ``kwargs``) or ``"generic"``
            (the structure as the first of ``kwargs``' ``structure=``).
        n: Number of nodes, for the temporal and unstructured models.
        graph: A kernellib graph, for the areal models.
        hyper: Hyperparameter priors by name (``{"tau": ..., "phi": ...}``;
            ``{"range_sigma": ...}`` for an SPDE).
        **kwargs: Passed through to the constructor (``scale_model=``,
            ``cyclic=``, ``alpha=``, ...).

    Returns:
        The component.

    Raises:
        ValueError: For an unknown model, a missing size or graph, or an
            unknown hyperparameter.

    Examples:
        >>> import kernellib as kl
        >>> import pyrox_lgm as lgm
        >>> g = kl.grid_graph((4, 4))
        >>> tau = lgm.PCPrecision(1.0, 0.01)
        >>> region = lgm.f("region", "bym2", graph=g, hyper={"tau": tau})
        >>> region.name, type(region).__name__
        ('region', 'BYM2')
        >>> lgm.f("week", "rw2", n=52, scale_model=True).n_index
        52
    """
    if model not in _MODELS:
        raise ValueError(f"unknown model {model!r}; choose from {sorted(_MODELS)}")
    build, sized_by, hyper_kw = _MODELS[model]
    args: list[Any] = []
    if sized_by == "n":
        if n is None:
            raise ValueError(f"model {model!r} needs n=")
        args.append(n)
    elif sized_by == "graph":
        if graph is None:
            raise ValueError(f"model {model!r} needs graph=")
        args.append(graph)
    elif model == "generic":
        if "structure" not in kwargs:
            raise ValueError("model 'generic' needs structure=")
        args.append(kwargs.pop("structure"))
    for key, prior in (hyper or {}).items():
        if key not in hyper_kw:
            raise ValueError(
                f"model {model!r} has no hyperparameter {key!r}; "
                f"it has {sorted(hyper_kw)}"
            )
        kwargs[hyper_kw[key]] = prior
    return build(*args, name=column, **kwargs)

Diagnostics

CPO, PIT, WAIC and DIC from one fit, with no refits.

diagnostics(result: INLAResult, *, order: int = 80) -> Diagnostics

DIC, WAIC, CPO and PIT of an INLAResult, with no refits.

Parameters:

Name	Type	Description	Default
result	INLAResult	What inla() returned.	required
order	int	Gauss–Hermite nodes per site and design point.	80

Returns:

Type	Description
Diagnostics	A Diagnostics.

Examples:

>>> import numpy as np
>>> import pyrox_lgm as lgm
>>> t = np.arange(25)
>>> y = np.sin(t / 4.0) + 0.2 * np.cos(5.0 * t)
>>> model = lgm.LGM((lgm.RW2(25, name="t"),), lgm.FixedEffects(("intercept",)))
>>> d = lgm.inla(model, {"y": y, "t": t}).diagnostics()
>>> d.cpo.shape, bool(((d.pit > 0) & (d.pit < 1)).all())
((25,), True)

Source code in packages/pyrox-lgm/src/pyrox_lgm/_diagnostics.py

def diagnostics(result: INLAResult, *, order: int = 80) -> Diagnostics:
    """DIC, WAIC, CPO and PIT of an `INLAResult`, with no refits.

    Args:
        result: What `inla()` returned.
        order: Gauss–Hermite nodes per site and design point.

    Returns:
        A `Diagnostics`.

    Examples:
        >>> import numpy as np
        >>> import pyrox_lgm as lgm
        >>> t = np.arange(25)
        >>> y = np.sin(t / 4.0) + 0.2 * np.cos(5.0 * t)
        >>> model = lgm.LGM((lgm.RW2(25, name="t"),), lgm.FixedEffects(("intercept",)))
        >>> d = lgm.inla(model, {"y": y, "t": t}).diagnostics()
        >>> d.cpo.shape, bool(((d.pit > 0) & (d.pit < 1)).all())
        ((25,), True)
    """
    model, data = result.model, result.data
    y = jnp.asarray(data["y"])
    nodes, gh = np.polynomial.hermite_e.hermegauss(order)
    log_g = jnp.log(jnp.asarray(gh / gh.sum()))
    z = jnp.asarray(nodes)
    lik = model.likelihood

    def nodes_of(m, v):  # (n, order) quadrature nodes of N(m, v)
        return m[:, None] + jnp.sqrt(v)[:, None] * z[None, :]

    def per_point(u, m, v):
        theta = model._likelihood_theta(model.unflatten(u))
        logp = lik.site_log_prob(y[:, None], nodes_of(m, v), theta, data)
        # Sites are independent given eta: the gradient of the sum is the
        # per-site gradient, and so on for the curvature.
        site = lambda e: jnp.sum(lik.site_log_prob(y, e, theta, data))
        g = jax.grad(site)(m)
        W = -jax.grad(lambda e: jnp.sum(jax.grad(site)(e)))(m)
        a = jnp.maximum(1.0 - W * v, 1e-12)
        c = v / a
        cav_mean = m - g * c
        # Z_ik on the marginal's own nodes, as E[p * cavity / marginal]: the
        # ratio is constant for a Gaussian likelihood and smooth for a
        # log-concave one, where the wider cavity's nodes under-resolve p.
        # Its log, expanded in the standard node z (eta = m + sqrt(v) z), has
        # no division by v, so a deterministic predictor (v = 0) gives p(y|m).
        zz = z[None, :]
        log_ratio = (
            0.5 * zz**2 * (1.0 - a)[:, None]
            - (jnp.sqrt(v) * g)[:, None] * zz
            - (0.5 * g**2 * c)[:, None]
            + 0.5 * jnp.log(a)[:, None]
        )
        log_z = logsumexp(log_g[None, :] + logp + log_ratio, axis=1)
        # F is bounded, so its cavity expectation takes the cavity's nodes.
        cdf = lik.site_cdf(y[:, None], nodes_of(cav_mean, c), theta, data)
        return logp, log_z, jnp.sum(gh / gh.sum() * cdf, axis=1)

    logp, log_z, cav_pit = jax.vmap(per_point)(
        result.theta_points, result.predictor_means, result.predictor_variances
    )  # (K, n, order), (K, n), (K, n)
    log_w = jnp.log(result.theta_weights)[:, None, None] + log_g[None, None, :]

    def expect_log(values):  # log E[exp(values)] over (k, node)
        return logsumexp(log_w + values, axis=(0, 2))

    # Design point k's leave-one-out weight is proportional to w_k / Z_ik.
    log_loo = jnp.log(result.theta_weights)[:, None] - log_z
    log_cpo = -logsumexp(log_loo, axis=0)
    cpo = jnp.exp(log_cpo)
    pit = jnp.exp(log_cpo + logsumexp(log_loo, axis=0, b=cav_pit))
    w = jnp.exp(log_w)
    mean_logp = jnp.sum(w * logp, axis=(0, 2))
    # Centred, so a large common log-density does not cancel (float32).
    var_logp = jnp.sum(w * (logp - mean_logp[None, :, None]) ** 2, axis=(0, 2))
    lppd = expect_log(logp)
    p_waic = jnp.sum(var_logp)
    waic = -2.0 * (jnp.sum(lppd) - p_waic)

    d_bar = -2.0 * jnp.sum(mean_logp)
    theta_hat = model._likelihood_theta(model.unflatten(result.theta_mode))
    eta_bar = result.linear_predictor.mean
    d_hat = -2.0 * jnp.sum(lik.site_log_prob(y, eta_bar, theta_hat, data))
    p_d = d_bar - d_hat
    return Diagnostics(
        cpo=cpo,
        pit=jnp.clip(pit, 0.0, 1.0),
        log_score=-jnp.mean(log_cpo),
        waic=waic,
        p_waic=p_waic,
        dic=d_bar + p_d,
        p_d=p_d,
    )

Diagnostics

Bases: NamedTuple

Leave-one-out and information criteria of an INLAResult.

Attributes:

Name	Type	Description
cpo	Float[Array, ' n']	Conditional predictive ordinates \(\pi(y_i\mid y_{-i})\).
pit	Float[Array, ' n']	Probability integral transforms \(P(Y_i \le y_i \mid y_{-i})\).
log_score	Float[Array, '']	\(-\frac1n\sum_i \log\mathrm{CPO}_i\) (lower is better).
waic	Float[Array, '']	\(-2\sum_i(\mathrm{lppd}_i - p_{\mathrm{WAIC},i})\).
p_waic	Float[Array, '']	Effective number of parameters of WAIC.
dic	Float[Array, '']	\(\bar D + p_D\).
p_d	Float[Array, '']	Effective number of parameters of DIC, \(\bar D - \hat D\).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_diagnostics.py

class Diagnostics(NamedTuple):
    """Leave-one-out and information criteria of an `INLAResult`.

    Attributes:
        cpo: Conditional predictive ordinates $\\pi(y_i\\mid y_{-i})$.
        pit: Probability integral transforms $P(Y_i \\le y_i \\mid y_{-i})$.
        log_score: $-\\frac1n\\sum_i \\log\\mathrm{CPO}_i$ (lower is better).
        waic: $-2\\sum_i(\\mathrm{lppd}_i - p_{\\mathrm{WAIC},i})$.
        p_waic: Effective number of parameters of WAIC.
        dic: $\\bar D + p_D$.
        p_d: Effective number of parameters of DIC, $\\bar D - \\hat D$.
    """

    cpo: Float[Array, " n"]
    pit: Float[Array, " n"]
    log_score: Float[Array, ""]
    waic: Float[Array, ""]
    p_waic: Float[Array, ""]
    dic: Float[Array, ""]
    p_d: Float[Array, ""]

Observation models

Gaussian

Bases: AbstractObservation

y ~ N(eta, 1/prec); prec under PCPrecision(1, 0.01).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_likelihood.py

class Gaussian(AbstractObservation):
    """``y ~ N(eta, 1/prec)``; ``prec`` under ``PCPrecision(1, 0.01)``."""

    name: str = eqx.field(static=True, default="lik")
    prec_prior: dist.Distribution = eqx.field(
        default_factory=lambda: PCPrecision(1.0, 0.01)
    )

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {"prec": (self.prec_prior, default_transform(self.prec_prior))}

    def build(self, y, theta, data):
        return gx.GaussianLikelihood(y, 1.0 / theta["prec"])

    def site_distribution(self, eta, theta, data):
        return dist.Normal(eta, jax.lax.rsqrt(theta["prec"]))

Poisson

Bases: AbstractObservation

y ~ Poisson(exp(eta)); put the log exposure in data["offset"].

Source code in packages/pyrox-lgm/src/pyrox_lgm/_likelihood.py

class Poisson(AbstractObservation):
    """``y ~ Poisson(exp(eta))``; put the log exposure in ``data["offset"]``."""

    name: str = eqx.field(static=True, default="lik")

    def build(self, y, theta, data):
        return gx.PoissonLikelihood(y)

    def site_distribution(self, eta, theta, data):
        return _Poisson(jnp.exp(eta))

Bernoulli

Bases: AbstractObservation

y ~ Bernoulli(sigmoid(eta)), y in {0, 1}.

Source code in packages/pyrox-lgm/src/pyrox_lgm/_likelihood.py

class Bernoulli(AbstractObservation):
    """``y ~ Bernoulli(sigmoid(eta))``, ``y`` in ``{0, 1}``."""

    name: str = eqx.field(static=True, default="lik")

    def build(self, y, theta, data):
        return gx.BernoulliLikelihood(y)

    def site_distribution(self, eta, theta, data):
        return _Binomial(jnp.ones_like(eta), eta)

Binomial

Bases: AbstractObservation

y ~ Binomial(n, sigmoid(eta)) with n = data["n_trials"].

Source code in packages/pyrox-lgm/src/pyrox_lgm/_likelihood.py

class Binomial(AbstractObservation):
    """``y ~ Binomial(n, sigmoid(eta))`` with ``n = data["n_trials"]``."""

    name: str = eqx.field(static=True, default="lik")

    def build(self, y, theta, data):
        return gx.BinomialLikelihood(y, jnp.asarray(data["n_trials"]))

    def site_distribution(self, eta, theta, data):
        return _Binomial(_per_site(data["n_trials"], eta), eta)

NegativeBinomial

Bases: AbstractObservation

y ~ NegBin(mean exp(eta), size); size under Gamma(1, 0.1).

Source code in packages/pyrox-lgm/src/pyrox_lgm/_likelihood.py

class NegativeBinomial(AbstractObservation):
    """``y ~ NegBin(mean exp(eta), size)``; ``size`` under ``Gamma(1, 0.1)``."""

    name: str = eqx.field(static=True, default="lik")
    size_prior: dist.Distribution = eqx.field(
        default_factory=lambda: dist.Gamma(1.0, 0.1)
    )

    def theta_spec(self) -> dict[str, tuple[dist.Distribution, Transform]]:
        return {"size": (self.size_prior, default_transform(self.size_prior))}

    def build(self, y, theta, data):
        return gx.NegativeBinomialLikelihood(y, theta["size"])

    def site_distribution(self, eta, theta, data):
        return _NegativeBinomial(eta, theta["size"])