GMRF Precision Builders

The components of a latent Gaussian model, as precision operators. Each is a quadratic form \(\tfrac{\tau}{2}\|Dx\|^2\) for a sparse difference operator \(D\), so \(Q = \tau D^\top D\) is sparse, and each builder returns the structured operator that sends solve, logdet, diag_inv and sampling to their cheapest exact path:

Builder	Returns	Exact path
iid_precision	lineax.DiagonalLinearOperator	elementwise
rw1_structure, rw2_structure, ar1_precision	BlockTriDiag (SparseOperator when cyclic)	block Cholesky and the block selected inverse, \(O(N d^3)\)
besag_structure, bym2_precision	SparseOperator	sparse Cholesky (dense below AutoSolver's threshold until it lands)
spde_precision (+ fem_matrices)	SparseOperator	sparse Cholesky, as above
spde_precision_grid	SpectralFunction of a KroneckerSum	factor eigenvectors, \(O(\sum_m n_m^3)\) once and \(O(N\sum_m n_m)\) per call

gaussx never builds graphs or meshes: a graph's structure matrix and null space arrive as an operator and an array (kernellib's Graph.laplacian_operator() and graph_null_space), and meshes come from fmesher, pygmsh or meshio.

The maths.

• RW1 / RW2. Increments \(x_{i+1}-x_i\) (second differences \(x_{i+1}-2x_i+x_{i-1}\)) are \(\mathcal N(0,\tau^{-1})\); the structure matrix has null space \(\{\mathbf 1\}\) (\(\{\mathbf 1, t\}\)). RW2 is stored with \(2\times 2\) blocks, so an odd \(n\) gets one decoupled unit-precision padding node; strip it from results. Their normalising constant \(\tfrac12\log|R|_+\) (and Besag's) is pseudo_logdet; the padding node's eigenvalue 1 adds nothing to it.
• AR(1). \(Q = \frac{\tau}{1-\rho^2}\operatorname{tridiag}(-\rho,\ 1+\rho^2,\ -\rho)\) with 1 in the corners: every marginal variance is \(1/\tau\).
• BYM2 (Riebler et al., 2016). The pair \((b, u^*)\) with \(b = (\sqrt{1-\phi}\,v + \sqrt\phi\,u^*)/\sqrt\tau\) has a sparse joint precision whose pattern does not depend on \((\tau, \phi)\). \(u^*\) is scaled by generalized_variance_scale, the geometric mean of the constrained marginal variances (Sørbye & Rue, 2014); it matches R-INLA's inla.scale.model (golden fixture, scripts/golden/inla/).
• SPDE (Lindgren, Rue & Lindström, 2011). \((\kappa^2-\Delta)^{\alpha/2}(\tau x) = \mathcal W\) has Matérn covariance with \(\nu = \alpha - d/2\). With P1 elements, \(K = \kappa^2\tilde C + G\) and \(Q_\alpha = \tau^2 K(\tilde C^{-1}K)^{\alpha-1}\); on a grid with spacing \(h\), \(Q_\alpha = \tau^2 h^d(\kappa^2 I + h^{-2}(L_1\oplus\cdots\oplus L_d))^\alpha\). matern_spde_params converts (range, \(\sigma\), \(\nu\)) into \((\kappa, \tau, \alpha)\).

Boundary effects and domain extension. The SPDE on a bounded domain, mesh or grid, has natural (Neumann) boundary conditions. They inflate the marginal variance within about one practical range of the boundary: up to about \(2\sigma^2\) on an edge and \(4\sigma^2\) in a corner of a raster. Extend the domain by at least one range beyond the region of interest and discard the extension: a larger raster, or a mesh with an outer ring of coarser triangles. Periodic axes of a grid (periodic=(False, True) for longitude on a global raster) and closed surfaces such as a sphere have no boundary. On the matching right-triangle mesh, spde_precision_grid equals spde_precision exactly at nodes at least \(\alpha\) cells from the boundary; nearer the boundary the mesh's lumped mass and half-weight boundary edges differ, which is the same boundary effect.

Example.

import jax.numpy as jnp
import numpy as np

import gaussx as gx

# Temporal: a daily RW2 trend and an AR(1) nuisance
R_trend = gx.rw2_structure(364)  # null space {1, t}
Q_ar = gx.ar1_precision(365, rho=0.8, tau=10.0)
sd_ar = jnp.sqrt(gx.diag_inv(Q_ar))  # block selected inverse, O(N)

# Areal: BYM2 on a graph (here the path 0 - 1 - 2 - 3; a county graph from
# kernellib in practice)
senders, receivers = np.array([1, 2, 3]), np.array([0, 1, 2])
degree = np.bincount(np.r_[senders, receivers], minlength=4).astype(float)
R = gx.besag_structure(
    gx.SparseOperator.from_coo(
        np.r_[np.arange(4), senders],
        np.r_[np.arange(4), receivers],
        jnp.asarray(np.r_[degree, -np.ones(3)]),
        (4, 4),
        symmetric=True,
    )
)
s = gx.generalized_variance_scale(R, jnp.ones(4))
Q_bym2 = gx.bym2_precision(s * R, tau=1.5, phi=0.7)  # sparse (b, u*) stack

# Continuous space: Matérn ν = 1 on a mesh, range 0.5, sd 2
vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [0.5, 0.5]])
triangles = np.array([[0, 1, 4], [1, 2, 4], [2, 3, 4], [3, 0, 4]])
C, G = gx.fem_matrices(vertices, triangles)
kappa, tau, alpha = gx.matern_spde_params(range=0.5, sigma=2.0, nu=1.0, d=2)
Q_spde = gx.spde_precision(C, G, kappa, tau, alpha)
stations = np.array([[0.2, 0.1], [0.7, 0.6]])
A = gx.fem_projector(vertices, triangles, stations)  # 3 non-zeros per row

# ...or on a global raster, with no mesh at all
Q_grid = gx.spde_precision_grid(
    (90, 180), kappa=0.3, tau=1.0, alpha=2, periodic=(False, True)  # wrap longitude
)
sd_grid = jnp.sqrt(gx.diag_inv(Q_grid))  # exact, two small matrix products per axis

Structured linear algebra and Gaussian primitives for JAX.

iid_precision(n: int, tau: Float[ArrayLike, '']) -> lx.DiagonalLinearOperator

Precision τ I of n independent effects with variance 1/τ.

Parameters:

Name	Type	Description	Default
n	int	Number of effects.	required
tau	Float[ArrayLike, '']	Precision τ > 0 (may be traced).	required

Returns:

Type	Description
DiagonalLinearOperator	A lineax.DiagonalLinearOperator of size n.

Examples:

import gaussx

Q = gaussx.iid_precision(5, tau=2.0)
Q.as_matrix()  # 2 I

Source code in src/gaussx/_gmrf/_temporal.py

def iid_precision(n: int, tau: Float[ArrayLike, ""]) -> lx.DiagonalLinearOperator:
    r"""Precision ``τ I`` of ``n`` independent effects with variance ``1/τ``.

    Args:
        n: Number of effects.
        tau: Precision ``τ > 0`` (may be traced).

    Returns:
        A `lineax.DiagonalLinearOperator` of size ``n``.

    Examples:
        ```python
        import gaussx

        Q = gaussx.iid_precision(5, tau=2.0)
        Q.as_matrix()  # 2 I
        ```
    """
    tau = _as_float(tau)
    return lx.DiagonalLinearOperator(jnp.full((n,), tau, dtype=tau.dtype))

