Structured operators — Dense, Diagonal, Kronecker, BlockDiag, LowRankUpdate

A linear operator is anything that turns a vector into another vector linearly. The dense matrix $A\in\mathbb{R}^{n\times n}$ is the universal one — it can represent every operator — but it is also the most expensive one. Once $n$ exceeds a few thousand, storing $A$ requires gigabytes and the cubic-cost primitives ($\mathrm{solve}$, $\mathrm{logdet}$, $\mathrm{chol}$) become impossibly slow.

Every Gaussian-process scaling story begins with the same realization: we rarely need a general operator. The kernel matrix on a tensor-product grid factorizes as $K_X\otimes K_Y$. The precision of a Markov chain is block tridiagonal. The posterior covariance after a low-rank measurement is $A + UDV^\top$. Structure is the rule, not the exception — and once we expose it as the type of the operator, every downstream primitive can dispatch to a cheaper algorithm.

This notebook is the catalog: we tour the five basic operators that show up in 90% of GP code (MatrixLinearOperator, DiagonalLinearOperator, Kronecker, BlockDiag, LowRankUpdate), compare their per-primitive costs, and end with a “which operator do I reach for” decision table.

What you will see by the end:

• The lineax.AbstractLinearOperator interface — matvec, transpose, structural tags — and how gaussx extends it.
• The arithmetic identities that make each structured class cheap (Kronecker vec-trick, block-wise solve, Sherman–Morrison–Woodbury).
• A side-by-side cost table that ranks operators by their $\mathrm{matvec}$, $\mathrm{solve}$, and $\mathrm{logdet}$ complexities.
• Theoretical FLOP curves plotted directly from the dispatch table (same approach as notebook 0.9).

By the end you should be able to look at a covariance and say “that’s a Kronecker(BlockDiag, Toeplitz)” — and know, without writing code, what the inference cost will be.

0. Where these structures show up in geoscience¶

The five operators in this notebook aren’t abstractions for their own sake — every one of them maps to a recurring shape in environmental data. A short, non-exhaustive map:

• Diagonal — heteroscedastic observation noise: each station / mooring / satellite footprint has its own measurement variance $\sigma_i^2$, and the noise covariance is $\mathrm{diag}(\sigma_1^2,\dots,\sigma_N^2)$. ARD lengthscales for multivariate inputs (latitude, longitude, depth, elevation) ride on a diagonal weighting too.
• Kronecker — anything on a regular tensor-product grid: sea-surface temperature on (lat × lon), CO₂ flux on (time × space), soil moisture on (depth × horizontal), seismic amplitude on (offset × azimuth). The kernel $K_{\mathrm{lat}}\otimes K_{\mathrm{lon}}$ replaces a dense $N\times N$ covariance with two factor matrices the size of one axis.
• BlockDiag — independent geographic regions (river basins, ocean basins, ecoregions) where cross-region correlation is negligible. Multi-output GPs with one block per pollutant species. Ensemble forecasts where members are conditionally independent given parameters.
• LowRankUpdate — assimilating $M$ observations into a high-dimensional state, with $M \ll N$. The Kalman update $P^{+} = P^{-} - K H P^{-}$ adds a rank-$M$ correction to the prior covariance; ensemble Kalman filters (EnKF) parameterize the ensemble covariance directly as $\frac{1}{M-1}EE^\top$, which is rank-$M$ low-rank by construction.
• Dense — calibration-scale inverse problems with $N \lesssim 10^3$ (gravity / magnetics inversion at coarse resolution, instrument calibration). The fallback when no structure is exploitable.

The composition rule matters too: a Kronecker of BlockDiags is a tensor-product grid where each output channel evolves independently — exactly the shape of multivariate spacetime fields like (temperature, salinity) on (lat × lon × depth).

1. The AbstractLinearOperator interface¶

lineax defines an operator $A$ by what it can do, not by how it is stored:

Three things follow:

• Storage is opaque. A dense MatrixLinearOperator stores all $n^2$ entries, but a Kronecker(A,B) stores only $\dim(A)^2 + \dim(B)^2$ entries. Both expose the same mv interface — primitives don’t care.
• Tags drive dispatch. When gaussx.solve sees an operator carrying the positive_semidefinite_tag, it dispatches to Cholesky; when it sees a Kronecker, it dispatches to the vec-trick. Tags are how structure becomes free speed.
• Operators compose. A BlockDiag of Kroneckers of LowRankUpdates is a perfectly valid operator — each layer adds its own structural shortcut. This is the whole point of gaussx’s Layer-1 stack.

