Derivation: isotropic elastic Born (linearised) propagator

Developer/reference documentation: the discrete derivation behind deepwave/elastic_born and its adjoint. See Born Modelling with the Elastic Wave Equation for the user-facing propagator page.

Implementation status (updated). The physics core (prepare_scatter_parameters, update_velocities_born, update_stresses_born), the full eager driver (elastic_born_eager_core), and the public deepwave.elastic_born(...) / ElasticBorn API (§8) are landed and validated on the eager backend against every physics gate of §10 that does not require the compiled backend: Born linearity (gate 1), Born == directional FD w.r.t. lamb/mu/buoyancy in 2D and 3D (gate 2), jvp == Born for the eager path (gate 3, machine-precision), FD-of-gradient / gradcheck + gradgradcheck (gates 4/5 — the eager backend gives correct second-order results through ordinary autograd), free-surface interaction (gate 7), and zero-scatter background byte-identity. The torch.func.jvp forward-mode rule on the compiled elastic propagator (§6) is now landed: ElasticForwardFunc.jvp (_elastic_forward_jvp) replaces the unsupported_jvp guard and routes the tangents through the Born relationship — the primal comes from the fast compiled forward, and the tangent from the compiled Born forward core (elastic_born_compiled_core; the eager Born driver is used only as the M-PML fallback) — so torch.func.jvp(elastic, …) and torch.autograd.forward_ad run entirely on the fast backend and match the elastic_born scattered output to near machine precision in 2D and 3D (gate 3, compiled path, well inside the ~1e-6 fast-math band), with the PML-band-tangent warning and the single-transform-level restriction. The second-order (HVP, §5) path is also landed: ElasticForwardFunc.backward drops @no_second_order for a nested ElasticBackwardFunc whose double-backward drives the compiled Born backend, so gradgradcheck/reverse-over-reverse Hessian-vector products work on the compiled backend (M-PML falls back to backend="eager"). The compiled elastic_born FORWARD kernel is now landed (elastic_born.c/.cu, symbols elastic_born_iso_*_forward_*; §9 ABI/CMake/backend_utils wiring done): a single coupled pass over the doubled 26/54-wavefield set (§2.5), forward-only (no gradient stores). elastic_born(..., backend="compiled"), and the default backend="auto", dispatches to it, validated to reproduce the eager Born driver to near machine precision on CPU (background AND scattered wavefields + both receiver blocks, 2D/3D, including the zero-scatter and pressure-source paths); the .cu nvcc-compiles in lockstep (GPU-untested on the CPU-only dev host — maintainer GPU-tests). It is forward-only: gradients on the compiled path are supplied by re-running the validated eager Born driver in the autograd Function’s backward (ElasticBornForwardFunc, a checkpoint-style bridge), so the default path stays differentiable. M-PML on the compiled path is deferred (compiled mpml_fraction > 0 raises, directing to backend="eager").

The compiled forward’s background gradient-store checkpointing is now landed (phase-5 prerequisite): when store_bg is true (a scatter/background model needs grad) the forward writes the eight/five/two background finite-difference terms the adjoint will correlate against — dv{i}dbuoyancy (= dt·W̄_i), dv{k}d{k} (= dt·Dv̄_k), dv{j}d{i}_plus_dv{i}d{j} (= dt·S̄_ij) — into a single packed bg_store buffer of shape [n_store, nt, n_shots, *grid] (a bespoke full-history buffer, not the shared step-based storage engine, matching the self-contained Born forward). The grad-free fast path (store_bg=false) is byte-identical (test test_forward_bg_store_does_not_perturb_and_is_populated). ABI grows additively by one pointer + one bool; the CPU kernel implements the writes (the CUDA symbol accepts the ABI, GPU store implementation deferred to the GPU adjoint pass).

Phase 6 landed the compiled backward_sc (scattered-only adjoint / imaging condition, §4.1). elastic_born.c/.cu now export elastic_born_iso_*_backward_sc_*: a self-contained reverse loop propagating a single scattered adjoint field with the plain elastic adjoint operator (fixed background models λ, μ, μ_ij, b_i), seeded by the scattered receiver residuals and/or scattered-final-wavefield gradients, reading the checkpointed bg_store in place of the shared-engine stores exactly as this document anticipated, and accumulating grad_δλ/grad_δμ/grad_δμ_ij/grad_δb_i with the same imaging forms as elastic.c’s nonlinear model gradients. Validated (tests/test_elastic_born.py::test_compiled_backward_sc_matches_eager_imaging) against the eager Born driver’s autograd gradient of a loss seeded on both the scattered final wavefields and the scattered receivers, in 2D and 3D, with and without PML (relative error ~1e-6, matching the fast-math §10 gate-3 band); abi_reference.json extended additively (+6 entries: elastic_born|backward_sc| {1,2,3}|{float,double}), no drift to existing symbols.

Phase 7 wired backward_sc into the public dispatch and extended bg_store with the scattered stores. ElasticBornForwardFunc.backward now dispatches to the compiled backward_sc (via _elastic_born_backward_sc_compiled) whenever no background-side quantity needs a gradient — no background model (lamb/mu/buoyancy), no background source amplitude, no background initial wavefield, and M-PML off (mirrors acoustic_born.py’s need_bg_adjoint condition) — instead of always falling back to the eager-recompute bridge; this is the first real, user-facing speedup on the compiled path (test_compiled_backward_sc_dispatch, which spies on the compiled core to prove it fires exactly once on the scattered-only path and never on the eager backend or a background-grad request, gradients matching the eager backend to <1e-6). The eager bridge remains the fallback for any case needing background-side gradients, until the full two-field backward below lands. Also extended the forward’s checkpointing with sc_store (store_sc-gated), the scattered-field analogue of bg_store (§2.5/§4.3 — W_sc_i/Dv_sc_k/S_sc_ij, written at the same 7 sites bg_store already writes, from the already-computed w_sum_sc/dv*_sc terms), the prerequisite the full two-field backward needs for its two-term background gradients. Validated (test_forward_sc_store_does_not_perturb_and_matches_fd) against a central finite difference of bg_store, taken through the public elastic_born() API (perturbing the physical lamb/mu/buoyancy) to ~1e-6 relative. A genuine, if surprising, finding surfaced while developing that test: comparing store_bg/store_sc on/off via strict torch.equal is not a valid gate once any sibling conditional exists in the same -Ofast-compiled SIMD loop, because -ffast-math permits floating-point reassociation and the compiler’s vectorisation can shift when new (even always-false-at-runtime) branches are added nearby — measured at ~2e-13 relative after 12 leapfrog steps, ULP-level and physically meaningless, confirmed by the independent FD check agreeing far more tightly (~1e-6, limited by truncation). Both store-perturbation tests now use a tight relative tolerance (<1e-9) instead of bitwise equality, documented inline in tests/test_elastic_born.py. backward_sc itself, the full two-field backward (§4.2’s 36-site coupling checklist), the HVP, and the jvp re-pointing remain future work, as detailed in the phase-6 plan below (still accurate for what remains — items 1 and the “extend bg_store” prerequisite of item 2 are now done; the 36-site coupling substitution itself is unstarted).

Phase 8 landed the full two-field backward (§4.2’s complete 36-site coupling checklist, §4.3’s two-term background gradients). elastic_born.c now exports elastic_born_iso_*_backward_*: a second (background) adjoint field runs alongside the scattered adjoint field of backward_sc, using elastic.c’s exact transpose macros with the + Dmodel·Q_sc coupling injected at every one of the 36 sites — all 9 velocity-adjoint macros (A), all 9 stress-adjoint macros (B), all 9 velocity memory-variable recursions (C), and all 9 stress memory-variable recursions (D), each in BOTH its PLAIN and _PML form (_BG-suffixed macros in the kernel: VZ_Z_BG/VZ_Z_BG_PML, …, SIGMAZZ_Z_BG/SIGMAZZ_Z_BG_PML, …, and the coupled m_v*_bg/m_sigma*_bg recursions) — no sub-family deferred. Per §5.3, only the m_sigma* family is double-buffered/ping-ponged (for both fields); m_v* updates in place for both fields. Background-model gradients (grad_lamb/grad_mu/grad_mu_ij/ grad_buoyancy_i) accumulate the exact two-term form of §4.3 (background adjoint × bg_store PLUS scattered adjoint × sc_store); scattering-model gradients (grad_dlamb/…) accumulate from the scattered adjoint alone, as in backward_sc. Wiring: get_elastic_born_backward_template in backend_utils.py (additive ABI, +1 pointer block for scattering models/ background residuals/background adjoint wavefields/sc_store/background grad buffers/background source-amplitude grads over backward_sc’s signature), elastic_born_backward_compiled in elastic_born.py (re-runs the forward once with store_bg=store_sc=True to regenerate both packed stores, then drives the kernel), and ElasticBornForwardFunc.backward’s dispatch now routes to this compiled kernel whenever a background-side quantity needs a gradient (the case backward_sc doesn’t cover) UNLESS M-PML is active, which still falls back to the eager bridge. Validated (test_compiled_backward_full_matches_eager, 2D+3D × pml∈{0,3} — pml=0 isolates the PLAIN sites, pml=3 exercises every _PML site and memory recursion) against the eager driver’s autograd gradient of a loss seeded on every forward output (both final wavefield blocks and both receiver blocks, so both adjoint fields are directly driven), requesting gradients for lamb/mu/buoyancy (background models), dlamb/dmu/dbuoyancy (scattering models), and the background source amplitude — relative error ~1e-7, all 4 cases green. abi_reference.json extended additively (+6 entries: elastic_born|backward|{1,2,3}|{float,double}); a real ABI-wiring bug was caught during verification (not a false green): backend_utils.py’s module-init loop that calls _assign_argtypes for every propagator/ndim/ accuracy/dtype combination was missing the plain elastic_born backward entry (only forward and backward_sc were registered), so the new kernel’s argtypes were never assigned and the first call raised ctypes.ArgumentError; fixed by adding the missing _assign_argtypes call. CUDA: elastic_born.cu still only exports forward; the full backward (like backward_sc) has no CUDA implementation yet (deferred, matching every prior phase’s compiled-adjoint scope — the maintainer GPU-tests before release). HVP (§5) and jvp re-pointing (§6) were NOT attempted this phase — budget went entirely to the 36-site kernel and its verification, per the task brief’s priority.

