test: faster CPU suite, FD checks that check

test/test_acoustic_tti.py

@@ -10,6 +10,8 @@
 radian angle, not the ~1° the bug produced.
 """
 
+from functools import lru_cache
+
 import numpy as np
 import pytest
 import torch
@@ -43,9 +45,11 @@ def _src(nz, nx):
     return torch.tensor(np.array([[nx // 2, nz // 2]], dtype=np.int64)[None])
 
 
+@lru_cache(maxsize=None)
 def _wavefront_tilt_deg(theta_rad, nz=NZ, nx=NX):
     """Propagate AcousticTTI with a constant ``theta`` (radians) and return the
-    principal-axis tilt of the final |h1| snapshot in degrees."""
+    principal-axis tilt of the final |h1| snapshot in degrees.  Cached: both
+    tests start from theta = 0, and the run is deterministic."""
     shape = (nz, nx)
     models = [
         torch.full(shape, VP_BG, dtype=torch.float32),

test/test_biot.py

@@ -9,6 +9,7 @@
 
 from sweep.equations import Biot, Elastic
 from sweep.equations.biot import biot_from_biot_constants, biot_from_rock, biot_wave_velocities
+from sweep.propagator.options import EagerOptions
 from sweep.propagator.torch import PropTorch
 from sweep.signal import ricker
 
@@ -23,9 +24,13 @@ def _wavelet(nt, dt, f0):
     return torch.from_numpy(ricker(t, f=f0))
 
 
-def _prop(eq, shape, dh, dt, src, rec, abcn=20, free_surface=False):
+def _prop(eq, shape, dh, dt, src, rec, abcn=20, free_surface=False, compiled=True):
+    # A cold torch.compile of the step costs seconds per grid, source field and gradient mode,
+    # as much as a few thousand uncompiled steps of these grids: only the runs with many steps
+    # per compiled graph (wave speeds, wavenumber integration) compile.
     return PropTorch(eq, shape=shape, dev=DEV, dh=dh, dt=dt, source_type=src, receiver_type=rec,
-                     abcn=abcn, free_surface=free_surface, backend="torch", impl="eager")
+                     abcn=abcn, free_surface=free_surface, backend="torch", impl="eager",
+                     eager_options=EagerOptions(use_compile=compiled))
 
 
 def _fields(coef, shape):
@@ -63,8 +68,10 @@ def test_no_relative_flow_limit_is_undrained_elastic():
     src = np.array([[35, 20]])
     rec = np.stack([np.arange(5, 65, 3), np.full(20, 40)], 1)[None]
     w = _wavelet(nt, dt, 12.0)
-    db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"])(w, src, rec, models=coef)
-    de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"])(w, src, rec, models=[vp, vs, rho])
+    db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], compiled=False)(
+        w, src, rec, models=coef)
+    de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], compiled=False)(
+        w, src, rec, models=[vp, vs, rho])
     err = float((db - de).abs().max() / de.abs().max())
     assert err < 1e-4, err
 
@@ -103,11 +110,16 @@ def test_gradient_matches_finite_difference_inside(free_surface):
     src = np.array([[24, 10]])
     rec = np.stack([np.arange(8, 40, 4), np.full(8, 30)], 1)[None]
     w = _wavelet(nt, dt, 10.0)
-    prop = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz", "p"], free_surface=free_surface)
+    prop = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz", "p"], free_surface=free_surface,
+                 compiled=False)
     obs = prop(w, src, rec, models=coef).detach()
+    # each receiver type in its own unit: as recorded, p (~1e-2 Pa) outweighs vx and vz
+    # (~1e-8 m/s) by 1e13 in the misfit, which left the solid-velocity adjoint unchecked
+    weight = 1.0 / obs.abs().amax(dim=(0, 1, 2))
 
     def J(mu):
-        return 0.5 * (prop(w, src, rec, models=[coef[0], mu, *coef[2:]]) - 1.05 * obs).double().square().sum()
+        r = (prop(w, src, rec, models=[coef[0], mu, *coef[2:]]) - 1.05 * obs) * weight
+        return 0.5 * r.double().square().sum()
 
     mu0 = coef[1].clone().requires_grad_()
     (g,) = torch.autograd.grad(J(mu0), mu0)
@@ -135,10 +147,10 @@ def test_free_surface_no_relative_flow_limit_is_elastic():
     rec = np.concatenate([np.stack([np.arange(5, 65, 3), np.zeros(20, int)], 1),
                           np.stack([np.arange(5, 65, 3), np.full(20, 12)], 1)])[None]
     w = _wavelet(nt, dt, 12.0)