rw1_structure(n: int, *, spacing: Float[ArrayLike, ' n-1'] | float | None = None, cyclic: bool = False) -> BlockTriDiag | SparseOperator

Structure matrix R = D₁ᵀ W D₁ of a first-order random walk.

The increments x_{i+1} − x_i ~ N(0, h_i/τ) give the precision τ R with R = D₁ᵀ diag(1/h) D₁: the weighted path-graph Laplacian with edge weights 1/h_i (unit weights for regular spacing). R is singular with null space span{1}.

Parameters:

Name	Type	Description	Default
n	int	Number of nodes.	required
spacing	Float[ArrayLike, ' n-1'] | float | None	Gaps h_i between consecutive nodes, shape (n − 1,) ((n,) with cyclic=True, the last one closing the cycle), or a scalar. None means unit spacing.	None
cyclic	bool	Join node n − 1 back to node 0 (a cycle-graph Laplacian, e.g. a seasonal effect). The corner entries leave the tridiagonal band, so the result is a SparseOperator.	False

Returns:

Type	Description
BlockTriDiag | SparseOperator	A positive-semidefinite BlockTriDiag with 1 × 1 blocks, or a
BlockTriDiag | SparseOperator	symmetric SparseOperator with cyclic=True.

Examples:

import jax.numpy as jnp
import gaussx

R = gaussx.rw1_structure(5)
R.mv(jnp.ones(5))  # zeros: constants are in the null space

Source code in src/gaussx/_gmrf/_temporal.py

def rw1_structure(
    n: int,
    *,
    spacing: Float[ArrayLike, " n-1"] | float | None = None,
    cyclic: bool = False,
) -> BlockTriDiag | SparseOperator:
    r"""Structure matrix ``R = D₁ᵀ W D₁`` of a first-order random walk.

    The increments ``x_{i+1} − x_i ~ N(0, h_i/τ)`` give the precision
    ``τ R`` with ``R = D₁ᵀ diag(1/h) D₁``: the weighted path-graph Laplacian
    with edge weights ``1/h_i`` (unit weights for regular spacing). ``R`` is
    singular with null space ``span{1}``.

    Args:
        n: Number of nodes.
        spacing: Gaps ``h_i`` between consecutive nodes, shape ``(n − 1,)``
            (``(n,)`` with ``cyclic=True``, the last one closing the cycle),
            or a scalar. ``None`` means unit spacing.
        cyclic: Join node ``n − 1`` back to node ``0`` (a cycle-graph
            Laplacian, e.g. a seasonal effect). The corner entries leave
            the tridiagonal band, so the result is a `SparseOperator`.

    Returns:
        A positive-semidefinite `BlockTriDiag` with ``1 × 1`` blocks, or a
        symmetric `SparseOperator` with ``cyclic=True``.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx

        R = gaussx.rw1_structure(5)
        R.mv(jnp.ones(5))  # zeros: constants are in the null space
        ```
    """
    n_edges = n if cyclic else n - 1
    if spacing is None:
        weights = jnp.ones(n_edges, dtype=jnp.result_type(float))
    else:
        h = _as_float(spacing)
        weights = jnp.broadcast_to(1.0 / h, (n_edges,))
    if cyclic:
        return _cyclic_structure(n, [weights], [1])
    degree = jnp.zeros(n, dtype=weights.dtype).at[1:].add(weights).at[:-1].add(weights)
    return BlockTriDiag(
        rearrange(degree, "(n a b) -> n a b", a=1, b=1),
        rearrange(-weights, "(n a b) -> n a b", a=1, b=1),
        tags=_PSD,
    )

rw2_structure(n: int, *, cyclic: bool = False) -> BlockTriDiag | SparseOperator

Structure matrix R = D₂ᵀ D₂ of a second-order random walk.

The second differences x_{i+1} − 2x_i + x_{i−1} ~ N(0, 1/τ) give the pentadiagonal precision τ R, the discrete cubic smoothing spline. Its null space is span{1, t} (span{1} for cyclic=True).

The band is stored as a BlockTriDiag with 2 × 2 blocks, which needs an even size: for odd n the operator has n + 1 rows, the last one a decoupled node with unit precision. That node does not interact with the others, so strip it from results (x[:n], diag_inv(R)[:n]); it adds nothing to log|R| and log 1 = 0 otherwise (log τ once R is scaled by τ).

Parameters:

Name	Type	Description	Default
n	int	Number of nodes (at least 3).	required
cyclic	bool	Wrap the second differences around (a seasonal RW2). The result is then a SparseOperator of size exactly n.	False

Returns:

Type	Description
BlockTriDiag | SparseOperator	A positive-semidefinite BlockTriDiag of size 2⌈n/2⌉, or a
BlockTriDiag | SparseOperator	symmetric SparseOperator with cyclic=True.

Raises:

Type	Description
ValueError	If n < 3.

Examples:

import jax.numpy as jnp
import gaussx

R = gaussx.rw2_structure(6)
t = jnp.arange(6.0)
R.mv(t)  # zeros: linear trends are in the null space

Source code in src/gaussx/_gmrf/_temporal.py

def rw2_structure(n: int, *, cyclic: bool = False) -> BlockTriDiag | SparseOperator:
    r"""Structure matrix ``R = D₂ᵀ D₂`` of a second-order random walk.

    The second differences ``x_{i+1} − 2x_i + x_{i−1} ~ N(0, 1/τ)`` give the
    pentadiagonal precision ``τ R``, the discrete cubic smoothing spline. Its
    null space is ``span{1, t}`` (``span{1}`` for ``cyclic=True``).

    The band is stored as a `BlockTriDiag` with ``2 × 2`` blocks, which needs
    an even size: for **odd** ``n`` the operator has ``n + 1`` rows, the last
    one a decoupled node with unit precision. That node does not interact
    with the others, so strip it from results (``x[:n]``,
    ``diag_inv(R)[:n]``); it adds nothing to ``log|R|`` and ``log 1 = 0``
    otherwise (``log τ`` once ``R`` is scaled by ``τ``).

    Args:
        n: Number of nodes (at least 3).
        cyclic: Wrap the second differences around (a seasonal RW2). The
            result is then a `SparseOperator` of size exactly ``n``.

    Returns:
        A positive-semidefinite `BlockTriDiag` of size ``2⌈n/2⌉``, or a
        symmetric `SparseOperator` with ``cyclic=True``.

    Raises:
        ValueError: If ``n < 3``.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx

        R = gaussx.rw2_structure(6)
        t = jnp.arange(6.0)
        R.mv(t)  # zeros: linear trends are in the null space
        ```
    """
    if n < 3:
        raise ValueError(f"rw2_structure needs n >= 3, got {n}.")
    dtype = jnp.result_type(float)
    if cyclic:
        ones = jnp.ones(n, dtype=dtype)
        return _cyclic_structure(n, [6.0 * ones, -4.0 * ones, ones], [0, 1, 2])
    # D₂ᵀD₂ for the open chain, padded to an even size with a decoupled node.
    size = n + (n % 2)
    D = np.zeros((n - 2, size))
    rows = np.arange(n - 2)
    D[rows, rows], D[rows, rows + 1], D[rows, rows + 2] = 1.0, -2.0, 1.0
    D = jnp.asarray(D, dtype=dtype)
    R = einsum(D, D, "k i, k j -> i j")
    if size != n:
        R = R.at[n, n].set(1.0)
    blocks = rearrange(R, "(N a) (M b) -> N M a b", a=2, b=2)
    num = size // 2
    diagonal = blocks[jnp.arange(num), jnp.arange(num)]
    sub = blocks[jnp.arange(1, num), jnp.arange(num - 1)]
    return BlockTriDiag(diagonal, sub, tags=_PSD)