Phase 9 landed the HVP (§5) and re-pointed the jvp (§6) — pure Python, no kernel change. ElasticForwardFunc.backward’s @no_second_order guard is removed and replaced by a nested differentiable ElasticBackwardFunc.apply(...) (mirroring AcousticForwardFunc/AcousticBackwardFunc): its forward runs the ordinary first-order compiled adjoint (via run_backward_core with the elastic SPEC, the m_sigma*-only aux swap of §5.3, callback_drop_last=1, cpu_tmp_batch, backend_extra="_mpml" when M-PML), and its @once_differentiable backward is the two-pass HVP of §5.2. The HVP drives the now-landed COMPILED Born backend, not the eager driver the phase-6 note (3) anticipated: Step 1 calls elastic_born_compiled_core (background = m, scatter = h in prepared-model space) for the scattered forward wavefields / scattered receivers (the grads w.r.t. grad_wavefields/grad_receivers), and Step 2 calls elastic_born_backward_compiled with the HVP seeding (background adjoint seeded at zero with empty background residual; scattered adjoint seeded by the original grad_wavefields and sourced by the original grad_receivers; sc_models = h, background models requires_grad), reading grad_bg as H·h, grad_f_bg as the source HVP, and the background adjoint final as the initial-wavefield HVP. model_gradient_sampling_interval must give an effective stride of one for compiled double backward; otherwise the sampled first gradient is not an exact gradient field and the Born directional operation is not its reverse Jacobian. h lives entirely in prepared-model space — the chain rule through prepare_parameters is handled by elastic’s differentiable front-end, exactly as acoustic’s HVP works in (K, B_d) space (so prepare_scatter_parameters is not called by the HVP). To let the double-backward recompute the forward field, ElasticForwardFunc.forward now passes save_wavefields=True (initial wavefields saved for backward). The jvp (_elastic_forward_jvp) now routes its tangent through elastic_born_compiled_core (the same fast kernel the elastic_born default uses) instead of the eager bridge, so torch.func.jvp/torch.autograd.forward_ad run entirely on the compiled engine (tangent and reference share the kernel, so they agree to near machine precision, well inside the ~1e-6 fast-math gate-3 band). M-PML fallback: the compiled Born backend has no _mpml variants, so the HVP backward raises NotImplementedError directing to backend="eager" under M-PML, and the jvp falls back to the eager Born driver under M-PML (first-order M-PML gradients are unaffected — the first-order adjoint keeps its _mpml dispatch). §5.3 aux-swap finding (confirmed): the memory-variable aux swap applies to elastic’s default (standard-PML) path, not only under M-PML — the m_sigma* family exists whenever pml_width > 0, and the first-order swap ping-pongs exactly that family (2D: wavefield indices 6:9 and 12; the m_v* family at 3:5, 11 is never swapped). Validation: compiled HVP == backend="eager" exact second-order to ~1e-11 (2D) / ~1e-12 (3D) for lamb/mu/buoyancy/source and cross-terms; gradgradcheck on compiled elastic succeeds (2D+3D) for lamb/mu/source (the buoyancy block’s finite-difference-of-forward numerical Hessian is unreliable even on the exact eager autograd path, so it is covered by the exact eager-comparison gate instead); HVP == FD-of-gradient ~2e-7. No .c/.cu/ABI change (Python-only, as the phase-8 assessment predicted).

Phase-6 plan (the remaining compiled fast path), with de-risking findings. (1) backward_sc is elastic.c’s FUNC(backward) almost verbatim — the scattered adjoint field obeys the plain elastic adjoint operator with the background models, seeded by the scattered data residuals; its imaging conditions (grad_δλ/δμ/δμ_ij/δb_i, §4.1) are identical in form to elastic’s grad_lamb/grad_mu/…, reading the bg_store terms in place of the shared-engine stores. So backward_sc is a self-contained reverse loop over bg_store with one adjoint field and the memory-variable double-buffer swap on the m_sigma* family only (§5.3). (2) The full two-field backward adds a second (background) adjoint field whose update is that same elastic adjoint plus the + Dmodel·Q_sc coupling at the §4.2 sites, and the two-term background gradients (§4.3), which additionally need the scattered stores (extend bg_store to also pack W_sc/Dv_sc/S_sc, written from the already-computed w_sum_sc/dv*_sc terms — the exact analogues of the bg store sites). (3) The HVP does not require the compiled Born backward: ElasticBackwardFunc.forward keeps the fast compiled first-order adjoint (first-order performance unchanged) and its @once_differentiable backward may compute H·h via the eager Born driver (one eager Born forward + one eager Born backward, §5.2); torch.is_grad_enabled() (the no_second_order discriminator) already isolates the second-order path. Re-pointing §6 jvp to the compiled backend is then a one-line swap (eager→compiled core), tightening gate 3 to the ~1e-6 fast-math band. The eager backend already covers all differentiability (first- and second-order) in the meantime. This document adapts the variable-density acoustic Born work (docs/acoustic_born_derivation.md) to the isotropic elastic velocity–stress staggered-grid propagator (elastic.c/.py). It derives the linearised (Born) system, its adjoint (imaging conditions + the full two-field background gradients), the second-order (Hessian-vector product) structure, and the forward-mode (jvp) rule; it fixes the planned public API, kernel filenames, symbols, and validation gates. Per the task’s hard requirement, NO .c/.cu/CMake code is written in this pass — the kernel implementation is a separate follow-up after this derivation is reviewed. The guards this work will eventually replace are elastic.py lines 1629–1633 (ElasticForwardFunc.jvp = unsupported_jvp and the @no_second_order decorator on ElasticForwardFunc.backward).*

Throughout, the isotropic elastic constitutive fields are the first Lamé parameter λ (lamb), the shear modulus μ (mu), and the buoyancy b = 1/ρ (buoyancy). The kernel additionally consumes the staggered derived models μ_zy, μ_zx, μ_yx (shear moduli harmonically averaged onto the shear-stress nodes) and b_z, b_y, b_x (buoyancy arithmetically averaged onto the velocity nodes), all produced by elastic.py::prepare_parameters(mu, buoyancy). The state is the particle velocity v = (v_z, v_y, v_x) and the symmetric stress tensor σ = (σ_zz, σ_yy, σ_xx, σ_yz, σ_xz, σ_xy). A background quantity carries a bar (v̄, σ̄, λ, μ, b); a first-order perturbation carries a δ (δv, δσ, δλ, δμ, δb).

---

1. The forward system (what elastic solves)

1.1 Continuous equations

The isotropic elastic first-order velocity–stress system is

velocity:   ρ ∂v_i/∂t = Σ_j ∂_j σ_ij   (+ f_i)      ⇔   ∂v_i/∂t = b (Σ_j ∂_j σ_ij) (+ b f_i)
stress:     ∂σ_ij/∂t = λ δ_ij (Σ_k ∂_k v_k) + μ (∂_i v_j + ∂_j v_i)   (+ source)

i.e. Newton’s law (buoyancy times the divergence of the stress row i) for each velocity component, and the isotropic Hooke’s law rate form for each stress component: the normal stresses σ_ii couple to the dilatation Σ_k ∂_k v_k through λ plus 2μ ∂_i v_i; the shear stresses σ_ij (i≠j) couple to the engineering shear strain rate ∂_i v_j + ∂_j v_i through μ.

1.2 Discrete staggered leapfrog (the reference backend)

