Skip to content

SEMPropagator

The time loop under PropSEM. It drives any SEM equation from its field_specs and takes a built equation, source and receiver stamps and a wavelet; PropSEM assembles those from the FD-shaped call. Use it for what PropSEM does not cover, such as an equation whose build takes keywords PropSEM does not supply, or stamps placed by hand.

sweep.sem.SEMPropagator 

SEMPropagator(
    eq: SEMEquation,
    dt: float,
    nt: int,
    sources: list,
    receivers: list,
    source_kinds: list,
    source_payloads: list,
    receiver_type: list = None,
    _compiled_step: object = None,
)

SEMPropagator(eq: 'SEMEquation', dt: 'float', nt: 'int', sources: 'list', receivers: 'list', source_kinds: 'list', source_payloads: 'list', receiver_type: 'list' = None, _compiled_step: 'object' = None)

build classmethod 

build(
    eq: SEMEquation,
    *,
    dt: float,
    nt: int,
    source_xz: list,
    receiver_xz: list,
    source_kinds: list | None = None,
    source_payloads: list | None = None,
    receiver_type: list | None = None,
    pml_guard: str = "warn"
) -> SEMPropagator

pml_guard ∈ {"warn", "error", "off"}: catch sources / receivers placed inside the C-PML absorbing band (where the field is artificially damped — a silent seismogram corruption). Default warns.

forward 

forward(
    source_wavelet: Tensor,
    *,
    record_field_every: int | None = None,
    snapshot_times=None,
    compile_step: bool = False,
    memory: str | None = None,
    bs_store_dtype=None
)

Run the time loop.

compile_step=True wraps eq.step in :func:torch.compile (cached on the propagator across calls) — the per-timestep tensor graph is fused, which on GPU is ~2× faster than eager for the elastic step. The first call pays a one-time compilation cost. With memory='boundary' it also fuses the reverse reconstruction and the per-step adjoint recompute (both call eq.step), ~2× on the BS rollout with bit-identical gradients.

memory='boundary' runs the backward pass with boundary-saving wavefield reconstruction (peak memory ~flat in nt) for any equation exposing the BS hooks (acoustic swap2nd / velocity-stress substep). Snapshots are unavailable in this mode (the forward state is discarded).

source_wavelet is (nt,) — one time function shared by every point source — or (P, nt) for the FD source-encoding mode: P = len(self. sources) superposed sources each with its own per-step amplitude s_k(t) (distinct polarity / amplitude / time-shift), injected as Σ_k stamp_k(x) · s_k(t).

Returns the seismogram tensor (nt, nrec, ncomp) -- the FD PropTorch layout minus its leading shot axis, since this propagator is single-shot. ncomp and the component ORDER come from receiver_type; it defaults to every receiver-capable field of the equation, so acoustic gives ncomp = 1 and elastic gives 2 (2-D) or 3 (3-D).

With record_field_every=k it returns (seis, snaps) instead -- the same shape as FD/JAX's (record, wavefield) pair -- where snaps is (nsnap, ncomp, n_global) sampled at time steps range(0, nt, k). snapshot_times names the steps instead (any order, duplicates collapse; snaps follows them ascending); the two are exclusive.

Point and moment stamps

A stamp holds the Lagrange weights of one point, and their derivatives, in the element that contains it: a source is spread over that element's GLL nodes with them, and a receiver reads the field through them. A moment-tensor source differentiates the basis, which jumps across element edges, so a source on an edge is stamped in every element that meets there (EdgeStamp).

sweep.sem.build_point_stamp 

build_point_stamp(mesh, *coords)

Locate a point and build the Lagrange/derivative weights.

Coordinate convention follows the mesh: 2-D meshes take (x, z), the 3-D :class:Hex8Mesh takes (x, y, z). Dispatches on mesh type: closed-form (ξ, η[, γ]) for axis-aligned :class:Quad4Mesh / :class:Hex8Mesh, Newton-Raphson on the Q2 map for a deformed :class:Quad9Mesh.

sweep.sem.PointStamp 

PointStamp(
    elem_id: int,
    w_eta: Tensor,
    w_xi: Tensor,
    dw_eta: Tensor,
    dw_xi: Tensor,
    ixx: float,
    ixz: float,
    iex: float,
    iez: float,
)

Lagrange weights + derivatives + local inv-Jacobian at one point.

A single representation that covers every source/receiver need: * scalar_stamp — acoustic body source, * force_stamp — elastic body force vector, * moment_tensor_stamp — elastic moment-tensor source, * gather — interpolate one field component at the point.

ixx, ixz, iex, iez are the local inverse-Jacobian components ∂ξ/∂x, ∂ξ/∂z, ∂η/∂x, ∂η/∂z at the point. For an axis-aligned Quad4 element ixz = iex = 0 so the moment-tensor formula reduces to the simple ∂ξ/∂x, ∂η/∂z chain rule.

gather_weights 

gather_weights() -> torch.Tensor

Flattened receiver weight stamp (ngll²,) aligned with ibool[elem].reshape(-1) for the vectorised receiver gather.

sweep.sem.PointStamp3D 

PointStamp3D(elem_id: int, w_gamma: Tensor, w_eta: Tensor, w_xi: Tensor)