The five basics below are the leaves of every operator tree you will build.

2. Dense — MatrixLinearOperator¶

The reference. No structure, no shortcuts: $A$ is stored as an $n\times n$ array, and every primitive falls back to the textbook algorithm.

Use it when $n \lesssim 10^3$, or when there is genuinely no structure to exploit. Everything else in this notebook is a strategy for avoiding the dense fallback.

3. Diagonal — the trivial structure¶

A diagonal operator $D = \mathrm{diag}(d_1,\dots,d_n)$ is the warm-up: every primitive collapses to an element-wise operation on the diagonal vector.

Cost is linear in $n$ across the board. Storage is $O(n)$ — the diagonal vector itself.

Diagonal operators show up everywhere: noise covariances $\sigma^2 I$, ARD lengthscale weights, mean-field variational families, preconditioners. lineax provides DiagonalLinearOperator and IdentityLinearOperator directly; gaussx recognizes them via the diagonal_tag.

4. Kronecker — separable structure¶

Whenever the covariance factorizes across independent input axes (a 2-D grid in space, time × feature), the kernel matrix becomes a Kronecker product:

The dense form has $(nm)^2$ entries, but the operator only needs $n^2 + m^2$. The savings come from Roth’s vec-lemma:

A matvec on a vector of length $nm$ becomes two matrix-multiplies of size $m\times n$ — total cost $O(nm(n+m))$ instead of $O((nm)^2)$. Solve and logdet inherit the same factorization:

For $n=m=64$ the dense matvec is $\sim 16\text{M}$ flops; the Kronecker matvec is $\sim 0.5\text{M}$ — a 30× speedup that grows with size.

5. BlockDiag — independent sub-problems¶

A block-diagonal operator partitions the input into $K$ independent groups, each handled by its own (possibly structured) sub-operator:

Every primitive decomposes block-wise:

If each block has size $b = n/K$, dense solve drops from $O(n^3)$ to $O(K b^3) = O(n^3/K^2)$ — embarrassingly parallel. Crucially, the blocks need not be dense: a BlockDiag of Kroneckers is the canonical multi-output / batched-GP structure.

6. LowRankUpdate — Sherman–Morrison–Woodbury¶

A low-rank update writes a large operator as a cheap base plus a rank-$k$ correction:

When $A$ is cheap (diagonal, identity, Kronecker) and $k \ll n$, every primitive on $M$ can be reduced to a cheap operation on $A$ plus a $k\times k$ work item. The key identity is Sherman–Morrison–Woodbury:

and the matching matrix determinant lemma:

Solve cost drops from $O(n^3)$ (dense) to $O(n k^2 + k^3)$ — linear in $n$ when $A$ is diagonal. Notebook 1.9 derives Woodbury in full; here we just verify the identities numerically. gaussx provides three constructors for the common shapes: LowRankUpdate(base, U, d, V) (general), low_rank_plus_diag(diag, U, d, V), and low_rank_plus_identity(U, d, V, scale).

7. The dispatch tables¶

We follow the 0.9 convention: separate the closed-form identity from the FLOP count from the storage footprint, so each axis of the cost trade-off is legible on its own.

7.1 Closed-form identities¶

7.2 FLOP count (leading order)¶

Let $N$ = ambient dimension, $k$ = rank, $K$ = number of blocks (each of size $b = N/K$), $A$ has size $n$ and $B$ has size $m$ so $N = nm$.

7.3 Storage¶

Two takeaways: (i) every structured operator is sub-quadratic in storage when its parameters are small, and (ii) cost gaps grow with $N$ — the scaling exponents in §7.2 are why we bother.

8. Theoretical FLOP curves¶

Plotting the closed-form FLOP table (§7.2) directly — same approach as notebook 0.9. We sweep $N$ from 32 to 4096 along a square Kronecker layout ($n=m=\sqrt N$), with rank-8 low-rank updates and $K=8$ blocks. Empirical timings would ride on these curves with a constant prefactor; what matters is the slope.

The y-axis is log-FLOPs and the x-axis is log-$N$, so a curve’s slope equals its asymptotic exponent: dense solve $=3$, Kronecker solve $=3/2$, low-rank solve $=1$, diagonal $=1$.

9. Which operator do I reach for?¶

A short flowchart for picking structure when you build a covariance:

The remaining 1.A notebooks fill in the rest of the catalog: 1.2 lazy algebra (sums/products/scaling), 1.3 KroneckerSum vs SumKronecker, 1.4 Toeplitz, 1.5 BlockTriDiag, 1.6 MaskedOperator. The next sections of this notebook (§10–§13) cover the structural-tag system that lets gaussx dispatch automatically across the entire catalog.

A reading hint: when you see a covariance in a paper, name its operator type out loud before reading the algorithm. Half the time the algorithm is the dispatch — and the structured cost is already in the table above.

10. How gaussx knows which algorithm to use — the two-tiered dispatch system¶

Every operator above carries a type (Diagonal, Kronecker, BlockDiag, LowRankUpdate, MatrixLinearOperator) and a frozenset of structural tags (positive_semidefinite_tag, symmetric_tag, diagonal_tag, kronecker_tag, low_rank_tag, …). gaussx’s primitives use these in two complementary ways:

1. isinstance dispatch in solve / logdet / cholesky — the class picks the algorithm. solve is essentially:

def solve(operator, b):
    if isinstance(operator, lx.DiagonalLinearOperator): return _solve_diagonal(operator, b)
    if isinstance(operator, BlockDiag):                 return _solve_block_diag(operator, b)
    if isinstance(operator, Kronecker):                 return _solve_kronecker(operator, b)
    if isinstance(operator, LowRankUpdate):             return _solve_low_rank(operator, b)
    ...
    return _solve_fallback(operator, b)   # generic Cholesky / LU

2. Tag predicates — the property picks the inner solver inside a recipe. is_positive_semidefinite(op) decides Cholesky vs LU; is_symmetric(op) decides eigh vs eig; is_diagonal(op) collapses to element-wise. Predicates are functools.singledispatch functions that each operator type registers via @lx.is_symmetric.register(MyType).

Class for which algorithm. Tags for which solver inside the algorithm. The two are deliberately complementary: class identity is the recipe; tags are the property assertions that recipes consult.

11. Tag inventory + what the basics advertise¶

The full tag catalog (re-exported from gaussx._tags):

Each comes with a matching predicate (is_symmetric, is_kronecker, etc.). Each operator class registers itself for the predicates it satisfies. The next cell builds one of every operator type from §2–§6 and queries the predicates that are relevant at this point in the tour.

12. Watching solve find the recipe — dispatch tracing¶

A small experiment: monkey-patch each _solve_* branch to record when it fires, then call gaussx.solve on each operator. The recorded branch is the recipe gaussx chose. Every structured operator routes to its own recipe; the dense PSD case falls through to _solve_fallback (which delegates to lineax’s Cholesky / LU autosolver).

13. Bring your own operator — extending the system¶

The dispatch system is open for extension. Three steps to plug a new operator type into gaussx:

1. Subclass lx.AbstractLinearOperator with the required methods (mv, as_matrix, transpose, in_structure, out_structure).
2. Register tag predicates for any properties your operator advertises (@lx.is_symmetric.register(MyType), @gx.is_low_rank.register(MyType), …).
3. Optionally provide a fast primitive recipe — a function the user calls directly, or a singledispatch registration on a custom dispatcher.

The example below builds a circulant operator — like Toeplitz but with periodic boundary conditions, common for fully-periodic domains in spectral GP regression (annual cycles on a daily grid, longitudinal kernels on a sphere). Circulant matrices are diagonalized by the DFT, so solve and logdet are pure FFT operations costing $O(n\log n)$. This is exactly the pattern a domain expert follows to add a Slepian-localized kernel, a stochastic-PDE precision, or a custom SDE-discretization for their corner of geophysics.

14. Class vs tag — when to use which¶

A short style guide for picking between isinstance(op, MyType) and is_my_property(op):

The reason gaussx uses isinstance for solve (not singledispatch) is the last point: keeping the dispatch in an if-chain inside a single function lets JAX trace through it cleanly without per-type closure overhead. Tag predicates are singledispatch because they’re called outside hot inner loops — during construction, structural inference, and user-facing introspection.

That closes the operator-basics tour. The next notebooks build on this same dispatch backbone: 1.2 introduces lazy algebra (Sum / Scaled / Product) whose tag-propagation rules are critical for getting the right primitives to fire on composites; 1.3–1.6 introduce the remaining structured leaves (KroneckerSum, SumKronecker, Toeplitz, BlockTriDiag, MaskedOperator), each of which registers its own class branch and tag predicates following exactly the pattern in §13.