On the staggered grid, λ, μ, and the normal stresses σ_ii live at the integer (“normal-stress”) nodes; each velocity v_i and its buoyancy b_i live half a cell from σ_ii in dimension i; each shear stress σ_ij and its shear modulus μ_ij live half a cell in both i and j. One internal step is (from elastic.c FUNC(forward); ∂_j denotes the appropriate half/integer staggered difference, optionally augmented by a CPML/M-PML memory variable m_*):

velocity:   W_i   = Σ_j ∂_j σ̄_ij  (+ PML memory)               (checkpointed: dv{i}dbuoyancy_store_1 = dt·W_i)
            v̄_i  ← v̄_i + b_i · dt · W_i

normal σ:   Dv̄_k = ∂_k v̄_k       (+ PML memory)               (checkpointed: dv{k}d{k}_store_1 = dt·Dv̄_k)
            σ̄_ii ← σ̄_ii + dt·( λ · Σ_k Dv̄_k  +  2μ · Dv̄_i )

shear σ:    S̄_ij = ∂_i v̄_j + ∂_j v̄_i  (+ PML memory)          (checkpointed: dv{j}d{i}_plus_dv{i}d{j}_store_1 = dt·S̄_ij)
            σ̄_ij ← σ̄_ij + dt · μ_ij · S̄_ij

The eight checkpointed stores (3D; 5 in 2D, 2 in 1D — elastic.c lines given for reference) are exactly the quantities the Born virtual sources need. Each store holds dt · its finite-difference quantity (the dt is baked in), and each is written only when the corresponding model requires grad:

store name (elastic.c)	holds dt ·	node	multiplies (fwd)	grad guard
dvzdbuoyancy_store_1 (l.994)	W_z = Σ_j ∂_j σ̄_zj	v_z	buoyancy_z	buoyancy_requires_grad
dvydbuoyancy_store_1 (l.1045)	W_y = Σ_j ∂_j σ̄_yj	v_y	buoyancy_y	buoyancy_requires_grad
dvxdbuoyancy_store_1 (l.1098)	W_x = Σ_j ∂_j σ̄_xj	v_x	buoyancy_x	buoyancy_requires_grad
dvzdz_store_1 (l.1168)	Dv̄_z = ∂_z v̄_z	σ_ii	λ (all σ_ii), 2μ (σzz)	lamb_or_mu_requires_grad
dvydy_store_1 (l.1183)	Dv̄_y = ∂_y v̄_y	σ_ii	λ (all σ_ii), 2μ (σyy)	lamb_or_mu_requires_grad
dvxdx_store_1 (l.1197)	Dv̄_x = ∂_x v̄_x	σ_ii	λ (all σ_ii), 2μ (σxx)	lamb_or_mu_requires_grad
dvydx_plus_dvxdy_store_1 (l.1244)	S̄_xy = ∂_y v̄_x + ∂_x v̄_y	σ_xy	mu_yx	mu_requires_grad
dvzdx_plus_dvxdz_store_1 (l.1283)	S̄_xz = ∂_z v̄_x + ∂_x v̄_z	σ_xz	mu_zx	mu_requires_grad
dvzdy_plus_dvydz_store_1 (l.1321)	S̄_yz = ∂_z v̄_y + ∂_y v̄_z	σ_yz	mu_zy	mu_requires_grad

Key discrete fact (used by the adjoint): on the staggered grid the difference operators satisfy the exact-transpose property (the same one that gives D_d = −G_dᵀ for acoustic), so the elastic adjoint kernel FUNC(backward) is a clean transpose: the adjoint velocity update differentiates the stress equation (carries λ/μ/μ_ij · σ_adj products) and the adjoint stress update differentiates the velocity equation (carries b_i · v_adj products). This is the machinery the Born adjoint reuses (§4).

Sources (elastic.py lines 1296–1328): each velocity/force source is pre-scaled by dt · b_i (buoyancy gathered at the source cell) and added into v̄_i; the “pressure” source is pre-scaled by −dt only, with no model factor, and added equally into all normal stresses σ_ii. (The pressure receiver reads −(Σ_i σ_ii)/ndim.) This asymmetry — velocity source carries a model factor, pressure source does not — drives the source-cell tangent difference in §5.

---

2. Born linearisation

Perturb the constitutive fields λ→λ+δλ, μ→μ+δμ, b_i→b_i+δb_i, and (as derived kernel models) μ_ij→μ_ij+δμ_ij. Split the state v = v̄+δv, σ = σ̄+δσ. Substitute into §1.1 and keep first order.

2.1 Velocity equation

∂(v̄_i+δv_i)/∂t = (b_i+δb_i) Σ_j ∂_j (σ̄_ij + δσ_ij)

Zeroth order (background): ∂v̄_i/∂t = b_i Σ_j ∂_j σ̄_ij. First order:

∂δv_i/∂t = b_i Σ_j ∂_j δσ_ij   +   δb_i Σ_j ∂_j σ̄_ij

2.2 Normal-stress equation

∂(σ̄_ii+δσ_ii)/∂t = (λ+δλ) Σ_k ∂_k(v̄_k+δv_k) + 2(μ+δμ) ∂_i(v̄_i+δv_i)

First order:

∂δσ_ii/∂t = λ Σ_k ∂_k δv_k + 2μ ∂_i δv_i   +   δλ Σ_k ∂_k v̄_k + 2δμ ∂_i v̄_i

2.3 Shear-stress equation

∂(σ̄_ij+δσ_ij)/∂t = (μ_ij+δμ_ij)(∂_i(v̄_j+δv_j) + ∂_j(v̄_i+δv_i))

First order:

∂δσ_ij/∂t = μ_ij(∂_i δv_j + ∂_j δv_i)   +   δμ_ij(∂_i v̄_j + ∂_j v̄_i)

2.4 The scattered field obeys the same operator, with virtual sources

Collecting §2.1–§2.3: the scattered field (δv, δσ) satisfies the identical isotropic elastic system with the background parameters (λ, μ, μ_ij, b_i), driven by three families of virtual sources — precisely the quantities the forward already checkpoints (§1.2):

velocity eq  (at v_i node):    virtual force      = δb_i · W̄_i           = δb_i · (dv{i}dbuoyancy_store_1 / dt)
normal-σ eq  (at σ_ii node):   virtual injection  = δλ · Σ_k Dv̄_k + 2δμ · Dv̄_i
shear-σ eq   (at σ_ij node):   virtual injection  = δμ_ij · S̄_ij         = δμ_ij · (dv{j}d{i}_plus_dv{i}d{j}_store_1 / dt)

with W̄_i = dv{i}dbuoyancy_store_1/dt, Dv̄_k = dv{k}d{k}_store_1/dt, S̄_ij = dv{j}d{i}_plus_dv{i}d{j}_store_1/dt. This is the exact elastic analogue of acoustic §2.3 (−δK ∇·v̄, −δB_d ∇_d p̄) and of scalar_born’s 2·v·scatter·dt²·w_sum coupling. Node alignment is automatic: δλ/δμ live at the normal-stress node (where Σ_k Dv̄_k/Dv̄_i are evaluated), δμ_ij at the shear node (where S̄_ij lives), δb_i at the v_i node (where W̄_i lives) — because the scattering models are δ of the same prepared model set the forward consumes (§3).

2.5 Discrete Born forward (one coupled pass)

The Born forward runs the background and scattered fields side by side in one pass, exactly like acoustic_born.c and scalar_born.c. The scattered updates reuse the stored background terms:

scattered velocity:   δv_i  ← δv_i  + b_i·dt·(Σ_j ∂_j δσ_ij + PML_sc)  +  δb_i · (dv{i}dbuoyancy_store_1)
scattered normal σ:   δσ_ii ← δσ_ii + dt·( λ·Σ_k ∂_k δv_k + 2μ·∂_i δv_i + PML_sc )
                                     +  δλ·(Σ_k dv{k}d{k}_store_1) + 2δμ·(dv{i}d{i}_store_1)
scattered shear σ:    δσ_ij ← δσ_ij + dt·μ_ij·(∂_i δv_j + ∂_j δv_i + PML_sc)
                                     +  δμ_ij·(dv{j}d{i}_plus_dv{i}d{j}_store_1)

So the Born virtual sources are literally the stored background terms multiplied by the scattering models — no new spatial derivatives of the background are needed in the Born forward.

Scattered stores, required for the full (two-field) backward (§4.3). When a background model (λ, μ, or any b_i) requires grad — the HVP path (§5) and any differentiable-background case — the full backward also needs the scattered analogues of the eight stores, evaluated on the scattered field:

W_sc_i  = Σ_j ∂_j δσ_ij   (+ PML_sc)          [scattered velocity-buoyancy store]
Dv_sc_k = ∂_k δv_k        (+ PML_sc)          [scattered normal store]
S_sc_ij = ∂_i δv_j + ∂_j δv_i  (+ PML_sc)     [scattered shear store]