3-D Lagrange weights at one point for a :class:Hex8Mesh element.

The axis order matches ibool[elem] = (kz, ky, kx) ↔ (z, y, x), so every stamp tensor is w_gamma ⊗ w_eta ⊗ w_xi (z ⊗ y ⊗ x). Covers the scalar (acoustic) source, the 3-component body force, and field gather. Moment-tensor sources are not yet surfaced in 3-D.

sweep.sem.build_moment_stamp 

build_moment_stamp(mesh, *coords)

:func:build_point_stamp for a source that differentiates the basis.

Inside an element it IS that function's stamp. On an element edge or corner it is an :class:EdgeStamp over every element whose closure holds the point. 3-D meshes have no moment-tensor stamp; they get the plain one.

sweep.sem.EdgeStamp 

EdgeStamp(parts: list, weights: list)

One source point stamped in every element that meets at it.

A moment-tensor load is -M : ∇φ_a(x_s), and the gradient of a C0 element's basis jumps across its edges. Stamped in the one element the locator returns, a source on an edge takes a one-sided derivative and radiates lopsided -- an explosion on a corner came out 7e-3 (v_z) and 2e-2 (v_x) off mirror symmetry, and integer grid sources land on edges every n_order cells. The limit of a symmetric source shrunk onto the point instead gives each element the share of a small disk it holds: weights are the sector angles over 2π -- 1/2 a side on an edge, the corner angles at a corner, the half-disk inside the domain on its boundary.

Force and scalar stamps need none of this: the Lagrange weights agree on the shared nodes and vanish off them, so one element carries them exactly.

Boundary saving

The eager boundary-saved rollout; PropSEM(memory=BoundarySaving()) is the usual way to ask for it. See Boundary saving.

sweep.sem.run_boundary_saved 

run_boundary_saved(
    eq,
    source_spec,
    wavelet,
    rec_G,
    rec_W,
    dt,
    nt,
    *,
    compile_step=False,
    store_dtype=None,
    frame_every=None,
    recv_idx=None
)

Boundary-saved single-shot rollout for any SEM equation with BS hooks.

frame_every=K additionally stores a full-state checkpoint every K steps and re-anchors the reverse march on them (memory: nt/K frames). Needed only by the coupled fluid–solid equation, whose reverse march accumulates ~2e-5 of the peak field without them and ~4e-11 with (the acoustic and elastic marches are flat and need none). Defaults to the equation's bs_frame_checkpoint_every class attribute when present.

Two source modes, selected by wavelet dimensionality:

  • shared (wavelet=(nt,)) — source_spec is a list of (n_global,) unit patterns, one per source field; the term is pattern_f(x) · wavelet[n] (one time function for all sources).

  • per-source (wavelet=(P, nt)) — FD source-encoding: source_spec is (src_G, src_W) from :meth:SEMPropagator._build_per_source_patterns (P superposed point sources), and the term is Σ_k stamp_k(x) · wavelet[k, n] (a distinct time function per source — polarity / amplitude / time-shift encoding). The model gradient is the true cross-talk source-encoded gradient (one rollout, superposed wavefield).

  • rec_G / rec_W — receiver gather matrices, sampled on every receiver field.

  • store_dtype — interface-ring storage precision (the dominant BS memory): None (compute dtype), torch.bfloat16 / torch.float16 (½), or torch.int8 (¼, per-step quantized). bf16 keeps the gradient within reconstruction noise; int8 trades a little accuracy for the smallest tape.
  • compile_step — wrap eq.step in :func:torch.compile. Because the reverse pass (reverse_step) and the per-step VJP (bs_recompute) both call self.step, compiling that one method fuses the forward recompute, the reverse reconstruction and the adjoint recompute — ~2× faster on GPU with bit-identical gradients. (eq is shallow-copied first so a cached caller equation is not mutated; the copy is kept alive by the per-step ctx until backward, so the compiled kernel is still in place then.)
  • returns {receiver_field_index: seis (nrec, nt)}.

Anisotropy

sweep.sem.thomsen_to_vti_cij 

thomsen_to_vti_cij(vp, vs, epsilon, delta, rho)

VTI (vertical transverse isotropy) Voigt stiffnesses from Thomsen (1986).

Returns (c11, c13, c33, c55) for the in-plane P-SV system, where the velocity-stress constitutive update is

σ_xx += dt (c11 ∂x v_x + c13 ∂z v_z)
σ_zz += dt (c13 ∂x v_x + c33 ∂z v_z)
σ_xz += dt  c55 (∂z v_x + ∂x v_z)

with vp / vs the vertical P / S velocities. Thomsen ε, δ are the standard anisotropy parameters:

c33 = ρ vp²            (vertical P modulus)
c55 = ρ vs²            (vertical S modulus = c44 in 2D)
c11 = c33 (1 + 2ε)     (horizontal P modulus)
c13 = √[2δ c33 (c33−c55) + (c33−c55)²] − c55

ε = δ = 0 reduces exactly to isotropy: c11 = c33 = λ+2μ, c13 = λ, c55 = μ — so a VTI equation degenerates bit-for-bit to the isotropic one (verified in the test-suite).