-    db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True)(
+    db = _prop(Biot(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True, compiled=False)(
         w, src, rec, models=coef)
-    de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True)(
-        w, src, rec, models=[vp, vs, rho])
+    de = _prop(Elastic(device=DEV), shape, dh, dt, ["sxx", "szz"], ["vx", "vz"], free_surface=True,
+               compiled=False)(w, src, rec, models=[vp, vs, rho])
     err = float((db - de).abs().max() / de.abs().max())
     assert err < 1e-4, err
 
@@ -160,7 +172,7 @@ def s_time(tt):
         return np.exp(-(math.pi * f0 * (tt - t0)) ** 2)
 
     w = torch.from_numpy(dt * np.gradient(s_time(t), t) / dx**2).float()
-    prop = _prop(Biot(device=DEV), shape, dx, dt, ["p"], ["p"], abcn=30, free_surface=True)
+    prop = _prop(Biot(device=DEV), shape, dx, dt, ["p"], ["p"], abcn=30, free_surface=True, compiled=False)
     with torch.no_grad():
         num = prop(-w, np.array([[xs, zs]]), rec_xz[None], models=coef)[0, :, :, 0].double().numpy()
 
@@ -245,7 +257,7 @@ def s_time(tt):
     src, rec = np.array([[pad, pad]]), np.array([[[pad + off, pad + off]]])
     num = 0.0
     for fields, scale in ((["sxx", "szz"], 1.0), (["p"], -qpart / phi)):
-        prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vz"], abcn=30)
+        prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vz"], abcn=30, compiled=False)
         with torch.no_grad():
             num = num + prop(w * scale, src, rec, models=coef)[0, :, 0, 0].double().numpy()
     ana = _analytic_line_source_vz(models, t, off * dx, off * dx + 0.5 * dx, s_time, qpart, phi)
@@ -331,10 +343,13 @@ def s_time(tt):
         prop = _prop(Biot(device=DEV), shape, dx, dt, fields, ["vx", "vz"], abcn=30, free_surface=True)
         with torch.no_grad():
             num = num + prop(w * scale, np.array([[xs, zs]]), np.array([[[xr, 0]]]), models=coef)[0, :, 0].double().numpy()
-    # vx sits at (x + dx/2, 0), vz at (x, dx/2) on the staggered grid
+    # vx sits at (x + dx/2, 0), vz at (x, dx/2) on the staggered grid.  One integration for both
+    # receivers, which share its per-frequency solves (most of its time), and only up to
+    # kx = 1 rad/m: past it the band's propagating modes are evanescent over the 20 m source
+    # depth (exp(-17) at 100 Hz), and the reference moves by < 1e-5 relative from kx_max = 2
     off = (xr - xs) * dx
-    ref_vx, _ = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off + 0.5 * dx, 0.0)], t, s_time)
-    _, ref_vz = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off, 0.5 * dx)], t, s_time)
-    for k, ref in enumerate((ref_vx[:, 0], ref_vz[:, 0])):
+    ref_vx, ref_vz = _dwn_halfspace(models, (1.0, qpart / phi), zs * dx, [(off + 0.5 * dx, 0.0), (off, 0.5 * dx)],
+                                    t, s_time, kx_max=1.0)
+    for k, ref in enumerate((ref_vx[:, 0], ref_vz[:, 1])):
         err = np.linalg.norm(num[:, k] - ref) / np.linalg.norm(ref)
         assert err < tol, (("vx", "vz")[k], err)

test/test_compiled_stencil_slices.py

@@ -30,11 +30,13 @@
 DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-@pytest.fixture(autouse=True)
+@pytest.fixture
 def _release_compiled_state():
     """Dynamo's cache keeps the compiled steps' propagators -- and their GPU
     buffers -- alive; release them so a later test that reads the absolute
-    peak memory (test_eager_boundary_saving) does not count them."""
+    peak memory (test_eager_boundary_saving) does not count them.  Taken by
+    the tests that compile only: the reset and a full gc.collect() cost
+    0.1-0.25 s a test, and the kernel-level tests build no compiled state."""
     yield
     if hasattr(torch, "_dynamo"):
         torch._dynamo.reset()
@@ -177,6 +179,7 @@ def _assert_matches(run, ref):
 
 
 @requires_compile
+@pytest.mark.usefixtures("_release_compiled_state")
 def test_the_conv_counter_sees_a_conv(monkeypatch):
     """With the slice form off the counter must find the convs, or the zero
     the tests below assert proves nothing."""
@@ -188,6 +191,7 @@ def test_the_conv_counter_sees_a_conv(monkeypatch):
 
 
 @requires_compile
+@pytest.mark.usefixtures("_release_compiled_state")
 @pytest.mark.parametrize("cls,shape,surface", EQUATIONS)
 def test_compiled_rollout_is_one_conv_free_graph_matching_uncompiled(cls, shape, surface):
     """A rollout compiles one graph (fullgraph=True raises on any break) with
