Add JopRBF: scattered-node compact-support RBF model parameterization

Project.toml

@@ -1,6 +1,6 @@
 name = "JetPack"
 uuid = "24ef3835-3876-54c3-8a7a-956cf69ca0b2"
-version = "2.2.1"
+version = "2.3.0"
 
 [deps]
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"

docs/JopRBF/JopRBF-demo.md

@@ -0,0 +1,158 @@
+# JopRBF: scattered-node model parameterization for JetPack
+
+`JopRBF` maps a small set of scattered control-node coefficients to a regular fine
+model grid using **normalized compact-support radial basis functions** (Wendland
+C2). It is a reduced (coarse) model parameterization for `Jets` inverse problems
+(FWI dimension reduction, regularization, preconditioning), with an exact adjoint
+and a matrix-free implementation that scales to very large grids (for example
+1000^3).
+
+This directory holds the operator's demo scripts and the figures below. The demos
+use a local project (`docs/JopRBF/Project.toml`) that adds PyPlot alongside a
+dev'd JetPack, so the plotting dependency stays out of JetPack itself:
+
+```
+julia --startup-file=no --project=docs/JopRBF -t 8 docs/JopRBF/JopRBF_waterbottom_demo.jl   # water-bottom comparison figure
+julia --startup-file=no --project=docs/JopRBF -t 8 docs/JopRBF/JopRBF_demo.jl               # kernel / smoothness figures
+```
+
+## The operator
+
+$$
+d(x) = \frac{\sum_j \varphi\!\left(r_j(x)\right)\, c_j}{\sum_j \varphi\!\left(r_j(x)\right)},
+\qquad
+r_j(x) = \sqrt{\sum_k \left(\frac{x_k - \xi_{j,k}}{\delta_{k,j}}\right)^2}
+$$
+
+$$
+\varphi(r) = (1-r)^4\,(4r+1)\ \text{ for } 0 \le r < 1, \qquad \varphi(r) = 0 \text{ otherwise (Wendland C2).}
+$$
+
+Here $c$ are the node coefficients (the optimization variables), $\xi_j$ the node
+locations (in fine-grid index coordinates), and $\delta$ the compact-support
+radius. It may be a scalar, a per-axis vector (for anisotropic grids), or a
+per-node $\delta_{k,j}$ matrix, so each node can carry its own support radius.
+Properties, each a unit test in `test/jop_RBF.jl`:
+
+- **Compact support, so it scales.** $\varphi$ vanishes for $r \ge 1$, so each
+  fine point only sees the $O(1)$ nodes within $\delta$. The forward and adjoint
+  are applied matrix-free with a node bucket index (no dense $N_\text{fine}\times M$
+  matrix is formed). The forward threads over fine points and gathers nearby
+  nodes; the adjoint threads over nodes and walks each node's fine-grid box. Both
+  enumerate the identical set of (fine point, node) pairs, so the adjoint is exact.
+- **C2 smooth, no hot zones.** $\varphi$, $\varphi'$ and $\varphi''$ all vanish at
+  $r = 1$, so the field is twice continuously differentiable (no bilinear-style
+  creases).
+- **Reproduces constants.** The Shepard normalization is a partition of unity.
+- **No overshoot.** $\varphi \ge 0$, so the field is a convex combination of
+  nearby coefficients and stays within their min and max (no ringing / halos).
+- **Build once, reuse.** The node bucket index and the normalization field
+  $1/\sum_j\varphi$ depend only on the node positions and $\delta$; they are
+  computed once at construction and reused for every apply. There is no linear
+  solve.
+
+### API
+
+```julia
+# nodes are a (ndims(range) x M) matrix in fine-grid index coordinates
+A = JopRBF(JetSpace(Float64, M), JetSpace(Float64, nz, nx), nodes; delta = 20.0)
+d = A * rand(domain(A))                       # coarse coefficients -> fine model
+g = A' * rand(range(A))                        # adjoint -> gather to nodes
+
+# anisotropic support (per axis)
+A = JopRBF(JetSpace(Float32, M), JetSpace(Float32, nz, nx, ny), nodes3; delta = [30.0, 50.0, 50.0])
+
+# per-node support radii (ndims x M): small support shallow, large support deep
+A = JopRBF(JetSpace(Float64, M), JetSpace(Float64, nz, nx), nodes; delta = deltas)
+```
+
+## Water-bottom "freeze" workflow
+
+The scattered nodes make it natural to parameterize only the sediment below a
+(possibly dipping) water bottom while freezing the water column.
+
+1. Get the water-bottom depth `wb(x)` (for example from `water_bottom_index`).
+2. Place the RBF centers below the water bottom with a depth-tapered scattered
+   cloud ([`tapered_rbf_nodes`](rbf_tapered_nodes.jl), pure base Julia, no external
+   mesher). The rows **conform to** (run parallel to) the water bottom: each row
+   sits at a fixed offset from the per-column WB and the offsets march downward by
+   the local spacing $h(z) = \lambda(z)/\text{ppw}(z)$, $\lambda(z) = v(z)/f$, so
+   the node spacing follows a points-per-wavelength rule (denser where the
+   wavelength is short), with `ppw` tapering linearly with depth (dense at the
+   water bottom, sparse at depth). Making the rows parallel to the WB avoids
+   horizontal rows cross-cutting a dipping boundary. One capped "ghost" row sits
+   just **above** the water bottom (masked out of the model, but its support
+   reaches below) so the near-WB fit has two-sided support and does not overshoot
+   at the boundary.
+3. Build the reduced parameterization operator as a composition:
+
+   ```julia
+   A  = JopRBF(JetSpace(T, M), JetSpace(T, nz, nx), nodes; delta)
+   S  = JopDiagonal(Float64.(below))      # 1 below the water bottom, 0 in the water
+   P  = S ∘ A                             # reduced parameterization (a JopLn)
+   ```
+
+4. The full model is the frozen water plus the RBF sediment:
+   $v(c) = v_\text{water} \cdot (\text{water mask}) + P\,c$.
+
+   In FWI you invert for $c$. The reduced gradient is $A'$ applied to the masked
+   fine-grid gradient ($P' = A' S'$), so the water column is never updated: the
+   freeze is automatic because there are no nodes in the water and `S` zeros the
+   water rows. To start from a given sediment model, fit the coefficients once,
+   `c = P \ vec(v_start .* below)`, then invert. Bounds and regularization now act
+   on the few hundred coefficients `c` instead of every fine sample, and the
+   sediment is smooth (C2) and free of overshoot by construction.
+
+**Figure 1**, produced by
+[`JopRBF_waterbottom_demo.jl`](JopRBF_waterbottom_demo.jl), shows this on a 101×201
+model with a sloping water bottom, comparing three inversion frequencies (lower
+frequency $\Rightarrow$ longer wavelength $\Rightarrow$ coarser node cloud). Row 1:
+the true model, and the partition-of-unity coverage below the water bottom (water
+shown white). Rows 2-4: the frozen-water + RBF parameterization at 6, 3 and 1.5 Hz
+(nodes overlaid) next to the signed relative error $\Delta v/v$ (red/blue, shared
+color scale). Node counts follow the points-per-wavelength taper: 6 Hz ≈ 608 nodes
+(25× reduction, 0.0008 RMS), 3 Hz ≈ 186 (83×, 0.0013), 1.5 Hz ≈ 65 (238×, 0.0032)
+below-WB relative error. Lower frequency = fewer nodes = a smoother, coarser
+representation.
+
+![Figure 1: water bottom freeze workflow](images/jopRBF_waterbottom_2d.png)
+
+## Kernel and smoothness
+
+**Figure 2** shows a single normalized basis function (a compact bump), the
+Wendland C2 kernel next to a C0 linear "tent" kernel, and a slice confirming the
+basis is exactly zero beyond $\pm\delta$.
+
+![Figure 2: basis and kernel](images/jopRBF_basis_2d.png)
+
+**Figure 3** explains why the C2 kernel matters. Both panels interpolate the *same*
+coefficients on the *same* coarse node cloud (white dots), differing only in the
+kernel. The linear "tent" kernel gives a **C0** field (continuous value but
+discontinuous slope) with visible creases/kinks between nodes; the Wendland kernel
+gives a smooth **C2** field (continuous value, slope, and curvature), which is what
+`JopRBF` uses. The transect along the red line (right) makes it explicit: the C0
+profile kinks at node boundaries, the C2 profile is smooth. Those C0 kinks are the
+"hot zones" a piecewise-linear (bilinear or barycentric) parameterization produces
+between nodes.
+
+![Figure 3: C0 vs C2 smoothness](images/jopRBF_smoothness_2d.png)
+
+## Why this design
+
+- **Global RBF** (thin-plate, multiquadric) has dense interpolation matrices and
+  global support, so it does not scale to 1000^3.
+- **Triangulation plus C1 cubic** (Clough-Tocher) is local and smooth but needs a
+  3D Delaunay mesh and fragile C1 macro-elements.
+- **Barycentric-linear on a mesh** is cheap but only C0, which reintroduces the
+  creases / hot zones.
+
+Normalized compact-support RBF keeps locality (scales), is C2 (no hot zones), is
+meshless (no 3D triangulation), reproduces constants, and does not overshoot.
+
+## Tests
+
+`test/jop_RBF.jl`: Wendland C2 kernel (value and first/second derivatives vanish
+at the support edge, non-negative); dot-product (adjoint) tests in 1D/2D/3D;
+partition of unity; compact support (influence vanishes beyond $\delta$); no
+overshoot (convex combination); anisotropic per-axis $\delta$; per-node
+(multiresolution) $\delta$; Float32.

docs/JopRBF/JopRBF_demo.jl

@@ -0,0 +1,143 @@
+# Kernel / smoothness figures for JopRBF (normalized compact-support Wendland-C2 RBF).
+#
+# Generates two PNGs under docs/JopRBF/images/:
+#   1. jopRBF_basis_2d.png       a single normalized basis, its C2 vs C0 kernel, compact support
+#   2. jopRBF_smoothness_2d.png  why C2: linear (C0) interpolation kinks between nodes; Wendland (C2) is smooth
+#
+# Run (global environment has JetPack dev'd + PyPlot):
+#   julia --startup-file=no docs/JopRBF/JopRBF_demo.jl
+#
+ENV["MPLBACKEND"] = "Agg"
+using JetPack, Jets, LinearAlgebra, Random, Printf, PyPlot
+
+Random.seed!(20260709)
+const OUT = joinpath(@__DIR__, "images")
+isdir(OUT) || mkpath(OUT)
+
+# scattered node cloud (denser shallow) used for the basis figure
+function make_nodes(nz, nx)
+    pts = Vector{Float64}[]
+    for x in range(1, nx, length = 12)
+        push!(pts, [1.0, x]); push!(pts, [Float64(nz), x])
+    end
+    for z in range(1, nz, length = 9)
+        push!(pts, [z, 1.0]); push!(pts, [z, Float64(nx)])
+    end
+    for _ = 1:120
+        push!(pts, [1 + (nz - 1) * rand()^0.65, 1 + (nx - 1) * rand()])
+    end
+    reduce(hcat, pts)
+end
+
+# normalized (Shepard) RBF evaluation on the full grid, for a chosen kernel
+function eval_normalized(nodes, c, nz, nx, delta, kernel)
+    d = zeros(nz, nx)
+    M = size(nodes, 2)
+    for iz = 1:nz, ix = 1:nx
+        num = 0.0; den = 0.0
+        for j = 1:M
+            r = sqrt((iz - nodes[1, j])^2 + (ix - nodes[2, j])^2) / delta
+            w = kernel(r)
+            num += w * c[j]; den += w
+        end
+        d[iz, ix] = den > 0 ? num / den : 0.0
+    end
+    d
+end
+
+wend_c2(r) = r < 1 ? (1 - r)^4 * (4r + 1) : 0.0     # Wendland C2 kernel (what JopRBF uses)
+linear0(r) = r < 1 ? (1 - r) : 0.0                  # linear "tent" kernel -> C0 (piecewise linear)
+
+# ---------------------------------------------------------------------------
+# Figure 1: a single normalized basis + C2 vs C0 kernel + compact support
+# ---------------------------------------------------------------------------
+nz, nx = 160, 220
+delta = 26.0
+nodes = make_nodes(nz, nx)
+M = size(nodes, 2)
+A = JopRBF(JetSpace(Float64, M), JetSpace(Float64, nz, nx), nodes; delta = delta)
+
+jc = argmin(vec(sum((nodes .- [80.0, 110.0]) .^ 2, dims = 1)))
+e = zeros(domain(A)); e[jc] = 1.0
+basis = A * e
+zc, xc = nodes[1, jc], nodes[2, jc]
+
+figure(figsize = (13, 4.2))
+subplot(1, 3, 1)
+im = imshow(basis, aspect = "auto", cmap = "viridis")
+scatter([xc], [zc], s = 45, c = "red", marker = "x")
+xlim(1, nx); ylim(nz, 1)
+title("normalized basis  A e_j  (compact bump)"); xlabel("x"); ylabel("z"); colorbar(im, fraction = 0.046)
+
+subplot(1, 3, 2)
+rr = range(0, 1.3, length = 400)
+plot(rr, wend_c2.(rr), lw = 2.4, label = "Wendland C2  (1-r)^4(4r+1)")
+plot(rr, linear0.(rr), lw = 1.8, ls = "--", label = "linear C0  (1-r)")
+axvline(1.0, color = "k", lw = 0.7, alpha = 0.5)
+xlabel("r = |x - ξ| / δ"); ylabel("φ(r)"); title("kernel shape: C2 vs C0"); legend(); grid(alpha = 0.3)
+
+subplot(1, 3, 3)
+row = clamp(round(Int, zc), 1, nz)
+plot(1:nx, basis[row, :], lw = 2.0)
+axvline(xc - delta, color = "r", ls = ":", lw = 1); axvline(xc + delta, color = "r", ls = ":", lw = 1)
+xlabel("x"); ylabel("basis value"); title("slice at z=$(row): zero beyond ±δ"); grid(alpha = 0.3)
+tight_layout()
+savefig(joinpath(OUT, "jopRBF_basis_2d.png"), dpi = 120); close()
+
+# ---------------------------------------------------------------------------
+# Figure 2: why C2. Same coarse nodes and coefficients, two interpolation
+# kernels. Linear weights give a piecewise-linear (C0) field that kinks between
+# nodes; the Wendland kernel (what JopRBF uses) is smooth (C2).
+# ---------------------------------------------------------------------------
+gz, gx = 130, 180
+h = 26.0
+snodes = let pts = Vector{Float64}[]
+    zs = collect(range(1.0, gz, step = h)); (gz - zs[end] > 1) && push!(zs, Float64(gz))
+    xs = collect(range(1.0, gx, step = h)); (gx - xs[end] > 1) && push!(xs, Float64(gx))
+    for z in zs, x in xs
+        jz = (1 < z < gz) ? 0.18h * (rand() - 0.5) : 0.0
+        jx = (1 < x < gx) ? 0.18h * (rand() - 0.5) : 0.0
+        push!(pts, [z + jz, x + jx])
+    end
+    reduce(hcat, pts)
+end
+sM = size(snodes, 2)
+sdelta = 1.15h
+crand = rand(sM)
+dC0 = eval_normalized(snodes, crand, gz, gx, sdelta, linear0)
+dC2 = eval_normalized(snodes, crand, gz, gx, sdelta, wend_c2)
+cmin = min(minimum(dC0), minimum(dC2)); cmax = max(maximum(dC0), maximum(dC2))
+row = 66
+
+figure(figsize = (15, 4.4))
+subplot(1, 3, 1)
+im = imshow(dC0, aspect = "auto", cmap = "viridis", vmin = cmin, vmax = cmax)
+scatter(snodes[2, :], snodes[1, :], s = 12, c = "white", edgecolors = "k", linewidths = 0.3)
+axhline(row, color = "r", lw = 1.5)
+xlim(1, gx); ylim(gz, 1)
+title("linear kernel → C0 field\n(kinks/creases between nodes)"); xlabel("x"); ylabel("z"); colorbar(im, fraction = 0.046)
+
+subplot(1, 3, 2)
+im = imshow(dC2, aspect = "auto", cmap = "viridis", vmin = cmin, vmax = cmax)
+scatter(snodes[2, :], snodes[1, :], s = 12, c = "white", edgecolors = "k", linewidths = 0.3)
+axhline(row, color = "r", lw = 1.5)
+xlim(1, gx); ylim(gz, 1)
+title("Wendland C2 kernel → C2 field\n(smooth; what JopRBF uses)"); xlabel("x"); ylabel("z"); colorbar(im, fraction = 0.046)
+
+subplot(1, 3, 3)
+plot(1:gx, dC0[row, :], lw = 1.8, ls = "--", color = "tab:orange", label = "C0 linear (kinks)")
+plot(1:gx, dC2[row, :], lw = 2.6, color = "tab:blue", label = "C2 Wendland (smooth)")
+xlabel("x"); ylabel("field value along red line"); title("transect at z=$(row)")
+legend(); grid(alpha = 0.3)
+tight_layout()
+savefig(joinpath(OUT, "jopRBF_smoothness_2d.png"), dpi = 120); close()
+
+# scaling note
+let x = rand(domain(A)), y = rand(range(A))
+    A * x; A' * y
+    t1 = @elapsed for _ = 1:5; A * x; end
+    t2 = @elapsed for _ = 1:5; A' * y; end
+    @printf("timing on %dx%d, M=%d, %d threads: forward %.1f ms, adjoint %.1f ms\n",
+        nz, nx, M, Threads.nthreads(), 1000t1 / 5, 1000t2 / 5)
+end
+println("wrote jopRBF_basis_2d.png and jopRBF_smoothness_2d.png in ", OUT)