ar1_precision(n: int, rho: Float[ArrayLike, ''], tau: Float[ArrayLike, '']) -> BlockTriDiag

Precision of a stationary AR(1) process with marginal precision τ.

x_t = rho x_{t−1} + ε_t, started from its stationary law, has

\[ Q = \frac{\tau}{1-\rho^2}\operatorname{tridiag}(-\rho,\ 1+\rho^2,\ -\rho) \]

with 1 in the two corners, so every marginal variance is 1/τ and the innovation precision is τ/(1 − rho²).

Parameters:

Name	Type	Description	Default
n	int	Number of time points (at least 2).	required
rho	Float[ArrayLike, '']	Lag-one correlation, |rho| < 1 (may be traced).	required
tau	Float[ArrayLike, '']	Marginal precision τ > 0 (may be traced).	required

Returns:

Type	Description
BlockTriDiag	A positive-definite BlockTriDiag with 1 × 1 blocks.

Examples:

import jax.numpy as jnp
import gaussx

Q = gaussx.ar1_precision(50, rho=0.8, tau=10.0)
gaussx.diag_inv(Q)  # all 0.1: the marginal variance 1/τ

Source code in src/gaussx/_gmrf/_temporal.py

def ar1_precision(
    n: int, rho: Float[ArrayLike, ""], tau: Float[ArrayLike, ""]
) -> BlockTriDiag:
    r"""Precision of a stationary AR(1) process with marginal precision ``τ``.

    ``x_t = rho x_{t−1} + ε_t``, started from its stationary law, has

    $$
    Q = \frac{\tau}{1-\rho^2}\operatorname{tridiag}(-\rho,\ 1+\rho^2,\ -\rho)
    $$

    with ``1`` in the two corners, so every marginal variance is ``1/τ``
    and the innovation precision is ``τ/(1 − rho²)``.

    Args:
        n: Number of time points (at least 2).
        rho: Lag-one correlation, ``|rho| < 1`` (may be traced).
        tau: Marginal precision ``τ > 0`` (may be traced).

    Returns:
        A positive-definite `BlockTriDiag` with ``1 × 1`` blocks.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx

        Q = gaussx.ar1_precision(50, rho=0.8, tau=10.0)
        gaussx.diag_inv(Q)  # all 0.1: the marginal variance 1/τ
        ```
    """
    if n < 2:
        raise ValueError(f"ar1_precision needs n >= 2, got {n}.")
    rho = _as_float(rho)
    tau = _as_float(tau)
    dtype = jnp.result_type(rho, tau)
    scale = tau / (1.0 - rho**2)
    main = jnp.full(n, 1.0 + rho**2, dtype=dtype).at[0].set(1.0).at[-1].set(1.0)
    off = jnp.full(n - 1, -rho, dtype=dtype)
    return BlockTriDiag(
        rearrange(scale * main, "(n a b) -> n a b", a=1, b=1),
        rearrange(scale * off, "(n a b) -> n a b", a=1, b=1),
        tags=_PSD,
    )

