Propagators¶

A propagator wraps an equation object, grid configuration, acquisition geometry, and model tensors into a callable solver. The two main entry points are:

sweep.propagator.torch.PropTorch — PyTorch-family runtime
sweep.propagator.jax.PropJax — JAX-family runtime

`backend` vs `impl`¶

PropTorch takes two related but distinct knobs:

backend — the array / autograd framework carrying tensors and gradients ("torch" for PyTorch eager / torch.compile; "jax" for JAX on PropJax).
impl — the implementation path under that framework ("eager" = pure-Python operators with PyTorch autograd; "c" = the prebuilt CUDA core (libsweep_core.so) driven through the pure-Python sweep._C ctypes layer).

The two axes are orthogonal:

Backend	Impl	Device	What runs
`torch`	`eager`	CPU or CUDA	Pure-Python operators + PyTorch autograd
`torch`	`c`	CUDA	Prebuilt CUDA core via the `sweep._C` ctypes layer
`jax`	—	CPU or CUDA	JAX implementation via `PropJax`

"cuda" is not a backend or impl value — it is a device choice driven by the tensors you pass in. The compiled path runs on CUDA tensors only; a host model tensor is refused with a clear error, and a CPU propagator resolves to impl="eager".

A minimal invocation for each path:

from sweep.propagator.torch import PropTorch

solver_eager = PropTorch(..., backend="torch", impl="eager")
solver_c     = PropTorch(..., backend="torch", impl="c")

from sweep.propagator.jax import PropJax

solver_jax   = PropJax(..., backend="jax")

Geometry conventions¶

sources — shape (nshots, ndim) for a single point source per shot, or (1, nsources, ndim) for one source-encoded supershot (all sources fire together; receivers is then (1, nreceivers, ndim) and the wavelet (nt,) or (nsources, nt))
receivers — shape (nshots, nreceivers, ndim)
2D coordinates are (x, z) in grid indices — horizontal first, depth last
3D coordinates are (x, y, z) in grid indices — horizontal axes first, depth last

Model and wavefield arrays themselves are stored as (nz, nx) / (nz, ny, nx) (depth first), which is the same axis order matplotlib.imshow expects when plotting with depth growing downward. The propagator reverses the user-supplied (x, z) / (x, y, z) coordinate tuples internally before indexing into the wavefield tensor.

Boundaries: free surface and PML¶

Every propagator takes free_surface= and abcn= (PML thickness). free_surface accepts several equivalent forms, all normalised to one canonical per-face boolean tuple (axis-major order — 2-D (z_lo, z_hi, x_lo, x_hi), 3-D (z_lo, z_hi, y_lo, y_hi, x_lo, x_hi)):

PropTorch(eq, ..., free_surface=False)                # absorbing everywhere (default)
PropTorch(eq, ..., free_surface=True)                 # historical top-only free surface
PropTorch(eq, ..., free_surface="top")                # single face name — same as True
PropTorch(eq, ..., free_surface=["top", "left"])      # list/set of face names
PropTorch(eq, ..., free_surface={"bottom": True,
                                 "right": True})      # dict: face name -> on/off
PropTorch(eq, ..., free_surface=(True, False,
                                 True, False))        # canonical per-face tuple: top + left
PropTorch(eq, ..., free_surface=(1, 1, 1, 1))         # same form with ints = closed box

Face names (also the order of the canonical tuple): 2-D top, bottom, left, right; 3-D top, bottom, front, back, left, right — top is the z-min face, matching the (nz, nx) depth-first array layout. An unknown name raises ValueError listing the valid names.

Each free face replaces its PML pad with the image-method boundary condition; the remaining faces stay absorbing. abcn takes a scalar or a canonical tuple of per-face PML widths (no face names or dict). Per-face widths are eager-only -- impl='c' raises NotImplementedError for them -- and, like per-edge free surfaces, 2-D Acoustic / Elastic / ViscoAcoustic only.

