feat(operators): compact flux-form diffusion leaf - #60
Open
kylebeggs wants to merge 5 commits into
Open
Conversation
The advertised `divergence(g) * scaling(κ) * gradient(g)` chains two centered
first differences, giving a wide 2h stencil whose row at cell I samples the flux
κ∇u only at I±e. No equation ever couples κ at adjacent cells, so a κ-inversion
decouples exactly into even/odd (i+j)-parity sublattices fit to disjoint halves of
the data. `Diffusion` is the finite-volume form instead: face fluxes
κ_{i+½}(u_{i+1}-u_i)/Δ differenced across the cell, with κ averaged to faces by
`ArithmeticMean` (default) or `HarmonicMean`.
Beyond removing the decoupling, the compact form is exactly symmetric for real κ
(the composed form is not) and declares `operator_diagonal`, which the composition
cannot — `operator_diagonal(::Composed)` needs both factors diagonal. That unlocks
Jacobi/Chebyshev multigrid for variable-coefficient diffusion, via `_rediscretize`
using the child mean rather than `Restriction`.
Boundary treatment is exact rather than approximate. The package's homogeneous
fills reflect about the face, so a Neumann face carries zero flux (its coefficient
is irrelevant) and a Dirichlet face carries κ_I(0-u_I)/(Δ/2), which the one-sided
face coefficient κ_f = κ_I reproduces exactly. The leaf stores its own copy of κ
with ghosts extended by `fill_coefficient_ghosts!` — an even mirror at walls and a
periodic wrap, deliberately not `apply_bc!`, which would antisymmetrize a material
coefficient at a Dirichlet wall.
Two implementation notes worth keeping. The diagonal weight of a face is s_f - 1
for s_f = ∂u[neighbour]/∂u_I, which covers the case a `_bc_sign`-style correction
misses: a periodic dimension of a single cell wraps onto itself. And dimensions are
unrolled by recursion, not `ntuple(Val(N)) do d` — with a second array in the
stencil that closure stops inlining and the broadcast loses vectorization, costing
8x with identical numerics (260 vs 33 µs per 256² sweep).
Scoped to CartesianGrid: an Interface face has no cross-block κ to average, so the
constructor rejects forests and slabs and `_distributable` stays false. Staged in
#54. Measured 62 µs vs the composition's 115 µs for a 256² prepared `mul!`.
Closes #48.
The example used three drive patterns because one was not enough: under the wide composed stencil a single-drive reconstruction speckled, and the extra drives shrank each parity sublattice's variance until the interleaved halves agreed. That was a workaround for the discretization, not a property of the inverse problem. With `diffusion(g, κ)` one excitation suffices — 3.09% relative error at n=48 with 1% noise, and a jaggedness of 0.069 against the truth's 0.056, so no checkerboard. The residual error is faint streaking where ∇u is small, which is a genuine resolution limit of one drive rather than a stencil artifact. The smoothness prior is now inert: sweeping λ from 0 to 5e-4 moves the recovery by under 0.01 percentage points, where before it was the difference between a usable answer and speckle. λ is kept only so the knob stays visible. Also reports relative error and jaggedness instead of asserting the improvement in prose, and records in DESIGN.md §10.2 that a staggered G/D pair cannot assume D = -Gᵀ: the compact operator factors as -Bᵀ diag(κ_f) B, but at a Dirichlet wall that row of B is √2/Δ, not 2/Δ.
|
Benchmark ResultsTime
Memory and allocations
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
Building the leaf inside the differentiated region stores the Const grid into the freshly built active coefficient Field. Enzyme's static activity analysis clears that on Julia 1.12 but not on 1.10 or 1.11, so all three κ-gradient tests through `diffusion` errored with EnzymeRuntimeActivityError there. Use set_runtime_activity at those call sites — the workaround already used at every other Enzyme call in the suite — and say so in the `diffusion` docstring and docs/pages/autodiff.qmd, since anyone inverting for κ hits the same wall.
The quoted pair labelled the leaf's own 115 µs as the composition's. Benchmark CI has the leaf at 115 µs, the composition at 337 µs, the Laplacian at 67.5 µs for the same 256² prepared mul!.
…tion The README and the scaling/divergence docstrings advertised divergence(g) * scaling(κ) * gradient(g) as the way to spell variable-coefficient diffusion — scaling's worked example handed it over labelled ∇·(κ∇u) with no caveat. Lead with diffusion(g, κ) everywhere; keep the composition as what it is, a demonstration that the algebra composes, with the wide-stencil cost stated.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stage 1 of #56: the compact flux-form
∇·(κ∇u)leaf on a uniformCartesianGrid.Closes #48.
What and why
divergence(g) * scaling(κ) * gradient(g)chains two collocated centered first differences, so its row at cellIsamples the fluxκ∇uonly atI±e. No equation ever couples κ at two adjacent cells, and a κ-inversion decouples exactly into the even and odd(i+j)-parity sublattices — fit to disjoint halves of the noisy data, tied together only by the regularizer.diffusion(g, κ)is the finite-volume form instead: face fluxesκ_{i+½}(u_{i+1} − u_i)/Δdifferenced across the cell, with κ averaged to faces byArithmeticMean(default) orHarmonicMean.The measured effect, in
examples/inverse_diffusion.jl(n=48, 1% noise, same optimizer):The prior is now inert — five orders of magnitude of λ move the answer by under 0.01 percentage points. Under the wide stencil, λ was the difference between a usable answer and speckle, and a 20× stronger prior still could not rescue it. The example accordingly drops from three drive patterns to one; the extra drives existed only to shrink each sublattice's variance until the interleaved halves agreed.
Beyond accuracy
test/algebra.jlcan only check its declared transpose against the dense transpose.operator_diagonal, which the composition structurally cannot (operator_diagonal(::Composed)needs both factors diagonal). This is the only spelling of variable-coefficient diffusion a Jacobi or Chebyshev multigrid smoother can touch, and_rediscretizecoarsens κ with the child mean rather thanRestriction.mul!— 2.9×. Against theLaplacian's 67.5 µs it costs 1.7× for 2× the memory traffic, so the second array rides largely free. (From this PR's benchmark comment, head1c365e4:diffusion mul!is the leaf,∇·(κ∇u) mul!the composition,laplacian mul!the baseline — three entries from the same run, so the comparison carries no cross-revision runner noise.)scaling/divergencedocstrings advertiseddivergence(g) * scaling(κ) * gradient(g)as the way to write ∇·(κ∇u) —scaling's worked example handed it over labelled# ∇·(κ∇u)with no caveat. All of them now lead withdiffusion(g, κ); the composition stays, framed as the algebra stress-test §8 says it is, with the wide-stencil cost stated.Boundary treatment is exact, not approximate
The package's homogeneous fills reflect about the face. A Neumann face therefore carries zero flux, so its coefficient is irrelevant; a Dirichlet face carries
κ_I(0 − u_I)/(Δ/2), which the one-sided face coefficientκ_f = κ_Ireproduces exactly. Worth knowing when inverting: that one-sided value is O(Δ) accurate in κ at a wall while interior faces are O(Δ²), so wall-adjacent κ cells have a different sensitivity structure. Documented in the docstring.The leaf stores its own copy of κ with ghosts extended by
fill_coefficient_ghosts!— an even mirror at walls, a periodic wrap — deliberately notapply_bc!, whose Dirichlet antisymmetric mirror would negate a material coefficient at a wall. Same hazard_average_to_coarsedocuments for coarsening.Two things worth a reviewer's attention
The diagonal has a case a
_bc_sign-style correction misses. A face's diagonal weight iss_f − 1fors_f = ∂u[neighbour slot]/∂u_I: −1 interior, −2 Dirichlet, 0 Neumann, and 0 for a periodic dimension of a single cell, where the wrap neighbour is the cell itself.Laplacianhas the same latent bug today — filed separately as #55, not fixed here. Theoperator_diagonaltest sweepsn ∈ {(5,4), (3,3), (2,2), (1,4), (4,1)}againstdiag(materialize(...))precisely to pin this.Dimensions are unrolled by recursion, not
ntuple(Val(N)) do d. With a second array in the stencil that closure stops inlining and the broadcast loses vectorization — 260 µs vs 33 µs for one 256² sweep, with identical numerics. The first version of this leaf was 10× slower than the composition it replaces because of it. Added to CLAUDE.md's gotchas; please keep_diff_axesclosure-free.Scope
CartesianGridonly. AnInterfaceface has no cross-block κ to average, so the constructor rejects forests and slabs and_distributablestaysfalse— staged as #57 / #58 / #59. The declared transpose gather forInterfacegrids is nevertheless implemented and tested here (through the inner constructor, with κ ghosts filled by hand), since that is what #58 will build on.AD stays rule-free on the κ path, asserted rather than assumed:
test/enzyme_rules.jlchecksrule_hits()is unchanged across aDiffusiongradient, and κ ghost entries get zero gradient in both Enzyme and the finite-difference reference.