Internals

This page is for someone opening src/ for the first time. It walks the main operations of the package — applying an operator, taking an adjoint, crossing into a Krylov solver, editing a forest, moving field data across a regrid — and for each one gives the entry point, the call sequence, and the concrete type of everything that exists at each step.

It is deliberately the only page written this way. API is generated from docstrings and has the signatures. DESIGN.md is the record of decisions and rationale — why the package is shaped like this, what was rejected, what is deferred. The sibling pages (AMR, Geometric Multigrid, Distributed Multi-GPU Solves, Automatic Differentiation) tell the user-facing story of each subsystem: what it means and when to reach for it. This page has call sequences and types, and nothing else.

How to read this page

Every section has the same five parts: a lede naming the one concrete example it traces, a metadata line pointing at the entry point and the source and test files, a call diagram (who calls whom), a type trace (what object exists at each step), then step commentary and a closing paragraph on what the operation throws and what it invalidates.

The call diagram and the type trace carry different information on purpose. The diagram is the control flow; the trace is the data. Read the diagram to learn the shape, the trace to learn the types.

No line numbers appear anywhere on this page — they rot faster than anything else, and a symbol name plus the Source line is enough to grep. Each exported symbol is linked to the API page once, on its first mention; everything else is in bare backticks. Where this page and the tests disagree, the tests are right and this page is a bug.

Pseudocode notation

Mark Meaning
name::ConcreteType = … a binding, with its concrete type named — never elided to an abstract one
f(x) -> T returns a freshly built T
f!(x) mutates x in place (Julia’s own convention; no extra mark)
# ALLOC allocates, naming what and how many
# MUTATES x mutates in place where the name does not end in !
# SAME hands back the object it was given — identity preserved, callers’ references stay valid
# THROWS T a guard, and the exception type it raises
# HOLDS: the invariant that is true once this step returns
# STALE: what this step invalidates

Vocabulary

Six terms recur in every section. They are defined once, here.

The two index spaces. Every field’s storage is halo-padded: its array has padded_size(g) entries, of which local_size(g) are owned cells and the rest are ghosts. interior(g) returns the CartesianIndices of the owned cells in padded coordinates, and interior(f) returns the matching SubArray. So a stencil at an owned cell on the low face reads I - δ, which is a ghost. Every operator body indexes the padded array and writes only the interior.

Ghost cells are scratch. They are filled by halo_update! and apply_bc! at the top of each apply!, so any apply! may clobber its input’s ghosts — that is in the contract, not a bug. Fields the package hands you ship with zeroed ghosts (allocate_output calls zero_ghosts!), and a nonzero ghost cotangent arriving at an adjoint is a bug upstream.

Interior-only flat vectors. The Krylov boundary is a flat AbstractVector spanning owned degrees of freedom only. flatten, flat_to_interior! and interior_to_flat! cross it. Ghost cells are never solver unknowns, and the flat side cannot address one: the length comes from local_size, and every copy goes through an interior view.

The apply! prologue. Each stencil leaf is two pieces: apply! fills ghosts (halo_update! then apply_bc!) and then calls _apply_raw!, the ghost-trusting kernel. The split exists because some callers must not re-run the prologue — boundary_rhs plants inhomogeneous ghost offsets that apply_bc! would immediately overwrite, and the distributed driver stages exchanged ghosts from its neighbours the same way. Those callers reach past apply! to _apply_raw!.

Traits are declared, never assumed. islinear, isconstant, isselfadjoint, isdiagonal and shares_exchange all default to false. Leaves opt in; combinators propagate explicitly. The same rule governs adjoints: there is no generic apply! and no generic apply_adjoint!, so an operator that forgot to declare one produces a MethodError or an ArgumentError. A forgotten declaration must degrade to an error, never to a wrong number. This single idea answers most “why is there no fallback here?” questions below.

Generations and staleness. A Forest carries a generation counter. Every BlockField and PreparedForest stamps the generation it was built against, and _require_current compares the stamp on use. Editing the topology bumps the counter, so “invalidated” here always means throws on next use, never silently reads wrong-shaped storage.

Where the code lives

File Owns Section
Foundation
src/Grids.jl AbstractGrid, CartesianGrid, interior, coarsen, the halo_update! seam §1
src/boundaries.jl AbstractBC and its four leaves, apply_bc!, fold_bc!, the ghost fill/fold primitives §1, §2
src/Fields.jl AbstractField, Field, the transfer-policy singletons, the flat-vector boundary §1, §3
Operator algebra
src/operators/abstract.jl AbstractOperator, the five traits, AdjointOp, the adjoint gather engine §1, §2
src/operators/algebra.jl Scaled, Added, Composed and their trait and adjoint propagation §2
src/operators/laplacian.jl, derivative.jl, gradient.jl, divergence.jl, scaling.jl, diffusion.jl, advection.jl one leaf each: the struct, _apply_raw!, the declared adjoint, the trait declarations §1
src/operators/diagonal.jl operator_diagonal per leaf and its combinator rules §9
src/operators/linearize.jl LinearizedOp, the JVP backend tags §4
Solver boundary
src/linalg.jl prepare, PreparedOperator, PreparedForest, the tree rewrite, flat mul!, boundary_rhs §3
Forest
src/topology.jl LeafKey, Forest, Morton ordering, refine!/coarsen!/balance! §5
src/BlockForest.jl BlockForest, per-leaf geometry, leaf_grid §5
src/blockfield.jl, src/packedfield.jl BlockField, PackedBlockField, the layout tags, pack/unpack §6
src/schedule.jl the exchange descriptor structs and their device twins §7
src/transfer.jl, src/transfer_kernels.jl schedule construction, the forest halo and BC sweeps, their transposes §7
src/operators/forest.jl, forest_packed.jl per-leaf and packed-kernel forest sweeps §7
src/amr.jl regrid! and the three transfer kernels §8
Solvers and partitioning
src/multigrid.jl, src/operators/restriction.jl, prolongation.jl the transfer pair, the level hierarchy, the V-cycle §9
src/partitioning.jl, src/distributed.jl slab partitioning, the distributability guards, the distributed tree §10
Extensions
ext/ Enzyme rules, the MDLA multi-GPU backend, Reactant. The core depends on none of them §7, §10

1 — apply!: fill the ghosts, then sweep the stencil

Everything else in the package composes this one operation. The example traced here is laplacian(g) * u where g::CartesianGrid{2,Float64,…} carries homogeneous Dirichlet faces and u::Field{Center,Matrix{Float64},…}.

Entry point apply(L, x) -> Field, apply!(y, L, x, g, α, β) -> y · Source src/operators/abstract.jl, src/Grids.jl, src/boundaries.jl, src/operators/laplacian.jl · Tests test/operators_abstract.jl, test/laplacian.jl, test/boundaries.jl

  L * x   or   L(x)
    │
    ▼
  ┌──────────────────────────────────────────────┐
  │ apply(L, x)                                  │  the allocating entry:
  │   y = allocate_output(L, x)                  │  similar(x) + zero_ghosts!
  └──────────────────────────────────────────────┘
    │
    ▼
  ┌──────────────────────────────────────────────┐
  │ apply!(y, L, x, g)  ->  (…, α=true, β=false) │  the in-place 6-arg form
  └──────────────────────────────────────────────┘
    │
    ├──▶ halo_update!(x, g)      no-op on a CartesianGrid (forest: §7)
    │
    ├──▶ apply_bc!(x)            homogeneous ghost fill, dims 1:N
    │
    └──▶ _apply_raw!(y, L, x, g, α, β)
           interior(y) .= α .* stencil.(x.data, interior(g)) [.+ β .* interior(y)]
apply(L::Laplacian, x::Field)
  y::Field{Center,Matrix{Float64},CartesianGrid{2,…}} = allocate_output(L, x)   # ALLOC 1
       similar(x) -> Field(similar(x.data), x.grid); zero_ghosts!(y)
  apply!(y, L, x, x.grid)
    apply!(y, L, x, g, true, false)                        # α=true, β=false
      halo_update!(x, g)                        # SAME — identity on a CartesianGrid
      apply_bc!(x)                              # MUTATES x.data ghost layers
        _apply_bc_dims!(x.data, g, Val(1), g.bc)  #   dims ascend 1:N
          _fill_ghost!(data, Val(D), ghost, bc, source)
                                    #   ghost .= _bc_sign(bc) .* source; Periodic +1 wrap,
                                    #   Dirichlet -1 mirror, Neumann +1 mirror,
                                    #   Interface skipped entirely
      # HOLDS: every ghost x needs now carries its homogeneous value
      _apply_raw!(y, L, x, g, true, false)
        inv_h2::NTuple{2,Float64} = _inv_spacing2(g)       # stack, no heap
        interior(y) .= _lap_at.(Ref(x.data), interior(g), Ref(inv_h2))
                                                           # MUTATES interior(y) only
  # HOLDS: y's own ghosts untouched; x's interior untouched
  return y

1. The allocating entry is a thin shell. L * x and L(x) both land on apply, which calls allocate_output and then the in-place form. allocate_output is where tensor rank is decided: the default is similar(x), Gradient returns similar(x, SVector{N,T}), and Divergence returns similar(x, T) — each throwing an ArgumentError if handed the wrong input rank. Operators carry no rank parameter; the element type carries it.