boundary_buffer= (default None) puts that many sigma=0 cells between the physical box and the PML ramp on every absorbing face. Boundary saving stores a shell M + 1 cells wide, which is all the reverse time step needs; an equation whose gradient stencil reaches further -- AcousticVRZ and AcousticVRZ3D take a divergence of a gradient, reach 2M -- would image cells the shell never restored. Such an equation declares BOUNDARY_BUFFER_REACH and None gives it REACH * M + 1 buffer cells under every memory strategy and backend, so full, checkpointing and boundary saving solve the same grid; the buffer is cropped with the pad, so shapes and coordinates do not change. Equations with pointwise imaging get no buffer. An explicit value is honoured, but boundary saving refuses one below the equation's need. A non-zero buffer cannot be combined with topography= yet (NotImplementedError).

Support matrix:

2-D Acoustic / Elastic — full per-edge free-surface support on the eager and compiled CUDA backends, with adjoint gradients across every backward memory mode (full / boundary-saving / checkpointing).
2-D ViscoAcoustic — full per-edge free-surface support on the eager and compiled CUDA backends (full / checkpointing gradients; boundary saving is refused — the dissipative spectral damping term cannot be reverse-time reconstructed, and the impl='c' default falls back to 'full').
3-D — top face only (free_surface=True); per-edge lists raise.
Topography (irregular surface) — a separate top-only feature; it cannot be combined with per-edge faces.
Anisotropic equations (AcousticVTI1st, AcousticVTI, AcousticTTI, ElasticTTI(SG), …) raise NotImplementedError for any free surface: the anisotropic stress-free condition couples through the stiffness tensor and is not the isotropic image method these solvers implement. Run with free_surface=False or use an isotropic equation.

See the per-edge free surface notebook for snapshots of each configuration and a closed-box energy check.

Implementation-specific options¶

Every option block lives in sweep.propagator.options and is documented in detail on the Propagator Options page. A quick reference:

Option block	Used with	Configures
`EagerOptions`	`impl="eager"`	`torch.compile` flags, debug knobs
`CUDAOptions`	`impl="c"`	Container for `impl="c"`-only options; its `memory=` is equivalent to `PropTorch(memory=...)`
`Full`	as `memory=`	Keep the whole forward wavefield. No parameters
`BoundarySaving`	as `memory=`	Boundary saving with GPU / CPU / disk storage
`Ckpt`	as `memory=`	Chunk- or recursive-mode checkpointing

MemoryOptions(strategy=..., boundary=..., ckpt=...) is the older way to say the same three things and is deprecated; BoundarySaving and Ckpt subclass the BoundaryOptions / CkptOptions it wrapped, so every field, default and validation rule is the one you already knew.

Pass each block through the matching kwarg. The memory strategy (Full, BoundarySaving, Ckpt) always goes through PropTorch's own memory=, whichever impl runs:

PropTorch(..., eager_options=EagerOptions(...))  # impl="eager"
PropTorch(..., cuda_options=CUDAOptions(...))    # impl="c"
PropTorch(..., memory=BoundarySaving(...))       # either impl

A full Option composition example¶

To FWI on a large model with the compiled CUDA path, boundary saving on pinned host memory, and asynchronous disk fallback at the largest scale:

import torch

from sweep.equations import Acoustic
from sweep.propagator.torch import PropTorch
from sweep.propagator.options import Full, BoundarySaving, Ckpt

dev = torch.device("cuda")
shape = (512, 2048)
dh, dt = 10.0, 0.002

solver = PropTorch(
    Acoustic(spatial_order=8, device=dev),
    shape=shape, dh=dh, dt=dt, device=dev,
    impl="c",
    memory=BoundarySaving(storage="cpu", pinned_memory=True),
)

Swap in memory=BoundarySaving(storage="disk", disk_async_read=True) for the disk fallback, memory=Ckpt(mode="chunk", chunks=100) for chunk-mode checkpointing, or memory=Full() to keep the whole forward wavefield. The strategy IS the object, so there is no separate strategy= string to keep in step with it.