This is the elastic counterpart of scalar_born.c/acoustic_born.c checkpointing both the background stores and the scattered stores (wsc_store_1 / div_sc, term_sc_d). Concretely: the forward must checkpoint all eight background stores AND all eight scattered stores whenever lamb.requires_grad or mu.requires_grad or any(buoyancy.requires_grad), roughly doubling per-step checkpoint storage relative to the scattered-only (backward_sc) path (which needs only the background stores). See §4.3 for how these scattered stores enter the background gradients.

---

3. Parameterisation choice: kernel scatters in the prepared model set

Decision: the kernel-level scattering variables are the perturbations of the exact model set elastic.c consumes — (δλ, δμ, δμ_zy, δμ_zx, δμ_yx, δb_z, δb_y, δb_x). The chain rule from the public scattering potentials (dlamb, dmu, dbuoyancy) lives in differentiable Python, as the forward-mode derivative of prepare_parameters.

3.1 prepare_scatter_parameters — forward-mode derivative of prepare_parameters

elastic.py::prepare_parameters(mu, buoyancy) maps (μ, b) to the six staggered derived models: μ_zy/μ_zx/μ_yx are 4-cell harmonic averages of μ in the two named dimensions’ plane, b_z/b_y/b_x are the reciprocal of the 2-cell arithmetic average of ρ = 1/b, all guarded by torch.where masks that zero any node whose stencil touches a vacuum cell (the IVM free surface, §7.1). λ and μ pass through untouched (they already live at the normal-stress node).

Mirroring acoustic_born.py::prepare_scatter (which is just torch.func.jvp over prepare_models), define:

def prepare_scatter_parameters(mu, buoyancy, dmu, dbuoyancy):
    # forward-mode derivative of prepare_parameters at (mu, buoyancy)
    _, tangents = torch.func.jvp(prepare_parameters, (mu, buoyancy), (dmu, dbuoyancy))
    return list(tangents)           # = [δμ_zy, δμ_zx, δμ_yx, δb_z, δb_y, δb_x]  (3D order)

with δλ = dlamb and δμ = dmu passing straight through (λ, μ are consumed at the normal-stress node without averaging). This yields, for the record (what the jvp reproduces automatically):

δμ_ij  = Σ_{four plane neighbours k} ( μ_ij_val² / (4 μ_k²) ) · δμ_k     on the mask; 0 where masked
         (derivative of the 4-cell harmonic mean 4/Σ_k 1/μ_k)
δb_i   = ½ · b_i_val² · ( ρ_a² δb_a + ρ_c² δb_c )                        on the mask; 0 where masked
         (from ρ = 1/b, ρ̄_i = (ρ_a+ρ_c)/2, b_i = 1/ρ̄_i, chain-ruled)

but the code never hand-derives these — torch.func.jvp differentiates the exact masked torch.where expressions, so we inherit node alignment, the correct averaging, AND the IVM vacuum masking “for free” (§7.1). Its VJP (the reverse mode used by backward to chain grad_δμ_*/grad_δb_* back to grad_dmu/grad_dbuoyancy) is handled automatically by autograd, because prepare_scatter_parameters is plain differentiable torch ops sitting in front of the kernel Function — there is no hand-written backward for the scatter step. (torch.func.jvp composes and differentiates fine under autograd, exactly as in acoustic_born.)

3.2 Why the prepared-model space (not physical space) at the kernel boundary

1. These are the perturbations that appear linearly in the wave operator (§2). The map (dmu, dbuoyancy) → (δμ_ij, δb_i) is a pointwise smooth reparameterisation containing no wave physics — the same kind of composition prepare_parameters (and, one level up, vpvsrho_to_lambmubuoyancy) already performs. Keeping it in Python keeps the kernel parameterisation-agnostic.

2. It matches how elastic already works. The nonlinear C kernel never sees vp/vs/ρ or even raw μ/b at shear/velocity nodes — elastic.py calls prepare_parameters and the kernel consumes lamb, mu, mu_ij, buoyancy_i. The Born kernel’s scattering models are then literally δ of the same models the forward consumes: forward and Born share model layout, node alignment, PML profiles, and IVM masking.

3. The adjoint is clean and reuses the elastic backward. Because Born is linear in the prepared scattering models, the kernel produces grad_δλ, grad_δμ, grad_δμ_ij, grad_δb_i with the same imaging condition as the nonlinear elastic model gradients (§4); the physical-parameter gradient is the Python VJP with no hand-derivation.

3.3 Public vs kernel parameterisation

The public elastic_born API takes the scattering potentials dlamb, dmu, dbuoyancy (Lamé/buoyancy space — the same space as elastic’s own public models lamb, mu, buoyancy). The Python front-end composes the eight prepared scatter models and hands them to the autograd Function/kernel; gradients w.r.t. dlamb/dmu/dbuoyancy flow back through the composition automatically. A user wanting physical (dvp, dvs, drho) scatter composes torch.func.jvp(vpvsrho_to_lambmubuoyancy, …) in front (its VJP is again autograd’s job) — no kernel change. (vpvsrho_to_lambmubuoyancy is common.py lines 3329–3356: λ=(vp²−2vs²)ρ, μ=vs²ρ, b=ρ/(ρ²+ε).)

---

4. Adjoint of the linearised system

The Born forward is a linear operator L : (δλ, δμ, δμ_ij, δb_i) ↦ δd (scattered receiver data) for fixed background and acquisition. We need Lᵀ for LSRTM, for the VJP backward, and as a building block of the HVP (§5).

4.1 Imaging conditions (scattered-only backward, backward_sc)

Run the elastic adjoint (elastic.c FUNC(backward), the transpose machinery) on a scattered adjoint field (q^sc_v, q^sc_σ) sourced by the data residuals grad_δd at the receiver locations (velocity-component residuals at velocity receivers, pressure residual −(Σ q^sc_σii)/ndim at pressure receivers), time-reversed. Then correlate with the stored forward background terms, using exactly the same forms as the nonlinear elastic model gradients (elastic.c grad accumulations; dt is already in each store so only step_ratio multiplies):

grad_δλ    = Σ_t ( Σ_i q^sc_σii ) · ( Σ_k Dv̄_k )        · step_ratio      (cf. nonlinear grad_lamb, l.2076–2091)
grad_δμ    = Σ_t 2·( Σ_i q^sc_σii · Dv̄_i )              · step_ratio      (cf. grad_mu,   l.2092–2108)
grad_δμ_yx = Σ_t q^sc_σxy · S̄_xy                        · step_ratio      (cf. grad_mu_yx, l.2173–2177)
grad_δμ_zx = Σ_t q^sc_σxz · S̄_xz                        · step_ratio      (cf. grad_mu_zx, l.2211–2215)
grad_δμ_zy = Σ_t q^sc_σyz · S̄_yz                        · step_ratio      (cf. grad_mu_zy, l.2248–2252)
grad_δb_z  = Σ_t q^sc_vz · W̄_z                          · step_ratio      (cf. grad_buoyancy_z, l.1907–1911)
grad_δb_y  = Σ_t q^sc_vy · W̄_y                          · step_ratio      (cf. grad_buoyancy_y, l.1959–1963)
grad_δb_x  = Σ_t q^sc_vx · W̄_x                          · step_ratio      (cf. grad_buoyancy_x, l.2014–2018)

where Dv̄_k = dv{k}d{k}_store_1, S̄_ij = dv{j}d{i}_plus_dv{i}d{j}_store_1, W̄_i = dv{i}dbuoyancy_store_1 are the checkpointed forward background terms. These are independent of the scatter (L linear), exactly as acoustic’s grad_δK is independent of δK. backward_sc propagates only the scattered adjoint field and accumulates only these eight gradients — the cheap fixed-background LSRTM path, mirroring scalar_born.c/acoustic_born.c backward_sc.

4.2 The adjoint coupling substitution checklist (full two-field backward)

The full backward propagates BOTH a background adjoint field (q^bg_v, q^bg_σ) and a scattered adjoint field (q^sc_v, q^sc_σ). It reuses elastic.c’s transpose adjoint unchanged, EXCEPT that every product of the form MODEL · Q (background model times a background adjoint field) gains a second coupling term + Dmodel · Q_sc (the scattering model times the scattered adjoint field) — the direct analogue of scalar_born.c embedding 2·VDT2·SCATTER·WFCSC into V2DT2_WFC, and acoustic_born.c’s K·P_adj + δK·P_adj_sc / B_d·V_adj_d + δB_d·V_adj_sc_d.

Symbolically, in every site below:

λ · Qσ      →  λ · Qσ      + δλ · Qσ_sc
2μ · Qσ     →  2μ · Qσ     + 2δμ · Qσ_sc            (normal-stress diagonal only)
μ_ij · Qσij →  μ_ij · Qσij + δμ_ij · Qσij_sc
b_i · Qv    →  b_i · Qv    + δb_i · Qv_sc