2. halo_update! is a seam, not a function. On any AbstractGrid it is halo_update!(x, ::AbstractGrid) = x — a literal identity. The package has exactly one overload, for BlockForest (§7). Even a partition_grid slab, whose cut faces are marked Interface, takes the identity: the distributed driver stages a neighbour’s planes into the prepared input buffer instead of calling through here (§10). The seam exists so operator bodies are written once and distributed and AMR change only what the grid is and what this function does.

3. apply_bc! fills only the homogeneous part. It walks dimensions in ascending order 1:N on a shrinking BC tuple, and for each halo layer writes ghost = sign · source: +1 and a wrapped source for Periodic, -1 and a mirrored source for Dirichlet, +1 and a mirrored source for Neumann, and nothing at all for Interface. Ascending order is what makes corner ghosts consistent ghost-of-ghost values. Because only the homogeneous part is applied, islinear(L) really does imply L(0) = 0; inhomogeneous boundary data enters the system separately, through boundary_rhs (§3).

4. The leaf body is an array-level broadcast. _lap_at reads x.data — the whole padded array — at interior(g), so the I ± δ taps at a face cell land on the ghosts step 3 just wrote. It writes through interior(y), a SubArray, so y’s ghosts are never touched and the β ≠ 0 accumulating form is safe. There is no device branch and no kernel: the same broadcast runs on a Matrix and on a CuArray, and Enzyme differentiates it without a rule. KernelAbstractions @kernel is the per-operator escape hatch, not the default.

5. Combinators are the same shape. Scaled folds its factor into α and costs nothing. Added runs both operands, the second with β = true so it accumulates. Composed allocates an intermediate, applies b into it, then applies a — passing tmp.grid rather than g, because a transfer chain changes grid between factors. On a single grid Added therefore runs the prologue twice over the same x; collapsing that to one exchange is a forest optimisation (§7), where it is worth something.

Throws and invalidates. Nothing here invalidates state. There is no generic six-argument apply! — every leaf and every combinator declares its own, so an operator that forgot is a MethodError at the call site rather than a silently wrong sweep. Field’s inner constructor throws DimensionMismatch unless size(data) == padded_size(grid). allocate_output for Gradient throws ArgumentError on a non-Number element type and for Divergence on a non-SVector one.

2 — Adjoints: a gather, then fold_bc!

apply_bc! is a linear map from interior cells to the padded array. Call it P, and call the stencil sweep S; then a leaf’s forward action on the interior is S·P, and its exact transpose is Pᵀ·Sᵀ. fold_bc! is Pᵀ. That identity is the whole of this section. The example traced is apply_adjoint!(x̄, L, ȳ, g) for the same Dirichlet Laplacian as §1.

Entry point adjoint(L) -> AbstractOperator, apply_adjoint!(x̄, L, ȳ, g, α, β) · Source src/operators/abstract.jl, src/operators/algebra.jl, src/boundaries.jl · Tests test/algebra.jl, test/operators_abstract.jl, test/boundaries.jl

  L'                              apply_adjoint!(x̄, L, ȳ, g, α, β)
    │                               │
    ▼                               ▼
  islinear(L)? ──no──▶ THROW      ┌─────────────────────────────────────┐
    │ yes                         │ adjoint_gather!(x̄, ȳ, gather, α, β) │
    ▼                             └─────────────────────────────────────┘
  adjoint_operator(L)               │
    │                               ├──▶ zero_ghosts!(ȳ)   only interior
    ├─ leaf with a closed form ──▶  │                       cotangents are inputs
    │    e.g. self-adjoint -> L     │
    │         Restriction -> 2⁻ᴺ·P  ├──▶ x̄.data .= gather.(ȳ.data, …)
    │                               │      sweeps the WHOLE padded array,
    └─ otherwise ──▶ AdjointOp(L)   │      so cotangents land in ghosts
                                    │
                                    └──▶ fold_bc!(x̄)      Pᵀ: scatter each ghost
                                           back into its mirror source,
                                           dims descend N:1, then zero it
adjoint(L::AbstractOperator)
  islinear(L) || throw(ArgumentError(...))       # THROWS ArgumentError — nonlinear
  adjoint_operator(L)                            # the gate-free internal hook
    -> L                 if the leaf is its own transpose  (isselfadjoint)
    -> Scaled(...)       if a closed form exists  (Restriction -> 2⁻ᴺ·Prolongation)
    -> AdjointOp(L)      otherwise               # ALLOC 1, lazy wrapper

apply_adjoint!(x̄::Field, L::Laplacian, ȳ::Field, g::CartesianGrid, α, β)
  adjoint_gather!(x̄, ȳ, gather, α, β)
    zero_ghosts!(ȳ)                              # MUTATES ȳ ghosts -> 0
    # HOLDS: only interior cotangents enter, matching the flat Krylov boundary
    if iszero(β)
      x̄.data .= gather.(Ref(ȳ.data), CartesianIndices(x̄.data))   # MUTATES all of x̄
      fold_bc!(x̄)
      isone(α) || (x̄.data .*= α)
    else
      zero_bc_ghosts!(x̄)                         # this call's ghosts only; Interface
                                                 # slabs kept — a sibling under the
                                                 # same Added still needs them
      x̄.data .= α .* gather.() .+ β .* x̄.data
      fold_bc!(x̄)
    end
  # HOLDS: ⟨Lx, y⟩ = ⟨x, Lᵀy⟩ exactly, boundary contributions included

1. The gate and the hook are two different functions. Base.adjoint checks islinear and refuses nonlinear operators with a message pointing at linearize (§4). adjoint_operator is the internal hook that skips the check; leaves override it when a cheaper closed form exists, and it is what every combinator recursion calls. AdjointOp is the lazy fallback: it forwards islinear, isconstant, isselfadjoint and isdiagonal to its operand. It explicitly declares shares_exchange = false: the forest adjoint needs a gather and a cross-block fold, which sharing a forward exchange would bypass. It swaps allocate_output with allocate_input, is its own involution, and swaps the two actions, so apply!(y, AdjointOp(L), x, …) is apply_adjoint!(y, L, x, …).

2. The gather sweeps the padded array, on purpose. A forward stencil reads ghosts, so its transpose must write them. adjoint_gather! therefore runs over CartesianIndices(x̄.data), not over the interior, using a bounds-masked read that returns zero outside the array. Cotangents pile up in ghost cells, and the next step is what turns them back into interior contributions.

3. fold_bc! is the transpose of apply_bc!, structurally. Where _fill_ghost! gathers ghost = sign · source, _fold_ghost! scatters source += sign · ghost and then zeroes the ghost. And where the fill recurses after doing its work, giving ascending dimensions 1:N, the fold recurses first, giving descending N:1 — necessary because corner ghosts are ghost-of-ghost values, so the fills compose and the transpose must reverse the composition. Interface faces are skipped on both sides: those ghosts belong to the forest exchange (§7) or the distributed reduction (§10), never to this pair. The trailing fill!(dst, 0) leaves the array in exactly the state the forest adjoint exchange requires as its precondition.

4. The accumulating branch is the subtle one. With β ≠ 0 the running total in already contains folded contributions from earlier terms. Handing that to fold_bc! again would double-count them, so zero_bc_ghosts! clears exactly the cells the fold touches — and deliberately not the Interface slabs, which carry neighbour-owned cotangents a sibling operator under the same Added still has to accumulate.

5. Trait propagation is written out, never inherited. Each combinator declares all five; nothing falls through to a default it did not earn.

Scaled Added Composed AdjointOp
islinear operand both both operand
isconstant operand both both operand
isdiagonal operand both both operand
isselfadjoint operand and isreal(α) both always false operand
shares_exchange operand both always false always false
adjoint_operator Scaled(opᵀ, conj α) Added(aᵀ, bᵀ) Composed(bᵀ, aᵀ) — order flips the operand

The hard falses encode limits on propagation. isselfadjoint(::Composed) cannot be true from its factors: (AB)ᵀ = BᵀAᵀ, which equals AB only when they commute, and nothing here can check that. shares_exchange(::Composed) is false because b(x) is a fresh intermediate that needs its own exchange before a reads it; AdjointOp needs the separate adjoint fold described above. isselfadjoint is also grid-aware rather than fixed: it asks _uniform_grid(L.grid), which for a BlockForest reads a live Ref — so the moment a regrid makes the forest non-uniform, the Laplacian stops claiming self-adjointness.

Throws and invalidates. adjoint(L) on a nonlinear L throws ArgumentError naming linearize as the alternative. apply_adjoint! has a six-argument fallback that throws ArgumentError("no adjoint action declared for …") — the mirror of §1’s missing generic apply!, and for the same reason. The Scaled/Added/Composed size query throws ArgumentError when a composition is bound to no grid at all.

3 — The solver boundary: prepare, and flat mul!

Krylov wants mul!(y::AbstractVector, A, x::AbstractVector) with no allocation per iteration. Operators want fields with ghosts. prepare is the seam: it walks the tree once, allocates every buffer the tree will ever need, and returns an object whose only job is to be that mul!. The example traced is P = prepare(scaling(σ) - laplacian(g), scalar_field(g)) followed by Krylov.cg(P, flatten(f)).