besag_structure(laplacian_op: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

Validate a graph Laplacian as a Besag (ICAR) structure matrix.

Checks that the operator is square and symmetric and, when its values are concrete, that its rows sum to zero (constants are in its null space), then tags it positive semidefinite. A SparseOperator or BlockTriDiag keeps its type (and so its sparse dispatch); any other operator is wrapped in a lineax.TaggedLinearOperator.

Parameters:

Name	Type	Description	Default
laplacian_op	AbstractLinearOperator	The weighted graph Laplacian R = D − W.	required

Returns:

Type	Description
AbstractLinearOperator	The same matrix, tagged symmetric and positive semidefinite.

Raises:

Type	Description
ValueError	If it is not square or symmetric, or a row sum is not zero.

Examples:

import jax.numpy as jnp
import numpy as np
import gaussx

# Path graph 0 - 1 - 2, each edge once
R = gaussx.SparseOperator.from_coo(
    np.array([0, 1, 2, 1, 2]),
    np.array([0, 1, 2, 0, 1]),
    jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
    (3, 3),
    symmetric=True,
)
R = gaussx.besag_structure(R)  # now tagged positive semidefinite

Source code in src/gaussx/_gmrf/_areal.py

def besag_structure(
    laplacian_op: lx.AbstractLinearOperator,
) -> lx.AbstractLinearOperator:
    r"""Validate a graph Laplacian as a Besag (ICAR) structure matrix.

    Checks that the operator is square and symmetric and, when its values
    are concrete, that its rows sum to zero (constants are in its null
    space), then tags it positive semidefinite. A `SparseOperator` or
    `BlockTriDiag` keeps its type (and so its sparse dispatch); any other
    operator is wrapped in a `lineax.TaggedLinearOperator`.

    Args:
        laplacian_op: The weighted graph Laplacian ``R = D − W``.

    Returns:
        The same matrix, tagged symmetric and positive semidefinite.

    Raises:
        ValueError: If it is not square or symmetric, or a row sum is not
            zero.

    Examples:
        ```python
        import jax.numpy as jnp
        import numpy as np
        import gaussx

        # Path graph 0 - 1 - 2, each edge once
        R = gaussx.SparseOperator.from_coo(
            np.array([0, 1, 2, 1, 2]),
            np.array([0, 1, 2, 0, 1]),
            jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
            (3, 3),
            symmetric=True,
        )
        R = gaussx.besag_structure(R)  # now tagged positive semidefinite
        ```
    """
    n = laplacian_op.in_size()
    if laplacian_op.out_size() != n:
        raise ValueError(
            f"A structure matrix must be square, got ({laplacian_op.out_size()}, {n})."
        )
    if not lx.is_symmetric(laplacian_op):
        raise ValueError(
            "A structure matrix must be symmetric (tag it lx.symmetric_tag)."
        )
    row_sums = laplacian_op.mv(jnp.ones(n, dtype=laplacian_op.in_structure().dtype))
    try:
        row_sums = np.asarray(row_sums)
        scale = float(np.max(np.abs(np.asarray(diag(laplacian_op)))))
    except jax.errors.TracerArrayConversionError:
        pass  # traced values: the check is skipped
    else:
        if np.max(np.abs(row_sums), initial=0.0) > _ROW_SUM_TOLERANCE * max(scale, 1.0):
            raise ValueError(
                "A Besag structure matrix is a graph Laplacian, whose rows sum "
                f"to zero; the largest row sum is {np.max(np.abs(row_sums)):.3e}."
            )
    if isinstance(laplacian_op, SparseOperator):
        return SparseOperator(
            laplacian_op.values,
            laplacian_op.pattern,
            tags=laplacian_op.tags | _PSD | {lx.symmetric_tag},
        )
    if isinstance(laplacian_op, BlockTriDiag):
        return BlockTriDiag(
            laplacian_op.diagonal,
            laplacian_op.sub_diagonal,
            symmetric=laplacian_op.symmetric,
            tags=laplacian_op.tags | _PSD,
        )
    return lx.TaggedLinearOperator(laplacian_op, _PSD | {lx.symmetric_tag})

generalized_variance_scale(structure: lx.AbstractLinearOperator, null_space: Float[ArrayLike, 'n k'] | Float[ArrayLike, ' n'], *, eps: float | None = None) -> Float[Array, '']

Generalized variance of an intrinsic GMRF (Sørbye & Rue, 2014).

The geometric mean of the marginal variances under the constraint Vᵀx = 0 (V the null space),

\[ s = \exp\Big(\frac1n\sum_i\log\Sigma_{ii}\Big),\qquad \Sigma = R^{+}\ \text{on}\ \operatorname{range}(R), \]

so that s · R has generalized variance one: the scaling that makes a precision τ mean the same thing for every graph (BYM2, scaled RW1 / RW2).

Two exact paths:

• Eigen-structured R (a KroneckerSum grid Laplacian, a DiagonalisedOperator or a SpectralFunction): the diagonal of the pseudo-inverse from the factor eigenvectors (gaussx.diag_inv with pinv=True). This assumes null_space spans exactly the zero eigenspace.
• Anything else (a SparseOperator graph Laplacian, a BlockTriDiag random walk): as R-INLA's inla.scale.model, the marginal variances of R + εI from gaussx.diag_inv (the block selected inverse for a BlockTriDiag; sparse Cholesky / Takahashi for a SparseOperator as that dispatch lands), then the kriging correction for the constraint, Σ = S − S V (Vᵀ S V)⁻¹ Vᵀ S with S = (R + εI)⁻¹, which needs k solves.

For a disconnected graph scale each connected component separately. structure may have one more row than null_space: the decoupled padding node of an odd-size gaussx.rw2_structure, which is then left out.

Parameters:

Name	Type	Description	Default
structure	AbstractLinearOperator	The structure matrix R (symmetric PSD).	required
null_space	Float[ArrayLike, 'n k'] | Float[ArrayLike, ' n']	Basis of its null space, shape (n, k) or (n,) (e.g. the constants for a connected graph).	required
eps	float | None	Ridge ε. Defaults to √(machine eps) · max diag(R), the value R-INLA uses.	None

Returns:

Type	Description
Float[Array, '']	The scalar s.

Raises:

Type	Description
ValueError	If the sizes disagree.

Examples:

import jax.numpy as jnp
import gaussx

R = gaussx.rw1_structure(20)
s = gaussx.generalized_variance_scale(R, jnp.ones(20))
# R_scaled = s * R has generalized variance 1

Source code in src/gaussx/_gmrf/_areal.py

def generalized_variance_scale(
    structure: lx.AbstractLinearOperator,
    null_space: Float[ArrayLike, "n k"] | Float[ArrayLike, " n"],
    *,
    eps: float | None = None,
) -> Float[Array, ""]:
    r"""Generalized variance of an intrinsic GMRF (Sørbye & Rue, 2014).

    The geometric mean of the marginal variances under the constraint
    ``Vᵀx = 0`` (``V`` the null space),

    $$
    s = \exp\Big(\frac1n\sum_i\log\Sigma_{ii}\Big),\qquad
    \Sigma = R^{+}\ \text{on}\ \operatorname{range}(R),
    $$

    so that ``s · R`` has generalized variance one: the scaling that makes a
    precision ``τ`` mean the same thing for every graph (BYM2, scaled
    RW1 / RW2).

    Two exact paths:

    - **Eigen-structured** ``R`` (a `KroneckerSum` grid Laplacian, a
      `DiagonalisedOperator` or a `SpectralFunction`): the diagonal of the
      pseudo-inverse from the factor eigenvectors (`gaussx.diag_inv` with
      ``pinv=True``). This assumes ``null_space`` spans exactly the zero
      eigenspace.
    - **Anything else** (a `SparseOperator` graph Laplacian, a `BlockTriDiag`
      random walk): as R-INLA's ``inla.scale.model``, the marginal variances
      of ``R + εI`` from `gaussx.diag_inv` (the block selected inverse for a
      `BlockTriDiag`; sparse Cholesky / Takahashi for a `SparseOperator` as
      that dispatch lands), then the kriging correction for the constraint,
      ``Σ = S − S V (Vᵀ S V)⁻¹ Vᵀ S`` with ``S = (R + εI)⁻¹``, which needs
      ``k`` solves.

    For a disconnected graph scale each connected component separately.
    ``structure`` may have one more row than ``null_space``: the decoupled
    padding node of an odd-size `gaussx.rw2_structure`, which is then
    left out.

    Args:
        structure: The structure matrix ``R`` (symmetric PSD).
        null_space: Basis of its null space, shape ``(n, k)`` or ``(n,)``
            (e.g. the constants for a connected graph).
        eps: Ridge ``ε``. Defaults to ``√(machine eps) · max diag(R)``, the
            value R-INLA uses.

    Returns:
        The scalar ``s``.

    Raises:
        ValueError: If the sizes disagree.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx

        R = gaussx.rw1_structure(20)
        s = gaussx.generalized_variance_scale(R, jnp.ones(20))
        # R_scaled = s * R has generalized variance 1
        ```
    """
    V = _as_float(null_space)
    if V.ndim == 1:
        V = rearrange(V, "n -> n 1")
    n = V.shape[0]
    size = structure.in_size()
    if size not in (n, n + 1):
        raise ValueError(
            f"null_space has {n} rows but the structure matrix has size {size}."
        )
    if _has_eigenbasis(structure) and size == n:
        variances = diag_inv(structure, pinv=True)
    else:
        if size == n + 1:
            V = jnp.concatenate([V, jnp.zeros((1, V.shape[1]), dtype=V.dtype)])
        dtype = jnp.result_type(structure.in_structure().dtype, V.dtype)
        if eps is None:
            eps_value = jnp.sqrt(jnp.finfo(dtype).eps) * jnp.max(diag(structure))
        else:
            eps_value = jnp.asarray(eps, dtype=dtype)
        S = _add_ridge(structure, eps_value)
        # diag(S⁻¹) and the correction are both O(1/ε) and cancel to O(1), so
        # they must come from the same factorisation for the rounding to cancel.
        if isinstance(S, SparseOperator):
            factor = sparse_cholesky(S)
            solve_S, diag_inv_S = factor.solve, factor.diag_inv
        else:
            solve_S, diag_inv_S = ft.partial(solve, S), ft.partial(diag_inv, S)
        W = jax.vmap(solve_S, in_axes=1, out_axes=1)(V)
        M = einsum(V, W, "i a, i b -> a b")
        correction = einsum(W, jnp.linalg.inv(M), W, "i a, a b, i b -> i")
        variances = diag_inv_S() - correction
    return jnp.exp(jnp.mean(jnp.log(variances[:n])))

bym2_precision(structure_scaled: lx.AbstractLinearOperator, tau: Float[ArrayLike, ''], phi: Float[ArrayLike, '']) -> SparseOperator

Joint precision of the BYM2 pair (b, u*) (Riebler et al., 2016).

b = (√(1−φ) v + √φ u*)/√τ with v ~ N(0, I) and u* the scaled ICAR field (structure R*) gives

\[ Q = \begin{pmatrix} \frac{\tau}{1-\phi}I & -\frac{\sqrt{\tau\phi}}{1-\phi}I\\ -\frac{\sqrt{\tau\phi}}{1-\phi}I & R^* + \frac{\phi}{1-\phi}I \end{pmatrix}, \]

and the marginal covariance of b is ((1−φ)I + φ R*⁺)/τ under u*'s sum-to-zero constraint. The pattern (R*'s, shifted, plus two diagonals) is built on the host once per pattern of R* and does not depend on (τ, φ).