This substitution MUST be applied at ALL of the following sites — the PLAIN form, the _PML form, AND the memory-variable recursion — because missing the PML/memory variant would silently zero the coupling inside the absorbing layer while looking correct in the interior. The exhaustive site list (from elastic.c; line numbers for the reviewer):

(A) Velocity-adjoint macros — λ/μ/μ_ij · σ_adj (transpose of the stress update). Each has a PLAIN and a _PML form; the _PML form also reads a m_v* memory variable.

macro (PLAIN / _PML)	products to augment with + Dmodel·Qσ_sc
VZ_Z (l.274 / l.280)	λ·σxx + λ·σyy + (λ+2μ)·σzz (PML adds P_BZ·m_vzz)
VZ_Y (l.278 / l.286)	μ_zy·σyz (PML P_BYH·m_vzy)
VZ_X (l.279 / l.289)	μ_zx·σxz (PML P_BXH·m_vzx)
VY_Z (l.296 / l.310) [3D]	μ_zy·σyz (PML P_BZH·m_vyz)
VY_Y (l.299/303 / l.315/321)	λ·σxx + (λ+2μ)·σyy [+ λ·σzz] (PML P_BY·m_vyy)
VY_X (l.308 / l.328)	μ_yx·σxy (PML P_BXH·m_vyx)
VX_Z (l.334 / l.348) [3D]	μ_zx·σxz (PML P_BZH·m_vxz)
VX_Y (l.337 / l.353) [2D+]	μ_yx·σxy (PML P_BYH·m_vxy)
VX_X (l.340/342 / l.358/362)	(λ+2μ)·σxx [+ λ·σyy] [+ λ·σzz] (PML P_BX·m_vxx)

(B) Stress-adjoint macros — b_i · v_adj (transpose of the velocity update). Each has a PLAIN and a _PML form (the _PML form reads a m_sigma* memory variable).

macro (PLAIN / _PML)	product to augment with + δb_i·Qv_sc
SIGMAZZ_Z (l.374 / l.371)	b_z·v_z (PML P_BZH·m_sigmazzz)
SIGMAYY_Y (l.380 / l.377)	b_y·v_y (PML P_BYH·m_sigmayyy)
SIGMAXX_X (l.385 / l.382)	b_x·v_x (PML P_BXH·m_sigmaxxx)
SIGMAXY_Y (l.391 / l.388)	b_x·v_x (PML P_BY·m_sigmaxyy)
SIGMAXY_X (l.395 / l.392)	b_y·v_y (PML P_BX·m_sigmaxyx)
SIGMAXZ_Z (l.402 / l.399) [3D]	b_x·v_x (PML P_BZ·m_sigmaxzz)
SIGMAXZ_X (l.406 / l.403) [3D]	b_z·v_z (PML P_BX·m_sigmaxzx)
SIGMAYZ_Z (l.411 / l.408) [3D]	b_y·v_y (PML P_BZ·m_sigmayzz)
SIGMAYZ_Y (l.415 / l.412) [3D]	b_z·v_z (PML P_BY·m_sigmayzy)

(C) Velocity-block memory-variable recursions — b_i · v_adj (the m_sigma*n updates, written when !DW_PML_SKIP). Each is of the form b_i·dt·P_A?·v_shot + P_A?·m_…; augment the b_i·…·v_shot term with + δb_i·…·v_sc_shot:

vz block:  m_sigmazzzn (l.1892–1894, b_z·P_AZH·vz),  m_sigmayzyn (l.1897–1899, b_z·P_AY·vz),  m_sigmaxzxn (l.1902–1904, b_z·P_AX·vz)
vy block:  m_sigmayzzn (l.1943–1945, b_y·P_AZ·vy),  m_sigmayyyn (l.1949–1951, b_y·P_AYH·vy), m_sigmaxyxn (l.1954–1956, b_y·P_AX·vy)
vx block:  m_sigmaxzzn (l.1996–1998, b_x·P_AZ·vx),  m_sigmaxyyn (l.2003–2005, b_x·P_AY·vx),  m_sigmaxxxn (l.2009–2011, b_x·P_AXH·vx)

(D) Stress-block memory-variable recursions — λ/(λ+2μ)/μ_ij · σ_adj (the m_v* updates, written when !DW_PML_SKIP). Augment every MODEL·…·σ_shot term with + Dmodel·…·σ_sc_shot:

m_vzz (l.2111–2117): (λ+2μ)·P_AZ·σzz + λ·P_AZ·σxx + λ·P_AZ·σyy   (+ P_AZ·m_vzz)
m_vyy (l.2121–2129): (λ+2μ)·P_AY·σyy + λ·P_AY·σxx [+ λ·P_AY·σzz]  (+ P_AY·m_vyy)
m_vxx (l.2132–2142): (λ+2μ)·P_AX·σxx [+ λ·P_AX·σyy] [+ λ·P_AX·σzz] (+ P_AX·m_vxx)
m_vxy (l.2180–2182): μ_yx·P_AYH·σxy   (+ P_AYH·m_vxy)
m_vyx (l.2185–2187): μ_yx·P_AXH·σxy   (+ P_AXH·m_vyx)
m_vxz (l.2218–2220): μ_zx·P_AZH·σxz   (+ …)          [3D]
m_vzx (l.2223–2225): μ_zx·P_AXH·σxz   (+ …)          [3D]
m_vyz (l.2255–2257): μ_zy·P_AZH·σyz   (+ …)          [3D]
m_vzy (l.2260–2262): μ_zy·P_AYH·σyz   (+ …)          [3D]

That is the complete enumeration: 9 velocity-adjoint macros (A) + 9 stress-adjoint macros (B) + 9 velocity memory recursions (C) + 9 stress memory recursions (D), each in its PLAIN and (for A,B) _PML variant, gains the + Dmodel · Q_sc coupling. Note the diagonal normal-stress coupling carries the factor 2 on δμ (from 2μ), and (λ+2μ) couplings split into a δλ part and a 2δμ part.

4.3 Full-backward background gradients (two terms each, no third term)

Using the scattered stores of §2.5, the background model gradients in the full backward are the background-adjoint × background-store term PLUS the scattered-adjoint × scattered-store term:

grad_λ    = Σ_t ( (Σ_i q^bg_σii)·(Σ_k Dv̄_k) + (Σ_i q^sc_σii)·(Σ_k Dv_sc_k) ) · step_ratio
grad_μ    = Σ_t 2·( (Σ_i q^bg_σii·Dv̄_i) + (Σ_i q^sc_σii·Dv_sc_i) )           · step_ratio
grad_μ_ij = Σ_t ( q^bg_σij·S̄_ij + q^sc_σij·S_sc_ij )                          · step_ratio
grad_b_i  = Σ_t ( q^bg_vi·W̄_i  + q^sc_vi·W_sc_i )                             · step_ratio

There is NO third “coupling-coefficient” cross-term (unlike scalar_born). Same argument as acoustic §4.3: in scalar_born the coupling 2·v·scatter·dt²·L(u) has explicit dependence on the background v (the leading factor), so differentiating the coupling itself w.r.t. v yields an extra term. Here every elastic Born coupling — δb_i·W̄_i, δλ·ΣDv̄_k + 2δμ·Dv̄_i, δμ_ij·S̄_ij — has no explicit dependence on the background models λ/μ/μ_ij/b_i: the background enters ONLY through the stored terms W̄_i/Dv̄_k/S̄_ij (the background state, not the coupling coefficient). That state-dependence is exactly what the q^sc · (scattered store) cross-term already captures, via the ordinary background-field adjoint recursion the plain nonlinear elastic backward already runs. So each grad_* above has exactly two terms.

Consequently the full backward requires the forward to also checkpoint the SCATTERED stores (W_sc_i, Dv_sc_k, S_sc_ij) whenever any background model needs grad (§2.5). backward_sc (fixed background) needs neither the coupling substitution of §4.2 nor the scattered stores — it propagates only q^sc and accumulates only the grad_δ* of §4.1.

---

5. Second-order (HVP) structure, mirroring AcousticBackwardFunc

Goal: for the nonlinear elastic propagator with a scalar loss J(m), m = (λ, μ, b) (or physical (vp, vs, ρ)), compute H·h without forming H, where h is a model-space perturbation.

5.1 Why Born powers the HVP

The first-order backward computes g(m) = ∂J/∂m as a correlation of the forward wavefield with the adjoint wavefield. H·h = d/dα|₀ g(m+αh) has two contributions — the perturbation of the forward field by h and the perturbation of the adjoint field by h — and both are Born propagations with scatter = h. Hence one Born forward + one Born backward per HVP, exactly as scalar/acoustic use their Born backend as the HVP engine.

5.2 The two-pass mechanism (elastic analogue of AcousticBackwardFunc.backward)