The validation rules are enforced in the dataclass __post_init__ methods, so incompatible combinations (e.g. disk_async_read=True with storage="cpu") fail loudly at construction time rather than during a long FWI run.

Memory-saving features¶

The gradient-memory mode is a three-way choice, identical for the eager and CUDA backends, and the choice is a type:

memory=Full()                                     # keep the forward wavefield
memory=BoundarySaving(storage="gpu")              # reconstruct from the boundary
memory=Ckpt(mode="chunk", chunks=100)             # rematerialise from checkpoints

Left unset, impl="c" defaults to BoundarySaving(storage="gpu") (fp32 ring; Full() for ViscoAcoustic and DASZhao3D, which have no compiled boundary saving) and the eager backend to Ckpt(). The modes are mutually exclusive, and conflicting requests raise a ValueError instead of one path silently winning.

The older spellings still work, and now warn

memory=MemoryOptions(strategy=..., boundary=..., ckpt=...) and the flat boundary_saving_config={...} dict are read exactly as before and emit a DeprecationWarning; the removal criterion is written next to the shim. use_ckpt= is not deprecated — it stays a plain supported keyword.

Three rules make that resolution predictable:

An off-switch means 'full', not "the other trick". use_ckpt=False (or boundary_saving_config={'enabled': False}) with nothing else selects full-wavefield storage on both backends — the long-standing meaning of impl='c', use_ckpt=False. The implicit backend default applies only when no gradient-memory knob is passed at all.
A request is honoured, not out-voted. boundary_saving_config={'enabled': True} now really runs the boundary backward; it used to lose silently to the use_ckpt=True default, so scripts that thought they were measuring boundary saving were checkpointing.
memory= may sit next to a legacy knob when they agree. memory=BoundarySaving(), use_ckpt=False states one intent twice and is accepted; memory=BoundarySaving(), use_ckpt=True contradicts it and raises. Where both carry detail, memory= wins.

Tail truncation has a legacy dict spelling too — boundary_saving_config= {'enabled': True, 'tail_steps': K} is equivalent to BoundarySaving(tail_steps=K).

Feature	Path	Configured by
Full storage (no reconstruction)	both	`Full()`
Boundary saving (GPU or CPU ring; + disk on `impl="c"`)	both	`BoundarySaving(storage=..., storage_dtype=..., ...)`
Asynchronous disk prefetch	`impl="c"`	`BoundarySaving(storage="disk", disk_async_read=True, ...)`
Boundary tail truncation (steady-state / freqsel objectives)	`impl="c"` acoustic	`BoundarySaving(tail_steps=...)`
Chunked checkpointing	both	`Ckpt(mode="chunk", chunks=...)`
Recursive (fixed-budget) checkpointing	`impl="c"`	`Ckpt(mode="recursive", count=...)`
`torch.compile` on the eager step	`impl="eager"`	`EagerOptions(use_compile=True, ...)`

A runnable comparison of these options lives in the Memory · strategies notebook, which exercises full-wavefield, boundary saving, and checkpointing on the same Marmousi shot and prints the per-mode peak GPU / host memory.

Boundary tail truncation (`BoundarySaving.tail_steps`)¶

For steady-state objectives — frequency-selection / DFT-comb FWI, where the loss reads only the last n_probe samples of the record and the adjoint source is therefore zero everywhere earlier — the reverse sweep does not need to walk the whole record. tail_steps=K makes the forward save only the last K steps' boundary strips and stops the backward after them:

memory=BoundarySaving(storage="gpu", tail_steps=n_probe + margin)