@@ -206,6 +210,7 @@ def test_compiled_rollout_is_one_conv_free_graph_matching_uncompiled(cls, shape,
 
 
 @requires_compile
+@pytest.mark.usefixtures("_release_compiled_state")
 def test_plain_warmup_prefix_changes_nothing(monkeypatch):
     """Before torch 2.5 a checkpointed rollout runs its first, compiling steps
     on the plain tape, because compiling inside the checkpoint breaks its

test/test_eager_compile_fallback.py

@@ -7,6 +7,7 @@
 uncompiled, warns once, and returns exactly what ``use_compile=False`` does.
 """
 import warnings
+from functools import lru_cache
 
 import numpy as np
 import pytest
@@ -46,6 +47,14 @@ def _run(use_compile, memory=None):
     return solver, out.detach(), vp.grad
 
 
+@lru_cache(maxsize=None)
+def _reference(boundary_saving):
+    """``(record, gradient)`` of the plain uncompiled run every fallback must
+    reproduce, per memory mode: deterministic, so one run serves every test."""
+    _, out, grad = _run(False, BoundarySaving() if boundary_saving else None)
+    return out, grad
+
+
 @pytest.fixture
 def compiles_on_cpu(monkeypatch):
     """These tests stand in their own torch.compile; on torch 2.3 the CPU step
@@ -57,7 +66,7 @@ def compiles_on_cpu(monkeypatch):
 @pytest.mark.usefixtures("compiles_on_cpu")
 @pytest.mark.parametrize("memory", [None, BoundarySaving()], ids=["full", "boundary_saving"])
 def test_falls_back_to_the_uncompiled_step(monkeypatch, memory):
-    _, ref_out, ref_grad = _run(False, memory)
+    ref_out, ref_grad = _reference(memory is not None)
 
     calls = []
     monkeypatch.setattr(torch, "compile", _compile_that_cannot_build(calls))
@@ -90,7 +99,7 @@ def test_a_torch_without_compile_runs_uncompiled_and_says_so_once(monkeypatch, m
     """torch < 2.0 has no torch.compile at all.  The propagator used to skip
     compiling there in silence, leaving users on the slow path unawares."""
     import sweep.propagator._torch_eager as TE
-    _, ref_out, ref_grad = _run(False, memory)
+    ref_out, ref_grad = _reference(memory is not None)
 
     monkeypatch.delattr(torch, "compile")
     monkeypatch.setattr(TE, "_NO_COMPILE_WARNED", False)
@@ -119,7 +128,7 @@ def test_torch_23_runs_uncompiled_on_cpu_and_says_so_once(monkeypatch, memory):
     while compiling the backward, where no fallback catches it), and nothing
     turns the faulty fusion off."""
     import sweep.propagator._torch_eager as TE
-    _, ref_out, ref_grad = _run(False, memory)
+    ref_out, ref_grad = _reference(memory is not None)
 
     calls = []
     monkeypatch.setattr(TE, "INDUCTOR_23", True)

test/test_eager_snapshot_buffer.py

@@ -11,6 +11,7 @@
 import torch
 
 from sweep.equations import Acoustic
+from sweep.propagator.options import EagerOptions
 from sweep.propagator.torch import PropTorch
 
 NZ, NX, NT = 24, 28, 12
@@ -24,7 +25,10 @@ def prop():
                      receiver_type=["h1"],
                      # snapshots need the full tape; chunk checkpointing, the
                      # eager default, refuses return_wavefield outright
-                     use_ckpt=False)
+                     use_ckpt=False,
+                     # the buffer is the rollout's, not the step's: a cold
+                     # torch.compile of a 12-step run is nearly all its time
+                     eager_options=EagerOptions(use_compile=False))
 
 
 def _shoot(prop):

test/test_elastic_curvilinear.py