Mirroring AcousticForwardFunc.backward (acoustic.py l.1259–1324), replace the @no_second_order C-backend call in ElasticForwardFunc.backward with a nested differentiable Function, ElasticBackwardFunc:

• ElasticBackwardFunc.forward runs the ordinary first-order elastic adjoint (via the shared run_backward / run_backward_core with the elastic SPEC), seeding the adjoint velocity/stress fields directly from the output-wavefield gradients grad_wavefields (no negation), with the elastic memory-variable aux swap (§5.3) as the only ping-pong. Returns the first-order grad_models, grad_sources, and final adjoint wavefields.

• ElasticBackwardFunc.backward (@once_differentiable) is the two-pass HVP:

  • Step 1 — Born forward. Run ElasticBornForwardFunc.forward (called directly, not .apply, so no graph is built) with background = m (models .detach().requires_grad_(True) to force both background and scattered store checkpointing) and scatter = h (i.e. (δλ, δμ, δμ_ij, δb_i) = chain-rule image of h via prepare_scatter_parameters). The background fields recompute the forward wavefield; the scattered fields give the effect of h on the forward wavefield. Record the scattered field at the original receivers (feeds the adjoint-source term of H·h); background receivers empty. h is zero-filled where a model/source/ wavefield direction is absent.

  • Step 2 — Born full backward. Run ElasticBornForwardFunc’s full backward (via run_backward_core with the Born SPEC, reusing Step-1’s storage manager) with the ORIGINAL adjoint source grad_r (velocity/stress data residuals) as the scattered adjoint source and the ORIGINAL grad_wavefields seeding the scattered adjoint field, while the background adjoint field is seeded at zero and carries h’s effect through the Born coupling (§4.2). Their combination accumulates the second-order model gradient H·h (state.grad_models[:n]), plus HVP source and initial-wavefield terms.

  Field/seed assignment mirrors acoustic exactly: scattered adjoint = ordinary backprop (sourced by grad_r, seeded by grad_wavefields); background adjoint = zero-seeded, coupling-driven. Return: None for the leading inputs, then grad wrt grad_receivers (= the Step-1 scattered receiver output), H·h for the models, HVP sources, grad wrt grad_wavefields (= Step-1 scattered forward wavefields), and HVP initial wavefields.

5.3 No negate hook — only the memory-variable aux swap

Verified from ElasticForwardFunc.backward: there is no grad_wfp-style negation hook (that scheme is specific to scalar/scalar_born’s two-time-level wfc/wfp storage). Elastic’s first-order backward seeds the adjoint (v_i, σ_ij) fields directly from grad_wavefields, and the only ping-pong is the PML memory-variable aux swap — but only on ONE of the two memory-variable families. Per the §4.2(C)/(D) site enumeration: the m_sigma* recursions (§4.2(C), the velocity-adjoint block’s memory variables — m_sigmazzzn, m_sigmayzyn, …) are the ones double-buffered m_*/m_*n and swapped by run_backward each odd-length chunk. The m_v* recursions (§4.2(D), the stress-adjoint block’s memory variables — m_vzz, m_vyy, …, elastic.c lines 2112–2262) are updated in place, with no swap. A future kernel implementer must not swap m_v* in the Born two-field HVP pass — only the m_sigma* blocks of both the background and scattered fields get the aux-swap treatment (mirroring _acoustic_born_swap, which swaps both fields’ psi blocks; here “psi” corresponds to the m_sigma* family only). No negate hook is added for either field.

5.4 Correspondence table (acoustic → elastic)

acoustic (AcousticBackwardFunc)	elastic (ElasticBackwardFunc)
models (K, B_z, B_y, B_x)	models (λ, μ, μ_zy, μ_zx, μ_yx, b_z, b_y, b_x)
h → (δK, δB_d) via jvp(prepare_models)	h → (δλ, δμ, δμ_ij, δb_i) via jvp(prepare_parameters) + λ/μ passthrough
couplings −dt·δK·div̄, −dt·δB_d·term̄	δλ·ΣDv̄+2δμ·Dv̄_i, δμ_ij·S̄_ij, δb_i·W̄_i
adjoint source = residuals for p, v_d	adjoint source = residuals for v_z,v_y,v_x + pressure
aux swap = psi block; NO negate	aux swap = m_v*/m_sigma* blocks; NO negate
Born fwd + Born bwd = H·h	Born fwd + Born bwd = H·h

5.5 Wiring change required (replaces the no_second_order guard) — DONE

ElasticForwardFunc.backward no longer uses @no_second_order; it builds ElasticBackwardFunc.apply(...) (as AcousticForwardFunc.backward builds AcousticBackwardFunc), whose @once_differentiable backward is the §5.2 double-backward. The no_second_order guard is removed from elastic. The double-backward drives the compiled Born backend (elastic_born_compiled_core + elastic_born_backward_compiled) rather than the eager driver, and raises NotImplementedError (→ backend="eager") under M-PML, where the compiled Born backend has no _mpml variant.

---

6. Forward-mode (jvp) via Born, mirroring AcousticForwardFunc.jvp

Replace ElasticForwardFunc.jvp = unsupported_jvp (l.1630) with a Born-routed rule mirroring _acoustic_forward_jvp. With the primal inputs saved via save_for_forward in setup_context, the tangent map is:

(dlamb, dmu, dbuoyancy) tangent → (δλ, δμ, δμ_ij, δb_i) via prepare_scatter_parameters → Born scatter
source tangents                → Born scattered sources (pre-scaling tangents per §5-src below)
initial-wf tangents            → Born scattered initial wavefields (all 13/27 components)

One ElasticBornForwardFunc.forward call — invoked directly on no_grad tensors at base dispatch level (_base_level_dispatch() clearing the functorch stack, so the ctypes backend sees real storage and builds no graph), storage_mode_str="none" — yields primal = elastic(...) (Born background outputs) and tangent = elastic_born scattered outputs together:

jvp(elastic, (λ,μ,b); (δλ,δμ,δb)).tangent = elastic_born(dlamb=δλ, dmu=δμ, dbuoyancy=δb).scattered_output
jvp(elastic, …).primal                     = elastic(λ,μ,b) = Born background output

Guards (mirroring acoustic): raise NotImplementedError on nested functorch depth > 1; warn (_model_tangent_in_pml) when a model tangent is nonzero inside the PML band.

Non-PML boundary-face guard — not applicable yet. Acoustic’s jvp additionally raises on any non-PML boundary face (boundary="free"/"rigid"/"periodic", Task G). elastic has no boundary argument at all today (Task G’s first-class boundary faces are a scalar/acoustic-only feature, not yet extended to elastic), so this guard is currently vacuous for elastic and is dropped from the initial jvp implementation. Note it here only as future-proofing: if/when elastic gains a boundary argument, ElasticForwardFunc.jvp must add the matching NotImplementedError guard at that time (Born does not linearise a non-PML boundary fill any more than it linearises the PML profile).

Source-cell tangents (§5-src) — the elastic asymmetry. Because the forward pre-scales sources (§1.2), differentiating the pre-scaling gives the scattered source:

velocity/force source tangent:  dt · ( δb_i · f_i  +  b_i · δf_i )      injected into δv_i
pressure source tangent:         − dt · δf_p                             injected into all δσ_ii

The velocity tangent carries the buoyancy model factor (exactly like acoustic’s dt·(δB_d·f + B_d·δf)); the pressure tangent has NO model factor (−dt · δf_p only) because the elastic pressure source is pre-scaled by −dt alone — unlike acoustic’s dt·(δK·s + K·δs). As in acoustic, these pre-scaled tangents are produced automatically by autograd through elastic’s differentiable source preprocessing, so they inject directly with no double counting.

PML caveat (identical to scalar/acoustic). Born does not linearise the PML profile’s dependence on the model (through max_vel/absorber design), so the jvp is exact only where the tangent is zero in the PML region (or pml_width = 0); a warning is emitted otherwise.

---

7. Elastic-specific complications

7.1 IVM free surface — scatterers automatically vanish in vacuum

The Immersed Vacuum Method encodes the free surface entirely in the models: lamb = mu = buoyancy = 0 in vacuum cells, with pml_width = 0 on that face (elastic docstring l.561–572). prepare_parameters turns those zeros into correct coefficients via torch.where masks (elastic.py l.187–268): a shear node whose 4-cell harmonic stencil touches any vacuum cell is masked to μ_ij = 0; a velocity node whose 2-cell density average is (near) zero is masked to b_i = 0, while a node straddling one solid + one vacuum cell gets b_i = 2/(ρ_solid + 0) = 2·b_solid (the arithmetic density average naturally halving the surface density).