Entry point prepare(L, x) -> PreparedOperator, mul!(y, P, x, α, β) · Source src/linalg.jl, src/Fields.jl · Tests test/prepare_linalg.jl, test/fields.jl

  prepare(L, x)                         mul!(y::Vector, P, x::Vector, α, β)
    │                                     │
    ├──▶ islinear(L)? ──no──▶ THROW       ├──▶ flat_to_interior!(P.xpad, x)
    │                                     │      reshape, no copy; ghosts LEFT ALONE
    ├──▶ xpad = similar(x)                │      (the distributed path stages
    │      zero_ghosts!(xpad)             │       neighbour ghosts here — §10)
    │                                     │
    ├──▶ ypad = allocate_output(L, x)     ├──▶ apply!(P.ypad, P.op, P.xpad, P.grid)
    │                                     │      ALWAYS α=true, β=false
    └──▶ op = _prepare_tree(L, x)         │
           │                              └──▶ interior_to_flat!(y, P.ypad, α, β)
           ├─ Composed  ─▶ PreparedComposed(a, b, tmp)      the caller's α and β
           ├─ AdjointOp ─▶ _push_adjoints, then              are fused into the
           │                PreparedAdjoint(op, scratch)     copy-OUT broadcast
           ├─ ScalingOp{BlockField} on a packed prototype
           │                ─▶ ScalingOp(pack(coeff))
           └─ leaves, Added, Scaled ─▶ structural / unchanged
prepare(L::AbstractOperator, x::Field) -> PreparedOperator
  islinear(L) || throw(ArgumentError(...))              # THROWS ArgumentError
  xpad::Field = similar(x); zero_ghosts!(xpad)          # ALLOC 1
  ypad::Field = allocate_output(L, x)                   # ALLOC 1 (maybe another eltype)
  op = _prepare_tree(L, x)                              # ALLOC 1 per Composed / AdjointOp
  PreparedOperator(op, x.grid, xpad, ypad)              # 4 fields, NOT an AbstractOperator

_prepare_tree(L::Composed, x)
  PreparedComposed(_prepare_tree(L.a, ), _prepare_tree(L.b, x),
                   allocate_output(L.b, x))             # tmp bound once, reused forever
_prepare_tree(L::AdjointOp, x)
  pushed = _push_adjoints(L)                            # (A*B)ᵀ  ->  Composed(Bᵀ, Aᵀ)
  pushed isa AdjointOp || return _prepare_tree(pushed, x)
  op = _prepare_tree(pushed.op, x)                      # normalise the leaf's coefficient layout
  scratch = allocate_output(pushed, x)                  # adjoint OUTPUT shape, not x's shape
  return PreparedAdjoint(op, scratch)

mul!(y::AbstractVector, P::PreparedOperator, x::AbstractVector, α, β)
  flat_to_interior!(P.xpad, x)     # MUTATES interior(P.xpad); NO zero_ghosts! here
  apply!(P.ypad, P.op, P.xpad, P.grid)                  # α=true, β=false — not the caller's
  interior_to_flat!(y, P.ypad, α, β)                    # vi .= α .* interior .+ β .* vi
  # HOLDS: zero allocation in steady state

1. A prepared operator is not an operator. PreparedOperator and PreparedForest are deliberately not AbstractOperator subtypes: they are the solver boundary, so they cannot re-enter +, *, adjoint or apply. They do implement the small interface Krylov consumes — size as (output DOFs, input DOFs), so rank changers report an honestly rectangular shape, and eltype unwrapping SVector{M,T} to T. Four traits are forwarded to the wrapped tree explicitly: islinear, isconstant, isselfadjoint and isdiagonal. Neither wrapper exposes shares_exchange; that trait is queried on nodes inside the tree. Calling it on a prepared solver wrapper raises MethodError, because there is no AbstractOperator supertype to supply the default.

2. The tree rewrite hoists every per-call allocation. Unprepared, Composed allocates its intermediate on every apply. _prepare_tree replaces it with PreparedComposed, which carries that intermediate as a field — and PreparedComposed is an AbstractOperator, so it re-enters ordinary apply! dispatch and the sweep code below it needs no changes. PreparedAdjoint does the same for a surviving leaf adjoint’s scratch. Both explicitly declare islinear, isconstant, isselfadjoint and isdiagonal, while shares_exchange uses the AbstractOperator default of false. PreparedAdjoint additionally declares adjoint_operator(L) = L.op, without which the generic forest adjoint walk would wrap it in another AdjointOp whose apply is this node’s transpose again, recursing forever.

3. _push_adjoints runs before any buffer is bound. It rewrites AdjointOp(A*B) into Composed(Bᵀ, Aᵀ), with a fixed-point guard so a leaf with no cheaper adjoint does not recurse. The reason is distributed: PreparedComposed carries one intermediate shaped for the forward pass, and a distributed adjoint walk would run aᵀ then bᵀ with no reduction between them, dropping the intermediate’s Interface cotangents. Normalising the tree so the composition is explicit puts the reduction back where it belongs (§10).

4. The flat boundary is interior-only by construction, not by assertion. The flat length is derived from local_size, flat_to_interior! reshapes the caller’s vector without copying and broadcasts into an interior view, and interior_to_flat! is the exact mirror. There is no index arithmetic that could reach a ghost. flatten is the one allocating member of the family — it materialises a copy of the interior, reinterpreting SVector{M,T} to component-fastest scalars for a vector field.

5. Two comments in mul! are load-bearing. The operator is applied with α = true, β = false always; the caller’s α and β are fused into the copy-out broadcast instead. And there is deliberately no zero_ghosts! on the way in: the distributed path stages exchanged Interface ghosts into P.xpad between applies, and zeroing would discard them. When β = 0 the output vector is never read, only written, so solvers may hand in an uninitialised buffer.

6. boundary_rhs is the affine lift, and it reaches past the prologue. apply! enforces homogeneous boundaries, so L is linear and inhomogeneous data has to enter the right-hand side: solve L·u = f − b where b is this lift. It builds a zeroed field, fills inhomogeneous ghosts into it — Dirichlet ghost 2·value, Neumann layer k ghost (2k−1)·Δ·flux, periodic and Interface left at zero — and then calls _apply_raw! rather than apply!, because the prologue would overwrite exactly the offsets it just planted. It recurses through the combinators on the algebra: Added sums the two lifts, Scaled scales, Composed uses a(b(x) + c_b) + c_a, and AdjointOp is identically zero because the adjoint action is built homogeneous.

Throws and invalidates. Nothing on a CartesianGrid goes stale — a PreparedOperator has no generation stamp because its grid cannot change shape.

Guard Condition Raises
prepare(L, x) !islinear(L) ArgumentError, naming linearize(L, u0)
prepare(L) one-argument form the operator is bound to no grid ArgumentError
boundary_rhs(L, x) !islinear(L) ArgumentError
_require_prepared_match x/y grid or eltype disagrees with the prepared buffers ArgumentError, four separate checks
mul!(::AbstractVector, ::AbstractOperator, …) called on a raw operator ArgumentError — “would allocate scratch every call; wrap with prepare
operator_diagonal(L) no diagonal declared for L ArgumentError — there is deliberately no fallback

The PreparedForest twin adds a generation stamp and one more guard; it is covered in §7, once the forest exists to explain it.

4 — linearize: a nonlinear operator becomes a Krylov-ready Jacobian

Nonlinear operators support apply! and differentiate fine, but they are not linear maps: adjoint and prepare both refuse them. linearize closes the gap by freezing a state u₀ and returning the Jacobian as an operator, which then satisfies islinear and can be prepared like anything else. The example traced is J = linearize(F, u0) for an F::Advection{…,SelfAdvection}, followed by apply!(y, J, v, g).

Entry point linearize(F, u0[, backend]) -> LinearizedOp, linearize!(J, u) · Source src/operators/linearize.jl, ext/MatrixFreeOperatorsEnzymeExt.jl · Tests test/linearize.jl

  linearize(F, u0)                 ──▶ FiniteDifferenceJVP()   (the default)
  linearize(F, u0, EnzymeJVP())    ──▶ provided by the Enzyme extension
    │                                   otherwise THROWS, naming `using Enzyme`
    ▼
  LinearizedOp(op, u0, shifted, fplus, fminus)      every buffer owned
    │
    ├── apply!(y, J, v, g, α, β)        TWO applies of F, central difference
    │     ε from ‖u₀‖ and ‖v‖
    │     y ← α·(F(u₀+εv) − F(u₀−εv)) / 2ε
    │
    └── adjoint(J) ──▶ THROWS           an FD JVP has no reverse mode

  EnzymeLinearizedOp(op, u0, ustate, ushadow, fout, fshadow)    (extension only)
    ├── apply!         Enzyme.Forward through apply!   — ONE apply, exact
    └── apply_adjoint! Enzyme.Reverse through apply!   — a real transpose
linearize(F::AbstractOperator, u0::Field, ::FiniteDifferenceJVP) -> LinearizedOp
  u0c::Field     = copy(u0)                       # ALLOC 1 — the frozen state is OWNED
  shifted::Field = similar(u0c)                   # ALLOC 1
  fplus::Field   = allocate_output(F, u0c)        # ALLOC 1
  fminus::Field  = allocate_output(F, u0c)        # ALLOC 1
  # HOLDS: every buffer is owned, so apply! allocates nothing in a Newton loop

apply!(y::Field, J::LinearizedOp, v::Field, g::AbstractGrid, α, β)
  T = _scalar_eltype(eltype(J.u0.data))
  normv = norm(v.data)
  ε = iszero(normv) ? sqrt(eps(T)) : sqrt(eps(T)) * (1 + norm(J.u0.data)) / normv
  J.shifted.data .= J.u0.data .+ ε .* v.data      # MUTATES J.shifted
  apply!(J.fplus,  J.op, J.shifted, g)            # MUTATES J.fplus
  J.shifted.data .= J.u0.data .- ε .* v.data
  apply!(J.fminus, J.op, J.shifted, g)            # MUTATES J.fminus
  interior(y) .=/ 2ε) .* (interior(J.fplus) .- interior(J.fminus)) [.+ β .* interior(y)]