The forward physics is unchanged — the wavefield still runs the full record so the steady state can ring up. Only the saved/reconstructed step range shrinks, so both the backward wall time and the one-shot boundary buffer drop by roughly 1 - K / nt (measured: 74 % backward time and 75 % buffer at K/nt = 25 %).
The restore at reverse step it consumes the strip saved at forward step it - 1, so one saved step is spent on alignment: the effective reverse depth is K - 1 — budget it inside margin.
margin is physical: the dropped gradient term is exactly the adjoint × ring-up-transient correlation that steady-state methods discard, and it decays as the adjoint field drains through the absorbing boundary. Sweep the margin once per setup: on a ramped-sine test the truncated gradient converges monotonically to the full one (cos 0.992 → 1.000000 for margin 0 → 800 steps on a 140×160 grid).
Do not use it with impulsive-source objectives: there the early adjoint–forward correlations are real gradient content and the truncated gradient is genuinely different (cos ≈ 0.1 in the same test).
Scope: impl="c" Acoustic 2-D/3-D with the boundary-saving backward, any storage/storage_dtype. Checkpointing, the eager backend and the other equations raise NotImplementedError/ValueError rather than silently ignoring the option. Unset (None, the default) is bit-exact legacy behaviour, and tail_steps >= nt degenerates to it bitwise.
Domain decomposition composes: ModelParallel inherits tail_steps from the wrapped propagator's memory config exactly like storage/storage_dtype, shrinks every tile's boundary ring to K steps, and stops the lockstep reverse halo loop at the same global step on every rank (the stop index is derived from (nt, tail_steps), which are identical across ranks by construction, so no rank can be left waiting in an exchange). The truncated DD gradient is bit-exact against the truncated single-domain gradient on fp32 boundaries (test/test_dd_tail_two_tile.py, test/dd_tail_nccl_check.py).

Environment toggles¶

A few knobs stay out of the API because the right value depends on the machine, not on the problem. All are read once per run. None changes the physics, but two change the numbers: a lossy SWEEP_BOUNDARY_DTYPE (fp16/bf16/int8) and SWEEP_VRZ_GRAD_SPLIT (a different summation, ~1 ULP).

Variable	Effect
`SWEEP_VRZ_GRAD_SPLIT=1`	`AcousticVRZ3D` backward: force the O(M) split gradient (materialise `c_d`/`e_d`, then one divergence) instead of the fused nested-stencil kernel that `order<=4` picks by default. The crossover is GPU-dependent — fused wins on RTX 6000 Ada, split is ~12 s/iter faster on V100 at production scale.
`SWEEP_DD_DISABLE_OVERLAP=1`	Domain decomposition: serial step-then-exchange instead of the overlapped forward (see Domain decomposition).
`SWEEP_BOUNDARY_DTYPE`	`impl="c"`: `storage_dtype` of the implicit default ring only (no memory knob passed). Any `BoundarySaving(...)` carries `storage_dtype="fp32"` unless set, so it always wins.
`SWEEP_DATASETS_CACHE`	Where `sweep.datasets` caches downloads (see Datasets).
`SWEEP_JIT_FULL=1`	Developer path: compile the pybind shim against your torch instead of using the ctypes layer over the prebuilt core.
`SWEEP_CORE=<path>/libsweep_core.so`	Use a custom CUDA core; its `core.json` sidecar gets the same fit check as the shipped one.
`SWEEP_JIT_ALLOW_OLD_CUDA=1`	Try an nvcc 12.0–12.3 for a local core build.
`TORCH_CUDA_ARCH_LIST`	Target arch(s) of a local core build; when set it overrides the visible GPU's arch, and it is required when no GPU is visible.
`TORCH_EXTENSIONS_DIR`	Where a locally built core is cached.
`CUDA_HOME` / `CUDA_PATH`	Toolkit for a local core build, version-checked against torch's CUDA major.

Consistency testing¶

The C memory modes are exercised by test/solver_gradient_mode_suite.py. The suite compares eager gradients against compiled full-wavefield, boundary saving, and checkpoint modes across interior, finite-difference edge, and free-surface source placements, and writes per-mode gradient figures to test/test_outputs/solver_gradient_mode_suite/.