Born inherits this for free. Because prepare_scatter_parameters is the forward-mode derivative of this same masked composition, the scattering models δμ_ij, δb_i are automatically zero wherever the corresponding prepared model is masked to zero (the mask is a piecewise-constant torch.where selector, so its jvp propagates the same indicator): a scatterer in a vacuum cell produces no scattered source. The “masked-average JVP” also correctly linearises the vacuum-adjacent nodes (the solid-side node straddling the surface, b_i = 2·b_solid), which is exactly where a naive fixed-stencil δ would get the surface physics wrong. So the free surface needs NO special Born handling beyond routing the scatter through prepare_scatter_parameters. Validity: Born is a first-order interior perturbation — perturbations must keep vacuum cells vacuum (i.e. dlamb/dmu/dbuoyancy are zero in vacuum); a perturbation that moves the free surface (creates/removes vacuum) is outside the linearisation (the mask indicator is not differentiated).

7.2 M-PML (CFS-M-PML) — keep the _mpml machinery; compile _mpml variants of all three entry points

The elastic kernel’s M-PML is selected by DW_MPML, which makes ELASTIC_VARIANT(name) append _mpml to the symbol and switches the P_* profile accessors to full-grid arrays and the DW_PML_SKIP region test to interior-only (elastic.c l.71–184). elastic_born.c must keep this machinery verbatim and compile _mpml variants of all three Born entry points (forward, backward, backward_sc): build each once with DW_MPML undefined and once defined; the _mpml suffix, full-grid P_* loads, and union-of-strips DW_PML_SKIP follow automatically from the macros. The coupling substitution of §4.2 applies inside the _PML macro forms regardless of the _mpml variant (the coupling reads the same adjoint/scatter fields; only the P_* accessor implementation differs). Dispatch is via backend_extra="_mpml" through run_forward_engine/run_backward (as elastic already does; mpml_fraction > 0 selects it). Validity: as with the free surface, Born is valid for interior perturbations only; a perturbation inside the M-PML band is not linearised (§6 PML caveat).

7.3 Doubled wavefields (26 in 2D, 54 in 3D) and the _sc initial-wavefield API

The elastic wavefield set is 4 (1D) / 13 (2D) / 27 (3D) components: velocities v_i, normal stresses σ_ii, shear stresses σ_ij, velocity-PML memory vars m_v* (9 in 3D), and stress-PML memory vars m_sigma* (9 in 3D) (elastic_wavefield_names). Born doubles this — a background field and a scattered field of every component — giving 26 (2D) / 54 (3D) wavefields, exactly as acoustic_born doubles the acoustic set. Only the physical velocity and stress components of each field are exposed as public initial-wavefield keyword arguments – BOTH a background set (vz_0, sigmazz_0, …) AND a scattered set (vz_sc_0, sigmazz_sc_0, …); the m_v*/m_sigma* PML memory variables of both fields are not keyword arguments at all, and are reachable only through the initial_state bundle (as part of the folded auxiliary_state). Outputs return the final background wavefields, one opaque auxiliary_state holding both fields’ PML memory variables, then the final scattered wavefields, then background receiver amplitudes, then scattered receiver amplitudes (the “background data before scattered data” convention). The eager core’s wavefields_flat = args_list[-27:]/[-13:]/[-4:] slicing must become the doubled slice for the Born spec.

---

8. Planned public API

Mirrors acoustic_born (background + scattered sources/receivers, both fields’ initial wavefields) and elastic (per-component velocity force sources + a pressure source, mpml_fraction).

def elastic_born(
    lamb, mu, buoyancy, grid_spacing, dt,
    dlamb=None, dmu=None, dbuoyancy=None,            # scattering potentials (Lamé/buoyancy space)
    # background + scattered share source *locations* (§5-src); amplitudes given once:
    source_amplitudes_z=None, source_locations_z=None,
    source_amplitudes_y=None, source_locations_y=None,
    source_amplitudes_x=None, source_locations_x=None,
    source_amplitudes_pressure_rate=None, source_locations_pressure_rate=None,
    # scattered-field receivers (the Born data) and, optionally, background receivers:
    sc_receiver_locations_z=None, sc_receiver_locations_y=None,
    sc_receiver_locations_x=None, sc_receiver_locations_p=None,
    bg_receiver_locations_z=None, bg_receiver_locations_y=None,
    bg_receiver_locations_x=None, bg_receiver_locations_p=None,
    accuracy=4, pml_width=20, pml_freq=None, mpml_fraction=0.0,
    max_vel=None, courant_number=None, survey_pad=None,
    # background initial wavefields (physical velocity/stress components only):
    vz_0=None, vy_0=None, vx_0=None, sigmazz_0=None, …, sigmaxx_0=None,
    # scattered initial wavefields (physical components only): vz_sc_0, …, sigmaxx_sc_0
    vz_sc_0=None, …, sigmaxx_sc_0=None,
    # the m_v*/m_sigma* PML memory variables of both fields are not keyword
    # arguments; they are reachable only through initial_state/auxiliary_state.
    initial_state=None, origin=None, nt=None,
    model_gradient_sampling_interval=1,
    freq_taper_frac=0.0, time_pad_frac=0.0, time_taper=False,
    forward_callback=None, backward_callback=None, callback_frequency=1,
    backend="compiled",
    storage_mode="device", storage_path=".", storage_compression=False,
    location_interpolation="auto", location_interpolation_halfwidth=4,
) -> PropagatorOutput:
    ...

Plus an ElasticBorn(torch.nn.Module) wrapper storing lamb, mu, buoyancy, grid_spacing (and dlamb, dmu, dbuoyancy if desired), matching Elastic/AcousticBorn.

Front-end composition (mirroring acoustic_born):

1. setup_propagator pads/survey-crops [lamb, mu, buoyancy, dlamb, dmu, dbuoyancy].

2. bg_models = [lamb_p, mu_p] + prepare_parameters(mu_p, buoyancy_p) = [λ, μ, μ_zy, μ_zx, μ_yx, b_z, b_y, b_x].

3. sc_models = [dlamb_p, dmu_p] + prepare_scatter_parameters(mu_p, buoyancy_p, dmu_p, dbuoyancy_p) = [δλ, δμ, δμ_zy, δμ_zx, δμ_yx, δb_z, δb_y, δb_x].

4. Shared-source scaling: background velocity amp = amp · gather(b_i) · dt, scattered virtual “source-cell” amp = amp · gather(δb_i) · dt, both at the shared sources_i; pressure background = −amp·dt, scattered = −amp_tangent·dt (no model factor). Where an amplitude is absent the scattered slot mirrors the background slot.

5. Call ElasticBornForwardFunc.apply(...); downsample receiver outputs.

PropagatorSpec (name="elastic_born", staggered=True):

• model_names(ndim) = background prepared models then scattered prepared models ([λ, μ, μ_ij…, b_i…, δλ, δμ, δμ_ij…, δb_i…]).

• wavefield_names(ndim) = elastic background wavefields (§7.3) followed by scattered ones (doubled: 26/54).

• n_source_types = ndim + 1; n_receiver_types = 2·(ndim + 1) (background + scattered) — or model background/scattered receiver sets like acoustic_born.

• output_names = doubled reordered wavefields + background then scattered receiver amplitudes {z,y,x,p}.

• Scatter-model flag indices analogous to elastic’s c_model_flag_indices=[0,1,-1] (the independently-differentiable owners: lamb, mu, raw buoyancy on the background side and dlamb, dmu, raw dbuoyancy on the scattered side).

Differentiability: outputs differentiable w.r.t. lamb, mu, buoyancy, dlamb, dmu, dbuoyancy, source amplitudes, and initial wavefields — all first-order gradients need no new hand-derivation beyond the kernel’s grad_δλ/grad_δμ/grad_δμ_ij/grad_δb_i (§4), because the prepare_scatter_parameters and prepare_parameters chain rules are differentiable Python (§3).

Note — this deliberately EXCEEDS the acoustic_born precedent. acoustic_born’s public API does not make background (v, rho) gradients a first-class public differentiability promise: background gradients there are produced only as a Function-level internal capability of the full two-field backward (consumed by AcousticBackwardFunc’s HVP, §5, and by the general autograd machinery when the background happens to require grad), not as a documented/tested public guarantee of acoustic_born(...). This document’s claim that elastic_born differentiability w.r.t. lamb, mu, buoyancy is public and first-class is thus an intentional extension beyond that precedent, not an oversight — flagged explicitly so the reviewer can weigh in. When implementing, ElasticBornForwardFunc.backward must mirror AcousticBornForwardFunc’s dispatch logic (acoustic_born.py lines ~843–857: dispatch to the full backward vs. the cheaper backward_sc based on whether any background-side input — model, source amplitude, or initial wavefield — needs grad) and its conditional scattered-store checkpointing flags (lines ~793–806: background stores saved whenever any model needs grad; scattered stores saved only when a background model needs grad), for consistency with the rest of the Born family.

---

9. Planned kernel filenames and symbols

• src/deepwave/elastic_born.c and src/deepwave/elastic_born.cu (kept in lockstep; .cu GPU-tested by the maintainer before release).