linearize!(J, u) = (copyto!(J.u0.data, u.data); J)   # SAME — J's identity survives

1. The Jacobian is an operator, so nothing downstream changes. LinearizedOp declares islinear = true, forwards allocate_output and allocate_input to the operand, reports a size built from its input and output component counts, and resolves its grid from the operand or the frozen state. That is exactly enough to pass prepare’s gate and be handed to Krylov. It declares nothing else, so isconstant, isselfadjoint, isdiagonal and shares_exchange stay false and operator_diagonal throws — which is correct, and means a linearized operator cannot be multigrid-smoothed.

2. The step size is scaled to both magnitudes. ε = √eps·(1 + ‖u₀‖)/‖v‖, falling back to √eps for a zero direction. The formula is a central difference, so the product costs two applications of F and is accurate to roughly √eps — good enough to drive a Newton–Krylov inner solve, not good enough to check a gradient against. Note both norms are taken over the padded array.

3. linearize! refreshes the state in place. A Newton loop wants a new Jacobian each outer iteration and the same buffers; linearize! copies the new state over J.u0.data and returns the identical object, so a PreparedOperator built around J stays valid. This is safe precisely because Krylov solves are never differentiated through.

4. The Enzyme backend is a different struct with a real transpose. EnzymeLinearizedOp carries six buffers instead of five: a primal input copy and an input shadow, a primal output and an output shadow. Forward mode seeds the input shadow with v and differentiates the four-argument apply!, giving an exact JVP in one application. Reverse mode zeroes the input shadow, copies the cotangent into the output shadow and calls zero_ghosts! on it — upholding §2’s “only interior cotangents are inputs” contract — then runs Enzyme.Reverse and blends interior(ushadow) into . It also declares adjoint_operator as the lazy AdjointOp wrapper, so J' works and routes back through that reverse pass.

Throws and invalidates. linearize! invalidates nothing. The rest:

Guard Condition Raises
linearize(F, u0, ::AbstractJVPBackend) EnzymeJVP() without the extension loaded ArgumentError — “run using Enzyme first”
adjoint_operator(::LinearizedOp) any call ArgumentError — FD has no reverse mode; differentiate the operator instead
apply_adjoint!(…, ::LinearizedOp, …) any call ArgumentError — no adjoint action declared
adjoint(F) on the nonlinear F itself !islinear(F) ArgumentError, pointing at linearize
prepare(F) on the nonlinear F itself !islinear(F) ArgumentError, pointing at linearize

Note which gate fires for adjoint(J) on a finite-difference Jacobian: J passes the islinear check, so the message is the specific one about reverse mode, not the generic nonlinear one.

5 — Forest topology: keys, the leaf vector, and the generation counter

Why the package refines whole blocks rather than cells, and why 2:1 balance is the rule everything rests on, is in AMR. This section is the data structure: what a leaf is, what orders the leaf vector, and what exactly the generation counter counts. The example traced is refine!(bf, x -> x[1] < 0.3) on a bf::BlockForest{2,Float64,…} built over a 2×2 root tiling.