@@ -10,6 +10,8 @@
 what the staggered chain-rule predecessor could not do.
 """
 
+from functools import lru_cache
+
 import numpy as np
 import pytest
 import torch
@@ -73,22 +75,28 @@ def _hill(height_cells=6.0, sigma_cells=12.0):
     return np.round(z.max() - z).astype(np.int64)
 
 
+@lru_cache(maxsize=None)
+def _flat_record():
+    """The flat-grid run of the default geometry; two tests read it (a run is
+    deterministic, so one serves both)."""
+    sources, receivers = _geometry()
+    _, prop = _make_prop(topography=None)
+    return prop(_wavelet(), sources, receivers, models=_models())
+
+
 # --------------------------------------------------------------------------
 
 
 def test_flat_topography_runs_and_is_stable():
-    sources, receivers = _geometry()
-    _, prop = _make_prop(topography=None)
-    syn = prop(_wavelet(), sources, receivers, models=_models())
+    syn = _flat_record()
     assert torch.isfinite(syn).all(), "flat curvilinear run produced NaN/Inf"
     assert syn.abs().max().item() > 0, "record is identically zero"
 
 
 def test_zero_topography_matches_no_topography():
     """A flat ``topography`` array must be the same problem as no array."""
     sources, receivers = _geometry()
-    _, prop_none = _make_prop(topography=None)
-    a = prop_none(_wavelet(), sources, receivers, models=_models())
+    a = _flat_record()
     _, prop_zero = _make_prop(topography=np.zeros(NX, dtype=np.int64))
     b = prop_zero(_wavelet(), sources, receivers, models=_models())
     assert torch.equal(a, b)
@@ -317,7 +325,10 @@ def misfit(vp):
     g = vp.grad.detach()
 
     iz, ix = np.unravel_index(int(g.abs().argmax()), g.shape)
-    h = 10.0
+    # that cell is next to the source, where the misfit is far from linear in vp: h = 10 m/s
+    # left a 4.8% truncation error against the 5% tolerance (1.1% at 5 m/s, 0.27% at 2.5;
+    # float64 keeps the round-off orders below)
+    h = 2.5
     js = []
     for sgn in (+1.0, -1.0):
         vpp = vp0.clone()
@@ -417,14 +428,22 @@ def _snap_run(prop, topo, **kw):
     return snaps.detach().cpu()
 
 
+@lru_cache(maxsize=None)
+def _smooth_hill_physical():
+    """A smooth-hill propagator and its default (physical-grid) snapshots; two
+    tests read them (a run is deterministic, so one serves both)."""
+    topo = _smooth_hill()
+    _, prop = _make_prop(topography=topo)
+    return prop, _snap_run(prop, topo)
+
+
 def test_snapshots_are_the_grid_rows_interpolated_per_column():
     """The default snapshot is the computational one, each Cartesian row
     interpolated from the grid's rows around its depth -- checked against a
     plain per-column np.interp, and against a second run that asks for it."""
     topo = _smooth_hill()
-    _, prop = _make_prop(topography=topo)
+    prop, phys = _smooth_hill_physical()
     comp = _snap_run(prop, topo, snapshot_grid="computational")
-    phys = _snap_run(prop, topo)
     mapped = prop._curvilinear_snapshots(comp.clone())
     scale = float(comp.abs().max())
     assert float((phys - mapped).abs().max()) <= 1e-6 * scale
@@ -444,9 +463,7 @@ def test_snapshots_are_the_grid_rows_interpolated_per_column():
 
 
 def test_air_cells_are_zero_in_every_field():
-    topo = _smooth_hill()
-    _, prop = _make_prop(topography=topo)
-    phys = _snap_run(prop, topo)
+    prop, phys = _smooth_hill_physical()
     surf = prop.curvilinear_grid.h_surface.numpy() / DH
     air = np.arange(NZ)[:, None] < surf[None, :]
     assert air.any()

test/test_kwarg_validation.py

@@ -16,6 +16,7 @@
 import torch
 
 from sweep.equations.acoustic import Acoustic
+from sweep.propagator.options import EagerOptions
 from sweep.propagator.torch import PropTorch
 
 DEV = "cuda" if torch.cuda.is_available() else "cpu"
@@ -60,7 +61,9 @@ def test_call_time_typo_raises_and_names_the_key():
 
 
 def test_call_time_recognised_extras_still_work():
-    p = _build()
+    # uncompiled: the extras reach init_abc, not the step, and a cold torch.compile
+    # of the step would cost ~30x this 50-step run
+    p = _build(eager_options=EagerOptions(use_compile=False))
     wavelet = torch.zeros(50, device=DEV)
     src = np.array([[[24, 28]]], dtype=np.int64)
     rec = np.array([[[10, 28]]], dtype=np.int64)