Parameters:

Name	Type	Description	Default
structure_scaled	AbstractLinearOperator	The scaled structure R* = s R (see generalized_variance_scale), a SparseOperator or a scalar multiple of one.	required
tau	Float[ArrayLike, '']	Precision τ > 0 of b (may be traced).	required
phi	Float[ArrayLike, '']	Mixing 0 ≤ φ < 1, the spatial share of the variance (may be traced).	required

Returns:

Type	Description
SparseOperator	A symmetric positive-semidefinite SparseOperator of size 2n
SparseOperator	acting on the stacked vector (b, u*); singular along u*'s
SparseOperator	null space, which the sum-to-zero constraint removes.

Raises:

Type	Description
TypeError	If structure_scaled is not (a multiple of) a SparseOperator.

Examples:

import jax.numpy as jnp
import numpy as np
import gaussx

R = gaussx.SparseOperator.from_coo(
    np.array([0, 1, 2, 1, 2]),
    np.array([0, 1, 2, 0, 1]),
    jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
    (3, 3),
    symmetric=True,
)
s = gaussx.generalized_variance_scale(R, jnp.ones(3))
Q = gaussx.bym2_precision(s * R, tau=1.5, phi=0.7)  # (6, 6)

Source code in src/gaussx/_gmrf/_areal.py

def bym2_precision(
    structure_scaled: lx.AbstractLinearOperator,
    tau: Float[ArrayLike, ""],
    phi: Float[ArrayLike, ""],
) -> SparseOperator:
    r"""Joint precision of the BYM2 pair ``(b, u*)`` (Riebler et al., 2016).

    ``b = (√(1−φ) v + √φ u*)/√τ`` with ``v ~ N(0, I)`` and ``u*`` the scaled
    ICAR field (structure ``R*``) gives

    $$
    Q = \begin{pmatrix}
    \frac{\tau}{1-\phi}I & -\frac{\sqrt{\tau\phi}}{1-\phi}I\\
    -\frac{\sqrt{\tau\phi}}{1-\phi}I & R^* + \frac{\phi}{1-\phi}I
    \end{pmatrix},
    $$

    and the marginal covariance of ``b`` is
    ``((1−φ)I + φ R*⁺)/τ`` under ``u*``'s sum-to-zero constraint. The
    pattern (``R*``'s, shifted, plus two diagonals) is built on the host
    once per pattern of ``R*`` and does not depend on ``(τ, φ)``.

    Args:
        structure_scaled: The scaled structure ``R* = s R`` (see
            `generalized_variance_scale`), a `SparseOperator` or a scalar
            multiple of one.
        tau: Precision ``τ > 0`` of ``b`` (may be traced).
        phi: Mixing ``0 ≤ φ < 1``, the spatial share of the variance (may be
            traced).

    Returns:
        A symmetric positive-semidefinite `SparseOperator` of size ``2n``
        acting on the stacked vector ``(b, u*)``; singular along ``u*``'s
        null space, which the sum-to-zero constraint removes.

    Raises:
        TypeError: If ``structure_scaled`` is not (a multiple of) a
            `SparseOperator`.

    Examples:
        ```python
        import jax.numpy as jnp
        import numpy as np
        import gaussx

        R = gaussx.SparseOperator.from_coo(
            np.array([0, 1, 2, 1, 2]),
            np.array([0, 1, 2, 0, 1]),
            jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
            (3, 3),
            symmetric=True,
        )
        s = gaussx.generalized_variance_scale(R, jnp.ones(3))
        Q = gaussx.bym2_precision(s * R, tau=1.5, phi=0.7)  # (6, 6)
        ```
    """
    values, pattern = _as_sparse(structure_scaled)
    tau = _as_float(tau)
    phi = _as_float(phi)
    n = pattern.shape[0]
    out_pattern, inverse = _bym2_plan(pattern)
    dtype = jnp.result_type(values, tau, phi)
    ones = jnp.ones(n, dtype=dtype)
    raw = jnp.concatenate(
        [
            tau / (1.0 - phi) * ones,
            -jnp.sqrt(tau * phi) / (1.0 - phi) * ones,
            values.astype(dtype),
            phi / (1.0 - phi) * ones,
        ]
    )
    out = jax.ops.segment_sum(raw, inverse, num_segments=out_pattern.nnz)
    return SparseOperator(out, out_pattern, tags=_PSD)

spde_precision(C_lumped: lx.DiagonalLinearOperator, G: SparseOperator, kappa: Float[ArrayLike, ''], tau: Float[ArrayLike, ''], alpha: int) -> SparseOperator

SPDE precision Q_α = τ² K (C̃⁻¹K)^{α−1} with K = κ²C̃ + G.