• Symbol convention (mirroring elastic_iso_* and acoustic_born_iso_*): elastic_born_iso_<ndim>d_<accuracy>_<dtype>_<name>[_mpml]_<device>, with name ∈ {forward, backward, backward_sc} — e.g. elastic_born_iso_2d_4_float_forward_cpu and its …_forward_mpml_cpu M-PML variant. The _mpml infix is inserted by ELASTIC_VARIANT before the ndim/accuracy suffix, exactly as in elastic.c.

  • forward: single coupled pass, background + scattered fields (§2.5); conditionally checkpoints the scattered stores when a background model needs grad (§2.5/§4.3).

  • backward: full two-field adjoint, coupling substitution at every §4.2 site, accumulating background and scattered model gradients (§4.3) — powers HVP and background gradients.

  • backward_sc: scattered-only adjoint, accumulates only grad_δλ/grad_δμ/ grad_δμ_ij/grad_δb_i (§4.1), no coupling, no scattered stores.

• Registered under backend name "elastic_born" in backend_utils (get_backend_function("elastic_born", ndim, name, …)), mirroring "acoustic_born".

• CMakeLists.txt gains the elastic_born translation units in the same ndim × accuracy × dtype × device specialisation matrix as elastic. (Done for the default (non-M-PML) variant; there is no compiled M-PML Born variant, which falls back to backend="eager".)

---

10. Planned validation gates

Complete only when all pass (CPU C backend + eager backend; .cu compiled, GPU-run by the maintainer). Names mirror tests/test_acoustic_born.py / tests/test_elastic.py.

1. Born linearity (per parameter) — the scattered output is exactly linear in each of (dlamb, dmu, dbuoyancy) separately and jointly: scaling a scatterer by α scales the scattered receivers by α; all-zero scatter gives zero scattered output and background output byte-identical to elastic(λ,μ,b).

2. Born == directional FD w.r.t. lamb / mu / dbuoyancy — scattered receivers equal (elastic(m+ε h) − elastic(m−ε h)) / (2ε) for interior perturbations h, for dlamb, dmu, dbuoyancy independently and mixed; background receivers equal elastic(λ,μ,b).

3. jvp == Born — torch.func.jvp(elastic, (λ,μ,b), (δλ,δμ,δb)) tangent equals the elastic_born scattered output (eager ~1e-12; compiled C ~1e-6 fast-math tol); legacy forward_ad path agrees; the PML-boundary caveat warns; nested-transform and non-PML-boundary cases raise. Status: DONE for both the eager path and the compiled elastic propagator’s ElasticForwardFunc.jvp (which currently computes the tangent through the eager Born driver): the two new tests/test_elastic_born.py::test_jvp_compiled_backend_matches_born[2,3] and ::test_jvp_compiled_backend_guards verify jvp == Born to machine precision on the compiled backend, the PML-band-tangent warning, and the nested-transform NotImplementedError. (When the compiled elastic_born backend lands the tolerance moves to the ~1e-6 fast-math band.)

4. HVP == FD-of-gradient — the second-order path gives H·h equal to (∇J(m+ε h) − ∇J(m−ε h)) / (2ε) for J a scattered/background data misfit and h an interior perturbation, for (λ,μ,b) and cross terms. Replaces the current elastic test_second_order_raises.

5. gradcheck + gradgradcheck — torch.autograd.gradcheck on elastic_born w.r.t. lamb, mu, buoyancy, dlamb, dmu, dbuoyancy, source amplitudes, and (background + scattered) initial wavefields (float64, accuracies 2/4/6/8, step_ratio 1 and >1, batched models, pml_width=0 and asymmetric); torch.autograd.gradgradcheck on elastic exercising the HVP path. Plus ≥1 absolute FD gradcheck on a tiny problem as ground truth; cross-backend (C vs eager) parity on forward, gradient, and HVP.

6. Reciprocity — source↔receiver reciprocity of the Born scattered response (accounting for velocity/stress/pressure source–receiver duality on the staggered grid).

7. Free-surface interaction — Born linearity + Born == FD hold with an IVM free surface present (vacuum cells + pml_width=0 on that face), for interior perturbations; scatterers in vacuum cells produce zero scattered output (§7.1).

8. M-PML interaction — Born linearity + Born == FD hold with mpml_fraction > 0; the _mpml Born variants match the default variants where the M-PML is inactive; the union-of-strips region behaviour is preserved (interior perturbation only).

9. ABI additive — abi_reference.json extended with the additive elastic_born_* symbols; no drift to existing elastic_* symbols.

10. CUDA lockstep — elastic_born.cu compiles clean under nvcc across all ndim × accuracy × dtype × {default, mpml} specialisations; GPU-run parity with the CPU C backend on forward, gradient, and HVP (maintainer-run).

Supporting checks: byte-identity of background outputs vs plain elastic when the scatter is zero; IGNORE_LOCATION handling; storage-mode (device/cpu/disk/none) and checkpointing parity (including the conditional scattered-store checkpointing); callback parity.

---

11. Summary of decisions (for the reviewer)

• Forward system: ∂v_i/∂t = b Σ_j ∂_j σ_ij (+bf_i), ∂σ_ii/∂t = λΣ_k∂_k v_k + 2μ∂_i v_i, ∂σ_ij/∂t = μ_ij(∂_i v_j + ∂_j v_i); staggered leapfrog with CPML/M-PML (elastic.c). IVM free surface encoded entirely in the models.

• Born: scattered (δv, δσ) obeys the same operator with background (λ, μ, μ_ij, b_i), virtual sources δb_i·W̄_i (velocity), δλ·ΣDv̄_k + 2δμ·Dv̄_i (normal σ), δμ_ij·S̄_ij (shear σ), where W̄/Dv̄/S̄ are exactly the eight forward stores dv{i}dbuoyancy_store_1 / dv{k}d{k}_store_1 / dv{j}d{i}_plus_dv{i}d{j}_store_1.

• Parameterisation (and why): the kernel scatters in δ of the prepared model set (δλ, δμ, δμ_ij, δb_i); the public API is Lamé/buoyancy space (dlamb, dmu, dbuoyancy), and prepare_scatter_parameters = torch.func.jvp(prepare_parameters, …) (+ λ/μ passthrough) does the chain rule in differentiable Python. Chosen because (1) these are the perturbations linear in the operator, (2) it matches how elastic already prepares its models (same layout, node alignment, PML profiles, IVM masking), (3) the adjoint is the same imaging condition as the nonlinear model gradient, and the physical/VJP chain rule is autograd’s job — it also gives node alignment AND IVM vacuum masking for free.

• Adjoint coupling checklist: every λ·Qσ, 2μ·Qσ, μ_ij·Qσij, b_i·Qv product gains +Dmodel·Q_sc, at all 36 sites: 9 velocity-adjoint macros (A) + 9 stress- adjoint macros (B), each in PLAIN and _PML forms, + 9 velocity memory recursions (C) + 9 stress memory recursions (D). Full-backward background gradients have exactly TWO terms each (bg adjoint × bg store + sc adjoint × sc store); NO third coupling- coefficient term (the couplings have no explicit λ/μ/b dependence).

• Elastic-specific complications: (1) IVM free surface — scatterers vanish in vacuum and the masked-average JVP handles vacuum-adjacent nodes automatically, because the scatter is the jvp of the masked prepare_parameters; Born valid for interior perturbations that keep vacuum vacuum. (2) M-PML — keep ELASTIC_VARIANT/P_*/ DW_PML_SKIP, compile _mpml variants of all three entry points; Born valid for interior perturbations only. (3) Doubled wavefields 26 (2D)/54 (3D) and a _sc initial-wavefield API for every component. (4) Pressure source tangent is −dt·δf_p ONLY (no model factor) — unlike acoustic — while velocity source tangent is dt·(δb_i·f + b_i·δf).

• HVP: one Born forward + one Born backward per H·h, mirroring AcousticBackwardFunc; requires ElasticForwardFunc.backward to build a nested differentiable ElasticBackwardFunc and removal of the no_second_order guard; NO negate hook (only the m_v*/m_sigma* memory-variable aux swap, applied to both fields).

• jvp: Born-based ElasticForwardFunc.jvp replacing unsupported_jvp, mirroring _acoustic_forward_jvp, with the same PML caveat and the pressure/velocity source- tangent asymmetry above.

• Files: elastic_born.c/.cu; symbols elastic_born_iso_<ndim>d_<accuracy>_<dtype>_<forward|backward|backward_sc>[_mpml]_<device>; backend name "elastic_born".

• Gates: Born linearity per parameter; Born == directional FD w.r.t. lamb/mu/dbuoyancy; jvp == Born; HVP == FD-of-gradient; gradcheck + gradgradcheck; reciprocity; free-surface interaction; M-PML interaction; ABI additive; CUDA lockstep.

File path: docs/elastic_born_derivation.md (this document).