Entry point refine!, coarsen!, balance! · Source src/topology.jl, src/BlockForest.jl · Tests test/topology.jl, test/blockforest.jl

  refine!(bf::BlockForest, predicate)          predicate takes a leaf CENTRE
    │                                          (regrid!'s takes a Field — §8)
    ▼
  refine!(forest::Forest, key -> …)
    │
    ├──▶ keys = Set(forest.leaves)
    │      marked leaf below maxlevel:  delete!(keys, key); push! its 2ᴺ children
    │      marked leaf AT maxlevel:     silently ignored
    │
    ├──▶ _set_leaves!(forest, keys)      the ONE commit point
    │      │
    │      ├─ sort by Morton code of the key's finest-level coords
    │      ├─ sorted == forest.leaves ?  ─yes─▶ return.  NO generation bump.
    │      ├─ empty!/append! forest.leaves,  empty!/rebuild forest.index
    │      ├─ forest.uniform[] = all leaves share a level
    │      └─ forest.generation[] += 1          # STALE: every field, every PreparedForest
    │
    └──▶ balance!(forest)                fixpoint; may commit several MORE times
struct LeafKey{N}                  # isbits. Geometry is NEVER stored.
    level::Int
    coords::NTuple{N,Int}
end
parent_key(k)  -> LeafKey(k.level - 1, k.coords .>> 1)
children(k)    -> Vector{LeafKey{N}}     # ALLOC 1, length 2ᴺ, coords .<< 1 .+ offsets

struct Forest{N}
    nroot::NTuple{N,Int}
    periodic::NTuple{N,Bool}
    maxlevel::Int
    leaves::Vector{LeafKey{N}}           # Morton-ordered; storage order of every field
    index::Dict{LeafKey{N},Int}          # key -> position in `leaves`; the exact inverse
    uniform::Base.RefValue{Bool}
    generation::Base.RefValue{Int}
end

struct BlockForest{N,T,BC<:Tuple,Dev} <: AbstractGrid{N}
    forest::Forest{N}
    extent::NTuple{N,Tuple{T,T}}
    spacing0::NTuple{N,T}                # ROOT-level spacing; per-leaf is spacing0 / 2^ℓ
    blocksize::NTuple{N,Int}
    halo::NTuple{N,Int}
    bc::BC                               # PHYSICAL domain BCs — leaves carry Interface
    device::Dev
    schedule::Base.RefValue{…}           # per-generation exchange plan cache (§7)
    schedule_device::Base.RefValue{…}
end

leaf_grid(bf, i) -> CartesianGrid{N,T,,Nothing}
  # every leaf is ONE concrete isbits grid type, bc = Interface on all faces,
  # geometry recomputed from level — never stored per leaf

1. A key is a level and integer coordinates, and nothing else. No extent, no spacing, no parent pointer. parent_key is a right shift, children is a left shift plus the 2ᴺ offsets, and every physical coordinate is recomputed from spacing0 / 2^level. That is what keeps LeafKey isbits and lets the whole topology live in a Vector and a Dict with no pointer chasing.

2. forest.index is the inverse of forest.leaves, and it is the addressing scheme. BlockField.blocks[i] and PackedBlockField.data[…, i] are indexed by that same position i, so the map from key to storage slot is forest.index. Lookup is a raw getindex, so a key that is not a leaf raises KeyError rather than a domain error. The ordering is the Morton code of the key’s coordinates projected to maxlevel, which is why refining one block renumbers its neighbours — and why §8’s transfer must key by LeafKey rather than by integer position.

3. _set_leaves! is the only function that writes either container, and the only one that bumps the generation. It sorts a fresh key vector by Morton code, compares it to the current leaves, and returns early if they are identical — no bump, no index rebuild. Otherwise it empties and refills both containers in place, refreshes uniform, and increments generation. The early return is what makes §8’s “nothing changed, so the same field objects come back and prepared operators stay valid” exact rather than approximate.

4. refine! and coarsen! are set edits; balance! is a fixpoint. refine! swaps each marked leaf for its children, ignoring marks already at maxlevel. coarsen! groups leaves by parent and collapses a family only when all 2ᴺ siblings are present and all are marked — an incomplete or partly marked family is untouched, and level-0 leaves are skipped. balance! then loops until no face violates 2:1: for each leaf, dimension and side it finds the face neighbour (wrapping when that dimension is periodic), resolves which leaf actually covers that region, and refines the cover if it is two or more levels coarser. Each pass is its own _set_leaves! commit, so a single refine! call can bump the generation several times — code that snapshots a generation must compare, not count.

5. Every leaf is one concrete grid type, and that is a performance decision. leaf_grid builds a CartesianGrid through the unchecked inner constructor with Interface() on all 2N faces and a local range of 1:blocksize[d]. Per leaf geometry is derived from the level, never stored. Because the type is identical for every leaf and isbits, the per-leaf sweep in §7 is type-stable and rebuilding the grid inside the loop is free. It also means per-leaf BC sweeps are no-ops — physical boundaries live on the forest and are applied by a separate face pass.

Throws and invalidates. Every commit that changes the leaf set invalidates every BlockField and every PreparedForest built on the old set; both carry a stamp and throw ArgumentError on next use (§6, §7).

Guard Condition Raises
Forest constructor the root tiling and maxlevel exceed 64-bit Morton capacity ArgumentError
BlockForest constructor blocksize < 1, base size not divisible by it, or maxlevel < 0 ArgumentError
BlockForest constructor a face BC that is not Periodic/Dirichlet/Neumann — this is what rejects a user-supplied Interface ArgumentError
leaf_index(forest, key) the key is not a current leaf KeyError

6 — Block field layouts: one contract, two storages

A field on a forest has to be two things at once: a list of per-leaf padded arrays, which is what the reference CPU path and every regrid want, and one contiguous buffer with leaves on a trailing dimension, which is what a single GPU kernel launch wants. Both satisfy the same contract, and which one you have is settled at compile time. The example traced is up = pack(u) for a u::BlockField{Center,Interpolated,Matrix{Float64},…}.

Entry point pack, unpack, with_transfer · Source src/blockfield.jl, src/packedfield.jl · Tests test/blockfield.jl, test/packedfield.jl, test/forest_packed.jl

  AbstractBlockField                    everything layout-agnostic is written
    │                                   ONCE against this: set!, zero_ghosts!,
    ├── BlockField                      flatten, block(), the five forest sweeps
    │     blocks::Vector{A}    ── _storage ──▶ Vector of padded arrays
    │                          ── _layout  ──▶ BlocksLayout()      singleton
    │
    └── PackedBlockField
          data::A             ── _storage ──▶ one (blocksize.+2h…, nleaves) array
          levels::V                            device-resident per-leaf level
                              ── _layout  ──▶ PackedLayout()       singleton

  the ONLY two functions that know the difference:
    _leaf_array(store, ::BlocksLayout, i)        = store[i]
    _leaf_array(store, ::PackedLayout,  i)       = view(store, :, …, i)
    _leaf_view(store, ::BlocksLayout, i, ranges) = view(store[i], ranges...)
    _leaf_view(store, ::PackedLayout,  i, ranges) = view(store, ranges..., i)
struct BlockField{L,P<:RegridTransferPolicy,A<:AbstractArray,G<:BlockForest}
        <: AbstractBlockField
    blocks::Vector{A}
    grid::G
    generation::Int          # stamped from grid.forest.generation[] at construction
end

struct PackedBlockField{L,P<:RegridTransferPolicy,A,V<:AbstractVector{Int},
                        G<:BlockForest} <: AbstractBlockField
    data::A                  # leaves on the TRAILING dimension, Morton order
    levels::V                # so a kernel can recompute spacing without the Forest
    grid::G
    generation::Int
end

pack(f::BlockField{L,P}) -> PackedBlockField{L,P}
  _require_current(f)                                   # THROWS ArgumentError if stale
  data   = similar(first(f.blocks), eltype(f), (psize..., nleaves(bf)))   # ALLOC 1
  levels = _leaf_levels(bf)                             # ALLOC 1, host Vector -> device
  for i in 1:nleaves(bf)
      _block_array(p, i) .= f.blocks[i]                 # PADDED storage verbatim —
  end                                                   # ghosts copied, not recomputed
  # HOLDS: L (location) and P (transfer policy) round-trip through pack/unpack

1. The layout is not a field, it is the type. There is no layout::Symbol anywhere. The information is carried two ways, both resolved at compile time: the concrete struct identity, and a zero-size tag value whose type travels through dispatch. That is a deliberate constraint, and the decisive reason is Enzyme: from Julia 1.12 a custom-rule argument may not mix GC-tracked pointers with inline floats, and a BlockField does exactly that through its embedded BlockForest. So the differentiated seams (§7) take raw storage plus a singleton tag, and a runtime layout flag would force the field struct back into the rule signature.

2. Three more things fall out of the same choice. The tag being a singleton means _leaf_array and _leaf_view each resolve to one concrete SubArray type, so the descriptor loops in §7 stay dispatch-free and allocation-free. Layout selection for the GPU kernel sweeps is a method table rather than a branch — _forest_sweep! has a PackedBlockField method per operator, falling back to the per-leaf path otherwise. And Adapt.adapt_structure rebuilds the same parameterised type, so the layout survives device adaptation because it is the type.

3. pack copies padded storage verbatim, ghosts included. That is not laziness: a Diffusion coefficient field has its ghosts filled at construction by fill_coefficient_ghosts!, and copying the padded array means prepare can normalise a coefficient’s layout without re-running an exchange. unpack is the mirror. The location parameter L and the transfer policy P survive both directions — a policy silently reset by a round trip would conserve on one regrid and not the next.

4. The transfer policy is a type parameter for the same reason. Interpolated, Conservative and SlopeLimited are empty singletons, and RegridTransferPolicy is a closed Union of the three rather than an abstract type. §8 resolves the policy once per field from P and dispatches the transfer kernel on it. with_transfer rebuilds the wrapper around the same storage and the same generation stamp, copying nothing.

5. One rule for contributors, and it is the one that bites. Never index block storage outside _leaf_array / _leaf_view. The two layouts produce SubArrays over parents of different rank; every descriptor sweep in src/transfer.jl and the coarse–fine rewrite in src/operators/diffusion.jl are written so that only those two functions know it. Reaching into f.blocks[i] or f.data[…, i] directly compiles, runs, and silently breaks the other layout.

Throws and invalidates. _require_current(f) compares f.generation against f.grid.forest.generation[] and throws ArgumentError — “allocated before the forest was regridded … allocate a fresh field on the current forest” — on any mismatch. It is called from block, from pack and unpack, from the packed flat transfers, and from all five forest sweeps, so there is no path that reads stale storage without passing a guard.

7 — Forest apply!: one exchange, one BC pass, then ordinary stencils

Operators barely change on a forest, and AMR says why. This section is the machinery that makes it true: a precomputed descriptor table, rebuilt only when the generation changes, that turns every inter-block ghost fill into a flat loop with no topology queries in it. The example traced is mul!(y, P, x) for a P::PreparedForest wrapping laplacian(bf) + scaling(σ) on a non-uniform forest.

Entry point halo_update!, apply_bc!, prepare(L, ::AbstractBlockField) -> PreparedForest · Source src/schedule.jl, src/transfer.jl, src/transfer_kernels.jl, src/operators/forest.jl, src/operators/forest_packed.jl · Tests test/exchange_schedule.jl, test/exchange_kernels.jl, test/forest_parity.jl, test/forest_prepare.jl

  apply!(y, L, x, bf, α, β)
    │
    ├─ shares_exchange(L) ? ──yes──▶ ONE exchange, then sweep both operands
    │                                 (Added/Scaled recurse BEHIND the fills)
    ▼
  ┌──────────────────────────────────────────────────┐
  │ halo_update!(x, bf)                              │  ONCE, whole forest
  │   _exchange_schedule(bf)   generation match?     │
  │     └─ miss ─▶ rebuild, cache in bf.schedule[]   │
  │   _exchange_storage!(_storage(x), _layout(x), s) │  ◀── the Enzyme seam
  │     1. copies    same-level face slabs           │
  │     2. interp    coarse ─▶ fine   (quadratic)    │  order is HARD:
  │     3. restrict  fine ─▶ coarse  (flux match)    │  restrict READS interp's output
  └──────────────────────────────────────────────────┘
    │
  ┌──────────────────────────────────────────────────┐
  │ apply_bc!(x, bf)                                 │  ONCE, whole forest
  │   _bc_storage!(store, layout, bc, bcfaces, h, n) │  ◀── the Enzyme seam
  │     per-(dim, side) leaf-index lists, dims 1:N   │  reuses the SINGLE-GRID
  └──────────────────────────────────────────────────┘  _fill_ghost! primitives
    │
    ├───────────────┬───────────────┐
    ▼               ▼               ▼
  per leaf:       per leaf:       per leaf:       ← the EXISTING single-grid
  stencil         stencil         stencil           apply! from §1, unchanged
struct ExchangeSchedule{N,T}
    copies::Vector{CopyDescriptor{N}}      # src/dst leaf + src/dst index ranges
    interp::Vector{GhostFill{N,T}}         # dst + Vector{SlabTerm} (block, ranges, weight)
    restrict::Vector{GhostFill{N,T}}
    cfflux::Vector{CFFluxDescriptor{N}}    # isbits: Diffusion's coarse–fine rewrite
    bcfaces::NTuple{N,NTuple{2,Vector{Int}}}   # leaf indices touching each physical face
    generation::Int                        # -1 sentinel: the first fetch always builds
end

prepare(L, x::AbstractBlockField) -> PreparedForest
  islinear(L) || throw(ArgumentError(...))
  xpad, ypad  = similar(x) |> zero_ghosts!, allocate_output(L, x)   # ALLOC 2
  op          = _prepare_tree(L, x)         # also packs coefficient layouts to match x
  _exchange_schedule(x.grid)                # ALLOC — warms the per-generation plan
  PreparedForest(op, x.grid, xpad, ypad,
                 similar(x),                # ALLOC — adjscratch: the accumulating
                 x.grid.forest.generation[])#         adjoint sweep's buffer

mul!(y::AbstractVector, P::PreparedForest, x::AbstractVector, α, β)
  _require_prepared_current(P)              # THROWS ArgumentError if the forest moved
  flat_to_interior!(P.xpad, x); _forest_capply!(...); interior_to_flat!(y, P.ypad, α, β)

1. The schedule is a cache keyed on the generation, and it never throws on a miss. _exchange_schedule(bf) compares the stored plan’s generation against the live one and rebuilds on mismatch, writing back into the forest’s Ref. A fresh forest holds a sentinel plan stamped -1, which can never match, so the first fetch always builds. prepare warms it eagerly so the first mul! in a Krylov loop is not the one paying for construction. There is a second, lazily built device twin, flattened into isbits rows plus a CSR-style term array, for the packed GPU path.

2. Building the schedule is one pass over (leaf, dimension, side). Each face falls into exactly one of four cases: no neighbour, so push this leaf’s index onto the physical-BC face list; a same-level leaf, so emit one CopyDescriptor over the face slabs; a coarser covering leaf, so emit coarse→fine GhostFills; or no covering leaf at all, meaning finer children, so emit fine→coarse GhostFills and the matching CFFluxDescriptors. The invariant that makes the adjoint simple is that each ghost region is the destination of exactly one descriptor across all three vectors — so the transpose is a plain scatter-add followed by a zero, with no double counting to reason about.

3. The two coarse–fine fills are different maps, and the order between them is not negotiable. Coarse→fine is a quadratic Martin–Cartwright interpolation: normal weights from a three-point Lagrange fit at the child’s quarter-cell offset, tangentially tensored with per-parity classes evaluated at ξ = ±1/4. Every term reads block interiors only, so it is order-independent and never touches a BC ghost. Fine→coarse is flux matching, not volume averaging — a plain 2ᴺ average leaves an O(1) truncation error at the interface and was rejected — and it reads the fine children’s interpolation-filled ghosts. That is why interp must run before restrict, always.

4. The physical BC pass reuses the single-grid primitives verbatim. It walks the per-face leaf-index lists and calls the same _fill_ghost!, _source_low, _source_high and _bc_sign helpers §1 uses, in the same ascending dimension order. That is what keeps layer indexing, corner order and mirror signs bitwise identical to a single grid — which is what makes the parity test in test/forest_parity.jl meaningful. Periodic dimensions have empty face lists; their wrap is the exchange’s business.

5. Every forward sweep has a declared transpose, and they are paired by construction.

Forward Transpose Shape of the transpose
_exchange_storage! _exchange_storage_adjoint! phases reversed: restrict, interp, copies
_run_copies! _run_copies_adjoint! descriptors reversed; src += ghost, then zero the ghost
_run_fills! _run_fills_adjoint! fills reversed; every term gets term += weight · dst, then zero
_bc_storage! _bc_storage_adjoint! recurses first, so dimensions run N:1 against the fill’s 1:N
_cf_flux_rewrite! _cf_flux_rewrite_adjoint! descriptors reversed, conjugated coefficient averages

The full forest adjoint is then: per-leaf transpose gathers leaving cotangents in ghosts, then fold_bc! over the face lists, then halo_update_adjoint! folding interface ghosts into neighbours’ interiors — §2’s single-grid order, lifted whole. Two shortcuts precede it: a self-adjoint operator on a uniform forest reuses the forward sweep, and a diagonal operator skips the fold entirely, because a pointwise transpose must not leak ’s ghost scratch into interiors. The adjoint exchange stays on host descriptor loops on every backend — its scatter-adds collide on shared source cells, so a kernel would need atomics and stop being bitwise reproducible.

6. shares_exchange is what turns two sweeps into one. An operator declares it when its forest action is one stencil sweep over an already-exchanged input: it reads x’s interiors and ghosts, writes only y’s interior, and never writes into x. When an Added reports it, apply! runs a single halo_update! and apply_bc! and then recurses into a _forest_sweep! combinator method behind the fills, so both operands sweep the one exchanged input. Those combinator sweep methods are reachable only through that gate — an ungated operand could rewrite x’s ghosts under its sibling. Diffusion declares it conditionally, true only on a uniform forest, because the coarse–fine flux rewrite writes into x’s ghosts and so violates the contract the moment the forest refines. A forgotten declaration costs a redundant exchange; it never costs correctness.

7. The differentiated seams are shaped by an Enzyme constraint, not by taste. _exchange_storage! and _bc_storage! take raw storage, a singleton layout tag, and isbits-or-Const descriptors — never a field, never a grid. Automatic Differentiation covers why they get custom rules at all and why those rules cannot swallow a parameter gradient. What matters here is the argument shape: _bc_storage! receives the grid already shredded into its BC tuple, face lists, halo and block size, because passing g would put a GC-pointer-and-inline-float mixture into a rule signature. Keeping the seam at this level rather than at halo_update! is also what lets one rule serve both layouts. CFFluxDescriptor was made fully isbits for the opposite reason — it must be taped, because coefficient gradients flow through it.

Throws and invalidates. A regrid invalidates a PreparedForest; a stale one throws on mul! and on apply!, the two public entries.

Guard Condition Raises
_require_prepared_current(P) P.generation != P.grid.forest.generation[] ArgumentError — re-run prepare
_require_current(f) field stamp mismatch, checked in all five sweeps ArgumentError
_throw_unbalanced(K) schedule build meets a face violating 2:1 ArgumentError — run balance!
_validate_coarse_fine(bf) non-uniform forest with a halo other than 1, or a block size too small or odd ArgumentError
diffusion(g, κ) built directly on a grid with Interface faces ArgumentError — build on the whole forest with a BlockField coefficient

8 — regrid!: a topology edit, then data motion keyed by identity

When to regrid, how to choose an indicator, and why a conserved field wants Conservative rather than the default are in AMR. This section is the sequence: what is snapshotted before the topology moves, why the transfer is keyed by LeafKey and never by position, and exactly which objects come out alive. The example traced is η, u = regrid!(η, u; refine = b -> maximum(interior(b)) > τ).

Entry point regrid! · Source src/amr.jl · Tests test/amr_driver.jl, test/forest_amr.jl

  regrid!(u, more...; refine, coarsen)
    │
    ├──▶ _regrid_marks       each leaf handed to the predicate as an ordinary Field
    │                         (a VIEW of the block — no copy)
    │
    ├──▶ gen0      = forest.generation[]          snapshot BEFORE mutating
    ├──▶ old_index = copy(forest.index)           _set_leaves! empties it in place
    ├──▶ old_blocks = f.blocks                    by REFERENCE — topology edits
    │                                              never touch field storage
    ├──▶ _regrid_topology!   ONE combined split+collapse commit, then balance!
    │
    ├──▶ generation unchanged? ──yes──▶ return the SAME field objects.
    │                                    prepared operators stay valid. DONE.
    ▼ no
  ┌───────────────────────────────────────────────────────────────┐
  │ per field:  _fresh_field   -> zeroed BlockField on the NEW    │
  │                               leaf set, policy P preserved    │
  │             _transfer!     -> per new leaf, look up old_index │
  │                unchanged leaf ─▶ copy the interior            │
  │                new child      ─▶ prolong from the old PARENT  │
  │                coarsened leaf ─▶ 2⁻ᴺ mean of the old CHILDREN │
  └───────────────────────────────────────────────────────────────┘
    │
    └──▶ ghosts left ZERO (scratch, refilled by §7 on the next apply)
         every old field and every PreparedForest now throws (§6, §7)
regrid!(u::BlockField, more::BlockField...; refine, coarsen = Returns(false))
  bf::BlockForest{N}  = u.grid
  forest::Forest{N}   = bf.forest
  _require_current(u); for f in more; f.grid === bf || THROW; _require_current(f); end
  all(iseven, bf.blocksize) || THROW               # _transfer_average! halves n
  refine_marks, coarsen_marks = _regrid_marks(u, bf, refine, coarsen)
                                      # -> Set{LeafKey{N}} x 2, DISJOINT by construction:
                                      #    refine wins, level-0 is never coarsen-marked
  gen0                = forest.generation[]
  old_index           = copy(forest.index)         # ALLOC 1 — Dict{LeafKey{N},Int}
  old_blocks          = map(f -> f.blocks, fields) # by reference, no copy
  _regrid_topology!(forest, refine_marks, coarsen_marks)
                                      # MUTATES forest: splits, collapses complete
                                      # families, ONE _set_leaves!, then balance!
  forest.generation[] == gen0 && return isempty(more) ? u : fields     # SAME
  out = map(fields, old_blocks) do f, blocks
      _transfer!(_fresh_field(f, bf), blocks, old_index, bf)           # ALLOC per leaf
  end
  # STALE: every field and every PreparedForest built on the old leaf set
  return isempty(more) ? out[1] : out

1. The predicates see a Field, not a coordinate. _regrid_marks builds block(u, i, leaf_grid(bf, i)) per leaf — an ordinary Field sharing the block’s storage, no copy — so a criterion is written as a reduction like b -> maximum(abs, interior(b)) > τ. This is the one place where a forest predicate differs from refine!’s, which takes a block centre. The two mark sets are disjoint without a check: a refine-marked leaf is never offered to coarsen, and level-0 leaves are never coarsen-marked at all.

2. Three things are snapshotted, for three different reasons. gen0 is the comparison that detects a no-op. old_index must be a copy because _set_leaves! empties the live dictionary in place. old_blocks is captured by reference and needs no copy at all, because a topology edit never touches field storage — the arrays outlive the leaf vector that used to name them.

3. The topology edit is one commit, not two. _regrid_topology! splits refine-marked leaves and collapses complete coarsen-marked families in a single pass into one key set, then does one _set_leaves!, then balance!. Calling refine! and coarsen! in sequence would work but would commit — and bump the generation — two or three times for no reason. Because refine-marked siblings are never coarsen-marked, the “all 2ᴺ siblings present and marked” count test already implements the refine-wins rule; there is no explicit conflict resolution anywhere.

4. The early return is exact, and it is the convergence signal. If the marks changed nothing, _set_leaves!’s identity check declined to bump the generation, and regrid! hands back the same objects===, not merely equal. So a prepared operator built before the call is still valid, and an adaptive loop can use “regrid! returned what I gave it” as its stopping test.

5. The transfer is keyed by LeafKey, and that is the whole design. Storage positions are Morton indices into a vector that was just re-sorted, so they mean nothing across the edit. For each new leaf K, _transfer! asks old_index three questions in order: was K itself a leaf before (copy its interior); was parent_key(K) a leaf before (prolong from the parent into K’s octant); are all of children(K) old leaves (average them in). Interiors only — old ghosts are never read, and the new field’s ghosts stay zero until §7’s next exchange. The transfer policy is read once per field from its type parameter, so the per-leaf loop dispatches on a singleton rather than branching.

6. The three kernels, and which weights each uses. Copy is a view-to-view assignment over the interior box. Prolong under the default Interpolated is per-dimension linear interpolation with weights (3/4, 1/4), flipping to a one-sided (5/4, −1/4) extrapolation where the quarter-weight tap would fall outside the parent’s interior; it is exact on linear data but not mean-preserving. Prolong under Conservative or SlopeLimited is instead a cell-conservative reconstruction u_child = u_parent + Σ_d ξ_d·σ_d with ξ_d = ∓1/4, and the reason it conserves is that all 2ᴺ children share one slope per parent cell per dimension, so their mean telescopes back to the parent exactly. The two differ only in the slope: a centred difference for Conservative, a minmod limiter for SlopeLimited — which also injects rather than reconstructs on every parent-block-face cell. Coarsening is always the 2⁻ᴺ arithmetic mean of the children, regardless of policy.

Throws and invalidates. Everything built on the old leaf set is now stale and throws through the §6 and §7 guards. regrid! itself:

Guard Condition Raises
field/forest check a field in more is on a different BlockForest ArgumentError
_require_current any input field is already stale ArgumentError
block-size check an odd block size in any dimension ArgumentError — the 2⁻ᴺ mean halves it
layout check any argument is a PackedBlockField ArgumentError — unpack, regrid, re-pack
_transfer! fallthrough a new leaf has no old leaf within one level ErrorException — a bug assert, not a user error

Because a real regrid invalidates every field on the old leaf set, anything needed afterwards must ride the same call. Compute indicator fields before regridding and pass them first, which is why the traced example puts η ahead of u.

9 — Multigrid: a tuple of levels, and one recursive V-cycle

Why multigrid breaks the link between grid size and iteration count, and how the transfer pair and the smoothers fall out of machinery that already exists, is in Geometric Multigrid. This section is the hierarchy’s representation and the cycle’s control flow. The example traced solves \(-\Delta u = f\) on a 256×256 grid with homogeneous Dirichlet boundaries. The minus sign makes the operator positive definite, as CG requires; both the prepared operator and its multigrid preconditioner are built from that same L.

using MatrixFreeOperators, Krylov

g = CartesianGrid(((0.0, 1.0), (0.0, 1.0)), (256, 256))  # homogeneous Dirichlet by default
L = -laplacian(g)
P = prepare(L)
M = MultigridPreconditioner(L)
b = flatten(set!(scalar_field(g), x -> 2π^2 * sinpi(x[1]) * sinpi(x[2])))
u, stats = Krylov.cg(P, b; M)

Entry point MultigridPreconditioner, mul!(z, M, r) · Source src/multigrid.jl, src/operators/restriction.jl, src/operators/prolongation.jl · Tests test/multigrid.jl

  MultigridPreconditioner(L)
    │
    ├──▶ _mg_grids(g, :auto)        coarsen while even, > 64 cells, < 10 levels
    │
    ├──▶ _build_levels              per level, recursively:
    │      A = _prepare_tree(L, proto)      the buffer-hoisted tree from §3
    │      R, P = restriction(…), prolongation(…)
    │      x, b, r, t = four scalar_fields
    │      smoother state from operator_diagonal(L)
    │      then recurse on _rediscretize(L, coarser)   ← NOT a Galerkin RAP
    │
    └──▶ _build_coarsest            materialise column by column, lu(), keep the
                                     factorization and a flat scratch vector

  mul!(z, M, r)                     Krylov's preconditioner convention: z ≈ L⁻¹r
    ├─ flat_to_interior!(top.b, r)
    ├─ _vcycle!(M.levels)           Tuple recursion — fully specialised per level
    └─ interior_to_flat!(z, top.x, α, β)

  _vcycle!((lvl, rest...))          _vcycle!((coarsest,))
    smooth FROM ZERO                  interior_to_flat! -> ldiv!(fact, bflat)
    r ← b − A·x                                        -> flat_to_interior!
    next.b ← R·r
    _vcycle!(rest)
    x ← x + P·next.x                (α=true, β=true — prolongate and ACCUMULATE)
    smooth again
struct MGLevel{O,RO<:Restriction,PO<:Prolongation,G<:CartesianGrid,F<:Field,S}
    A::O          # _prepare_tree'd rediscretised operator — NOT a PreparedOperator
    R::RO; P::PO  # this level -> coarser, coarser -> this level
    g::G
    x::F; b::F; r::F; t::F      # correction, RHS, residual, spare — all preallocated
    smoother::S
end
struct MultigridPreconditioner{Lv<:Tuple}
    levels::Lv    # (MGLevel, MGLevel, …, MGCoarsest) — a TUPLE, not a Vector
end

1. The hierarchy is a Tuple on purpose. _vcycle! recurses with first/Base.tail, so with a tuple every level is a distinct type and the whole cycle specialises into straight-line code with no dynamic dispatch and no allocation. A Vector{MGLevel} would make the recursion type-unstable at every step, in the hottest loop the preconditioner has.

2. Restriction and Prolongation are a declared adjoint pair, not two independent stencils. Prolongation is per-dimension linear interpolation on the cell-centred 2:1 pair, weights (3/4, 1/4). Restriction’s forward action literally is the scaled transpose gather 2⁻ᴺ·Pᵀ — the same kernel, called from the other side — so ⟨Px, y⟩ = ⟨x, Pᵀy⟩ holds to machine precision by construction rather than by agreement. Their declared adjoints return each other, scaled: Rᵀ = 2⁻ᴺ·P and Pᵀ = 2ᴺ·R. Note Restriction’s forward path runs no prologue: the gather engine from §2 zeroes its input’s ghosts first, so only fine interior values contribute and boundary_rhs(R, ·) is identically zero.

3. Coarse operators are rediscretised, not assembled. _rediscretize walks the tree and rebuilds each leaf on the coarse grid — Laplacian(gc), Diffusion with an averaged coefficient, ScalingOp with an averaged coefficient, structural recursion through Scaled/Added/Composed — and throws on anything it does not know. Coefficient averaging uses a plain 2ᴺ-child arithmetic mean, deliberately not Restriction: full weighting folds Dirichlet’s −1 mirror through the boundary, which is right for a residual and wrong for a material coefficient at a wall.

4. Smoothers are built from the operator’s own diagonal. operator_diagonal returns a Number where the diagonal is uniform and a Field otherwise, with the boundary mirror fills’ contribution included, and throws where no diagonal is declared — there is no generic fallback, so an operator without one errors rather than smoothing with a wrong diagonal. Jacobi stores ω, a sweep count and the inverted diagonal. Chebyshev additionally runs ten power iterations on D⁻¹A at setup, seeded with a checkerboard, to bracket the spectrum; the two scratch fields that costs are discarded afterwards.

5. The default Jacobi smoother skips a separate zeroing pass. The pre-smooth runs “from zero”: its first sweep writes x rather than updating it, which both skips a residual evaluation and removes the need for a separate zeroing pass. Chebyshev’s from-zero path explicitly clears x before its recurrence. The coarse correction comes back through apply!(lvl.x, lvl.P, nxt.x, lvl.g, true, true) — prolongate and accumulate, using §1’s β ≠ 0 form.

Pre- and post-smoothing use the same fixed operator, making the cycle symmetric for symmetric L. Symmetry alone does not establish CG compatibility: Krylov’s CG contract requires both the system and preconditioner to be positive definite. The Dirichlet \(-\Delta\) example with the default damped Jacobi smoother satisfies those conditions; test/multigrid.jl checks the preconditioner’s symmetry and positive eigenvalues on a small grid.

6. Everything is preallocated at construction; mul! allocates nothing. The construction cost includes the coarsest level’s dense Matrix and its lu factorization — the only place in the package where a matrix is assembled, and only ever at a level small enough that levels=:auto stopped coarsening. The standalone MultigridSolver wraps a real PreparedOperator plus two flat vectors; its solve allocates only the returned iterate, and on exhausting maxiter it warns and returns the current iterate rather than throwing.

Throws and invalidates. Nothing here goes stale; the hierarchy is immutable once built.

Guard Condition Raises
MultigridPreconditioner(L) !islinear(L) ArgumentError
MultigridPreconditioner(L) the grid is a BlockForest ArgumentError — forest multigrid is not supported
MultigridPreconditioner(L) cycle other than :V ArgumentError
_mg_grids :auto found nothing coarsenable, or levels < 2 ArgumentError
_build_coarsest the coarse operator is singular — an all-Neumann or all-periodic Poisson has a constant nullspace ArgumentError — add a Dirichlet face or a zeroth-order term
_rediscretize an unsupported leaf ArgumentError, listing the supported set
_validate_transfer grids not 2:1, extents differing, Interface faces, mismatched BC kinds, different devices five separate ArgumentErrors

10 — Distributed: slabs, and one exchange in the middle of the tree

What a partition model is, why a cut face is a boundary condition, and why the list of distributable operators is deliberately short are in Distributed Multi-GPU Solves. This section is where the guards run, how a global operator becomes a per-slab one, and why the walk over the tree exists at all. The example traced is P = prepare_distributed(laplacian(g) - scaling(σ), 4).

Entry point partition_grid, prepare_distributed, local_grids · Source src/partitioning.jl, src/distributed.jl, ext/MatrixFreeOperatorsMDLAExt.jl · Tests test/partitioning.jl, test/multigpu/

  prepare_distributed(L0, nparts)
    │
    ├──▶ L = _push_adjoints(L0)         normalise AdjointOp DOWN to leaves FIRST,
    │                                   so the guards only ever see leaf adjoints
    ├──▶ _check_distributable(L)        ◀── on the GLOBAL tree
    ├──▶ _check_one_grid(L, g)          ◀── on the GLOBAL tree
    │      both MUST run here: after localisation a ScalingOp reports the SLAB
    │      grid while its Laplacian sibling still reports the global one, so a
    │      re-check of a localised tree would reject a valid operator
    │
    ├──▶ partition_grid(g, nparts)  -> Vector{CartesianGrid}
    │
    ├──▶ per partition:  prepare(adapt(CuArray, _slab_op(Lh, locals[d])), …)
    │                     ◀── localisation happens HERE, after the guards
    │
    └──▶ _dist_tree(…)              a DistNode mirror of the prepared tree

  mul!(y, P, x, α, β)
    ├──▶ _root_scatter!                 fill every slab's Interface ghosts
    ├──▶ _dist_capply!(ypads, tree, xpads, ctx, true, false)
    │      DistComposed ─▶ apply b, THEN _dist_scatter!, THEN apply a
    │                      the exchange fires in the MIDDLE of the tree
    └──▶ interior_to_flat! per device   α and β applied at the flat boundary,
                                        never across a collective
partition_grid(g::CartesianGrid{N,T}, nparts) -> Vector{CartesianGrid{N}}
  # A slab is a plain CartesianGrid. Per slab:
  extent      = g.extent            # the GLOBAL extent, verbatim
  spacing     = g.spacing           # copied, never recomputed from the slab span
  size        = length(zr) in dim N, global elsewhere
  bc[N]       = (p > 1 || periodic ? Interface() : bc_global[N][1], )
  local_range = ntuple(d -> d == N ? zr : g.local_range[d], Val(N))
  topology    = nothing
  # HOLDS: extent describes the DOMAIN, local_range describes OWNERSHIP.
  #        cell_center evaluates at the GLOBAL cell index:
  #          z = first(local_range[d]) - 1 + I[d] - halo[d]
  #          extent[d][1] + (z - 0.5) * spacing[d]

_slab_op(L, lg)                     # localisation: only coefficient-bearing leaves move
  ScalingOp{<:Field} -> ScalingOp(_slab_field(S.coeff, lg))
  Diffusion          -> Diffusion(lg, _slab_field(D.κ, lg), D.avg)   # INNER constructor
  Added/Scaled/Composed/AdjointOp -> structural recursion
  anything else      -> SAME object, every partition                 # narrows DistLeaf.ops

_slab_field(f::Field{L}, lg) -> Field{L}
  win = one PADDED window per dim, offset by the slab's local_range
  Field{L}(copy(view(f.data, win...)), lg)      # ALLOC 1 — a RESTRICTION, never a
                                                # re-evaluation, so no coordinate drift

1. A slab keeps the global extent, and that is not an oversight. Extent describes the domain; local_range describes ownership; cell_center reconstructs a global cell index from the two and evaluates there. Rewriting a slab’s extent to describe its own span would round twice — once for the slab origin and once for the cell offset — and drift by an ulp, which would make a coordinate-assembled right-hand side depend on how many partitions you happened to use. The same reasoning is why spacing is copied rather than recomputed.

2. A cut face is an Interface, and halo_update! stays a no-op. Each partition’s cut-dimension faces are marked Interface(), which §1’s BC sweep skips and §2’s fold skips. Filling them is the driver’s job, not the grid’s: the exchange stages a neighbour’s planes into the prepared input buffer between applies. Under a periodic cut both cut faces are Interface on every partition, and a two-partition periodic split therefore needs at least 2h planes per slab rather than h.

3. The guards live in core and run once, on the global tree. _distributable is a whitelist predicate whose AbstractOperator fallback is false, so an operator that was never considered is refused rather than silently mishandled. _check_one_grid covers what a per-operator predicate structurally cannot: a mismatched pair. Both run before _slab_op, and that ordering is the point — after localisation the tree is deliberately inconsistent, with localised leaves reporting slab grids and grid-free leaves reporting the global one. Both guards live in src/distributed.jl rather than the extension so they are tested on CI, which has no GPU.

4. _push_adjoints runs before the guards, for the same reason it runs inside prepare (§3). It folds AdjointOp down to leaves, so _distributable only ever has to reason about leaf-level adjoints, and the composition structure a distributed adjoint needs its reduction inside is explicit in the tree.

5. The distributed tree exists because of one node type. DistLeaf, DistAdded and DistScaled are bookkeeping; DistComposed is the reason for the walk. A composition’s intermediate lives on each slab with unfilled Interface ghosts, so the exchange has to fire between the two factors — mid-tree, not at the boundary. _reads_ghosts decides whether that exchange is needed at all, and it must gate the forward scatter and the adjoint reduce together, or the adjoint identity breaks. DistAdjoint carries its own input copy, and that is not an optimisation: §2’s gather opens with zero_ghosts! on its input, so sharing the enclosing segment’s field would destroy exchanged ghosts a sibling under the same Added still needs — making the result depend on term order, silently.

6. Two shortcuts that hold elsewhere are deliberately not taken here. The distributed adjoint walk skips the isselfadjoint and isdiagonal fast paths, because _selfadjoint_grid reports true for any CartesianGrid — including a slab whose cut faces are Interface, where the claim is false. And _dist_boundary_rhs! calls _apply_raw! rather than apply!, the same reason as §3: the prologue would overwrite the inhomogeneous ghost offsets. Finally, local_grids returns host grids on purpose — a field allocated on a device grid would land on whichever device happened to be current, and assemble_rhs would then read it from a different one.

Throws and invalidates. Nothing goes stale; every refusal is up front.

Guard Condition Raises
partition_grid the grid is already distributed, nparts < 1, or a slab is thinner than the halo ArgumentError
_check_distributable any operator outside the whitelist, or a complex / already-distributed coefficient ArgumentError, naming the culprit and why
_check_one_grid a grid-bearing operator is on a different grid from the one being partitioned ArgumentError
prepare_distributed the grid is not a CartesianGrid, or nparts exceeds the available devices ArgumentError
_node_exchange a non-scalar intermediate — a rank changer would need its own partition spec ArgumentError
assemble_rhs wrong number of source fields, or one sized against the wrong slab ArgumentError

Rules for changing this code

Each of these is an obligation the type system cannot enforce, paired with the test that catches you if you skip it.

  1. A new leaf owes four checks, and its adjoint owes the boundary. Action against an analytic solution, the adjoint identity ⟨Lx,y⟩ = ⟨x,Lᵀy⟩, composition laws, and AD gradients against finite differences — for the field and for operator parameters. Boundary conditions break self-adjointness even for the Laplacian, so the adjoint must be declared including its boundary contribution; do not infer one. test/test_utils.jl has fd_gradient and materialize, which densifies a prepared operator on a small grid so structure can be checked exactly. Caught by the per-operator file, e.g. test/laplacian.jl, test/diffusion.jl.

  2. A new wrapper must preserve the traits its interface exposes. The lazy combinators declare all five traits, including explicit false declarations where propagation is unsafe. The internal PreparedComposed and PreparedAdjoint nodes declare the four algebraic traits and use the inherited shares_exchange = false default. The solver wrappers PreparedOperator and PreparedForest forward only the four algebraic traits; they have no AbstractOperator fallback and expose no shares_exchange method. Forgetting islinear on an internal node hides a linear tree from consumers; forgetting a method on a solver wrapper raises MethodError. Check test/algebra.jl, test/operators_abstract.jl and test/prepare_linalg.jl when changing these declarations.

  3. A new forest-level sweep owes a declared transpose and an isbits descriptor. The transpose is what keeps the adjoint identity exact on a forest; the descriptor shape is what keeps the Enzyme seam viable. If the sweep must be differentiated through rather than ruled, its descriptor has to be fully isbits — see CFFluxDescriptor, which carries no Vector for exactly that reason. Caught by test/exchange_schedule.jl, test/exchange_kernels.jl, test/enzyme_rules.jl.

  4. Anything that reads ghosts owes a shares_exchange declaration, and it gates forward and adjoint together. Declaring it when the operator writes into its input, or into an intermediate needing its own exchange, is silently wrong rather than merely slow. Caught by test/forest_parity.jl and test/forest_diffusion.jl.

  5. Never index block storage outside _leaf_array / _leaf_view. Reaching into f.blocks[i] or f.data[…, i] compiles, runs, passes on the layout you tested, and breaks the other one. Caught by test/forest_packed.jl — but only if the new code path has a packed test.

  6. An Enzyme rule body must not allocate. A Dict lookup and a closure in an augmented-primal body segfaulted on Linux x86_64 while passing on macOS/aarch64. Enzyme behaves differently across platforms generally: a green local run is not evidence about CI, so push and read the matrix. Caught by test/enzyme_rules.jl, on CI.

  7. Measure every allocation claim under --check-bounds=yes. Pkg.test and therefore CI always pass it, and it blocks the SROA that elides a Ref or a view escaping a broadcast body — the same code measures 0 B without the flag and hundreds of bytes with it. A CI-only allocation failure is this before it is anything exotic. Related: a stencil reading two arrays must unroll dimensions by recursion over Val(D), not ntuple(Val(N)) — with one array the closure inlines, with a coefficient array it stops, the broadcast loses vectorisation, and a 256² sweep goes from 33 µs to 260 µs with identical numerics. See _diff_axes in src/operators/diffusion.jl.

  8. prepare is stateful and single-threaded. Call it once per concurrent solve, not once globally. And note that a green suite does not mean the GPU paths ran: test/device_gpu.jl is gated behind MFO_TEST_GPU=true and an available CUDA.jl.

One more, about this page. If you change a call sequence described here, update the section that describes it. The only thing this page is worth is being true.