The pattern of Q_α (the (α−1)-ring neighbourhood of G's) and the index triples of each sparse product are computed once on the host per pattern and cached; only the values depend on κ and τ, so jit, grad and vmap over them never redo the symbolic work. The result goes to sparse Cholesky through gaussx.solve, gaussx.logdet and gaussx.diag_inv as those dispatch for a SparseOperator.

Parameters:

Name	Type	Description	Default
C_lumped	DiagonalLinearOperator	Lumped (diagonal) mass matrix C̃, from gaussx.fem_matrices.	required
G	SparseOperator	Stiffness matrix, from gaussx.fem_matrices.	required
kappa	Float[ArrayLike, '']	Inverse range parameter κ > 0 (may be traced).	required
tau	Float[ArrayLike, '']	Scale τ > 0 (may be traced).	required
alpha	int	Integer order α ≥ 1; nu = α − d/2.	required

Returns:

Type	Description
SparseOperator	A symmetric positive-definite SparseOperator.

Raises:

Type	Description
ValueError	If alpha is not a positive integer or the sizes disagree.

Examples:

import numpy as np
import gaussx

# Two triangles forming the unit square
vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
triangles = np.array([[0, 1, 2], [0, 2, 3]])
C, G = gaussx.fem_matrices(vertices, triangles)
kappa, tau, alpha = gaussx.matern_spde_params(
    range=0.5, sigma=1.0, nu=1.0, d=2
)
Q = gaussx.spde_precision(C, G, kappa, tau, alpha)  # (4, 4), alpha = 2

Source code in src/gaussx/_gmrf/_spde.py

def spde_precision(
    C_lumped: lx.DiagonalLinearOperator,
    G: SparseOperator,
    kappa: Float[ArrayLike, ""],
    tau: Float[ArrayLike, ""],
    alpha: int,
) -> SparseOperator:
    r"""SPDE precision ``Q_α = τ² K (C̃⁻¹K)^{α−1}`` with ``K = κ²C̃ + G``.

    The pattern of ``Q_α`` (the ``(α−1)``-ring neighbourhood of ``G``'s) and
    the index triples of each sparse product are computed once on the host
    per pattern and cached; only the values depend on ``κ`` and ``τ``, so
    ``jit``, ``grad`` and ``vmap`` over them never redo the symbolic work.
    The result goes to sparse Cholesky through `gaussx.solve`,
    `gaussx.logdet` and `gaussx.diag_inv` as those dispatch for a
    `SparseOperator`.

    Args:
        C_lumped: Lumped (diagonal) mass matrix ``C̃``, from
            `gaussx.fem_matrices`.
        G: Stiffness matrix, from `gaussx.fem_matrices`.
        kappa: Inverse range parameter ``κ > 0`` (may be traced).
        tau: Scale ``τ > 0`` (may be traced).
        alpha: Integer order ``α ≥ 1``; ``nu = α − d/2``.

    Returns:
        A symmetric positive-definite `SparseOperator`.

    Raises:
        ValueError: If ``alpha`` is not a positive integer or the sizes
            disagree.

    Examples:
        ```python
        import numpy as np
        import gaussx

        # Two triangles forming the unit square
        vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
        triangles = np.array([[0, 1, 2], [0, 2, 3]])
        C, G = gaussx.fem_matrices(vertices, triangles)
        kappa, tau, alpha = gaussx.matern_spde_params(
            range=0.5, sigma=1.0, nu=1.0, d=2
        )
        Q = gaussx.spde_precision(C, G, kappa, tau, alpha)  # (4, 4), alpha = 2
        ```
    """
    if not isinstance(alpha, int | np.integer) or alpha < 1:
        raise ValueError(f"alpha must be a positive integer, got {alpha!r}.")
    if not isinstance(G, SparseOperator):
        raise TypeError(f"G must be a SparseOperator, got {type(G).__name__}.")
    c = lx.diagonal(C_lumped)
    n = G.pattern.shape[0]
    if c.shape != (n,):
        raise ValueError(f"C_lumped must have size {n}, got {c.shape[0]}.")
    kappa = _as_float(kappa)
    tau = _as_float(tau)
    K = G.add_diagonal(kappa**2 * c)
    Q = K
    for _ in range(int(alpha) - 1):
        Q = _sandwich(K, 1.0 / c, Q)
    return SparseOperator(
        tau**2 * Q.values,
        Q.pattern,
        tags=frozenset({lx.positive_semidefinite_tag}),
    )

spde_precision_grid(shape: tuple[int, ...], kappa: Float[ArrayLike, ''], tau: Float[ArrayLike, ''], alpha: int, *, spacing: float = 1.0, periodic: bool | tuple[bool, ...] = False) -> SpectralFunction

SPDE precision on a regular grid: a function of a Kronecker sum.

With spacing h the right-triangle mesh has C̃ = h^d I and G = h^{d−2}(L_1 ⊕ … ⊕ L_d), the (2d+1)-point Laplacian, so

\[ Q_\alpha = \tau^2 h^d\big(\kappa^2 I + h^{-2}(L_1\oplus\cdots\oplus L_d)\big)^\alpha, \]

which is τ²h²(κ² + λ/h²)^α on the eigenvalues λ of the Kronecker sum for a raster (d = 2). Each L_m is the path-graph Laplacian (natural boundary) or, on a periodic axis, the cycle-graph Laplacian. Its eigendecomposition is computed once on the host, so gaussx.solve, gaussx.logdet, gaussx.diag_inv and exact sampling (SpectralFunction.sqrt_matmul with inverse=True) cost O(Σ_m n_m³) once and O(N Σ_m n_m) per call, with no mesh and no Cholesky. In the interior it equals spde_precision on the matching right-triangle mesh; at a non-periodic boundary the mesh's lumped mass and half-weight boundary edges differ (see the module notes on boundary effects and domain extension).

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Grid shape (n_1, …, n_d); vectors are flattened row-major.	required
kappa	Float[ArrayLike, '']	Inverse range κ > 0 in units of the coordinates (may be traced).	required
tau	Float[ArrayLike, '']	Scale τ > 0 (may be traced).	required
alpha	int	Integer order α ≥ 1.	required
spacing	float	Grid spacing h (the same on every axis).	1.0
periodic	bool | tuple[bool, ...]	Wrap all axes (True) or some of them, e.g. (False, True) for a global latitude-longitude raster.	False

Returns:

Type	Description
SpectralFunction	A positive-definite gaussx.SpectralFunction.

Raises:

Type	Description
ValueError	If alpha is not a positive integer or periodic has the wrong length.

Examples:

import gaussx

# Matérn nu = 1 on a 30 x 60 raster, wrapping the second axis
Q = gaussx.spde_precision_grid(
    (30, 60), kappa=0.3, tau=1.0, alpha=2, periodic=(False, True)
)
sd = gaussx.diag_inv(Q) ** 0.5  # exact, through the factor eigenvectors

Source code in src/gaussx/_gmrf/_spde.py

def spde_precision_grid(
    shape: tuple[int, ...],
    kappa: Float[ArrayLike, ""],
    tau: Float[ArrayLike, ""],
    alpha: int,
    *,
    spacing: float = 1.0,
    periodic: bool | tuple[bool, ...] = False,
) -> SpectralFunction:
    r"""SPDE precision on a regular grid: a function of a Kronecker sum.

    With spacing ``h`` the right-triangle mesh has ``C̃ = h^d I`` and
    ``G = h^{d−2}(L_1 ⊕ … ⊕ L_d)``, the ``(2d+1)``-point Laplacian, so

    $$
    Q_\alpha = \tau^2 h^d\big(\kappa^2 I
        + h^{-2}(L_1\oplus\cdots\oplus L_d)\big)^\alpha,
    $$

    which is ``τ²h²(κ² + λ/h²)^α`` on the eigenvalues ``λ`` of the Kronecker
    sum for a raster (``d = 2``). Each ``L_m`` is the path-graph Laplacian
    (natural boundary) or, on a periodic axis, the cycle-graph Laplacian.
    Its eigendecomposition is computed once on the host, so `gaussx.solve`,
    `gaussx.logdet`, `gaussx.diag_inv` and exact sampling
    (`SpectralFunction.sqrt_matmul` with ``inverse=True``) cost
    ``O(Σ_m n_m³)`` once and ``O(N Σ_m n_m)`` per call, with no mesh and no
    Cholesky. In the interior it equals `spde_precision` on the matching
    right-triangle mesh; at a non-periodic boundary the mesh's lumped mass
    and half-weight boundary edges differ (see the module notes on
    boundary effects and domain extension).

    Args:
        shape: Grid shape ``(n_1, …, n_d)``; vectors are flattened row-major.
        kappa: Inverse range ``κ > 0`` in units of the coordinates (may be
            traced).
        tau: Scale ``τ > 0`` (may be traced).
        alpha: Integer order ``α ≥ 1``.
        spacing: Grid spacing ``h`` (the same on every axis).
        periodic: Wrap all axes (``True``) or some of them, e.g.
            ``(False, True)`` for a global latitude-longitude raster.

    Returns:
        A positive-definite `gaussx.SpectralFunction`.

    Raises:
        ValueError: If ``alpha`` is not a positive integer or ``periodic``
            has the wrong length.

    Examples:
        ```python
        import gaussx

        # Matérn nu = 1 on a 30 x 60 raster, wrapping the second axis
        Q = gaussx.spde_precision_grid(
            (30, 60), kappa=0.3, tau=1.0, alpha=2, periodic=(False, True)
        )
        sd = gaussx.diag_inv(Q) ** 0.5  # exact, through the factor eigenvectors
        ```
    """
    if not isinstance(alpha, int | np.integer) or alpha < 1:
        raise ValueError(f"alpha must be a positive integer, got {alpha!r}.")
    shape = tuple(int(n) for n in shape)
    if isinstance(periodic, bool):
        periodic = (periodic,) * len(shape)
    if len(periodic) != len(shape):
        raise ValueError(
            f"periodic must have one entry per axis ({len(shape)}), got {periodic}."
        )
    kappa = _as_float(kappa)
    tau = _as_float(tau)
    h = jnp.asarray(spacing, dtype=jnp.result_type(kappa, tau))
    factors = [
        _cast(
            EigenFactorization.from_matrix(_laplacian_1d(n, wrap), symmetric=True),
            h.dtype,
        )
        for n, wrap in zip(shape, periodic, strict=True)
    ]
    symbol = _MaternSymbol(kappa, tau, h, int(alpha), len(shape))
    return SpectralFunction.from_eigen_factorizations(
        factors, symbol, tags=frozenset({lx.positive_semidefinite_tag})
    )

matern_spde_params(range: Float[ArrayLike, ''], sigma: Float[ArrayLike, ''], nu: float, d: int) -> tuple[Array, Array, int]

SPDE parameters (κ, τ, α) of a Matérn field.

κ = √(8 nu)/rho for the practical range rho (correlation ≈ 0.13 at distance rho), α = nu + d/2, and τ from the marginal variance

\[ \sigma^2 = \frac{\Gamma(\nu)}{\Gamma(\alpha)(4\pi)^{d/2}\kappa^{2\nu}\tau^2}. \]

Parameters:

Name	Type	Description	Default
range	Float[ArrayLike, '']	Practical range rho > 0 (may be traced).	required
sigma	Float[ArrayLike, '']	Marginal standard deviation σ > 0 (may be traced).	required
nu	float	Smoothness nu > 0, concrete.	required
d	int	Spatial dimension, concrete.	required

Returns:

Type	Description
tuple[Array, Array, int]	(kappa, tau, alpha) with an integer alpha.

Raises:

Type	Description
ValueError	If nu + d/2 is not a positive integer (rational non-integer α is not supported).

Examples:

import gaussx

kappa, tau, alpha = gaussx.matern_spde_params(
    range=50.0, sigma=2.0, nu=1.0, d=2
)  # alpha == 2

Source code in src/gaussx/_gmrf/_spde.py

def matern_spde_params(
    range: Float[ArrayLike, ""],
    sigma: Float[ArrayLike, ""],
    nu: float,
    d: int,
) -> tuple[Array, Array, int]:
    r"""SPDE parameters ``(κ, τ, α)`` of a Matérn field.

    ``κ = √(8 nu)/rho`` for the practical range ``rho`` (correlation ≈ 0.13 at
    distance ``rho``), ``α = nu + d/2``, and ``τ`` from the marginal variance

    $$
    \sigma^2 = \frac{\Gamma(\nu)}{\Gamma(\alpha)(4\pi)^{d/2}\kappa^{2\nu}\tau^2}.
    $$

    Args:
        range: Practical range ``rho > 0`` (may be traced).
        sigma: Marginal standard deviation ``σ > 0`` (may be traced).
        nu: Smoothness ``nu > 0``, concrete.
        d: Spatial dimension, concrete.

    Returns:
        ``(kappa, tau, alpha)`` with an integer ``alpha``.

    Raises:
        ValueError: If ``nu + d/2`` is not a positive integer (rational
            non-integer ``α`` is not supported).

    Examples:
        ```python
        import gaussx

        kappa, tau, alpha = gaussx.matern_spde_params(
            range=50.0, sigma=2.0, nu=1.0, d=2
        )  # alpha == 2
        ```
    """
    alpha_float = float(nu) + d / 2
    alpha = round(alpha_float)
    if alpha < 1 or not math.isclose(alpha, alpha_float):
        raise ValueError(
            f"alpha = nu + d/2 = {alpha_float} must be a positive integer; "
            "non-integer alpha is not supported."
        )
    rho = _as_float(range)
    sigma = _as_float(sigma)
    kappa = jnp.sqrt(8.0 * nu) / rho
    log_tau2 = (
        gammaln(nu)
        - gammaln(float(alpha))
        - 0.5 * d * math.log(4.0 * math.pi)
        - 2.0 * nu * jnp.log(kappa)
        - 2.0 * jnp.log(sigma)
    )
    return kappa, jnp.exp(0.5 * log_tau2).astype(kappa.dtype), alpha

fem_matrices(vertices: Float[ArrayLike, 'V D'], triangles: Int[ArrayLike, 'T 3']) -> tuple[lx.DiagonalLinearOperator, SparseOperator]

Lumped mass C̃ and stiffness G of P1 elements on a triangle mesh.

The local matrices (module docstring) are computed for all triangles at once with einx and scattered with a segment_sum into G's pattern, which is built on the host from triangles (two vertices are coupled iff they share a triangle). vertices may be traced; only triangles must be concrete.

Parameters:

Name	Type	Description	Default
vertices	Float[ArrayLike, 'V D']	Vertex coordinates, shape (V, 2) for a planar mesh or (V, 3) for a surface mesh.	required
triangles	Int[ArrayLike, 'T 3']	Vertex indices of each triangle, shape (T, 3) (concrete host integers).	required

Returns:

Type	Description
DiagonalLinearOperator	(C_lumped, G): the diagonal lumped mass matrix as a
SparseOperator	lineax.DiagonalLinearOperator and the stiffness matrix as a
tuple[DiagonalLinearOperator, SparseOperator]	symmetric positive-semidefinite SparseOperator (its null space is
tuple[DiagonalLinearOperator, SparseOperator]	the constants on each connected component).

Raises:

Type	Description
ValueError	If the shapes are wrong or an index is out of range.

Examples:

import numpy as np
import gaussx

# The unit right triangle: |T| = 1/2
vertices = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
C, G = gaussx.fem_matrices(vertices, np.array([[0, 1, 2]]))
C.as_matrix()  # I / 6
G.as_matrix()  # [[1, -1/2, -1/2], [-1/2, 1/2, 0], [-1/2, 0, 1/2]]

Source code in src/gaussx/_gmrf/_fem.py

def fem_matrices(
    vertices: Float[ArrayLike, "V D"],
    triangles: Int[ArrayLike, "T 3"],
) -> tuple[lx.DiagonalLinearOperator, SparseOperator]:
    r"""Lumped mass ``C̃`` and stiffness ``G`` of P1 elements on a triangle mesh.

    The local matrices (module docstring) are computed for all triangles at
    once with einx and scattered with a ``segment_sum`` into ``G``'s
    pattern, which is built on the host from ``triangles`` (two vertices are
    coupled iff they share a triangle). ``vertices`` may be traced; only
    ``triangles`` must be concrete.

    Args:
        vertices: Vertex coordinates, shape ``(V, 2)`` for a planar mesh or
            ``(V, 3)`` for a surface mesh.
        triangles: Vertex indices of each triangle, shape ``(T, 3)``
            (concrete host integers).

    Returns:
        ``(C_lumped, G)``: the diagonal lumped mass matrix as a
        `lineax.DiagonalLinearOperator` and the stiffness matrix as a
        symmetric positive-semidefinite `SparseOperator` (its null space is
        the constants on each connected component).

    Raises:
        ValueError: If the shapes are wrong or an index is out of range.

    Examples:
        ```python
        import numpy as np
        import gaussx

        # The unit right triangle: |T| = 1/2
        vertices = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
        C, G = gaussx.fem_matrices(vertices, np.array([[0, 1, 2]]))
        C.as_matrix()  # I / 6
        G.as_matrix()  # [[1, -1/2, -1/2], [-1/2, 1/2, 0], [-1/2, 0, 1/2]]
        ```
    """
    X = _as_float(vertices)
    tri = _host_triangles(triangles, X.shape[0])
    if X.ndim != 2 or X.shape[1] not in (2, 3):
        raise ValueError(f"vertices must have shape (V, 2) or (V, 3), got {X.shape}.")
    gram, area = _local_geometry(X[tri])
    G_local = einx.divide("t i j, t -> t i j", gram, 4.0 * area)
    values = rearrange(G_local[:, _LOCAL_ROWS, _LOCAL_COLS], "t k -> (t k)")
    n = X.shape[0]
    G = SparseOperator.from_coo(
        tri[:, _LOCAL_ROWS].ravel(),
        tri[:, _LOCAL_COLS].ravel(),
        values,
        (n, n),
        symmetric=True,
        tags=frozenset({lx.positive_semidefinite_tag}),
    )
    mass = jax.ops.segment_sum(
        repeat(area / 3.0, "t -> (t k)", k=3), tri.ravel(), num_segments=n
    )
    return lx.DiagonalLinearOperator(mass), G

fem_projector(vertices: Float[ArrayLike, 'V D'], triangles: Int[ArrayLike, 'T 3'], points: Float[ArrayLike, 'n D'], *, triangle_index: Int[ArrayLike, ' n'] | None = None) -> SparseOperator

Observation matrix A of P1 interpolation: (A w)_k = Σ_i ψ_i(s_k) w_i.

Row k holds the three barycentric weights of point s_k in its triangle, so A has three non-zeros per row and, in a Laplace Hessian Q + AᵀWA, never couples vertices that are not already neighbours.

Point location (on the host, blocked brute force over all triangles):

• Planar meshes (V × 2): the barycentric test; a point outside every triangle raises.
• Surface meshes (V × 3) that are star-shaped about the centroid of their vertices (spheres, icospheres; every triangle must face away from the centroid, which also requires a consistent orientation): the ray from the centroid through each point is intersected with the triangles, and the weights are taken at the intersection, i.e. the point is projected radially onto the mesh (its distance to the mesh, the chord error, is ignored).
• Any other surface mesh needs triangle_index.

With triangle_index no location is done and the weights are those of the point's orthogonal projection onto the triangle's plane; points may then be traced.

Parameters:

Name	Type	Description	Default
vertices	Float[ArrayLike, 'V D']	Vertex coordinates, shape (V, 2) or (V, 3).	required
triangles	Int[ArrayLike, 'T 3']	Vertex indices of each triangle, shape (T, 3) (concrete host integers).	required
points	Float[ArrayLike, 'n D']	Observation locations, shape (n, D); concrete unless triangle_index is given.	required
triangle_index	Int[ArrayLike, ' n'] | None	The triangle containing each point, shape (n,) (concrete host integers).	None

Returns:

Type	Description
SparseOperator	A SparseOperator of shape (n, V).

Raises:

Type	Description
ValueError	If a point is not on the mesh, or a surface mesh is not star-shaped about its centroid and triangle_index is missing.

Examples:

import numpy as np
import gaussx

vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
triangles = np.array([[0, 1, 2], [0, 2, 3]])
A = gaussx.fem_projector(vertices, triangles, np.array([[0.5, 0.25]]))
A.as_matrix()  # [[0.5, 0.25, 0.25, 0]]: barycentric weights

Source code in src/gaussx/_gmrf/_fem.py

def fem_projector(
    vertices: Float[ArrayLike, "V D"],
    triangles: Int[ArrayLike, "T 3"],
    points: Float[ArrayLike, "n D"],
    *,
    triangle_index: Int[ArrayLike, " n"] | None = None,
) -> SparseOperator:
    r"""Observation matrix ``A`` of P1 interpolation: ``(A w)_k = Σ_i ψ_i(s_k) w_i``.

    Row ``k`` holds the three barycentric weights of point ``s_k`` in its
    triangle, so ``A`` has three non-zeros per row and, in a Laplace
    Hessian ``Q + AᵀWA``, never couples vertices that are not already
    neighbours.

    Point location (on the host, blocked brute force over all triangles):

    - **Planar meshes** (``V × 2``): the barycentric test; a point outside
      every triangle raises.
    - **Surface meshes** (``V × 3``) that are star-shaped about the
      centroid of their vertices (spheres, icospheres; every triangle must
      face away from the centroid, which also requires a consistent
      orientation): the ray from the centroid through each point is
      intersected with the triangles, and the weights are taken at the
      intersection, i.e. the point is projected radially onto the mesh (its
      distance to the mesh, the chord error, is ignored).
    - Any other surface mesh needs ``triangle_index``.

    With ``triangle_index`` no location is done and the weights are those of
    the point's orthogonal projection onto the triangle's plane; ``points``
    may then be traced.

    Args:
        vertices: Vertex coordinates, shape ``(V, 2)`` or ``(V, 3)``.
        triangles: Vertex indices of each triangle, shape ``(T, 3)``
            (concrete host integers).
        points: Observation locations, shape ``(n, D)``; concrete unless
            ``triangle_index`` is given.
        triangle_index: The triangle containing each point, shape ``(n,)``
            (concrete host integers).

    Returns:
        A `SparseOperator` of shape ``(n, V)``.

    Raises:
        ValueError: If a point is not on the mesh, or a surface mesh is not
            star-shaped about its centroid and ``triangle_index`` is
            missing.

    Examples:
        ```python
        import numpy as np
        import gaussx

        vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
        triangles = np.array([[0, 1, 2], [0, 2, 3]])
        A = gaussx.fem_projector(vertices, triangles, np.array([[0.5, 0.25]]))
        A.as_matrix()  # [[0.5, 0.25, 0.25, 0]]: barycentric weights
        ```
    """
    X = _as_float(vertices)
    tri = _host_triangles(triangles, X.shape[0])
    if X.ndim != 2 or X.shape[1] not in (2, 3):
        raise ValueError(f"vertices must have shape (V, 2) or (V, 3), got {X.shape}.")
    P = jnp.asarray(points, dtype=X.dtype)
    if P.ndim != 2 or P.shape[1] != X.shape[1]:
        raise ValueError(
            f"points must have shape (n, {X.shape[1]}) like vertices, got {P.shape}."
        )
    centre = None
    if triangle_index is None:
        try:
            X_host, P_host = np.asarray(X), np.asarray(P)
        except jax.errors.TracerArrayConversionError as err:
            raise TypeError(
                "Point location needs concrete vertices and points; pass "
                "triangle_index to build the projector from traced values."
            ) from err
        if X.shape[1] == 2:
            index = _locate_planar(X_host, tri, P_host)
        else:
            index = _locate_radial(X_host, tri, P_host)
            centre = einx.mean("v d -> d", X)
    else:
        index = np.asarray(triangle_index)
        if index.shape != (P.shape[0],) or not np.issubdtype(index.dtype, np.integer):
            raise ValueError(
                f"triangle_index must be {P.shape[0]} integers, got {index.shape}."
            )
        if index.size and (index.min() < 0 or index.max() >= tri.shape[0]):
            raise ValueError("triangle_index out of range.")
    corners = X[tri[index]]  # (n, 3, D)
    if centre is not None:
        P = _radial_projection(corners, P, centre)
    weights = _barycentric(corners, P)
    n = P.shape[0]
    return SparseOperator.from_coo(
        np.repeat(np.arange(n), 3),
        tri[index].ravel(),
        rearrange(weights, "n k -> (n k)"),
        (n, X.shape[0]),
    )