Utilities reference

The shared machinery every propagator uses or returns: the output and callback types, the type aliases, checkpointing, storage, planning and diagnostics, the discretisation helpers, off-grid source and receiver locations, wavelets, transducer arrays, misfit functions, wavefield-space tools, and the parameter converters. Usage explains when to reach for each of them.

Anisotropy parameter converters

Each anisotropic propagator takes the parameterisation it actually solves in, which is not the one models are usually stored or inverted in. deepwave.elastic_vti() takes the VTI stiffnesses c11, c13, c33, c55 (and c66 in 3-D) together with the buoyancy \(1/\rho\); deepwave.acoustic_vti() takes the same stiffnesses less the shear ones (c11, c13, c33) and the buoyancy. Practitioners, by contrast, reason in Thomsen’s \((v_{P0}, v_{S0}, \rho, \epsilon, \delta, \gamma)\) parameters, or – when the model came from surface P-wave velocity analysis, which cannot resolve \(\epsilon\) and \(\delta\) separately – in Alkhalifah and Tsvankin’s \((v_{P0}, v_{nmo}, \eta)\). The deepwave.anisotropy module holds the exact (not weak-anisotropy linearised) conversions between them, plus the near-source taper deepwave.acoustic_vti() sometimes needs. Every one of them is ordinary differentiable tensor arithmetic, so inverting in Thomsen or Alkhalifah-Tsvankin space is just a matter of building the propagator’s inputs inside the loss and letting autograd chain through (see Model parameterisations and Thomsen-space FWI with the pseudo-acoustic VTI propagator).

thomsen_to_cij returns the buoyancy alongside the stiffnesses, so its whole return value unpacks straight into a propagator call, and a vacuum cell (\(\rho = 0\)) becomes buoyancy = 0, the sentinel the propagators expect, rather than the infinity a plain 1 / rho would give.

Converter	From	To
deepwave.anisotropy.thomsen_to_cij()	vp, vs, rho, epsilon, delta, gamma	c11, c13, c33, c55, c66, buoyancy – the inputs of elastic_vti (and, dropping c55/c66, of acoustic_vti, for which vs = 0)
deepwave.anisotropy.cij_to_thomsen()	c11, c13, c33, c55, c66, buoyancy	vp, vs, rho, epsilon, delta, gamma – the exact inverse, taking and returning the tuples in the orders that let the two compose directly
deepwave.anisotropy.alkhalifah_to_thomsen()	vp, vnmo, eta	epsilon, delta – feed to thomsen_to_cij with a chosen vs, since P-wave data constrains no shear information
deepwave.anisotropy.thomsen_to_alkhalifah()	vp, epsilon, delta	vnmo, eta – the inverse, for comparing with or initialising from a conventional NMO-velocity model
deepwave.anisotropy.source_taper()	c11, c13, c33, source_locations, radius	a modified c13 – not a parameterisation change, but the near-source ellipticity taper that suppresses acoustic_vti’s pseudo-shear artefact

The isotropic pair deepwave.vpvsrho_to_lambmubuoyancy() / deepwave.lambmubuoyancy_to_vpvsrho() stays at the package root: it is not anisotropic, and it shipped at that spelling.

deepwave.anisotropy.thomsen_to_cij(vp: Tensor, vs: Tensor, rho: Tensor, epsilon: Tensor, delta: Tensor, gamma: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None, Tensor]

Converts Thomsen (1986) parameters to VTI stiffnesses.

Uses the exact (not weak-anisotropy linearised) relations. With the symmetry axis as the Voigt “3” direction:

c33 = rho * vp**2
c55 = rho * vs**2
c11 = c33 * (1 + 2 * epsilon)
c66 = c55 * (1 + 2 * gamma)
c13 = sqrt(2 * delta * c33 * (c33 - c55) + (c33 - c55)**2) - c55
buoyancy = 1 / rho

where vp and vs are the P and S wave speeds along the symmetry axis. The c13 expression inverts Thomsen’s exact definition delta = ((c13 + c55)**2 - (c33 - c55)**2) / (2 c33 (c33 - c55)), taking the positive square-root branch, which corresponds to c13 + c55 > 0. This branch holds for virtually all real materials; “anomalous” media with c13 + c55 < 0 cannot be represented in Thomsen parameters unambiguously and are not supported. The argument of the square root is non-negative iff delta >= -(c33 - c55) / (2 c33); more negative delta values are rejected with an error. Cells with rho = 0 (or vp = vs = 0), the vacuum convention of the elastic propagators, map to all-zero stiffnesses (the square root is masked there so no NaN gradients are produced) and to buoyancy = 0 exactly, which is the vacuum sentinel the propagators themselves expect.

The buoyancy (the reciprocal of the density) is returned alongside the stiffnesses because it, not rho, is what deepwave.elastic_vti() and deepwave.acoustic_vti() take, so the returned tuple can be passed straight to them. It is computed with the same zero-guarded reciprocal as deepwave.vpvsrho_to_lambmubuoyancy(), so a vacuum cell yields exactly 0 rather than inf.

The function is differentiable, so gradients with respect to the Thomsen parameters may be obtained by chaining autograd through it (inverting in Thomsen space).

Parameters:

• vp – P-wave speed along the symmetry axis.

• vs – S-wave speed along the symmetry axis.

• rho – Density.

• epsilon – Thomsen’s epsilon, (c11 - c33) / (2 c33).

• delta – Thomsen’s delta (exact definition, see above).

• gamma – Thomsen’s gamma, (c66 - c55) / (2 c55). Required to build c66, which is only used in 3-D. Optional, default None.

Returns:

(c11, c13, c33, c55, c66, buoyancy), with c66 set to None when gamma is None. In 2-D, where c66 does not enter the equations, unpack it into _; in 3-D pass it to the propagator’s c66 keyword argument.

Return type:

Tuple[torch.Tensor, …]

Raises:

DeepwaveValueError – If vp, vs or rho contains negative values, if c33 < c55 (vp < vs) anywhere, or if delta is more negative than -(c33 - c55) / (2 c33) anywhere.

deepwave.anisotropy.cij_to_thomsen(c11: Tensor, c13: Tensor, c33: Tensor, c55: Tensor, c66: Tensor | None, buoyancy: Tensor) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor | None]

Converts VTI stiffnesses to Thomsen (1986) parameters.

The exact inverse of deepwave.anisotropy.thomsen_to_cij() (not the weak-anisotropy linearised approximation). With the symmetry axis as the Voigt “3” direction, this inverts the relations documented in deepwave.anisotropy.thomsen_to_cij():

rho = 1 / buoyancy
vp = sqrt(c33 * buoyancy)
vs = sqrt(c55 * buoyancy)
epsilon = (c11 - c33) / (2 * c33)
delta = ((c13 + c55)**2 - (c33 - c55)**2) / (2 * c33 * (c33 - c55))
gamma = (c66 - c55) / (2 * c55)

The arguments are in the order deepwave.anisotropy.thomsen_to_cij() returns them, and the return values are in the order deepwave.anisotropy.thomsen_to_cij() takes them, so the two compose directly in both directions: cij_to_thomsen(*thomsen_to_cij(vp, vs, rho, epsilon, delta, gamma)) round-trips to the inputs (to floating-point precision) whenever vp > vs (so c33 != c55 and delta is not degenerate) and rho > 0, and thomsen_to_cij(*cij_to_thomsen(...)) likewise returns the stiffnesses it was given. Note that this function takes buoyancy, not rho: that is what the propagators use, and what deepwave.anisotropy.thomsen_to_cij() now returns.

A vacuum cell (buoyancy == 0, equivalently rho == 0, the vacuum convention of the elastic propagators) has no recoverable Thomsen parameters at all: deepwave.anisotropy.thomsen_to_cij() maps every such cell to all-zero stiffnesses whatever its vp/vs/epsilon/delta/ gamma were. It is mapped here to vp = vs = rho = epsilon = delta = gamma = 0, which is the canonical vacuum representative: feeding it back to deepwave.anisotropy.thomsen_to_cij() reproduces exactly the all-zero stiffnesses and buoyancy = 0 it came from. The masking also keeps the gradients there finite.

The vp == vs (isotropic, c33 == c55) case is also singular – deepwave.anisotropy.thomsen_to_cij() maps every delta to the same c13 there – but, unlike vacuum, it has no canonical representative to collapse to, so it is deliberately not special-cased: the formula above is evaluated directly and produces a division by zero (nan/inf) for such cells, loudly, rather than inventing a value. The same applies to gamma when c55 == 0 (a pseudo-acoustic, vs = 0 model).

Parameters:

• c11 – The c11 stiffness.

• c13 – The c13 stiffness.

• c33 – The c33 stiffness.

• c55 – The c55 stiffness (shear modulus along the symmetry axis).

• c66 – The c66 stiffness (horizontal-plane shear modulus), or None if it is not available (in 2-D it does not enter the equations). gamma is then None too.

• buoyancy – The reciprocal of density, as the propagators take it.

Returns:

(vp, vs, rho, epsilon, delta, gamma), with gamma set to None when c66 is None.

Return type:

Tuple[torch.Tensor, …]

Raises:

DeepwaveValueError – If buoyancy contains negative values.

deepwave.anisotropy.alkhalifah_to_thomsen(vp: Tensor, vnmo: Tensor, eta: Tensor) → tuple[Tensor, Tensor]

Converts Alkhalifah-Tsvankin (1995) parameters to Thomsen (1986) epsilon/delta.

P-wave surface seismic data alone cannot resolve the full Thomsen parameter set: time processing (NMO, DMO, time migration) in a VTI medium is controlled only by the P-wave NMO velocity and the anellipticity parameter eta, not by epsilon/delta individually (Alkhalifah and Tsvankin, “Velocity analysis for transversely isotropic media”, Geophysics 60, 1995, 1550-1566). This helper converts that P-wave-derivable parameterisation back to Thomsen’s epsilon/delta using the exact relations:

delta = (vnmo**2 / vp**2 - 1) / 2
epsilon = eta * (1 + 2 * delta) + delta

which invert Alkhalifah and Tsvankin’s defining pair vnmo = vp * sqrt(1 + 2*delta) and eta = (epsilon - delta) / (1 + 2*delta).

The function is plain differentiable PyTorch arithmetic (no in-place ops), so gradients with respect to vp/vnmo/eta may be obtained by chaining autograd through it – e.g. inverting directly in Alkhalifah-Tsvankin parameter space, then continuing the chain through deepwave.anisotropy.thomsen_to_cij() to build a VTI stiffness model (see that function, and deepwave.elastic_vti()/deepwave.acoustic_vti() for the propagators that consume the resulting Cij).

Parameters:

• vp – P-wave velocity along the symmetry axis (the vertical P-wave velocity).

• vnmo – The short-spread P-wave normal-moveout (NMO) velocity, the quantity conventional velocity analysis of surface seismic data actually estimates.

• eta – Alkhalifah and Tsvankin’s anellipticity parameter, which controls P-wave time processing beyond the NMO velocity (dip moveout, time migration); zero for elliptical anisotropy (epsilon == delta).

Returns:

(epsilon, delta), Thomsen’s parameters, suitable as input to deepwave.anisotropy.thomsen_to_cij() (with a chosen vs, since shear information is not recoverable from P-wave-only vnmo/eta).

Return type:

Tuple[torch.Tensor, torch.Tensor]

deepwave.anisotropy.thomsen_to_alkhalifah(vp: Tensor, epsilon: Tensor, delta: Tensor) → tuple[Tensor, Tensor]

Converts Thomsen (1986) epsilon/delta to Alkhalifah-Tsvankin (1995) parameters.

The exact inverse of alkhalifah_to_thomsen(), using Alkhalifah and Tsvankin’s defining pair directly:

vnmo = vp * sqrt(1 + 2 * delta)
eta = (epsilon - delta) / (1 + 2 * delta)

so alkhalifah_to_thomsen(vp, *thomsen_to_alkhalifah(vp, epsilon, delta)) returns (epsilon, delta) and thomsen_to_alkhalifah(vp, *alkhalifah_to_thomsen(vp, vnmo, eta)) returns (vnmo, eta), both to floating-point precision. It is the conversion to use when a model is held in Thomsen parameters but the P-wave-derivable quantities are wanted – to compare with, or initialise from, a conventional NMO-velocity/anellipticity model, or to invert in the parameterisation that surface P-wave data actually constrains (see alkhalifah_to_thomsen() for why that is the NMO velocity and eta, not epsilon and delta individually).

The function is plain differentiable PyTorch arithmetic (no in-place ops), so gradients with respect to vp/epsilon/delta may be obtained by chaining autograd through it.

Parameters:

• vp – P-wave velocity along the symmetry axis (the vertical P-wave velocity).

• epsilon – Thomsen’s epsilon, (c11 - c33) / (2 c33).

• delta – Thomsen’s delta (exact definition, see deepwave.anisotropy.thomsen_to_cij()). Must be greater than -0.5 everywhere, since 1 + 2 delta is the squared ratio of the NMO to the vertical velocity.

Returns:

(vnmo, eta), the short-spread P-wave normal-moveout velocity and the anellipticity parameter, suitable as input to alkhalifah_to_thomsen().

Return type:

Tuple[torch.Tensor, torch.Tensor]

Raises:

DeepwaveValueError – If delta <= -0.5 anywhere, which would make 1 + 2 delta non-positive and vnmo imaginary.

References

T. Alkhalifah, I. Tsvankin, “Velocity analysis for transversely isotropic media,” Geophysics, vol. 60, no. 5, 1995.

deepwave.anisotropy.source_taper(c11: Tensor, c13: Tensor, c33: Tensor, source_locations: Tensor, radius: float) → Tensor

Locally elliptical c13 around sources to suppress pseudo-shear.

This returns a modified model, not a mask or a window. The c13 it gives back has been blended towards sqrt(c11 c33) within radius cells of every source, so a propagator run on it is running on a medium that differs from the one you passed in, near the sources: the anisotropy there has been made elliptical. That is the point – it is what removes the pseudo-shear artefact – but it means the simulated wavefield is not the wavefield of your own model in that neighbourhood, and a gradient computed with respect to the tapered c13 is not the gradient with respect to yours there either. Keep radius small (a few cells), and if the near-source medium matters to you, do not use this.

A point source in a region where eps != delta (i.e. c13 != sqrt(c11 c33)) emits a diamond/rhombic pseudo-shear wave, the V_S = 0 kinematic remnant. Following Duveneck et al. (2008), this helper blends c13 towards the locally elliptical value sqrt(c11 c33) (delta -> eps, for which no pseudo-shear artifact is emitted) inside a radius (in grid cells) of each source location, with a raised-cosine taper:

w(r) = 0.5 (1 + cos(pi r / radius))   for r < radius, else 0
c13_tapered = w sqrt(c11 c33) + (1 - w) c13

so c13 equals sqrt(c11 c33) exactly at a source (r = 0, w = 1) and is returned unchanged at and beyond radius (w = 0), varying smoothly in between. The blend perturbs the model near the source; keep radius small (a few cells).

The taper is the union over every location in source_locations (the nearest source governs each cell), and is applied on the model’s trailing ndim spatial dimensions (broadcasting over any leading shot dimension of the models). Distances are measured in grid cells (Euclidean).

Parameters:

• c11 – The c11 stiffness model.

• c13 – The c13 stiffness model (the one that is tapered).

• c33 – The c33 stiffness model.

• source_locations – Integer source locations, shape [..., ndim] (e.g. [n_shots, n_sources, ndim]); the leading dimensions are flattened and every listed location contributes. This helper does not interpret deepwave.IGNORE_LOCATION: pass only the locations you want tapered.

• radius – The taper radius in grid cells (> 0).

Returns:

A new c13 tensor (same shape/dtype/device as the input c13) with the near-source blend applied. Differentiable with respect to c11, c13, and c33.

PropagatorOutput, AuxiliaryState and RunInfo

class deepwave.PropagatorOutput(names: Sequence[str], values: Sequence[Any] | None = None, *, info: RunInfo | None = None)

Named, attribute-accessible view of a propagator’s outputs.

Returned by every propagator. Each final wavefield and receiver-amplitude tensor is accessible by name (e.g. out.receiver_amplitudes), which removes the most common indexing mistake: the position of receiver_amplitudes in the positional return shifts with the number of spatial dimensions, so code written for 2-D silently grabs the wrong tensor in 3-D.

The object is a tuple (it subclasses tuple rather than merely implementing collections.abc.Sequence): iterating, unpacking, and integer or slice indexing yield the values in the current positional return order. Positional unpacking therefore remains available, though it should be reviewed against the compatibility policy when upgrading:

out = scalar(...)
out.receiver_amplitudes        # by name (ndim-independent)
*state, receivers = out        # by position, as before

For a fixed propagator and dimensionality, the current field names and tuple order follow the stable compatibility policy (see Stability policy). Future metadata is added as attributes, not tuple elements. Named access is recommended: a positive index can silently mean a different field on another propagator or in another number of dimensions. For example, out[-2] is the opaque auxiliary state on scalar, but a receiver field on scalar_born.

Exactly one element is not a torch.Tensor: auxiliary_state, a AuxiliaryState holding the absorbing boundary’s memory variables and any other auxiliary state, which sits at the index the first of those fields used to occupy. It is a real tuple element (so len(out), iteration and unpacking see it) and also a typed property, and it is what initial_state= consumes to continue a simulation. Code that iterates the whole tuple expecting tensors – sum(o.sum() for o in out) and the like – must skip it or iterate its tensors.

Subclassing tuple (not abc.Sequence) is deliberate and load-bearing: torch.autograd.Function and torch.utils.checkpoint propagate requires_grad to an output only when it is a Tensor or a genuine tuple. An abc.Sequence that is not a tuple is silently detached (the root cause of issue #120), so returning a PropagatorOutput straight through torch.utils.checkpoint would break gradients if it were not a real tuple. Being a real tuple makes the named return compose with checkpointing exactly like the plain tuple return would; the non-Tensor auxiliary_state element does not change that for the tensor elements, and outside a checkpoint the tensors it holds keep their grad_fn.

torch.utils.checkpoint(..., use_reentrant=True) is the exception, and the reason to prefer use_reentrant=False (which PyTorch also recommends) when a checkpointed call’s state is fed onward. The reentrant implementation runs the wrapped function under torch.no_grad() and re-attaches only the Tensor elements of what it returns, so:

• the return is degraded to a plain tuple, losing the field names and info – pre-existing, and independent of auxiliary_state; and

• every tensor held inside auxiliary_state comes back detached, so continuing a simulation from it gives no gradient through the absorbing boundary. use_reentrant=False preserves both the subclass and those gradients. A function wrapped in a reentrant checkpoint that needs the auxiliary state downstream must return its tensors as top-level elements of its own return value, which is what deepwave.checkpointing.run_segments() does internally.

The public state fields use the same names as the propagator’s initial-state arguments (wavefield_0, wavefield_m1, …), so an output can be fed straight back in to continue a simulation. Fields that do not exist for the given number of dimensions (e.g. the z components in a 2-D simulation) are simply absent. final_state is the mapping of exactly the fields initial_state= consumes.

info is a RunInfo carrying non-differentiable metadata about how the call actually ran (resolved step_ratio, internal dt, Courant number, nt, PML width, PML frequency, survey origin/extent and more). It is plain attribute metadata, not a tuple element: it does not appear in iteration, unpacking, len(), or indexing. A propagator always supplies one. It is unreachable in exactly two situations: under torch.utils.checkpoint(..., use_reentrant=True), where the return is degraded to a plain tuple that has no attributes at all, and under torch.nn.DataParallel.gather, which does not reach the point of building an output.

Registered as a pytree node, so torch.func and torch.compile see the tensors inside it rather than one opaque object: a function that returns a whole output can be wrapped in torch.func.vmap(), carried as the auxiliary value of torch.func.grad() or torch.func.jacrev() (has_aux=True), compiled, or walked with torch.utils._pytree’s tree_map. The children are the tuple elements in return order (with auxiliary_state, itself a pytree node, expanding into the tensors it holds, so the leaves are exactly the output’s tensors); the field names and info travel as the node’s context, because the names vary per instance and the object cannot be rebuilt without them. A rebuilt output therefore still carries the names and the info of the run it came from, and its auxiliary_state still carries the recorded run parameters – format_version included – that make it checkable as an initial_state=.

Not usable as a full ``torch.nn.DataParallel`` return value on more than one GPU (DistributedDataParallel is unaffected). A forward() wrapping a propagator should return a single tensor (e.g. out.receiver_amplitudes) or a small tuple of plain tensors instead; see Distributed (multi-GPU) execution for the pattern and PropagatorOutput: DataParallel and pytrees for why this is unsupported by design.

Builds the tuple of values and records the field names.

as_dict() → dict[str, Any]

Returns an ordered mapping of field name to value.

property auxiliary_state: AuxiliaryState

The opaque absorbing-boundary and auxiliary state of this run.

Raises AttributeError on a propagator that has no auxiliary-state field. Declared as a property, rather than left to __getattr__, because it is the one element of the tuple that is not a torch.Tensor and so needs a static type of its own.

property bg_receivers: Mapping[str, Tensor]

The recorded background data of a Born propagator, keyed by channel.

As receivers, for the background (unperturbed) wavefield’s receivers. Raises AttributeError on the eight propagators that have no background receivers, rather than returning an empty mapping: an empty mapping would read as “the background data was not requested”, which is a different thing from “this propagator has none”.

detach() → PropagatorOutput

A copy of this output with every field detached from the graph.

Truncated backpropagation through a long record – running many segments but backpropagating through only the most recent ones – means cutting the autograd graph at a segment join deliberately. Every state field has to be cut for the graph to actually be cut: detaching only the auxiliary_state leaves the physical wavefields attached, so the join still carries the whole history and the gradient is the untruncated one:

out = scalar(v, dx, dt, source_amplitudes=chunk)
loss(out.receiver_amplitudes).backward()
state = out.detach()          # the whole join, cut
out = scalar(v, dx, dt, source_amplitudes=next_chunk,
             initial_state=state)

The auxiliary_state of the copy is detached through AuxiliaryState.detach(), so it records that the detachment was deliberate and the next call does not warn about a lost gradient. info and the field names are carried over unchanged.

property final_state: Mapping[str, Tensor | AuxiliaryState]

The fields that initial_state= consumes, as a read-only mapping.

The public state fields plus auxiliary_state, excluding the receiver amplitudes. prop(..., initial_state=out) and prop(..., initial_state=out.final_state) are equivalent; the mapping form is the one to build by hand when only part of the state is being replaced. Mutating it raises.

property names: tuple[str, ...]

The field names, in return order.

property receivers: Mapping[str, Tensor]

The recorded receiver data, keyed by channel, as a read-only mapping.

The same vocabulary the receivers= argument uses, so a wrapper can round-trip receivers={channel: locations} into out.receivers[channel] without knowing how this propagator spells its output fields. "u" on the two scalar propagators; "p", "y", "x" (and "z" in 3-D) on the acoustic and elastic families; "sigma_v"/"sigma_h" as well on acoustic_vti.

Every channel of the propagator is present, whether or not it was requested: a channel with no receiver locations holds the same empty tensor the returned tuple holds for it. An unrequested channel raising KeyError would be a worse interface than a documented empty tensor, because a caller could not tell it apart from a typo. That empty tensor is shaped like the data it would have held – [n_shots, 0, nt], the same shot count and the same number of time samples as a requested channel – so torch.cat(list(out.receivers.values()), dim=1) is a valid multichannel gather however few channels were asked for.

On a Born propagator this is the scattered data, matching the sc_receiver_* field names; the background data is bg_receivers.

Mutating the mapping raises.

class deepwave.AuxiliaryState(_tensors: Mapping[str, Tensor], propagator: str | None = None, ndim: int | None = None, accuracy: int | None = None, pml_width: tuple[int, ...] | None = None, boundary: tuple[str, ...] | None = None, n_mechanisms: int | None = None, mpml: bool | None = None, mpml_fraction: float | None = None, grid_spacing: tuple[float, ...] | None = None, surface_method: str | None = None, q_freq: float | None = None, q_band: tuple[float, float] | None = None, fit_convention: str | None = None, field_names: tuple[str, ...] | None = None, field_shapes: tuple[tuple[int, ...], ...] | None = None, shape: tuple[int, ...] | None = None, dtype: dtype | None = None, device: device | None = None, origin: tuple[int, ...] | None = None, pml_freq: float | None = None, scheme: str | None = 'cfs-cpml', detached_deliberately: bool = False, max_vel: float | None = None, internal_dt: float | None = None, format_version: int = 2)

Opaque absorbing-boundary and auxiliary state of one propagation.

A propagator’s absorbing boundary (a CFS-CPML) carries memory variables from one time step to the next, and continuing a simulation exactly requires carrying them across the join. They are implementation state, not physics: this object holds them so that a continuation is a matter of passing one value back (prop(..., initial_state=out)) rather than of naming fields whose number and shape are properties of the absorbing-boundary implementation.

The bundle is opaque as an output element: do not index it positionally, and do not build on how many tensors it holds, on their order or on their shapes (see Stability policy). Its keys are named – they are the provisional internal diagnostic names listed in Internal callback names. The mapping surface exists for inspection and debugging; the supported continuation use is to hand an AuxiliaryState back to a propagator unchanged.

It is a read-only mapping from name to torch.Tensor (names(), bs[name], name in bs, len(bs), iteration, keys(), items(), values(), as_dict()), plus to() for moving it between devices and dtypes and detach() for cutting the autograd graph at a join on purpose. It is not a torch.Tensor and never becomes one; the tensors it holds keep their grad_fn, so a continuation is differentiable end to end.

== is identity comparison, deliberately: a generated __eq__ would compare the tensors and return a Tensor, making bs1 == bs2 raise inside any if or in. Use compatible_with() to ask whether one auxiliary state may be reused by another call.

Registered as a pytree node, so torch.func and torch.compile reach the tensors it holds (it is one element of every propagator’s output, so without that the output could not be flattened to tensors at all). The children are the tensors; the tensor names and every recorded run parameter, format_version included, travel as the node’s context and come back unchanged, so a state that has been through a transform is still one a continuation checks itself against rather than an unrecorded bundle.

propagator

The name of the propagator that produced it (e.g. "scalar"), or None for a state that was assembled by hand rather than returned by a run.

Type:

str | None

ndim

The number of spatial dimensions of the run that produced it.

Type:

int | None

accuracy

The finite-difference order of accuracy that run used.

Type:

int | None

pml_width

The resolved per-face PML widths, length 2 * ndim.

Type:

tuple[int, …] | None

boundary

The resolved per-face boundary conditions, length 2 * ndim, from {"pml", "free", "rigid", "periodic"}.

Type:

tuple[str, …] | None

n_mechanisms

The number of relaxation mechanisms (1 except for the viscoacoustic propagator).

Type:

int | None

mpml

Whether multiaxial PML cross-damping was in use.

Type:

bool | None

shape

The shape of the tensors, [n_shot, *spatial], where the spatial part is the extracted survey region plus the PML. Checked: a continuation whose resolved simulation domain is a different shape is not continuing the same simulation.

Type:

tuple[int, …] | None

dtype

The dtype of the tensors.

Type:

torch.dtype | None

device

The device the tensors are on.

Type:

torch.device | None

origin

The index, in the full model, of cell 0 of the region the run extracted, length ndim. Used (not checked) to restore the same survey region on a continuation. The region’s size is checked, through shape.

Type:

tuple[int, …] | None

pml_freq

The PML frequency, in Hz, that run resolved.

Type:

float | None

max_vel

The maximum wavespeed, in m/s, that run used to build the PML profiles and choose its internal time step. A changed resolved value is incompatible with checked continuation.

Type:

float | None

internal_dt

The internal (sub-stepped) time step, in seconds, that run took. Checked: it is the condition every propagator’s initial_state documentation names for a continuation that exactly reproduces a single combined run.

Type:

float | None

scheme

Names the absorbing-boundary implementation the tensors belong to, so that a future change of scheme cannot silently reuse state from the old one.

Type:

str | None

detached_deliberately

Whether the tensors were detached from the autograd graph on purpose, by detach(). A state that arrived detached by accident (see detach()) is warned about; one marked this way is accepted silently. It is not part of compatible_with(): it says nothing about what the tensors mean.

Type:

bool

format_version

The version of the tensor bundle’s layout and numerical interpretation. A correction can require a new version, including in a patch release. Saved states require a compatible implementation; they are not a long-term interchange format. Handing back a state whose format_version differs from the one this release writes raises deepwave.DeepwaveValueError, naming both versions, rather than reading the tensors as something they are not.

Type:

int

A field recorded as None was not supplied by a run and is neither reported nor checked.

as_dict() → dict[str, Tensor]

A fresh dict of the tensors; mutating it does not affect this.

compatible_with(other: AuxiliaryState) → bool

Whether a call described by other may reuse this auxiliary state.

Compares format_version – a bundle written to a different layout is not readable as this one whatever else matches – and then the recorded run parameters that determine what the tensors mean: propagator, ndim, accuracy, pml_width, boundary, n_mechanisms, mpml, scheme, shape and internal_dt, as well as the resolved grid, attenuation and absorbing-boundary profile. dtype and device are deliberately excluded (see to()). A changed pml_freq or max_vel is incompatible with checked continuation.

shape is included – although to() cannot change it – because it is what pins the size of the survey region the tensors describe. origin pins where that region starts; without shape a continuation whose sources or receivers have moved outside the previous region can have the state re-read at a larger extent and silently zero-padded, displacing every field.

detach() → AuxiliaryState

A copy whose tensors are detached from the autograd graph.

Cutting the graph at a segment join deliberately – truncated backpropagation through a long simulation – is indistinguishable, from inside a propagator, from having lost the gradient by accident, so this records the intent: a state detached this way is accepted silently, while one that arrived detached is warned about.

This detaches the auxiliary state alone, which is not by itself enough to truncate anything: the physical wavefields of the same join are separate fields and stay attached, so the graph is not cut, no memory is freed, and the gradient is still the full one. Use PropagatorOutput.detach(), which cuts every field of the join, unless you are assembling the continuation state field by field and detaching the rest yourself.

items() → ItemsView[str, Tensor]

The (name, tensor) pairs, as a mapping view.

keys() → KeysView[str]

The names of the tensors, as a mapping view.

property names: tuple[str, ...]

The names of the tensors, in order.

to(device: str | device | None = None, dtype: dtype | None = None) → AuxiliaryState

A copy with the tensors moved to device and/or cast to dtype.

The recorded device/dtype are updated to match. Neither is part of the compatibility check (compatible_with()), so a round trip through another device – which is what deepwave.checkpointing.run_segments()’s checkpoint_device does – is legitimate.

values() → ValuesView[Tensor]

The tensors, as a mapping view.

Every PropagatorOutput’s .info attribute is a deepwave.RunInfo carrying non-differentiable metadata about how the call actually ran: the user dt and grid_spacing, the accuracy, the per-face boundary, pml_width, pml_freq, mpml_fraction and surface_method, the attenuation model’s q_freq/q_n_mechanisms/q_band, the resolved internal step_ratio, the internal dt, the targeted Courant number, nt, max_vel, and the survey origin/extent. Together these are the numerical configuration of the run, so out.info is enough to set up an identical one. pml_freq is the frequency the run actually used – the argument you gave, or, if you gave none, the value inferred from your source amplitudes (see Choosing pml_freq); pass it back in as pml_freq= to pin a later run to the same absorbing boundary. It is plain attribute metadata, not a tuple element, so it never affects iteration, unpacking, len(), or indexing.

class deepwave.RunInfo(step_ratio: int, internal_dt: float, courant_number: float, nt: int, pml_width: tuple[int, ...], pml_freq: float, origin: tuple[int, ...], extent: tuple[int, ...], dt: float = 0.0, grid_spacing: tuple[float, ...] = (), accuracy: int = 0, boundary: tuple[str, ...] = (), mpml_fraction: float = 0.0, surface_method: str = '', q_freq: float | None = None, q_n_mechanisms: int = 1, q_band: tuple[float, float] | None = None, max_vel: float | None = None, ndim: int = 0, dtype: dtype | None = None, device: device | None = None, model_batched: tuple[bool, ...] = (), wavefield_names: tuple[str, ...] = (), receiver_names: tuple[str, ...] = (), location_interpolation_used: str = '', backend: str = '', backward_backend: str = '', storage: Storage | None = None, lossy_compression_used: bool = False, model_gradient_sampling_interval: int = 1, model_gradient_stride: int | None = None)

Non-differentiable metadata about how a propagator call ran.

Every propagator attaches one to its return, as out.info. It records what the call resolved – the values the engine actually used, which are not always the ones the caller asked for (dt may be sub-stepped to satisfy the CFL condition, pml_freq may be inferred from the source spectrum, the survey region may be cut down around the sources and receivers) – so that a later call can be pinned to the same choices.

Existing fields follow the stable compatibility policy (see Stability policy). New resolved metadata may be added as attributes; an existing field is changed only through that policy.

dt

The user-level time step, in the caller’s own time unit, exactly as passed. internal_dt is this divided by step_ratio.

Type:

float

grid_spacing

The cell size along each dimension, in the caller’s own length unit, one entry per dimension (a scalar argument is expanded). Resolved.

Type:

tuple[float, …]

accuracy

The finite-difference order of accuracy in space.

Type:

int

boundary

The per-face boundary condition, length 2 * ndim, in the face order (begin dim 0, end dim 0, begin dim 1, ...), spelled as the boundary= argument spells them ("pml", "free", "rigid", "periodic"). Resolved, including the faces a zero pml_width turned into something else.

Type:

tuple[str, …]

mpml_fraction

The multiaxial-PML cross-damping fraction used (see deepwave.elastic()). 0.0 on a standard C-PML run.

Type:

float

surface_method

How the free surface implied by elastic vacuum cells was discretised (see deepwave.elastic()).

Type:

str

q_freq

The reference frequency of the attenuation model, or None on a propagator without one.

Type:

float | None

q_n_mechanisms

The number of relaxation mechanisms in the attenuation model (1 where there is none).

Type:

int

q_band

The frequency band the constant-Q fit spans, or None.

Type:

tuple[float, float] | None

step_ratio

The number of internal time steps taken per user-level time sample. 1 unless the user dt exceeded the CFL limit. Resolved.

Type:

int

model_gradient_sampling_interval

The positive integer requested by the caller, including when no model gradient is active.

Type:

int

model_gradient_stride

The effective native history/correlation stride in internal steps (step_ratio * model_gradient_sampling_interval), 1 for active eager model differentiation, or None when no model gradient is active. It describes sampling, not the accuracy of the complete parameter derivative.

Type:

int | None

internal_dt

The internal time step, in the time units of dt: the user dt divided by step_ratio. Resolved.

Type:

float

courant_number

The Courant number the internal time step was chosen to satisfy. The user’s courant_number argument when one was given, otherwise the propagator’s conservative default. Resolved.

Type:

float

nt

The number of user-level time samples propagated. The caller’s nt when one was given (the source amplitudes are padded or truncated to it), otherwise the length of the source amplitudes. Resolved.

Type:

int

pml_width

The per-face PML widths, in cells, length 2 * ndim, in the face order (begin dim 0, end dim 0, begin dim 1, ...). Resolved (a scalar argument is expanded to every face).

Type:

tuple[int, …]

pml_freq

The PML frequency, in the reciprocal of dt’s time unit (Hz when dt is in seconds): the pml_freq argument when one was given, otherwise the value inferred from the source amplitudes’ spectrum (or 0.5 / dt / 5 when there were no sources). Pass it back in as pml_freq= to pin a later run to the same absorbing boundary. Resolved.

Type:

float

origin

The index, in cells, in the full model, of cell 0 of the region the run extracted, one entry per dimension. Resolved (from the caller’s origin, from survey_pad, or from the source and receiver locations).

Type:

tuple[int, …]

extent

The spatial shape of the region extracted from the model, in cells, before FD and PML padding, one entry per dimension. Resolved.

Type:

tuple[int, …]

max_vel

The conservative speed bound, including density/stiffness weighting of staggered operators, in grid_spacing units per dt unit, used to choose the internal time step and to build the PML profiles: the max_vel argument when one was given, otherwise the maximum over the model. Resolved. None if the propagator did not record it.

Type:

float | None

ndim

The number of spatial dimensions. Resolved (from the model and location shapes when it was not given explicitly). 0 if the propagator did not record it.

Type:

int

dtype

The floating-point dtype the simulation ran in. Resolved from the models. None if the propagator did not record it.

Type:

torch.dtype | None

device

The device the simulation ran on. Resolved from the models. None if the propagator did not record it.

Type:

torch.device | None

model_batched

One flag per model, in the propagator’s model order, saying whether that model was given per shot (an extra leading batch dimension) rather than shared by every shot. Resolved. Empty if the propagator did not record it.

Type:

tuple[bool, …]

wavefield_names

The public state field names of this call, in return order – exactly the keys of PropagatorOutput.final_state, and therefore the names initial_state accepts. Depends on ndim and, for the attenuating propagators, on the number of relaxation mechanisms. Empty if the propagator did not record it.

Type:

tuple[str, …]

receiver_names

The receiver-amplitude field names of this call, in return order. Empty if the propagator did not record it.

Type:

tuple[str, …]

location_interpolation_used

The location interpolation actually used for this call’s source and receiver coordinates: "none" for the exact (on-cell) path, "hicks" for Kaiser-windowed sinc interpolation, or "mixed" when one kind of location was split across the two paths – some of a propagator’s source channels interpolated and others not, say. Sources are compared with sources and receivers with receivers, never with each other, so an off-node source with on-node receivers is "hicks", not "mixed"; only a propagator with several channels of one kind can report "mixed". A channel with no locations (None, an empty Tensor, or nothing but deepwave.IGNORE_LOCATION) has taken neither path and is not counted. "auto" resolves to one of those. The empty string if the propagator did not record it.

This is how to tell whether location_interpolation="auto" routed your coordinates through Hicks – in particular for on-node positions with requires_grad=True, which are always interpolated (only that path carries gradients back to the positions), and so give different forward values from the same positions without requires_grad.

Type:

str

backend

The backend that ran the forward pass, as one of "compiled", "eager" or "torch_compile". Resolved, so it is never "auto": that is a request, and this is the answer to it. This is how to find out what backend="auto" chose – including inside a distributed worker, where a warning would not have been seen. The empty string if the propagator did not record it.

Type:

str

backward_backend

The backend that would run the backward pass, in the same vocabulary. It is reported separately because a propagator is permitted to pair a compiled forward with an adjoint that runs on a different backend, but no configuration reachable today does: on every propagator and every backend it is currently equal to backend. (The one near-miss is deepwave.elastic_born() under M-PML, whose adjoint would use an eager-recompute bridge – but its compiled forward raises for mpml_fraction > 0 before any RunInfo is built.) Compare the two rather than assuming they agree; do not rely on their being equal. The empty string if the propagator did not record it.

Type:

str

storage

The fully resolved storage policy, as a deepwave.Storage. path is the directory that was actually chosen (mode="disk" only), so a run that fell back to the system temporary directory can be read back off here rather than guessed at. mode is the resolved mode, which records the "cpu"-to-"device" downgrade the engine performs when the compute device is already the CPU – otherwise invisible. None if the propagator did not record it.

Type:

deepwave.Storage | None

lossy_compression_used

Whether the wavefield history was actually stored with lossy compression. Resolved, not requested: it is False whenever no history was kept at all (mode="none"), even if compression was asked for, matching what the storage layer does.

Type:

bool

Attenuation response

deepwave.attenuation.response(q: Tensor, frequencies: Tensor, *, q_freq: float | None = None, q_n_mechanisms: int = 1, q_band: tuple[float, float] | None = None) → Tensor

Return the dimensionless complex modulus M(omega)/K.

The output has shape (*q.shape, frequencies.numel()) and complex64 or complex128 dtype matching q. Its real and imaginary parts are storage and positive loss. At DC the result is one. For multiple mechanisms, q is the weak-loss band target; real/imag is the actual effective Q and generally differs from that target, especially at strong loss. Gradients with respect to q are preserved. Positive infinite q gives the exact lossless response and the limiting zero Q derivative.

deepwave.checkpointing

deepwave.checkpointing.run_segments(propagator: Callable[[...], Any], model_args: Sequence[Tensor], grid_spacing: float | int | Sequence[float], dt: float, source_amplitudes: Mapping[str, Tensor], n_segments: int | None = None, *, seg_lengths: Sequence[int] | None = None, initial_state: PropagatorOutput | Mapping[str, Tensor | AuxiliaryState] | None = None, segment_callback: Callable[[int, PropagatorOutput], None] | None = None, storage: str | Storage | None = None, storage_mode: Literal['device', 'cpu', 'disk', 'none'] | None = None, storage_path: str | PathLike[str] | None = None, storage_compression: bool | None = None, combine: bool = True, optional_model_args: Mapping[str, Tensor] | None = None, checkpoint: bool = False, checkpoint_device: str | device | None = None, use_reentrant: bool = True, _slab_pml: bool = True, **kwargs: Any) → PropagatorOutput

Runs a propagator segment-by-segment, threading state via initial_state.

Stability: provisional (see Stability policy).

This is the general driver underlying snapshots() and deepwave.wavefields.measure(); it is also reusable on its own (e.g. to emit a wavefield snapshot at each segment boundary via segment_callback, or to bound backward-pass memory with checkpoint=True). Apart from the optional per-segment torch.utils.checkpoint wrapping it performs no gradient bookkeeping of its own – whether the result is differentiable is governed entirely by whether the inputs require grad and by storage_mode (as for a normal propagator call).

This takes model_args as a positional sequence and source_amplitudes as a name->tensor mapping. That generic-wrapper convention is what lets one implementation drive any current or future Deepwave propagator, but it is provisional and it is not the ergonomic way to segment a run: most users should reach for checkpoint_segments= on the propagator itself, which builds this call for them.

Internal sub-stepping (``step_ratio > 1``) is handled automatically. When the user dt does not satisfy the CFL condition, Deepwave runs internally at a smaller time step, resampling the source amplitudes up and the receiver amplitudes down with a global FFT. Because that resampling does not decompose across a segment boundary, a run of two or more segments at step_ratio > 1 instead resamples the whole trace once: every source tensor is upsampled a single time, every segment is driven at the internal time step (so each segment’s own step_ratio is 1), model_gradient_sampling_interval is scaled so the model gradient is sampled at the same steps a direct call samples it, and the time-concatenated receiver output is downsampled once at the end. Forward values and gradients then match an unsegmented direct call within rounding, given the same numerical setup. An omitted max_vel is resolved once through the propagator’s own preparation, including its heterogeneous and attenuation CFL bounds, before the record is split. This is behaviour, not an option: a caller wanting raw segments at the internal rate simply passes the internal dt itself, making step_ratio 1.

Combining with ``model_gradient_sampling_interval > 1``: the model-gradient match above holds only if every segment boundary falls on a multiple of the sampling interval, because each segment’s sampling phase restarts at 0 and a boundary that splits a sampling period would sample the gradient at the wrong steps. Rather than leave that to the caller, a misaligned split raises DeepwaveValueError up front, in the forward pass, naming the offending segment and a segment count that would work. The check is skipped when no model gradient will be computed, since there is then nothing for the interval to sample. This is unrelated to step_ratio and applies whether or not internal sub-stepping is happening.

Parameters:

• propagator – A Deepwave propagator function – any propagator that deepwave.spec() accepts, currently every propagator deepwave provides (e.g. deepwave.scalar, acoustic, elastic, elastic_vti, acoustic_vti, scalar_born, viscoacoustic, acoustic_born, elastic_born). Every propagator (including the Born family) takes its required models positional (before grid_spacing/dt) and its optional scattering models (dv/drho/dlamb/dmu/dbuoyancy/c66) keyword-only (after dt); since run_segments calls the propagator directly (not through a torch.autograd.Function boundary), those optional models may simply be passed in **kwargs like any other keyword argument and remain fully differentiable – unless checkpoint=True, which does cross such a boundary, and so needs them in optional_model_args instead (see below).

• model_args – The model tensors, in the order the propagator takes them positionally (e.g. (v,) for scalar; (lamb, mu, buoyancy) for elastic).

• grid_spacing – The grid spacing (as passed to the propagator).

• dt – The user time step.

• source_amplitudes – Mapping from source-amplitude argument name (e.g. "source_amplitudes", or "source_amplitudes_y") to its tensor of shape [shot, source, time]. All must share the same time length, which defines nt.

• n_segments – The number of near-equal contiguous time segments to split into. Exactly one of n_segments or seg_lengths must be given.

• seg_lengths – Explicit per-segment step counts (must sum to nt), for splits that are not near-equal (e.g. snapshot boundaries). Exactly one of n_segments or seg_lengths must be given.

• initial_state – An optional starting state (a PropagatorOutput or a name->tensor mapping), passed to the first segment.

• segment_callback –

  An optional segment_callback(cumulative_step, output) invoked after each segment completes, with cumulative_step the number of user time steps run so far (running total of the segment lengths, ending at nt) and output that segment’s PropagatorOutput. Segment boundaries always fall on user steps, so this count is unaffected by automatic resampling; the per-segment output’s receiver fields, however, are at the internal sample rate while resampling is active (they are downsampled only once, after the last segment) – detect that case by comparing the per-segment output.info.internal_dt with the dt you passed.

  This is distinct from a per-step forward_callback/backward_callback forwarded in kwargs, which kwargs passes through to each segment’s own propagator call. Two things about those are easy to get wrong. callback_frequency counts the steps of each segment’s call. When whole-trace resampling is active, these are internal steps, with interval output.info.internal_dt; otherwise they are user steps. Multiply the desired user-step frequency by the whole-run step_ratio when resampling is active. And each segment is a fresh propagator call, so CallbackState.step restarts at 0 in every segment: a callback that indexes a whole-run array by state.step will overwrite the first segment’s entries with every later segment’s. During the original forward pass, a whole-run time is the preceding segment’s cumulative_step * dt plus state.step * state.dt. Do not use a mutable segment offset for backward callbacks or checkpoint recomputation: they execute after that offset has advanced. Forward callbacks can run again during checkpoint recomputation, so their side effects must be safe to repeat. Use segment_callback for observations made once per segment in the original forward pass.

• storage – The storage policy passed to each segment, as one object: either a mode string or a deepwave.Storage. An alternative spelling for the three individual storage_* arguments below, which are permanent and not deprecated; passing both spellings raises DeepwaveValueError. Resolved once, here, so that every segment runs the same policy.

• storage_mode – The storage_mode passed to each segment. Defaults to "device", matching the propagators’ own default (so a segmented run is differentiable by default, the same as a direct call).

• storage_path – The storage_path passed to each segment (used only by storage_mode="disk"). Defaults to the system temporary directory (tempfile.gettempdir()).

• storage_compression – The storage_compression passed to each segment. Lossy – changes the model gradient by about 0.2-4% in 2-D and 0.2-8% in 3-D. See deepwave.elastic(). Defaults to False.

• combine – Which sample rate the returned receiver amplitudes are at, and how much is retained to produce them. If True (default) return a combined PropagatorOutput for the whole run – final wavefield state plus time-concatenated receiver amplitudes, downsampled once as a whole to the user sample rate if the run was internally sub-stepped. If False return the last segment’s output directly, at the internal sample rate, and do not retain the per-segment receiver traces, which keeps a purely streaming driver (used by deepwave.wavefields.measure(), whose segment_callback reduces each segment on the fly). Because the downsampling FFT is global in time and cannot be applied per segment, combine=False together with active resampling (two or more segments at step_ratio > 1) raises DeepwaveValueError rather than quietly returning internal-rate data.

• optional_model_args – The keyword-only optional model tensors (dv, drho, dlamb, dmu, dbuoyancy, c66) as a name->tensor mapping. Without checkpoint these are simply merged into **kwargs (passing them in kwargs directly remains supported and equivalent). With checkpoint=True they must be given here rather than in kwargs, because they are then threaded positionally through each checkpoint boundary so they receive gradients (closure-captured tensors get none under reentrant checkpointing). A name given both here and in kwargs is a ValueError.

• checkpoint – If True, wrap each segment except the last in torch.utils.checkpoint.checkpoint, so that segment’s forward activations (the stored wavefield history) are freed after it runs and recomputed on demand during the backward pass. This is the standard way to bound backward-pass memory: peak backward memory becomes roughly one segment’s history plus the auxiliary states, instead of the whole run’s history, at the cost of one extra forward pass per checkpointed segment. The last segment is deliberately not checkpointed (recomputing it would cost time and save no peak memory, since its history is needed immediately). A pass with nothing to differentiate – no model, source, optional-model, initial-state or keyword tensor requires grad, or grad mode is off – skips the wrapper entirely and calls each segment directly, since recomputation buys nothing without a backward pass. That is bit-exact (the wrapper runs the same code, only under torch.no_grad()), marginally cheaper, and free of the per-segment “None of the inputs have requires_grad=True” warning torch would otherwise emit, so checkpoint=True is safe to leave set on a forward-only run. The PML memory-variable fields of the state threaded between checkpointed segments are stored as their two per-axis boundary slabs (thickness pml_width + accuracy per side) rather than full-size, and rebuilt inside the segment: those fields are exactly zero outside the slabs, so this is bitwise-exact in both the forward values and every gradient, and it shrinks the n_segments x auxiliary-state term that dominates checkpointed peak memory by roughly 2-3x, growing with grid size. It applies to the threaded state only – segment_callback and every returned output still see full-size tensors – and falls back silently and exactly to full-size states for propagators that do not declare PropagatorSpec.pml_slab_axes, for backends other than the compiled one, for mpml_fraction > 0, for a non-default boundary, and for grids too small for a slab to be smaller than the field. Take the gradient of a checkpointed run with loss.backward(): the checkpointing is reentrant by default, and torch.utils.checkpoint refuses torch.autograd.grad() (and an inputs= argument to .backward()) in that mode, raising rather than returning a wrong number. Set use_reentrant=False if torch.autograd.grad is needed. Defaults to False.

• checkpoint_device – If given (with checkpoint=True), the state tensors threaded between segments are moved to this device after each checkpointed segment and moved back to the compute device inside the next segment, freeing compute-device memory between segments (e.g. "cpu"). Device transfers are differentiable, so this does not change gradients. It is not applied after the final segment, so the returned state is on the compute device exactly as without it. Giving it without checkpoint is a ValueError.

• use_reentrant – The use_reentrant value passed to every torch.utils.checkpoint.checkpoint call. Defaults to True, which saves native wavefield history between segments. Set False to use torch.autograd.grad or backward(inputs=...) with checkpointing; this can retain a history buffer per segment and substantially increase memory use.

• _slab_pml – Set False to thread the PML memory-variable fields of the checkpointed state full-size instead of as their two per-axis boundary slabs. Underscored and outside the stability policy, for Deepwave’s own tests.

• **kwargs – Remaining keyword arguments forwarded unchanged to every segment (source_locations, receiver_locations, pml_freq, max_vel, accuracy, origin, …). The survey may be given in any call form the propagator itself accepts: a sources=/receivers=/bg_receivers= mapping is resolved into the flat keywords before the first segment runs, and a flat source_amplitudes* keyword is lifted into the source_amplitudes mapping this function slices per segment. Neither is forwarded whole, which would give every segment the entire record. nt, if supplied, pads or truncates the whole source record before splitting it, as on a direct propagator call. Giving one channel’s amplitudes both ways raises. storage_mode, storage_path, and storage_compression must not appear here (use the parameters above). Under checkpoint=True a tensor here that requires grad is rejected, since kwargs are captured by the checkpointed closure and would silently receive no gradient (only top-level tensor values are checked). If pml_freq is absent, it is resolved once, from the whole source record (not per segment, which would give the segments different PML profiles while one set of boundary memory variables is threaded through them all) and then passed to every segment explicitly.

Returns:

A combined PropagatorOutput (combine=True) or the last segment’s output (combine=False).

deepwave.checkpointing.snapshots(propagator: Callable[[...], Any], model_args: Sequence[Tensor], grid_spacing: float | int | Sequence[float], dt: float, source_amplitudes: Mapping[str, Tensor], snapshot_steps: Sequence[int], *, initial_state: PropagatorOutput | Mapping[str, Tensor | AuxiliaryState] | None = None, max_vel: float | None = None, courant_number: float | None = None, **kwargs: Any) → SnapshotResult

Differentiable intermediate wavefield snapshots of a propagation.

Stability: provisional (see Stability policy).

Runs propagator once, from step 0 to nt, but breaks the run into contiguous segments at the requested snapshot_steps (outer-step indices) and returns the full wavefield state at each of those steps inside the autograd graph. A loss built on any returned snapshot is differentiable w.r.t. the model and source tensors, exactly as a loss on the final output is (the segments are chained through initial_state, so autograd threads the adjoint state across every boundary).

This is the read-write counterpart to the read-only forward_callback wavefield surface: instead of observing a detached wavefield mid-run, the returned snapshots are live graph nodes you can put a loss on (e.g. to match a measured wavefield at an intermediate time, to regularise the field, or to build a wavefield-space objective for WRI-style inversion).

Implemented as a thin, propagator-agnostic layer over run_segments() (each segment is a normal, differentiable propagator call with storage_mode="device"). It therefore inherits that function’s exactness behaviour, and the same model_args (positional)/source_amplitudes (name->tensor mapping) generic-wrapper convention: what lets one implementation drive any propagator, and provisional along with the rest of this module. Users who only want to bound backward-pass memory, rather than to differentiate a mid-run wavefield, should reach for checkpoint_segments= on the propagator itself instead.

Snapshots are supported at any step_ratio. When the user dt requires internal sub-stepping (step_ratio > 1), the driver resolves max_vel from the whole call if omitted and resamples the whole propagation once, so a snapshot still equals the wavefield a single unsegmented run would have reached; see run_segments().

The snapshot at outer step s is the wavefield state after s outer (user) steps have been applied – i.e. after s * step_ratio internal steps, the same wavefield a forward_callback observes when it is invoked at step == s (the forward callback fires before the chunk starting at step). Snapshots are returned in the propagator’s “pml” view (model region plus PML, finite-difference padding removed), matching CallbackState.get_wavefield(name, view="pml").

Parameters:

• propagator – A Deepwave propagator function (deepwave.scalar etc.).

• model_args – The model tensors in positional order ((v,) for scalar; (lamb, mu, buoyancy) for elastic). Those with requires_grad receive a gradient from any loss on the snapshots or final output.

• grid_spacing – Grid spacing, as passed to the propagator.

• dt – The user time step.

• source_amplitudes – Mapping from source-amplitude argument name to tensor ({"source_amplitudes": src} for scalar), defining nt.

• snapshot_steps – The outer-step indices at which to return a snapshot. Each must satisfy 1 <= step <= nt. Duplicates are ignored and the order does not matter (the returned mapping is keyed by step).

• initial_state – Optional starting state (PropagatorOutput or mapping), threaded into the first segment.

• max_vel – Maximum wave speed, used to resolve the internal time step for the whole-trace resampling and passed to the propagator. If omitted, the propagator’s preparation resolves its conservative bound.

• courant_number – Optional Courant number, affecting the resolved step_ratio and passed to the propagator.

• **kwargs – Remaining propagator keyword arguments forwarded to every segment (source_locations, receiver_locations, pml_freq, accuracy, origin, …). storage_mode and run_segments()’s checkpointing parameters must not appear.

Returns:

A SnapshotResult, a named (output, snapshots) tuple. output is the whole propagation’s PropagatorOutput (receiver amplitudes over all nt steps and the final wavefield state), as for a direct propagator call. snapshots is a dict mapping each requested step to a {wavefield_name: tensor} dict of that step’s wavefield state (differentiable).

class deepwave.checkpointing.SnapshotResult(output: PropagatorOutput, snapshots: dict[int, dict[str, Tensor]])

The result of deepwave.checkpointing.snapshots().

A named two-tuple, so combined, snaps = snapshots(...) and result.snapshots both read naturally.

Stability: provisional (see Stability policy).

output

The whole propagation’s PropagatorOutput, as a direct call would return it.

Type:

deepwave.PropagatorOutput

snapshots

A mapping from each requested outer step to that step’s {wavefield_name: tensor} wavefield state.

Type:

dict[int, dict[str, torch.Tensor]]

Create new instance of SnapshotResult(output, snapshots)

CallbackState

Every accessor below takes an internal name – the kernel’s own spelling of a model or wavefield, not the argument you passed the propagator. Internal callback names lists them all, per propagator and per dimensionality, and distinguishes stable physical fields from provisional internal diagnostics.

class deepwave.CallbackState(dt: float, step: int, wavefields: dict[str, Tensor], models: dict[str, Tensor], gradients: dict[str, Tensor], fd_pad: list[int], pml_width: list[int], forward_wavefields: dict[str, Tensor] | None = None, *, direction: Literal['forward', 'backward'], n_steps: int, internal_dt: float | None = None)

State provided to user callbacks.

What you may write to

Three separate promises, deliberately not one:

1. get_gradient() returns a live view into the gradient accumulator, and mutating it in place is supported. This accumulator is intermediate, even at the last callback. To scale the complete model gradient, use a PyTorch tensor gradient hook or transform model.grad after .backward(); see Custom imaging condition. The internal accumulator is available only on the compiled backend, which is the only backend that populates the gradient dictionary at all: the eager backward path and every forward callback are given an empty one, so there is nothing to mutate there. Ask writable() rather than assuming.

2. get_wavefield() and get_model() return, in the default mode="write", the tensor Deepwave itself holds, and writing into it in place changes the running simulation in inference calls. During differentiable runs a forward callback may observe these tensors but may not mutate them: an arbitrary edit has no defined adjoint and raises deepwave.DeepwaveRuntimeError. set_wavefield() is the checked route into the wavefields – it validates shape, dtype and device instead of letting a mismatch become a silent broadcast or a host-to-device copy in the middle of the time loop – and is what to prefer where it applies. There is no set_model; a model is written in place. A name that is derived from the internal state rather than being part of it is read-only, and writable() is how to ask which a name is.

3. Nothing you keep past the end of the callback is stable. "read" mode gives a detached view, not a snapshot: the buffer it points at is overwritten at the next time step. If you keep a tensor past the end of the callback, ``.clone()`` it – Visualising wavefields and gradients with callbacks is the modelled idiom.

When PyTorch invokes a backward node with an undefined adjoint, the callback receives a zero tensor for that field rather than None. PyTorch can also prune an unused backward node entirely, in which case its callback does not fire. A callback is not a finalisation hook.

step

User time-step index. Forward callbacks observe 0, f, 2*f, ... < n_steps for callback_frequency=f; backward callbacks traverse that same set in reverse at the matching state boundary. There is no extra callback at n_steps; the return value carries the final state.

direction

"forward" or "backward".

n_steps

The number of user time steps in this call. PyTorch may prune backward callback nodes that are not on the requested gradient path. In particular, eager torch.autograd.grad can prune the step-zero hook when initial wavefields are not requested; a backward callback is not a guaranteed finalisation hook.

dt

The caller’s time step. time == step * dt is measured from the start of this call.

internal_dt

The CFL-resolved time step used inside propagation.

time

Time of this state boundary in the caller’s units. Pressure and stress fields are at integer boundaries while staggered velocity components have their usual half-step offsets; time does not erase those component-specific offsets.

In a forward_callback the wavefields are the forward wavefields; in a backward_callback they are the adjoint wavefields, and the time-coincident forward wavefields may additionally be available through role="forward" – see get_wavefield(), has_forward_wavefield and available_roles.

Initialise the state.

Deepwave constructs this; a caller normally only receives one. A test double that constructs its own must supply every argument below.

param dt:

The user time step size.

param step:

The current user time-step index.

param wavefields:

A dictionary of wavefields.

param models:

A dictionary of models.

param gradients:

A dictionary of gradients.

param fd_pad:

The padding used for the finite difference stencil.

param pml_width:

The width of the PML.

param forward_wavefields:

Optionally, in a backward_callback, a dictionary of the forward wavefields reconstructed at the same time step as the adjoint wavefields in wavefields.

param direction:

Which pass this state belongs to.

param n_steps:

The number of user time steps in this pass.

param internal_dt:

The CFL-resolved time step.

property available_roles: tuple[str, ...]

The role values get_wavefield() will accept here.

("adjoint",), or ("adjoint", "forward") when the backend supplied the forward wavefield. None, which means the same as "adjoint", is not listed: it is not a distinct capability to branch on.

get_gradient(name: str, view: Literal['full', 'pml', 'inner'] = 'inner', mode: Literal['read', 'write'] = 'write') → Tensor

Get a gradient.

Parameters:

• name – The name of the gradient.

• view – The part of the gradient to return. One of ‘inner’ (the model region, default), ‘pml’ (the model region plus the PML), or ‘full’ (the entire gradient, including finite difference padding).

• mode – "write" (the default) returns the live accumulator; "read" returns .detach() of it. Unlike get_wavefield() and get_model(), this one is never a detached snapshot of a buffer that is about to be overwritten: mutating the gradient is the point. As on those two, "write" is the permanent default.

Returns:

The specified part of the gradient, as a live view into the intermediate gradient accumulator: mutating it in place changes the gradient Deepwave ultimately returns, and the changed result need not be a derivative of the unchanged forward loss. This is not the completed gradient with respect to the caller’s model. Only the compiled backend populates gradients; writable() reports whether this one does.

get_model(name: str, view: Literal['full', 'pml', 'inner'] = 'inner', mode: Literal['read', 'write'] = 'write') → Tensor

Get a model.

Parameters:

• name – The name of the model.

• view – The part of the model to return. One of ‘inner’ (the model region, default), ‘pml’ (the model region plus the PML), or ‘full’ (the entire model, including finite difference padding).

• mode – "write" (the default) returns exactly the tensor Deepwave holds; "read" returns .detach() of it. As for get_wavefield(), "read" is a detached view, not a copy – ``.clone()`` it if you keep it past the end of the callback – and "write" is the permanent default, not scheduled to flip in any release (see Stability policy).

Returns:

The specified part of the model. For an internal model name, writing into it in place is supported and changes the running simulation, as it did in v0.0.27; there is no set_model because writing in place is the route. writable(name, "model") reports whether that holds for this name.

get_wavefield(name: str, view: Literal['full', 'pml', 'inner'] = 'inner', role: Literal['adjoint', 'forward'] | None = None, mode: Literal['read', 'write'] = 'write') → Tensor

Get a wavefield.

Parameters:

• name – The name of the wavefield.

• view – The part of the wavefield to return. One of ‘inner’ (the model region, default), ‘pml’ (the model region plus the PML), or ‘full’ (the entire wavefield, including finite difference padding).

• role – Which wavefield to return, when both are available inside a backward_callback. None (default) or "adjoint" returns the primary wavefield of the callback (the forward field in a forward callback, the adjoint field in a backward callback). "forward" returns the reconstructed forward wavefield time-coincident with the adjoint field, for building custom imaging conditions; it is available only when the backend supplied forward wavefields, which has_forward_wavefield and available_roles report. Under storage_compression=True it is reconstructed from the lossy 8-bit history, so it is an approximation of the forward field (of order 0.1% of the local block maximum) rather than the field the forward pass propagated.

• mode –

  "write" (the default) returns exactly the tensor Deepwave holds. "read" returns .detach() of it, removing the autograd edge and stating that you do not intend to write through it.

  "read" is a detached view, not a copy: it does not make the tensor immutable and it does not stop the buffer being overwritten at the next time step. If you keep it past the end of the callback, ``.clone()`` it.

  "write" is the permanent default: it is not scheduled to flip in any release (see Stability policy). Writing into the returned tensor in place is supported for an internal wavefield name, exactly as in v0.0.27; writable() reports which names those are, and set_wavefield() is the checked route that validates shape, dtype and device for you.

Returns:

The specified part of the wavefield.

Raises:

NotImplementedError – If role="forward" is asked for where the backend did not supply the forward wavefield.

property has_forward_wavefield: bool

Whether get_wavefield(..., role="forward") will work here.

The branch-on-able form of the NotImplementedError get_wavefield() raises for role="forward": True exactly when the backend supplied the reconstructed forward wavefield for this callback. Lets one callback be written that computes an imaging condition where the forward field is available and falls back to inspecting the adjoint field where it is not, without a try/except.

set_wavefield(name: str, tensor: Tensor, view: Literal['full', 'pml', 'inner'] = 'inner') → None

Write into a wavefield of the running simulation.

The supported way to alter the simulation state from a callback. view mirrors get_wavefield()’s default, so state.set_wavefield(n, state.get_wavefield(n) + 1) round-trips.

The value is copied into Deepwave’s own tensor; tensor itself is not retained, and no broadcasting, casting or device transfer is performed. A mismatch raises rather than silently doing the conversion, because a silent host-to-device copy in the middle of a time loop is a performance trap and a silent cast loses precision where it is least visible.

Parameters:

• name – The name of the wavefield to write into.

• tensor – The values to write. Must match the target’s shape, dtype and device exactly.

• view – The part of the wavefield tensor corresponds to, as for get_wavefield().

Raises:

• DeepwaveValueError – If tensor’s shape or device differs from the target’s.

• DeepwaveTypeError – If tensor’s dtype differs from the target’s.

• KeyError – If name is not a wavefield of this callback, listing the names that are.

writable(name: str, kind: Literal['wavefield', 'model', 'gradient'] = 'wavefield') → bool

Whether writing to name is supported in this callback.

The branch-on-able form of the class docstring’s three promises, so one callback can be written against all nine propagators and both passes without a try/except around every write. It reports whether a write is supported, not merely whether one would happen to land somewhere today.

Parameters:

• name – The name to ask about.

• kind – Which name space name belongs to. The same string can be both a model and a gradient name ("v" is both), so the space has to be named for the question to be well posed.

Returns:

For kind="gradient", whether that gradient is present and so can be mutated in place through get_gradient() (true only on the compiled backend’s backward pass). For kind="wavefield", whether set_wavefield() will accept that name, which is also whether get_wavefield(name) may be written into in place. For kind="model", whether get_model(name) may be written into in place: true for an internal model name this callback holds, false for a name that is merely derived from the internal state (there is nothing behind such a name to write into). An unknown name returns False rather than raising – this is a query.

A deprecated name that the accessors still resolve (the DEPRECATED_CALLBACK_NAMES aliases) answers for the name it resolves to, so this agrees with the accessor it predicts. Being a query, it does not warn about the alias; the accessor does.

During a differentiable forward pass, model and wavefield writes are unsupported and return False even when the name is present.

MemoryEstimate

class deepwave.MemoryEstimate

The dict returned by estimate_memory().

Type aliases

The closed sets of values Deepwave’s arguments accept, named so that they can be used in your own annotations. boundary entries are deliberately plain str rather than a Literal; see Stability policy for why.

deepwave.StorageMode

alias of Literal[‘device’, ‘cpu’, ‘disk’, ‘none’]

deepwave.WavefieldView

alias of Literal[‘full’, ‘pml’, ‘inner’]

deepwave.CallbackRole

alias of Literal[‘adjoint’, ‘forward’]

deepwave.CallbackMode

alias of Literal[‘read’, ‘write’]

deepwave.CallbackKind

alias of Literal[‘wavefield’, ‘model’, ‘gradient’]

deepwave.CallbackDirection

alias of Literal[‘forward’, ‘backward’]

deepwave.LocationInterp

alias of Literal[‘auto’, ‘hicks’, ‘none’]

deepwave.Reduction

alias of Literal[‘mean’, ‘sum’, ‘none’]

deepwave.Accuracy

alias of Literal[2, 4, 6, 8]

deepwave.InitialState

alias of PropagatorOutput | Mapping[str, Tensor | AuxiliaryState]

deepwave.ON_NODE_TOLERANCE

How close to a grid node a coordinate must be, in cells, for Deepwave to treat it as being on that node.

This one number decides three things:

• location_interpolation="auto" builds a deepwave.Hicks interpolator only if some coordinate is further than this from its component’s node (but see that argument’s documentation for the requires_grad exception, which always interpolates);

• location_interpolation="none" rejects a locations Tensor that has any coordinate further than this off-node;

• deepwave.to_cells() with snap=True cleans up floating-point error out to this distance.

It is the single source for all three: Hicks’ default eps is this value, and so is the propagators’ internal node test. Changing it changes which calls are interpolated, and therefore changes numbers.

Stability: stable (see Stability policy).

boundaries Function

deepwave.boundaries(ndim: int, *, default: str = 'pml', width: int | Sequence[int] = 20, top: str | None = None, dim0_begin: str | None = None, dim0_end: str | None = None, dim1_begin: str | None = None, dim1_end: str | None = None, dim2_begin: str | None = None, dim2_end: str | None = None) → tuple[list[str], list[int]]

Builds a matching (boundary, pml_width) pair.

The two arguments have to agree – a named boundary face must have pml_width 0, and a zero-width face that is not named is an implicit reflector (see Reflecting edges). Building them together makes it impossible to get that wrong:

boundary, pml_width = deepwave.boundaries(2, top="free")
out = deepwave.scalar(
    v, dx, dt, ..., boundary=boundary, pml_width=pml_width
)

In 3-D, and with a per-face width:

boundary, pml_width = deepwave.boundaries(
    3, top="free", width=[0, 20, 20, 20, 20, 20]
)

Faces are named by axis, in the same order and with the same words the pml_width argument uses:

• beginning of first (slowest) dimension – dim0_begin

• end of first dimension – dim0_end

• beginning of second dimension – dim1_begin

• end of second dimension – dim1_end

• …

top is an alias for dim0_begin, provided because the beginning of dimension 0 is conventionally the free surface in seismic models; Deepwave itself gives no axis a physical meaning, so use dim0_begin…``dim2_end`` if your model is oriented differently. There are deliberately no bottom/left/right/front/ back spellings: they cannot be defined without inventing a handedness convention that the rest of the library refuses to fix (see Units and source conventions).

Parameters:

• ndim – The number of spatial dimensions: 1, 2 or 3.

• default – The condition for every face not named by a keyword. One of "pml" (the default), "free", "rigid", "periodic".

• width – The PML width in cells, applied only to "pml" faces: a single non-negative int, or a sequence of 2 * ndim of them. Every non-"pml" face gets width 0 automatically, which is what keeps the two returned lists consistent.

• top – An alias for dim0_begin. Giving both raises.

• dim0_begin – The condition for the beginning of dimension 0.

• dim0_end – The condition for the end of dimension 0.

• dim1_begin – The condition for the beginning of dimension 1.

• dim1_end – The condition for the end of dimension 1.

• dim2_begin – The condition for the beginning of dimension 2.

• dim2_end – The condition for the end of dimension 2.

Returns:

a list of 2 * ndim condition names and a list of 2 * ndim widths, ready to pass straight to any propagator’s boundary= and pml_width= arguments. This helper does not know which propagator the pair is for, so it builds any combination of names; a propagator that does not accept one of them (see Which faces each propagator accepts) raises when it is given the pair.

Return type:

(boundary, pml_width)

Raises:

DeepwaveValueError – If ndim is not 1, 2 or 3; if a dim{k}_* keyword is given for a k that ndim does not have; if both top and dim0_begin are given; if a condition name is not recognised; if width is negative or the wrong length; or if a "pml" face would end up with width 0, which is the implicit reflector this helper exists to replace.

Stability: provisional (see Stability policy).

Callback

class deepwave.Callback(*args, **kwargs)

A propagator callback: called with the live CallbackState.

Anything callable with one positional argument satisfies it – a def, a lambda, a bound method, a functools.partial, or an object with a __call__. The parameter is positional-only (the / in the signature below), so your callback’s parameter may be named anything; a protocol without it would silently reject every callback whose parameter is not called state.

runtime_checkable is set, so isinstance(cb, deepwave.Callback) works – but be aware that, as for every runtime-checkable protocol, it checks only for the presence of __call__, not the signature. It will therefore accept a callable of the wrong arity; use it as a sanity check, not as validation.

Example:

def report(state: deepwave.CallbackState) -> None:
    print(state.step, state.get_wavefield("wavefield_0").norm())

out = deepwave.scalar(v, dx, dt, ..., forward_callback=report)

Stability: stable (see Stability policy).

__call__(state: CallbackState, /) → None

Receives the live state for one time step.

cosine_taper_end Function

deepwave.discretisation.cosine_taper_end(signal: Tensor, n_taper: int) → Tensor

Tapers the end of the final dimension of a torch.Tensor using a cosine.

A half period, shifted and scaled to taper from 1 to 0, is used. The taper is dimensionless, so the result is in the same units as signal, whatever they are, and the sample interval of the final dimension is unchanged.

Parameters:

• signal – The torch.Tensor that will have its final dimension tapered. Its units are the caller’s own and are carried through.

• n_taper – The length of the cosine taper, in number of samples of the final dimension.

Returns:

The tapered signal, in the units of signal.

upsample Function

deepwave.discretisation.upsample(signal: Tensor, step_ratio: int, freq_taper_frac: float = 0.0, time_pad_frac: float = 0.0, time_taper: bool = False) → Tensor

Upsamples the final dimension of a torch.Tensor by a factor.

Low-pass upsampling is used to produce an upsampled signal without introducing higher frequencies than were present in the input. The Nyquist frequency of the input will be zeroed.

Units are the caller’s own and are carried through: the result is in the same units as signal, and its sample interval is signal’s divided by step_ratio. step_ratio is a count and the two fractions are dimensionless.

Parameters:

• signal – The torch.Tensor that will have its final dimension upsampled. Its final dimension is time, uniformly sampled; its units are the caller’s own.

• step_ratio – The integer factor by which the signal will be upsampled. The input signal is returned if this is 1 (freq_taper_frac, time_pad_frac, and time_taper will be ignored).

• freq_taper_frac – A float specifying the fraction of the end of the signal amplitude in the frequency domain to cosine taper. This might be useful to reduce ringing. A value of 0.1 means that the top 10% of frequencies will be tapered before upsampling. Defaults to 0.0 (no tapering).

• time_pad_frac – A float specifying the amount of zero padding that will be added to the signal before upsampling and removed afterwards, as a fraction of the length of the final dimension of the input signal. This might be useful to reduce wraparound artifacts. A value of 0.1 means that zero padding of 10% of the length of the signal will be used. Defaults to 0.0.

• time_taper – A bool specifying whether to force the upsampled signal to zero at both edges (with a Hann window) after its inverse FFT. The FFT-based resampling treats the time axis as periodic. A signal that does not already decay to (near) zero at its edges – e.g. real/field data cropped from a longer recording – otherwise produces wraparound ringing in the resampled result. Leave as False if the signal already tapers to (near) zero at its edges, since the window also attenuates genuine signal there.

Returns:

The signal after upsampling, in the units of signal, at a sample interval of signal’s divided by step_ratio.

downsample Function

deepwave.discretisation.downsample(signal: Tensor, step_ratio: int, freq_taper_frac: float = 0.0, time_pad_frac: float = 0.0, time_taper: bool = False, shift: float = 0.0) → Tensor

Downsamples the final dimension of a torch.Tensor by a factor.

Frequencies higher than or equal to the Nyquist frequency of the downsampled signal will be zeroed before downsampling.

The input time length must be divisible by step_ratio. Pad or trim the signal explicitly when a different output length is needed.

Units are the caller’s own and are carried through: the result is in the same units as signal, and its sample interval is signal’s multiplied by step_ratio. step_ratio is a count, the two fractions are dimensionless, and shift is in samples of signal.

Parameters:

• signal – The torch.Tensor that will have its final dimension downsampled. Its final dimension is time, uniformly sampled; its units are the caller’s own. Its length must be divisible by step_ratio.

• step_ratio – The integer factor by which the signal will be downsampled. The input signal is returned if this is 1 and shift is 0 (freq_taper_frac, time_pad_frac, and time_taper will be ignored).

• freq_taper_frac – A float specifying the fraction of the end of the signal amplitude in the frequency domain to cosine taper. This might be useful to reduce ringing. A value of 0.1 means that the top 10% of frequencies will be tapered after downsampling. Defaults to 0.0 (no tapering).

• time_pad_frac – A float specifying the amount of zero padding that will be added to the signal before downsampling and removed afterwards, as a fraction of the length of the final dimension of the output signal. This might be useful to reduce wraparound artifacts. A value of 0.1 means that zero padding of 10% of the length of the output signal will be used. Defaults to 0.0.

• time_taper – A bool specifying whether to force the signal to zero at both edges (with a Hann window) before this function’s FFT-based resampling, which treats the time axis as periodic. A signal that does not already decay to (near) zero at its edges – e.g. real/field data cropped from a longer recording – otherwise produces wraparound ringing in the resampled result. Leave as False if the signal already tapers to (near) zero at its edges, since the window also attenuates genuine signal there.

• shift – Amount, in samples of the input signal, to shift the data in time before downsampling. Defaults to 0.0.

Returns:

The signal after downsampling, in the units of signal, at a sample interval of signal’s multiplied by step_ratio.

cfl_condition Function

deepwave.discretisation.cfl_condition(dy: float, dx: float, *args: Any, **kwargs: Any) → tuple[float, int]

Calculates the time step interval for 2D models.

This is a convenience wrapper around cfl_condition_n for 2D models, and takes its units: any length unit and any time unit, as long as the grid spacings, dt and max_vel share the same pair.

Parameters:

• dy – The grid spacing in the y-dimension, in the length unit max_vel is expressed in.

• dx – The grid spacing in the x-dimension, in the same unit as dy.

• *args – Additional positional arguments to pass to cfl_condition_n.

• **kwargs – Additional keyword arguments to pass to cfl_condition_n.

Returns:

• inner_dt: A time step interval that obeys the CFL condition, in the same time unit as dt.

• step_ratio: The integer dt / inner_dt, dimensionless.

Return type:

Tuple

cfl_condition_n Function

deepwave.discretisation.cfl_condition_n(grid_spacing: float | Sequence[float], dt: float, max_vel: float, eps: float = 1e-15, courant_number: float | None = None) → tuple[float, int]

Calculates the time step interval to obey the CFL condition.

The output time step will be a factor of the input time step.

We use the maximum dt to calculate the integer factor (step_ratio) we need to divide the provided dt by to ensure it is less than or equal to this value.

Units are the caller’s own: any length unit and any time unit will do, as long as grid_spacing, dt and max_vel are expressed in the same pair. step_ratio and the Courant number are dimensionless, and inner_dt comes back in the unit dt was given in.

Parameters:

• grid_spacing – The grid spacing in each spatial dimension, in the length unit max_vel is expressed in. Can be a single float (applied as the only entry) or a sequence of floats.

• dt – The time step interval, in the time unit max_vel is expressed in.

• max_vel – The maximum absolute wavespeed in the model, in grid_spacing units per dt unit.

• eps – Retained for call compatibility. Zero speed is handled explicitly, so this parameter no longer regularises the speed or changes the selected time step.

• courant_number – An optional Courant number to use instead of the conservative default (0.6). If None (the default), 0.6 is used. If given, it is the actual dimensionless Courant number the time step will target; a larger value gives a larger inner_dt (fewer internal steps). This function does not itself enforce an upper bound: a caller that needs to guarantee stability should compare it against deepwave.discretisation.critical_courant_number() first and raise if it is exceeded, which is what the propagators do while setting up a call.

Returns:

• inner_dt: A time step interval that obeys the CFL condition, in the same time unit as dt.

• step_ratio: The integer dt / inner_dt, dimensionless.

Return type:

Tuple

critical_courant_number Function

deepwave.discretisation.critical_courant_number(accuracy: Literal[2, 4, 6, 8], *, propagator: str | Callable[[...], Any] | type | None = None) → float

Returns the maximum stable Courant number for a leapfrog scheme.

These are the homogeneous critical Courant numbers C_crit re-derived from the actual finite-difference stencil coefficients (see the appendix of Discretisation: the exact scheme). A propagation is stable when dt * v_max * sqrt(sum_i 1/dx_i^2) <= C_crit. Staggered propagators also check the density-weighted material operator; v_max must cover that conservative bound for heterogeneous models.

The critical number depends on whether the propagator uses a staggered grid, so name the propagator and let Deepwave answer that:

>>> import deepwave
>>> deepwave.discretisation.critical_courant_number(
...     4, propagator=deepwave.elastic
... )
0.857142857142857...

Parameters:

• accuracy – The spatial finite-difference accuracy order (2, 4, 6, or 8).

• propagator – The propagator the number is for, as the functional (deepwave.elastic), the Module class (deepwave.Elastic) or the name ("elastic").

Returns:

The critical (maximum stable) Courant number.

Raises:

• DeepwaveValueError – If propagator is not given, or if accuracy is not one of 2, 4, 6, 8.

• DeepwaveTypeError – If propagator does not name a propagator.

suggest_discretisation Function

deepwave.suggest_discretisation(min_vel: float, max_freq: float, ndim: int, accuracy: Literal[2, 4, 6, 8] | None = None, *, propagator: str | Callable[[...], Any] | type | None = None, courant_number: float | None = None, max_vel: float | None = None, dispersion_error: float = 0.01, min_points_per_wavelength: float = 6.0) → Discretisation

Suggests a grid spacing and time step for a wave you want to model.

Given the slowest wavespeed in your model and the highest frequency you need to propagate, this returns a cell spacing fine enough to carry that wave without excessive numerical dispersion, and the largest time step that is then stable on it. It is a starting point to sanity-check your own numbers against, not an oracle: it knows nothing about your geometry, your PML, or how far your waves have to travel.

Where the numbers come from, none of which is re-derived here:

• The shortest wavelength is \(\lambda_\mathrm{min} = v_\mathrm{min}/f_\mathrm{max}\) and the cells per wavelength are \(\lambda_\mathrm{min}/\Delta x\) (see Accuracy and grid design).

• The minimum points per wavelength for each order is the tabulated “Minimum points-per-wavelength (axial)” table in Discretisation: the exact scheme, which is the smallest sampling at which that stencil’s axial phase-velocity error falls to dispersion_error. Axial is the conservative direction: dispersion is not isotropic, and at 5 points per wavelength the regular 2-D order-2 axial phase velocity is \(0.936\,v\) while the diagonal is \(0.967\,v\), so a diagonally travelling wave is slightly better resolved than this suggestion promises.

• dt is the CFL limit from deepwave.discretisation.cfl_condition_n(), at DEFAULT_COURANT_NUMBER unless you pass courant_number. The number of dimensions enters only through that condition’s \(\sqrt{\sum_i 1/\Delta x_i^2}\) term.

Two criteria, and which one binds. Deepwave’s own runtime check warns when a propagator is given fewer than six cells per wavelength, and that single number is calibrated for the default accuracy=4. It is stricter than the 1% axial-dispersion criterion for orders 4, 6 and 8 (which ask for only 5.3, 3.9 and 3.4), because it is also standing in for everything the plane-wave dispersion analysis leaves out: how a point-source wavefront, a sharp model contrast and the PML are sampled. The suggestion therefore respects both, taking whichever asks for more cells, so that what it hands you never immediately provokes Deepwave’s own warning. note says which one bound. The consequence at the default 1% target is that orders 4, 6 and 8 all land on the same six-cell grid, and the coarser grids the higher orders could otherwise run on are not offered; lower min_points_per_wavelength to see them, but expect the warning when you then call a propagator.

Choosing the accuracy. With accuracy=None the order that minimises a rough estimate of the total work is chosen. Cells scale as \((1/\Delta x)^\mathrm{ndim}\) and time steps as \(1/\Delta t \propto 1/\Delta x\), so the work goes as points_per_wavelength ** (ndim + 1) times a relative cost per cell of (ndim * accuracy + 3) / (ndim * 4 + 3) – the multiply-adds a stencil of that order needs, normalised to 1 at accuracy=4. That weight is a rough relative figure and nothing more: it counts arithmetic and ignores cache behaviour, vector width and the wider halo, all of which make high orders dearer in practice than it says. Treat the choice as a suggestion to time, not a measurement. The chosen order and the runner-up, with the ratio between their estimates, are reported in note.

Parameters:

• min_vel – The slowest (non-zero) wavespeed in the model, in the same length unit as the grid spacing you want back, per second.

• max_freq – The highest frequency, in Hz, that must be propagated without excessive dispersion. This is the top of the wavelet’s significant band, not its dominant frequency.

• ndim – The number of spatial dimensions (1, 2 or 3).

• accuracy – The spatial finite-difference order (2, 4, 6 or 8), or None (the default) to choose the cheapest by the cost model above.

• propagator – The propagator the suggestion is for, as the functional (deepwave.elastic), the Module class (deepwave.Elastic) or the name ("elastic"). It selects the staggered or regular-grid column of the dispersion table and the matching critical Courant number. With None (the default) the conservative value of each is used: the regular-grid dispersion numbers, which ask for slightly more cells, and the staggered critical Courant number, which is slightly lower.

• courant_number – The Courant number to compute dt at. None (the default) uses DEFAULT_COURANT_NUMBER. A value above the critical Courant number for the order is rejected; when accuracy is None, orders whose critical number it exceeds are simply not considered.

• max_vel – The fastest wavespeed in the model, in the same length unit per second as min_vel, which is what the CFL condition is set by. Defaults to min_vel, i.e. to assuming a homogeneous model – pass it for any model with faster material in it, or the suggested dt will be too large for the grid and a propagator will have to sub-step it.

• dispersion_error – The axial phase-velocity error to design the grid for: 0.01 (1%, the default) or 0.001 (0.1%). The error is per wavelength travelled and accumulates, so a 1% error slips a full cycle over ~100 wavelengths of propagation; use 0.001 for long models.

• min_points_per_wavelength – The floor, in cells per wavelength, applied to the dispersion criterion, defaulting to the six cells per wavelength that Deepwave’s own runtime check wants (see above). Pass 0 to follow the dispersion table alone.

Returns:

A Discretisation naming the suggested grid_spacing, accuracy, dt, the points_per_wavelength and courant_number they correspond to, and a note.

Raises:

• DeepwaveValueError – If min_vel, max_freq, max_vel or courant_number is not a finite positive number, if min_points_per_wavelength is not a finite non-negative number, if max_vel is below min_vel, if ndim is not 1, 2 or 3, if accuracy is not 2, 4, 6, 8 or None, if dispersion_error is neither 0.01 nor 0.001, or if courant_number exceeds the critical Courant number of the requested order (or of every order, when one is being chosen).

• DeepwaveTypeError – If one of the numeric arguments is not a number, or propagator does not name a propagator.

Examples

>>> import deepwave
>>> suggestion = deepwave.suggest_discretisation(1500, 25, 2)
>>> suggestion.grid_spacing
(10.0, 10.0)
>>> suggestion.accuracy
4

Discretisation

class deepwave.Discretisation(grid_spacing: tuple[float, ...], accuracy: Literal[2, 4, 6, 8], dt: float, points_per_wavelength: float, courant_number: float, note: str)

What suggest_discretisation() returns.

Stability: provisional (see Stability policy).

grid_spacing

The suggested cell spacing, one entry per dimension (equal in every dimension), in the model’s length unit. Pass it straight to a propagator’s grid_spacing.

Type:

tuple[float, …]

accuracy

The spatial finite-difference order the suggestion is for (the one that was requested, or the one that was chosen).

Type:

Literal[2, 4, 6, 8]

dt

The largest time step, in seconds, that satisfies the CFL condition on this grid at courant_number.

Type:

float

points_per_wavelength

Cells per shortest wavelength that grid_spacing actually achieves.

Type:

float

courant_number

The Courant number dt was computed at.

Type:

float

note

A human-readable sentence saying which constraint was binding, what the runner-up accuracy was, and what was assumed.

Type:

str

Create new instance of Discretisation(grid_spacing, accuracy, dt, points_per_wavelength, courant_number, note)

DEFAULT_COURANT_NUMBER

deepwave.DEFAULT_COURANT_NUMBER

The Courant number Deepwave targets when courant_number is not given.

Every propagator, and deepwave.discretisation.cfl_condition_n(), uses this value to choose the internal time step when the caller leaves courant_number=None. It is deliberately conservative – the same number for every accuracy order and every dimensionality, well below the critical Courant number deepwave.discretisation.critical_courant_number() reports for any of them – because it is the value Deepwave has always used, and a run that does not set courant_number is reproduced bit-for-bit by passing this constant explicitly.

It is exported so that the default can be named rather than hard-coded: courant_number=1.2 * deepwave.DEFAULT_COURANT_NUMBER says what it means, and RunInfo.courant_number can be compared against it to tell a defaulted run from an explicit one. Raising courant_number above it permits a larger internal time step (fewer steps, faster) at the cost of stability margin; the value is validated against the critical Courant number and rejected if it would be unstable.

Stability: stable (see Stability policy).

check_finite Function

deepwave.check_finite(x: Tensor, name: str = 'tensor') → Tensor

Raises if x contains a NaN or Inf, otherwise returns x unchanged.

A cheap (one reduction over x) sanity check for a diverged simulation, for the tensors Deepwave does not check itself. Every propagator now checks its own outputs – the final wavefields and the receiver amplitudes – before returning them, and raises this same error type if any is non-finite; that is what check_output_finite=True (the default) does. What Deepwave still does not check is every step, which would cost a synchronisation per step and is prohibitively slow, nor any tensor that is not a propagator output: a gradient, a loss, or a model an optimiser has just updated. Call this on those (or, for early detection during a long run, from inside a forward_callback/ backward_callback on the wavefield – see the “Detecting divergence (NaN) early” recipe in Troubleshooting) to fail fast with a clear message instead of discovering the NaN much later.

Parameters:

• x – The tensor to check.

• name – A label for x, used in the error message (e.g. "receiver data", "v.grad").

Returns:

x, unchanged, so the call can be chained inline.

Raises:

deepwave.DeepwaveFloatingPointError – If x contains any non-finite (NaN or Inf) value. This is also a FloatingPointError (the type this function raised historically), so an existing except FloatingPointError still catches it, and it is a deepwave.DeepwaveError, so except deepwave.DeepwaveError catches it alongside Deepwave’s other errors.

Example

>>> import deepwave
>>> out = deepwave.scalar(v, grid_spacing, dt, ...)
>>> receiver_data = deepwave.check_finite(
...     out.receiver_amplitudes, "receiver data"
... )

Stability: provisional (see Stability policy).

estimate_memory Function

deepwave.estimate_memory(propagator: Callable[[...], Any] | str, model_shape: Sequence[int], nt: int, n_shots: int, dtype: dtype = torch.float32, *, accuracy: int = 4, pml_width: int | Sequence[int] = 20, model_gradient_sampling_interval: int = 1, storage: str | Storage | None = None, storage_mode: Literal['device', 'cpu', 'disk', 'none'] | None = None, storage_path: str | PathLike[str] | None = None, storage_compression: bool | None = None, n_models_requiring_grad: int | None = None, models_requiring_grad: Sequence[str] | None = None, q_n_mechanisms: int = 1, device: str | device = 'cpu', n_segments: int | None = None) → MemoryEstimate

Rough byte-size estimate for a propagator call, before running it.

This is meant to answer “will this run fit in memory” for a proposed model_shape/nt/n_shots/storage_mode before spending the time to build the tensors and start propagating. It mirrors, as closely as is practical without actually allocating anything, the two dominant sources of memory use identified for the backward pass: StorageManager (the sampled wavefield history kept for the model gradient) and allocate_model_gradients (the per-thread/per-shot scratch buffers used to accumulate that gradient without a race). It also gives a (much rougher) figure for the wavefields alive during the forward time loop itself.

Approximations (read before trusting the numbers):

• The number of same-shaped buffers alive at once during the forward time loop is equation- and dimensionality-specific (each block below is the total wavefield count returned by that propagator, both PML memory variables included); this is the single biggest approximation. Pass propagator to fill it in from the table below (using len(model_shape) as the number of dimensions); without it, a placeholder of 3 is used:

  • scalar / scalar_born background: 2 + 2 * ndim (2-D: 6, 3-D: 8); scalar_born (background + scattered): double that (2-D: 12, 3-D: 16).

  • acoustic: 1 + 3 * ndim (2-D: 7, 3-D: 10).

  • elastic: 13 (2-D), 27 (3-D).

• The number of stored wavefield histories is not the number of models you ask for a gradient for: the engine allocates one StorageManager slot per stored spatial derivative, which for acoustic and elastic is several slots per model. Name the models with models_requiring_grad and this is filled in exactly, from the propagator’s own slot function (the one its forward pass uses). Give only n_models_requiring_grad and a worst case over that many models is used instead; give neither a propagator nor names and the old one-slot-per-model assumption remains, which under-predicts acoustic’s density gradient by 1 + ndim times and elastic’s three-parameter gradient by 9/3 in 3-D.

• scalar with a backward_callback keeps a second full wavefield history (the raw forward field the callback is given); that extra copy is not modelled here.

• Like StorageManager itself, any dtype other than torch.float32 is costed as 8 bytes/element (i.e. as if it were torch.float64); this matches current Deepwave behaviour but is a simplification for exotic dtypes.

• model_shape should be your unpadded model’s spatial shape; this function adds the finite-difference (accuracy // 2 per side) and PML padding itself, the same way extract_survey does. The staggered-grid propagators (acoustic, elastic) really pad one cell less on the far side of each axis, so for them every per-model term here is slightly (about 0.7% per axis at 140 padded cells) generous.

• storage_compression=True is modelled from StorageManager’s own block-quantisation formula (8-per-side blocks, 1 byte/element plus one dtype_size-wide scale per block): it shrinks the sampled-history term, but StorageManager always keeps two extra full-size single-shot buffers alive alongside the compressed history (for building the next block before it is written), so compression’s net saving is smaller than the raw quantisation ratio, especially at small n_shots. Disk-mode host/device splits are not modelled; storage_mode="disk" is costed the same as "device", an upper bound on total (device + host/disk) footprint.

• This is a static estimate: it does not know about survey_pad shrinking the model per shot, per-shot-batched models, or checkpointing schemes external to Deepwave.

• The unmodelled fixed per-call overhead is the propagator’s other full-size buffers: the backward pass’s own set of padded wavefields (a second copy of the forward set counted above), the padding-removed wavefields returned to the caller (alive across .backward()), the padded/staggered internal model copies, and one gradient buffer per internal model requiring a gradient. It scales with the propagator’s wavefield and model counts, not with nt, so it matters most for 3-D elastic at small nt: measured at 96^3, 1 shot, nt=40, it is about 0.9 GiB (of which 0.27 GiB is the backward wavefield set, 0.24 GiB the returned wavefields, 0.19 GiB the internal model copies, 0.09 GiB the ping-pong buffers and 0.08 GiB the model gradients).

• Measured against real propagator runs (see Troubleshooting’s “Measured accuracy of estimate_memory” table), this estimate under-predicts real peak memory at every configuration tested, never over-predicts: by roughly 4-20% for the default storage_mode="device" uncompressed with the storage slots named exactly (models_requiring_grad), growing markedly worse (up to ~70%) as the modelled terms (sampled history, gradient scratch) shrink relative to that fixed per-call overhead – i.e. storage_compression=True, a large model_gradient_sampling_interval, and storage_mode="cpu"/ "disk" are all directions where the gap widens, as is a short nt on a many-wavefield propagator. Treat the returned figure as a lower bound and leave headroom, not a ceiling.

• n_segments= is part of that guarantee rather than an exception to it. Every term it adds is a floor on what it stands for: the boundary states are sized field by field, without the finite-difference halo and with the PML memory variables costed as the boundary slabs they are stored as, and they are counted once per segment only where a model gradient is being computed – a forward-only segmented run holds one state, whatever the segment count. Measured against real peaks on a 64-cubed scalar run at segment counts from 2 to 120, with and without a model gradient, the figure stayed below the measured peak in every case.

Parameters:

• propagator – The propagator this call would use: its name, its function, or its Module class. Any of the nine is accepted. It is required, because every number below depends on it: the forward wavefield count comes from its spec, and the backward-storage slot count from its own storage-slot rule. (It used to be optional, substituting a hard-coded guess of three forward wavefields – a figure no propagator’s spec agreed with.) Raises ValueError if it is not one of the recognised propagators.

• model_shape – The spatial shape of one (unpadded) model, e.g. v.shape.

• nt – The number of user-facing time steps (as passed to the propagator, i.e. matching source_amplitudes.shape[-1]).

• n_shots – The number of shots propagated together.

• dtype – The dtype of the model/computation.

• accuracy – The finite-difference accuracy order.

• pml_width – The PML width (int, or one value per face as accepted by the propagators).

• model_gradient_sampling_interval – As passed to the propagator.

• storage – The storage policy, as one object: a mode string or a deepwave.Storage. An alternative spelling for the three individual storage_* arguments, which are permanent and not deprecated; passing both spellings raises.

• storage_mode – One of "device", "cpu", "disk", or "none" ("none" estimates 0 for both backward-storage and gradient-scratch, matching what StorageManager actually allocates in that mode).

• storage_path – Accepted so that one dict of storage arguments can be passed to both this function and the propagator. It does not affect any byte count: disk mode is costed identically to device mode (see the approximations above).

• storage_compression – As passed to the propagator (device/cpu/disk modes only; ignored, matching StorageManager, when storage_mode="none"). Lossy – changes the model gradient. See deepwave.elastic().

• n_models_requiring_grad – How many model tensors require a gradient, when you do not want to say which (defaults to 1). With propagator, the storage-slot and gradient-buffer counts are the worst case over any choice of that many of that propagator’s models, i.e. an upper bound. Prefer models_requiring_grad, which is exact. Raises ValueError if negative, if it exceeds the propagator’s model count, or if it contradicts models_requiring_grad.

• models_requiring_grad – Which models require a gradient, by their user-facing argument names: "v" for scalar; "v" and "dv" for scalar_born; "v" and "rho" for acoustic; "lamb", "mu" and "buoyancy" for elastic. Any propagator’s names are deepwave.PropagatorSpec.model_names(). This is the accurate option on the four propagators that carry storage-slot tables (scalar, scalar_born, acoustic, elastic): the number of stored wavefield histories is then taken from the propagator’s own storage-slot function, so it matches what the engine really allocates (e.g. ["v", "rho"] in 3-D acoustic stores four histories, not two). On the other seven the same conservative one-slot-per-model fallback as n_models_requiring_grad is used. Raises ValueError for an unrecognised or repeated name, or a repeated name.

• storage_compression – If True, the stored wavefield history is compressed, which changes the arithmetic this function models. A deepwave.Storage passed here contributes its mode and its lossy_compression, which change the arithmetic. path has no bearing on a byte count.

• q_n_mechanisms – The number of GSLS relaxation mechanisms, for the viscoacoustic propagator, whose wavefield count depends on it. Ignored by the other families. Defaults to 1.

• device – The device the propagation would run on (only its type, cpu vs cuda, matters here).

• n_segments –

  If given, the number of time segments the run would be split into by deepwave.checkpointing.run_segments() (or by the propagator’s own checkpoint_segments=). The stored history is then divided by n_segments and an auxiliary-state term added for the states held alive between segments. Without it, the estimate has no model of checkpointing at all and over-predicts a checkpointed run’s peak substantially. Held to the same standard as the propagator’s own spelling of it: a bool or a non-int raises DeepwaveTypeError, and a value below 1 or above nt raises DeepwaveValueError (a run cannot have more segments than time steps, and the propagator refuses that too).

  How many auxiliary states this costs depends on whether a model gradient is being computed. With one, each segment’s state is held from the forward pass until that segment is recomputed, so the term scales with n_segments. Without one – no model named in models_requiring_grad, or n_models_requiring_grad=0 – there is nothing to checkpoint: the segmented driver threads a single state forward and drops the previous one, so the term is the size of one state however many segments there are.

Returns:

a dict (see that TypedDict) with the following keys:

• padded_model_shape: model_shape after adding finite-difference and PML padding.

• padded_model_elements: number of elements in one padded model.

• n_threads: the thread count used for the gradient-scratch estimate (torch.get_num_threads() on CPU, 1 on CUDA).

• forward_wavefields_bytes: estimated bytes for the forward time-loop wavefields.

• backward_storage_bytes: estimated bytes for the sampled wavefield history kept for the model gradient (StorageManager).

• gradient_scratch_bytes: estimated bytes for the per-thread/per-shot gradient accumulation buffers (allocate_model_gradients), one per internal model requiring a gradient (again more than one per user-facing model for acoustic/elastic; the gradient buffers themselves are part of the unmodelled overhead above).

• auxiliary_state_bytes: estimated bytes for the segment auxiliary states held alive between checkpointed segments. Zero unless n_segments was given and is greater than 1. Each state is sized field by field, with the finite-difference halo removed and the PML memory variables costed as the boundary slabs they are actually stored as, so it is a lower bound on the state rather than a full-grid over-estimate of it.

• total_bytes_lower_bound: the sum of the four components above.

Return type:

MemoryEstimate

Example

>>> import deepwave
>>> est = deepwave.estimate_memory(
...     model_shape=(2301, 751), nt=750, n_shots=115, accuracy=4,
...     pml_width=20, propagator="scalar",
...     models_requiring_grad=["v"],
... )
>>> est["total_bytes_lower_bound"] / 1e9
648.41...

(That figure is dominated by backward_storage_bytes, since it stores 750 full-model snapshots per shot for the gradient – exactly the cost that model_gradient_sampling_interval or a smaller storage_mode lets you trade against recompute or disk I/O; see Troubleshooting.)

Storage

Where and how the wavefield history that the model gradient is computed from is kept. Every propagator accepts the policy as one storage= object, or as the individual storage_mode, storage_path and storage_compression arguments, which remain supported and are not deprecated. Passing both spellings in one call raises, because there would be no honest way to merge them. storage_compression=True (Storage.lossy_compression) is lossy: it changes the model gradient, by about 0.2-4% in 2-D and 0.2-8% in 3-D in relative L2 error, with the largest errors in the model’s outermost cells (measured on scalar_born’s dv gradient, normalised by its own maximum: 0.03-0.07 in 1-D, about 1e-4 in 2-D and about 1e-5 in 3-D at the edge, against below 1e-3, 2e-5 and 1e-5 in the interior), because an edge cell’s gradient sums large, largely cancelling contributions from the absorbing layer. Without compression the stored history is lossless; exact model gradients also require internal-step sampling and fixed numerical parameters. The keyword’s name does not say any of this, which is why its documentation does:

out = deepwave.scalar(
    v, dx, dt, source_amplitudes=s, source_locations=sl,
    storage=deepwave.Storage(mode="disk", path="/scratch/history"),
)

class deepwave.Storage(mode: _StorageModeStr = 'device', path: StoragePath | None = None, lossy_compression: bool = False)

Where and how the wavefield history for the model gradient is kept.

One object carrying the whole storage policy, accepted as storage= by every propagator (functional and Module, on both the constructor and forward), by deepwave.checkpointing.run_segments(), by deepwave.estimate_memory(). It is an alternative spelling, never a requirement: the individual storage_mode/storage_path/storage_compression arguments are permanent and are not deprecated. Passing both spellings in one call raises, because there would be no honest way to merge them.

It is plain data – it carries no device, no time-step count and no shapes, so one object can be reused across a sweep. Derive a variant with dataclasses.replace(storage, mode="disk").

Stability: provisional (see Stability policy) – fields may be added as more of a storage policy becomes worth recording. The fields it has are not renamed or withdrawn without a release’s notice under deepwave.DeepwaveFutureWarning, and the individual storage_mode / storage_path / storage_compression arguments it is an alternative spelling for are stable, as is passing a plain mode string to storage=.

Parameters:

• mode – One of "device" (keep the history on the compute device, the default and the fastest), "cpu" (offload it to host memory), "disk" (write it to files under path), or "none" (keep nothing, which makes the call forward-only with respect to the model gradient). On a CPU device "cpu" and "device" are the same thing and the engine resolves the former to the latter; deepwave.RunInfo.storage records which was used.

• path – The directory mode="disk" writes into. None (the default) means the system temporary directory, resolved when you call, so TMPDIR set after import deepwave is honoured. Accepts str or any os.PathLike (a pathlib.Path). Ignored, with a warning, by every other mode. Each run makes its own deepwave_tmp_<pid>_<random hex> subdirectory here when it allocates its wavefield history, and removes the whole subdirectory once nothing refers to that history any more (after backward(), or when the output and its graph go out of scope). Only a process killed outright leaves one behind, since no destructor runs then.

• lossy_compression – If True, store the history as 8-bit block max-abs quantisation. This is lossy and it changes the model gradient – see deepwave.scalar()’s storage_compression entry for the measured error and what is unaffected. Defaults to False.

plan Function

deepwave.plan() answers “what will this run do, and what will it cost” without running it. It is not a model of a propagator call: it is the call, stopped before propagation and gradient-history allocation, so plan(...).info is the very deepwave.RunInfo the run would return and cannot drift from it:

p = deepwave.plan(
    deepwave.scalar, v, dx, dt,
    source_amplitudes=s, source_locations=sl, receiver_locations=rl,
)
print(p.summary())

Plan.warnings contains only warnings emitted by Deepwave itself. It does not capture or suppress warnings from PyTorch or other dependencies.

deepwave.plan(propagator: str | Callable[[...], Any], *model_args: Any, **call_kwargs: Any) → Plan

Resolves a propagator call without running it, and reports what it costs.

Call it exactly as you would call the propagator:

p = deepwave.plan(
    deepwave.scalar, v, dx, dt,
    source_amplitudes=s, source_locations=sl, receiver_locations=rl,
)
print(p.summary())

The propagator’s name works too (deepwave.plan("scalar", ...)), so a plan can be driven from a configuration file.

This is not a model of a propagator call: it is the call, stopped at the moment before propagation and gradient-history allocation. The resolution – the internal time step, the step ratio, the inferred PML frequency, the survey extraction, the backend choice, the storage policy – is done by the propagator’s own code, so Plan.info is the very deepwave.RunInfo the run would return, and cannot drift from it.

What it raises. Every error your call would raise about its arguments: a bad model, mismatched source and location shapes, receivers outside the model, a courant_number above critical, a model_gradient_sampling_interval that does not divide the time steps, an unknown storage_mode, storage_mode="none" with a model requiring a gradient, and an unsupported option for the requested backend.

What it cannot raise, because these depend on the run rather than on its arguments: an out-of-memory condition; a compiled-backend failure (including a disk write failure, and the missing-symbol error for an (ndim, accuracy, dtype, device) combination that was not built, since which kernel variant is needed is settled after the point this stops at); and the refusals that a torch.func transform or a double backward triggers on the result.

What it costs. It does the propagator’s setup, so it pads the models and resamples the source amplitudes just as the run would. Front-end preprocessing can allocate tensors proportional to the source time length; planning never allocates the propagation or gradient history.

Warnings. Plan.warnings collects Deepwave’s own warnings on this execution context. It does not alter Python’s process warning filters, so unrelated threads emit warnings normally. Third-party warnings raised during setup follow the caller’s ordinary warning policy and are not added to the plan. A Deepwave warning promoted to an error by the caller’s filter is still collected as text during planning; the real call raises as usual.

Parameters:

• propagator – The propagator to plan: the function itself (deepwave.scalar), its name as a string ("scalar"), or a constructed Module instance (Scalar(v, dx)), which is planned exactly as it would run. The Module class is refused, because it is not something that can be called and the plan for it would differ from the plan for any instance built from it – see _resolve_callable.

• *model_args – The positional arguments of that propagator, unchanged.

• **call_kwargs – Its keyword arguments, unchanged. checkpoint_segments= is understood here as well: the arguments a segmented call refuses are refused here too, and then the resolution reported describes the whole run while the memory estimate accounts for the segmentation.

Returns:

the resolved call and its estimated cost.

Return type:

Plan

Raises:

DeepwaveRuntimeError – If the propagator returned without reaching the point a plan stops at, which means it is not instrumented.

class deepwave.Plan(info: RunInfo, memory: MemoryEstimate, resampling: bool, warnings: tuple[str, ...] = (), checkpoint_segments: int | None = None)

What a propagator call would resolve to, and what it would cost.

Returned by deepwave.plan(). Every field is produced by the propagator’s own resolution code, so it is what the call would really do rather than a second opinion about it.

Stability: provisional (see Stability policy) – fields may be added, and summary() may be reworded, as more of a run’s resolved configuration becomes worth reporting. The fields it has are not renamed or withdrawn without a release’s notice under deepwave.DeepwaveFutureWarning.

info

Exactly the deepwave.RunInfo the call would attach to its output, built by the same code. Read the resolved internal time step, step ratio, PML width and frequency, survey origin and extent, backend and storage policy off it.

Type:

deepwave.RunInfo

memory

The call’s estimate of its own footprint, in the same shape deepwave.estimate_memory() returns. It is more accurate than the parametric form: the padded shape, the sampling interval, the storage-slot count and the wavefield count are read off the resolved run rather than derived from arguments. It remains a lower bound, including under checkpoint_segments= – see deepwave.estimate_memory() for what it does not model.

Type:

deepwave.MemoryEstimate

resampling

Whether the call would sub-step, i.e. whether info.step_ratio is greater than 1. The receiver amplitudes are resampled back to the requested rate, and info.internal_dt gives the rate they were computed at.

Type:

bool

warnings

Deepwave warnings the call would emit, as text. Collected on this execution context without changing process warning filters or spending the once-per-process warning budget.

Type:

tuple[str, …]

checkpoint_segments

The number of checkpointed time segments this plan costed, or None for an unsegmented call. The resolution in info always describes the whole run; only memory reflects the segmentation.

Type:

int | None

summary() → str

A short human-readable block describing the planned call.

Returns:

A multi-line block covering the resolved grid and padded shape, the timing (the user’s nt and dt, and the internal step count and internal_dt they resolve to), the backend, the storage policy, the memory breakdown, and any warnings the call would emit.

deepwave.diagnostics

deepwave.diagnostics.report() → Report

Collects the environment facts a Deepwave bug report should carry.

Print the result:

>>> import deepwave
>>> print(deepwave.diagnostics.report())
Deepwave 0.1.0
PyTorch  2.5.1
...

Returns:

A Report. It is a dataclass with named fields rather than a plain dict so that the shape of the answer is documented and so that nothing is tempted to parse the printed text.

Stability: provisional (see Stability policy).

class deepwave.diagnostics.Report(deepwave_version: str, torch_version: str, python_version: str, machine: str, required_baseline: str | None, simd_features: dict[str, bool | None], cuda_available: bool, cuda_devices: tuple[str, ...])

A snapshot of the environment Deepwave is running in.

Returned by report(). Print it (or str() it) to get the text to paste into a bug report; read the individual fields if you need one fact programmatically.

deepwave_version

Deepwave’s version string.

Type:

str

torch_version

PyTorch’s version string.

Type:

str

python_version

The running Python’s version string.

Type:

str

machine

The machine architecture, from platform.machine().

Type:

str

required_baseline

The instruction-set baseline the compiled backend was built to require. An empty string reports no fixed ISA requirement; None means the backend did not report it.

Type:

str | None

simd_features

Per-feature CPU SIMD support. True/False where the feature is meaningful and could be probed, None where it is not applicable (any non-x86 machine).

Type:

dict[str, bool | None]

cuda_available

Whether torch.cuda.is_available().

Type:

bool

cuda_devices

The name of each visible CUDA device, in device order. Empty when CUDA is unavailable.

Type:

tuple[str, …]

Stability: provisional (see Stability policy).

IGNORE_LOCATION

deepwave.IGNORE_LOCATION

Sentinel value to place in a source or receiver location tensor to mark that point as unused, allowing a single rectangular [shot, point, ndim] locations tensor to describe shots that have differing numbers of active sources or receivers.

A point whose coordinates are all IGNORE_LOCATION is ignored: a source injects nothing, a receiver records zeros. A point with only some coordinates set to the sentinel raises, because that is far more likely to be a mistake than an intention.

Its numeric value is unspecified and may change. Always compare against, and assign from, deepwave.IGNORE_LOCATION itself; never hard-code the number it happens to hold.

The *_mask arguments are the modern spelling of the same mechanism – a [shot, point] boolean tensor, True = active – and produce bit-identical results. Supplying both for one component raises.

Type aliases

deepwave.LocationsLike

alias of Tensor | Hicks

deepwave.SourceSpec

alias of tuple[torch.Tensor, LocationsLike] | tuple[torch.Tensor, LocationsLike, Optional[torch.Tensor]]

deepwave.ReceiverSpec

alias of LocationsLike | tuple[LocationsLike, Optional[torch.Tensor]]

Introspection

The machine-readable description of what a propagator offers, for code that must drive any of the nine without hard-coding one of them. Each function accepts the functional (deepwave.elastic), the Module class (deepwave.Elastic) or the name ("elastic").

deepwave.spec(propagator: Callable[[...], Any] | type | str) → PropagatorSpec

What a propagator offers, as data.

The machine-readable description a generic wrapper needs in order to drive any of the nine propagators without hard-coding one of them: the channels it has, the options it implements, and the names of everything it returns. See deepwave.PropagatorSpec.

Parameters:

propagator – The functional (deepwave.elastic), the Module class (deepwave.Elastic) or the name ("elastic").

Returns:

The propagator’s frozen specification. Cached, so the same propagator always yields the same object.

Raises:

• deepwave.DeepwaveTypeError – If the argument does not name a propagator.

• deepwave.DeepwaveValueError – If a string is not a propagator’s name.

Examples

>>> import deepwave
>>> deepwave.spec(deepwave.elastic).source_channels(2)
('pressure_rate', 'y', 'x')
>>> deepwave.spec("elastic").source_amplitudes_kwarg("pressure_rate")
'source_amplitudes_pressure_rate'

class deepwave.PropagatorSpec(name: str, status: str, ndims: tuple[int, ...], supported_options: frozenset[str], synthetic_source_channels: frozenset[str], synthetic_receiver_channels: frozenset[str], has_bg_receivers: bool, has_forward_wavefield: bool)

What one propagator offers, as data.

Returned by deepwave.spec(); not constructed directly. Every field is plain data and every method is a pure function of the number of spatial dimensions, so instances compare equal, repr usefully, and pickle.

Stability: provisional (see Stability policy). Fields and methods may be added as more of a propagator’s interface becomes worth publishing, and what a field reports changes when the thing it describes changes. The fields and methods it has are not renamed or withdrawn without a release’s notice under deepwave.DeepwaveFutureWarning.

name

The propagator’s name, e.g. "elastic".

Type:

str

status

Its stability tier, one of "stable", "provisional", "experimental", "deprecated".

Type:

str

ndims

The numbers of spatial dimensions it supports.

Type:

tuple[int, …]

supported_options

The shared options it implements ("mpml_fraction", "surface_method", "q_freq", "q_n_mechanisms", "q_band"). Every propagator accepts every shared option at its default value; a non-default value of one that is not listed here raises.

Type:

frozenset[str]

synthetic_source_channels

Source channels the propagator offers that have no backend slot of their own because they are expanded into other channels. {"pressure_rate"} on acoustic_vti, empty everywhere else.

Type:

frozenset[str]

synthetic_receiver_channels

Receiver channels the propagator offers that have no backend slot of their own because it computes them from other channels’ recordings. {"p"} on acoustic_vti, whose pressure is -(sigma_v + (ndim - 1) sigma_h) / ndim, empty everywhere else. Such a channel still occupies its own output tuple slot.

Type:

frozenset[str]

has_bg_receivers

Whether the propagator takes a bg_receivers= mapping (equivalently, whether it is a Born propagator).

Type:

bool

has_forward_wavefield

Whether the propagator implements backward_callback_needs_forward_wavefield=True, and so can serve deepwave.CallbackState.get_wavefield()’s role="forward" inside a backward_callback. True on scalar alone today. All nine accept the argument at its default of False; passing True where this is False raises deepwave.DeepwaveNotImplementedError.

Type:

bool

bg_receiver_channels(ndim: int) → tuple[str, ...]

The background-receiver channels; empty on a non-Born propagator.

bg_receiver_locations_kwarg(channel: str) → str

The flat keyword carrying this channel’s background locations.

bg_receiver_mask_kwarg(channel: str) → str

The flat keyword carrying this channel’s background mask.

bg_receiver_output_name(channel: str) → str

The output field holding this channel’s background data.

channel_aliases() → Mapping[str, str]

The permanent alternative spellings of this propagator’s channels.

Maps each accepted alias to the channel it names, on both the source and the receiver side. Empty on all nine propagators: every channel has exactly one spelling. Read-only; any entry would be permanent. A spelling that is going away is in deprecated_source_channels() instead.

deprecated_source_channels() → Mapping[str, str]

The deprecated source channel spellings, and what replaces each.

{"p": "volume_rate"} on the acoustic family and {"p": "pressure_rate"} on the elastic family and acoustic_vti; empty on the two scalar propagators. The old spelling keeps working, warning deepwave.DeepwaveFutureWarning, until it is removed. It is a source-side deprecation only: "p" is the canonical receiver channel name on every propagator that records a pressure.

model_names(ndim: int) → tuple[str, ...]

The propagator’s model tensors, in the order it takes them.

The user-facing models – the arguments the propagator’s public function takes, in signature order – not the internal prepared models the kernel consumes: acoustic reports ("v", "rho"), the two tensors a caller passes, rather than the bulk modulus and per-component buoyancies built from them. A Born propagator’s perturbation models are included (("v", "rho", "dv", "drho") for acoustic_born), as is any optional model tensor the propagator accepts at this ndim (elastic_vti’s c66, which is a model in 3-D only).

These are the names deepwave.estimate_memory()’s models_requiring_grad takes, and the names a generic wrapper needs in order to assemble a call.

Parameters:

ndim – The number of spatial dimensions.

Returns:

The model argument names.

Raises:

deepwave.DeepwaveValueError – On an unsupported ndim.

output_names(ndim: int, **options: Any) → tuple[str, ...]

Every field of the returned tuple, in order.

The state fields of wavefield_names() followed by the receiver fields of receiver_names(); the same tuple as deepwave.PropagatorOutput.names.

receiver_channels(ndim: int) → tuple[str, ...]

The receiver channels, in the order of the returned tuple.

receiver_locations_kwarg(channel: str) → str

The flat keyword carrying this channel’s receiver locations.

On a Born propagator this is the scattered receivers’ keyword, spelled sc_receiver_locations*; the background ones are bg_receiver_locations_kwarg().

receiver_mask_kwarg(channel: str) → str

The flat keyword carrying this channel’s receiver mask.

receiver_names(ndim: int) → tuple[str, ...]

The receiver-amplitude output fields, in return order.

On a Born propagator the background block comes first, then the scattered block, matching the returned tuple.

receiver_output_name(channel: str) → str

The output field holding this channel’s recorded data.

The inverse of receiver_channels(): the name under which channel’s data appears in the returned tuple and in deepwave.PropagatorOutput.receivers.

source_amplitudes_kwarg(channel: str) → str

The flat keyword carrying this channel’s source amplitudes.

source_channels(ndim: int) → tuple[str, ...]

The source channels, in canonical order.

source_locations_kwarg(channel: str) → str

The flat keyword carrying this channel’s source locations.

source_mask_kwarg(channel: str) → str

The flat keyword carrying this channel’s source mask.

wavefield_names(ndim: int, **options: Any) → tuple[str, ...]

The state fields the propagator returns, in return order.

Exactly the keys of deepwave.PropagatorOutput.final_state: the public wavefields plus "auxiliary_state", and no receiver fields. These are the names initial_state= consumes, so they are what a wrapper needs in order to continue a simulation.

Parameters:

• ndim – The number of spatial dimensions.

• **options – Options that change the list. q_n_mechanisms on viscoacoustic is the only one today; any other key raises rather than being ignored.

Returns:

The state field names.

Raises:

deepwave.DeepwaveValueError – On an unsupported ndim or an option key this propagator’s names do not depend on.

deepwave.wavefield_names(propagator: Callable[[...], Any] | type | str, ndim: int, **options: Any) → tuple[str, ...]

The state fields a propagator returns, in return order.

Exactly the keys of deepwave.PropagatorOutput.final_state – the public wavefields plus "auxiliary_state", and no receiver fields – so this is the list an initial_state= consumer needs.

Parameters:

• propagator – The functional, the Module class or the name.

• ndim – The number of spatial dimensions.

• **options – Options that change the list; q_n_mechanisms on the viscoacoustic propagator is the only one today, and any other key raises rather than being silently ignored.

Returns:

The state field names.

Raises:

deepwave.DeepwaveValueError – On an unsupported ndim, or an option key this propagator’s names do not depend on.

Examples

>>> import deepwave
>>> deepwave.wavefield_names(deepwave.scalar, 2)
('wavefield_0', 'wavefield_m1', 'auxiliary_state')

deepwave.receiver_names(propagator: Callable[[...], Any] | type | str, ndim: int) → tuple[str, ...]

The receiver-amplitude fields a propagator returns, in return order.

On a Born propagator the background block comes first and the scattered block second, matching the returned tuple.

Parameters:

• propagator – The functional, the Module class or the name.

• ndim – The number of spatial dimensions.

Returns:

The receiver output field names.

Raises:

deepwave.DeepwaveValueError – If ndim is one the propagator does not support.

Examples

>>> import deepwave
>>> deepwave.receiver_names(deepwave.acoustic, 2)
('receiver_amplitudes_p', 'receiver_amplitudes_y', 'receiver_amplitudes_x')

vpvsrho_to_lambmubuoyancy Function

deepwave.vpvsrho_to_lambmubuoyancy(vp: Tensor, vs: Tensor, rho: Tensor, eps: float = 1e-15) → tuple[Tensor, Tensor, Tensor]

Converts vp, vs, rho to lambda, mu, buoyancy.

All input Tensors must have the same shape.

Parameters:

• vp – A torch.Tensor containing the p wavespeed.

• vs – A torch.Tensor containing the s wavespeed.

• rho – A torch.Tensor containing the density.

• eps – An optional float to avoid division by zero. Default 1e-15.

Returns:

• lambda: A torch.Tensor containing the first Lamé parameter.

• mu: A torch.Tensor containing the second Lamé parameter.

• buoyancy: The buoyancy (1/density) model. A cell with buoyancy = 0 is a vacuum, not an infinitely dense cell; see the vacuum sentinel in Boundary conditions, which also explains how the sentinel is averaged onto the staggered nodes. rho = 0 is mapped to buoyancy = 0 rather than to infinity, so a vacuum built this way carries the sentinel.

Return type:

Tuple

lambmubuoyancy_to_vpvsrho Function

deepwave.lambmubuoyancy_to_vpvsrho(lamb: Tensor, mu: Tensor, buoyancy: Tensor, eps: float = 1e-15) → tuple[Tensor, Tensor, Tensor]

Converts lambda, mu, buoyancy to vp, vs, rho.

All input Tensors must have the same shape.

Parameters:

• lamb – A torch.Tensor containing the first Lamé parameter.

• mu – A torch.Tensor containing the second Lamé parameter.

• buoyancy – The buoyancy (1/density) model. A cell with buoyancy = 0 is a vacuum, not an infinitely dense cell; see the vacuum sentinel in Boundary conditions, which also explains how the sentinel is averaged onto the staggered nodes. Such a cell is mapped back to rho = 0 rather than to an infinite density.

• eps – An optional float to avoid division by zero. Default 1e-15.

Returns:

• vp: A torch.Tensor containing the p wavespeed.

• vs: A torch.Tensor containing the s wavespeed.

• rho: A torch.Tensor containing the density.

Return type:

Tuple

to_cells Function

deepwave.to_cells(coordinates: Tensor, grid_spacing: float | Sequence[float], grid_origin: float | Sequence[float] | None = None, *, snap: bool | float = False) → Tensor

Converts physical-unit coordinates to float grid-cell coordinates.

Computes (coordinates - grid_origin) / grid_spacing elementwise, with grid_spacing/grid_origin broadcasting over the last (ndim) dimension of coordinates. The result is suitable as float source_locations/receiver_locations for the propagators, which auto-detect off-grid-node coordinates and Hicks-interpolate them (see deepwave.Hicks); on-node results are used directly by the exact (integer) path.

With the default snap=False this function does not round: an off-node result is returned as-is, and even a physically on-node coordinate may not come back as an exact integer due to floating-point error. Snapping to the nearest integer grid cell is the propagators’ documented behaviour, not this function’s; call .round().long() yourself first if you specifically want integer cell indices, or pass snap=True to clean up the floating-point error only (below).

Parameters:

• coordinates – A [..., ndim] float Tensor of physical-unit coordinates.

• grid_spacing – The grid cell spacing: a single float (applied to every dimension) or a sequence of ndim floats.

• grid_origin – The physical-unit coordinate of grid cell 0 in each dimension: a single float, a sequence of ndim floats, or None (default), meaning zero in every dimension. Named for the grid, not for the propagators’ origin, which is the cell index at which a model region begins.

• snap –

  Whether to remove floating-point error from results that should have landed exactly on a node. False (the default) returns the division’s result untouched, bit-for-bit. True replaces any result within deepwave.ON_NODE_TOLERANCE cells of an integer by that integer. A float instead of True uses that tolerance, in cells, in place of ON_NODE_TOLERANCE.

  This is a cleanup of floating-point error, not a rounding of genuinely off-node coordinates: a coordinate that is really 0.4 cells off a node is returned unchanged, whatever snap is, and will still be Hicks-interpolated. Its purpose is the case where (coordinates - grid_origin) / grid_spacing is arithmetically an integer but comes back as 6.999999999999999 – which, without snapping, routes an exactly-on-node survey through the interpolated path and changes the numbers. Keyword-only, because the other three parameters are positional in released code.

Returns:

A [..., ndim] float Tensor of cell coordinates, the same shape, dtype and device as the unsnapped result. The dtype is coordinates’ own if that is a floating type, and the default float dtype otherwise (an integer coordinates tensor is accepted: the conversion is done in floating point, so a fractional grid_spacing/origin is not truncated).

Raises:

DeepwaveValueError – If snap is a negative number.

Hicks Class

class deepwave.Hicks(locations: Tensor, halfwidth: int = 4, *, free_surfaces: Sequence[bool] | None = None, free_surface_locs: Sequence[float] | None = None, model_shape: Sequence[int] | None = None, monopole: Tensor | bool = True, dipole_dim: Tensor | int = 0, dtype: dtype | None = None, eps: float = 1e-05, offset: Sequence[float] | None = None, rigid_surfaces: Sequence[bool] | None = None, differentiable: bool = False)

Location interpolation onto grid using method of Hicks.

Hicks (2002, https://doi.org/10.1190/1.1451454) proposed using a Kaiser windowed sinc function to interpolate source and receiver locations onto a grid when they are not centred on a grid point. This class implements this approach. It can also be used to create dipole sources and receivers.

The interpolator supports 1, 2, and 3 spatial dimensions (matching the Deepwave propagator being used) and is staggered-grid aware via the offset argument, so a field component that does not live on the regular-grid (pressure) node can still be interpolated about its true location.

Parameters:

• locations – A three dimensional torch.Tensor [shot, per_shot, ndim] specifying the locations to interpolate, provided as float values in units of cells relative to the origin of the grid, so [[[0.5, 0.5]]] corresponds to the point half a cell from the origin in both dimensions. ndim may be 1, 2, or 3.

• halfwidth – An integer specifying the halfwidth of the window that will be used to interpolate each point onto the grid. A halfwidth of 4 (the default) means that an 8x8 window will be used. Possible values are in [1, 10].

• free_surfaces –

  A sequence of 2 * ndim booleans specifying whether the corresponding edge of the grid is a free surface, in the order: [beginning of first dimension, end of first dimension, beginning of second dimension, end of second dimension, …]. For example, in 2D [True, False, False, True] means that the beginning of the first dimension and end of the second dimension are free surfaces, while the other two edges are not. Defaults to no edges being free surfaces.

  A free surface is a point at which the field is zero (an odd/antisymmetric mirror, also called a Dirichlet or pressure-release boundary). Part of the interpolation window that extends past the surface is folded back into the model with a sign flip (the method of images).

  Note

  Matching a Deepwave boundary. Which condition a face with pml_width = 0 produces, and where its mirror plane sits, depends on the propagator and the field component; the measured table is at Reflecting edges, and is the single source for it. What Hicks needs from that table is only the surface location, which is what free_surface_locs defaults to: one cell outside the model edge for a free surface, half a cell outside for a rigid one. Choose the surface type and location to match the propagator and field component you are placing.

  scalar and scalar_born need no Hicks of their own for this: their source and receiver points are on the regular grid and their boundary="free"/"rigid" faces mirror about exactly these default planes, so a plain float location next to such a face is folded automatically, giving bit for bit what building this class with those surfaces gives. The other seven propagators carry their components on staggered nodes, where the image is component dependent, and refuse a float location whose window would cross such a face.

• free_surface_locs – An optional sequence of 2 * ndim floats giving the location (in cells) of each face named in free_surfaces/rigid_surfaces, in the same order. Defaults to the location Deepwave’s own boundaries use: one cell outside the model edge (-1/ model_shape) for a free surface, half a cell outside (-0.5/model_shape - 0.5) for a rigid surface.

• model_shape – A sequence of ndim integers specifying the size of the grid. This is only used when the model contains free or rigid surfaces.

• monopole – A boolean or torch.Tensor of booleans (of shape [shot, per_shot]) specifying whether the source/receiver is a monopole. If False, the point will be a dipole. Default True.

• dipole_dim – An integer or torch.Tensor of integers (of shape [shot, per_shot]) in [0, ndim) specifying the dimension in which the dipole is oriented. This is only used for points that are dipoles (not monopoles). Default 0.

• dtype – The datatype to use for the interpolation weights. None (the default) means torch.get_default_dtype(), matching the other dtype parameters in the package. Under the usual float32 default this is the dtype the weights were always built in, so a propagator’s results are bit-identical. Under a float64 default the weights are built in float64 and then promoted (or demoted) to the propagator’s own dtype when they are used, which is a different rounding path from building them in float32: measured differences of order 1e-6 relative in the weights themselves and in receiver amplitudes from a float32 model, and of order 5e-8 relative from a float64 model. Pass dtype=torch.float to build the weights in float32 regardless of the default.

• eps – A small value to prevent division by zero. Points closer to a grid cell centre than this will be rounded to the grid cell. Default 1e-5.

• offset –

  An optional sequence of ndim floats giving the half-cell offset of the field component being interpolated relative to the regular-grid (pressure) node, used to make interpolation staggered-grid aware. Defaults to zero offset in every dimension (regular grid).

  Note

  Staggered surfaces. The default surface locations and the offset shift applied to them are validated for the regular grid (zero offset), which covers the scalar free surface and the acoustic pressure (rigid) surface. For a staggered component (non-zero offset) the image symmetry is component dependent (Robertsson, 1996): different components reflect with different signs and about different planes. Out-of-model cells are always dropped so the result is safe to use, but to be physically exact for a staggered component you should pass the correct free_surfaces/rigid_surfaces and explicit free_surface_locs for that component.

• rigid_surfaces –

  A sequence of 2 * ndim booleans specifying whether the corresponding edge of the grid is a rigid surface, in the same order as free_surfaces. Defaults to no edges being rigid surfaces. A face must not be both a free and a rigid surface.

  A rigid surface is a point at which the field has zero normal gradient (an even/symmetric mirror, also called a Neumann or zero-normal-velocity boundary). The part of the interpolation window past the surface is folded back into the model with the same sign (no flip); a point exactly on the surface has its weight doubled.

• differentiable –

  If True, the interpolation weights are built as differentiable torch expressions of locations (with no float-keyed weight cache), so gradients flow from source()/receiver() back to the source/receiver positions. This enables gradient-based survey design and source localisation. Default False (the fast, cached, bit-for-bit-stable path).

  Stability: provisional (see Stability policy).

  Weights are differentiable within a cell. At integer crossings the support changes: monopole values remain continuous, but their position gradients can jump; dipole values can jump too. The support indices are detached from autograd.

  float64 is recommended for gradient work. IGNORE_LOCATION sentinel points are handled as usual (their gradient is zero). Free and rigid surfaces are not supported in this mode (a NotImplementedError is raised); dipole sources/receivers are supported. See the “Differentiable positions” section of the usage documentation for the recommended recipe.

  Note

  A differentiable=True instance holds its interpolation weights in the autograd graph (they are torch expressions of locations). Calling backward() frees that graph, so a second backward through the same instance raises the usual “backward through the graph a second time” error. In an optimisation loop either construct a fresh Hicks instance each iteration (the normal pattern: build it from the current positions inside the loop), or pass retain_graph=True to backward() if you deliberately reuse one instance across steps.

Initialises the Hicks interpolation class.

See the class docstring for a description of the arguments.

get_locations(shot_idxs: Tensor | None = None) → Tensor

Get the interpolated locations.

The interpolated locations can be provided to a Deepwave propagator as the source or receiver locations.

Parameters:

shot_idxs – A 1D torch.Tensor containing the indices of the shots that amplitudes are being provided for, within the locations provided at initialisation. This is useful if shots are being used in batches, so only some of the shots are needed now. Default (None) means all shots.

Returns:

A torch.Tensor containing the interpolated locations.

receiver(amplitudes: Tensor, shot_idxs: Tensor | None = None) → Tensor

Convert receiver amplitudes from interpolated to original locations.

Parameters:

• amplitudes – A torch.Tensor containing the amplitudes recorded at interpolated receiver locations.

• shot_idxs – A 1D torch.Tensor containing the indices of the shots that amplitudes are being provided for, within the locations provided at initialisation. This is useful if shots are being used in batches, so only some of the shots are needed now. Default (None) means all shots.

Returns:

The amplitudes of receivers at the original locations.

source(amplitudes: Tensor, shot_idxs: Tensor | None = None) → Tensor

Calculate the amplitudes of the interpolated sources.

Parameters:

• amplitudes – A torch.Tensor containing the amplitudes of the sources at the original locations.

• shot_idxs – A 1D torch.Tensor containing the indices of the shots that amplitudes are being provided for, within the locations provided at initialisation. This is useful if shots are being used in batches, so only some of the shots are needed now. Default (None) means all shots.

Returns:

The amplitudes of the interpolated sources, ready to be provided to a Deepwave propagator.

ricker Function

deepwave.wavelets.ricker(freq: float, nt: int = <required>, dt: float = <required>, peak_time: float | None = None, dtype: dtype | None = None, *, length: int | None = None, shape: int | Sequence[int] | None = None) → Tensor

Return a Ricker wavelet with the specified central frequency.

Parameters:

• freq – The central frequency (Hz, if dt is in seconds). Must be positive.

• nt – The number of time samples. Required: omitting it (and its deprecated length alias), or passing None, raises TypeError.

• dt – The time sample spacing (seconds). Required: omitting it, or passing None, raises TypeError.

• peak_time – The time (in secs) of the peak amplitude. Optional; if not provided the wavelet is auto-centred in the record (peak_time = nt * dt / 2), the same record-centring convention as tone_burst(), gaussian_pulse() and ormsby() (which also default to peak_time = nt * dt / 2), so the (symmetric) wavelet is not truncated. Passing an explicit value is unchanged.

• dtype – The PyTorch datatype to use. Optional, defaults to PyTorch’s default (float32).

• length – Deprecated alias for nt, retained for backward compatibility. Removed in Deepwave 0.4.0. Providing both raises ValueError. Keyword-only, so a positional call such as ricker(freq, 100, dt) is unaffected by this deprecation.

• shape – Keyword-only. The leading dimensions of the returned tensor, with nt always last. None (the default) returns a one-dimensional [nt] tensor. (n_shots, n_sources) returns [n_shots, n_sources, nt], ready to pass straight to a propagator’s source_amplitudes. An int n is shorthand for (n,) and returns [n, nt]; an empty sequence () returns [nt], as None does. Every entry must be a positive integer. The leading dimensions are added with torch.Tensor.expand(), so no data is copied and the result is a view: every shot and source shares one copy of the wavelet, and each added dimension has stride zero. Call .contiguous() or .clone() on it before modifying it in place. Deepwave’s propagators do not modify their inputs, so passing the view to one directly is safe.

Returns:

A PyTorch tensor representing the Ricker wavelet, of shape [nt], or [*shape, nt] if shape was given.

Raises:

• TypeError – If nt (or its deprecated length alias) or dt is not provided, or is provided as None, or if shape is not None, an int, or a sequence of ints.

• ValueError – If freq is not positive, if dt is zero, if both nt and length are provided, or if any entry of shape is not positive.

tone_burst Function

deepwave.wavelets.tone_burst(freq: float, n_cycles: float, *, nt: int, dt: float, peak_time: float | None = None, dtype: dtype | None = None, window: str = 'hann', shape: int | Sequence[int] | None = None) → Tensor

Return a windowed sinusoid (“tone burst”).

A tone burst is a sinusoid of frequency freq multiplied by a smooth envelope of finite duration. It is the standard excitation in ultrasound and non-destructive testing, where the number of cycles n_cycles sets the trade-off between temporal resolution (few cycles, broad band) and spectral purity (many cycles, narrow band). The envelope duration is n_cycles / freq seconds, centred on peak_time.

Concretely, with \(T = n_{cycles} / f\) the burst duration and \(\tau = t - t_{peak}\), the wavelet is

\[s(t) = w(\tau)\,\sin(2\pi f\,\tau), \qquad |\tau| \le T/2,\]

and zero outside that window, where \(w\) is the chosen taper (a Hann window \(w(\tau) = \tfrac12\,[1 + \cos(2\pi\tau/T)]\) by default).

The record must be long enough to hold the burst. With the default, record-centred peak_time that means nt * dt of at least n_cycles / freq + 2 * dt. A shorter record cuts the burst off part-way through its oscillation, so the wavelet steps from zero to a non-zero value in a single sample; that step is broadband, which is exactly what choosing a freq and an n_cycles was meant to avoid, and it would otherwise pass unnoticed because nothing about it is an error. Deepwave therefore emits a UserWarning naming the end that was cut, the burst duration, the record length and an nt that would fit. Lengthen the record or shorten the burst; if the truncation is deliberate, filter the warning.

Parameters:

• freq – The centre frequency of the sinusoid (Hz, if dt is in seconds).

• n_cycles – The number of cycles of the sinusoid spanned by the window (may be fractional). The window duration is n_cycles / freq, in the same time units as dt.

• nt – The number of time samples in the returned wavelet.

• dt – The time sample spacing (seconds).

• peak_time – The time (seconds) of the centre of the window. Optional; if not provided the burst is centred in the record (peak_time = nt * dt / 2), the same record-centring convention as ricker(), gaussian_pulse() and ormsby(), so that swapping one wavelet for another does not move the source in time. Deepwave previously defaulted instead to peak_time = n_cycles / freq / 2, which centred the burst window rather than the record, placing the start of the burst at t = 0; pass peak_time=n_cycles / freq / 2 explicitly to reproduce that waveform.

• dtype – The PyTorch datatype to use. Optional, defaults to PyTorch’s default (float32).

• window – Keyword-only. The envelope taper; see deepwave.transducers.window_weights() for the recognised names and their aliases (default "hann").

• shape – Keyword-only. The leading dimensions of the returned tensor, with nt always last. None (the default) returns a one-dimensional [nt] tensor. (n_shots, n_sources) returns [n_shots, n_sources, nt], ready to pass straight to a propagator’s source_amplitudes. An int n is shorthand for (n,) and returns [n, nt]; an empty sequence () returns [nt], as None does. Every entry must be a positive integer. The leading dimensions are added with torch.Tensor.expand(), so no data is copied and the result is a view: every shot and source shares one copy of the wavelet, and each added dimension has stride zero. Call .contiguous() or .clone() on it before modifying it in place. Deepwave’s propagators do not modify their inputs, so passing the view to one directly is safe.

Returns:

A PyTorch tensor of shape [nt] containing the tone burst, or of shape [*shape, nt] if shape was given.

Raises:

• TypeError – If shape is not None, an int, or a sequence of ints.

• ValueError – If dt is zero, if freq or n_cycles is not positive, if window is not a recognised name, or if any entry of shape is not positive.

gaussian_pulse Function

deepwave.wavelets.gaussian_pulse(freq: float, bandwidth_frac: float, *, nt: int, dt: float, peak_time: float | None = None, dtype: dtype | None = None, shape: int | Sequence[int] | None = None) → Tensor

Return a Gaussian-modulated sinusoid with a prescribed bandwidth.

This is a cosine of frequency freq multiplied by a Gaussian envelope, the excitation used to model broadband ultrasound/NDT transducers. The envelope width is set by the fractional -6 dB bandwidth bandwidth_frac following the convention of scipy.signal.gausspulse (F. Harris, “On the use of windows…”, Proc. IEEE 66(1), 1978, for the windowing background): the magnitude spectrum is Gaussian and falls to \(-6\,\text{dB}\) (a factor \(10^{-6/20} = 1/2\)) of its peak at the frequencies \(f_c(1 \pm b/2)\), where \(f_c\) is freq and \(b\) is bandwidth_frac. Equivalently the wavelet is

\[s(t) = e^{-a\,\tau^2}\,\cos(2\pi f_c\,\tau), \qquad a = \frac{(\pi f_c b)^2}{4\ln 2}, \qquad \tau = t - t_{peak},\]

since a Gaussian time envelope \(e^{-a\tau^2}\) has the Gaussian spectrum \(\propto e^{-\pi^2 f^2 / a}\), whose half-amplitude (\(-6\,\text{dB}\)) half-width is \(\Delta f = f_c b / 2\).

Parameters:

• freq – The centre frequency of the sinusoid (Hz, if dt is in seconds).

• bandwidth_frac – The fractional -6 dB bandwidth \(b = \Delta f_{-6dB} / f_c\) (e.g. 0.5 for a 50%-bandwidth transducer). Must be positive.

• nt – The number of time samples in the returned wavelet.

• dt – The time sample spacing (seconds).

• peak_time – The time (seconds) of the envelope peak. Optional; if not provided the pulse is centred in the record (peak_time = nt*dt/2), the same record-centring convention as ricker(), tone_burst() and ormsby().

• dtype – The PyTorch datatype to use. Optional, defaults to PyTorch’s default (float32).

• shape – Keyword-only. The leading dimensions of the returned tensor, with nt always last. None (the default) returns a one-dimensional [nt] tensor. (n_shots, n_sources) returns [n_shots, n_sources, nt], ready to pass straight to a propagator’s source_amplitudes. An int n is shorthand for (n,) and returns [n, nt]; an empty sequence () returns [nt], as None does. Every entry must be a positive integer. The leading dimensions are added with torch.Tensor.expand(), so no data is copied and the result is a view: every shot and source shares one copy of the wavelet, and each added dimension has stride zero. Call .contiguous() or .clone() on it before modifying it in place. Deepwave’s propagators do not modify their inputs, so passing the view to one directly is safe.

Returns:

A PyTorch tensor of shape [nt] containing the Gaussian pulse, or of shape [*shape, nt] if shape was given.

Raises:

• TypeError – If shape is not None, an int, or a sequence of ints.

• ValueError – If dt is zero, if freq or bandwidth_frac is not positive, or if any entry of shape is not positive.

ormsby Function

deepwave.wavelets.ormsby(f1: float, f2: float, f3: float, f4: float, *, nt: int, dt: float, peak_time: float | None = None, dtype: dtype | None = None, shape: int | Sequence[int] | None = None) → Tensor

Return an Ormsby wavelet with the specified four corner frequencies.

The Ormsby wavelet is the standard seismic wavelet whose amplitude spectrum is an exact trapezoid: zero below f1, ramping linearly up to full amplitude at f2, flat (unity) between f2 and f3, then ramping linearly back down to zero at f4 and beyond. This gives more direct control over the excited passband than ricker() (a single central frequency) or gaussian_pulse() (a fractional bandwidth around one centre frequency).

With \(\tau = t - t_{peak}\), the wavelet is

\[w(\tau) \propto \frac{\pi f_4^2\,\mathrm{sinc}^2(\pi f_4 \tau) - \pi f_3^2\,\mathrm{sinc}^2(\pi f_3 \tau)}{f_4 - f_3} - \frac{\pi f_2^2\,\mathrm{sinc}^2(\pi f_2 \tau) - \pi f_1^2\,\mathrm{sinc}^2(\pi f_1 \tau)}{f_2 - f_1},\]

where \(\mathrm{sinc}(x) = \sin(x)/x\) (the mathematical, not normalised, sinc), normalised so the returned wavelet’s peak absolute amplitude is exactly 1.0. (Note: torch.sinc(x) computes the normalised sinc, \(\sin(\pi x)/(\pi x)\); since \(\mathrm{sinc}(\pi f \tau) = \sin(\pi f \tau) / (\pi f \tau) = \mathrm{torch.sinc}(f\tau)\), the mathematical sinc above is computed as torch.sinc(f * tau) directly, with no extra factor of \(\pi\).)

Parameters:

• f1 – The low-cut corner frequency (amplitude zero below this). All four corner frequencies are in Hz, if dt is in seconds, and must satisfy 0 < f1 < f2 < f3 < f4.

• f2 – The low-pass corner frequency (amplitude reaches full scale here).

• f3 – The high-pass corner frequency (amplitude starts rolling off here).

• f4 – The high-cut corner frequency (amplitude zero above this).

• nt – The number of time samples.

• dt – The time sample spacing (seconds).

• peak_time – The time (in secs) of the peak amplitude. Optional; if not provided the wavelet is auto-centred in the record (peak_time = nt * dt / 2), the same record-centring convention as ricker(), tone_burst() and gaussian_pulse(), so the (symmetric) wavelet is not truncated. Passing an explicit value is unchanged.

• dtype – The PyTorch datatype to use. Optional, defaults to PyTorch’s default (float32).

• shape – Keyword-only. The leading dimensions of the returned tensor, with nt always last. None (the default) returns a one-dimensional [nt] tensor. (n_shots, n_sources) returns [n_shots, n_sources, nt], ready to pass straight to a propagator’s source_amplitudes. An int n is shorthand for (n,) and returns [n, nt]; an empty sequence () returns [nt], as None does. Every entry must be a positive integer. The leading dimensions are added with torch.Tensor.expand(), so no data is copied and the result is a view: every shot and source shares one copy of the wavelet, and each added dimension has stride zero. Call .contiguous() or .clone() on it before modifying it in place. Deepwave’s propagators do not modify their inputs, so passing the view to one directly is safe.

Returns:

A PyTorch tensor of shape [nt] containing the Ormsby wavelet, or of shape [*shape, nt] if shape was given.

Raises:

• TypeError – If shape is not None, an int, or a sequence of ints.

• ValueError – If dt is zero, if the corner frequencies do not satisfy 0 < f1 < f2 < f3 < f4, or if any entry of shape is not positive.

Transducer array helpers

Stability: provisional (see Stability policy) – this module is expected to gain element types, functions and constructor keywords. The unit convention has settled: positions in grid cells, speeds in cells per unit time, delays and dt in one time unit of your choosing. What it has today is not respelled or withdrawn without a release’s notice under deepwave.DeepwaveFutureWarning.

The deepwave.transducers module builds extended transducer elements (lines, arcs, discs, focused bowls) as tessellated point sets with apodisation weights and focusing delay laws, for ultrasound and NDT modelling. See Ultrasound and NDT (coming from k-Wave) for an overview and worked usage.

deepwave.transducers.line_element(centre: Sequence[float] | Tensor, direction: Sequence[float] | Tensor, length: float, spacing: float = 1.0, window: str = 'rect', dtype: dtype | None = None) → tuple[Tensor, Tensor]

Tessellate a straight line-segment aperture (2-D or 3-D).

The element is a line segment of the given length (in grid cells), centred on centre and oriented along direction. It is sampled by n equally spaced points whose spacing is <= spacing cells; the endpoints are symmetric about centre.

Parameters:

• centre – The centre of the segment, shape [ndim] (grid cells).

• direction – The orientation vector, shape [ndim] (need not be normalised).

• length – The total length of the segment (grid cells); must be non-negative (a length of 0 gives a single point at centre).

• spacing – The maximum spacing between adjacent points (grid cells). Optional, default 1 cell.

• window – Window (apodisation) type passed to window_weights(). Optional, default "rect" (uniform) – note that this differs from window_weights()’s own "hann" default, because an element should not be apodised unless you ask for it.

• dtype – The PyTorch datatype. Optional, defaults to PyTorch’s default.

Returns:

(locations, weights) where locations has shape [n, ndim] (grid cells) and weights has shape [n]. Prefer passing both to element_source(), which applies the weights (and any delay law) and returns the (source_amplitudes, source_locations) pair a propagator takes; the raw pair is the low-level path, and forgetting to apply weights silently changes the beam pattern.

deepwave.transducers.arc_element(centre_of_curvature: Sequence[float] | Tensor, radius: float, angles: Sequence[float], spacing: float = 1.0, window: str = 'rect', dtype: dtype | None = None) → tuple[Tensor, Tensor]

Tessellate a circular-arc aperture (2-D only).

Models a curved (focused) 2-D element: points lie on a circle of the given radius about centre_of_curvature, spanning the angular range angles = (theta_start, theta_end) in radians. The arc’s centre of curvature is not a point on the arc. The point at angle \(\theta\) is at centre_of_curvature + radius * [cos(theta), sin(theta)], i.e. theta is measured from the first (slow) grid dimension towards the second. Adjacent points are spaced <= spacing cells along the arc.

A curved array focused at a point is obtained by placing an arc of the desired radius about that focal point.

For a full circle (an angular span that is an integer multiple of \(2\pi\)) the final sample would coincide with the first; the duplicated seam point is dropped so every returned cell is distinct.

Parameters:

• centre_of_curvature – The centre of the circle the arc lies on, shape [2] (grid cells). This is the arc’s geometric focus, not a point on the arc itself.

• radius – The arc radius (grid cells).

• angles – (theta_start, theta_end) in radians.

• spacing – Maximum arc spacing between adjacent points (grid cells).

• window – Window (apodisation) type passed to window_weights(). Optional, default "rect" (uniform) – note that this differs from window_weights()’s own "hann" default, because an element should not be apodised unless you ask for it.

• dtype – The PyTorch datatype. Optional.

Returns:

(locations, weights) with locations of shape [n, 2]. Prefer passing both to element_source(), which applies the weights (and any delay law) and returns the (source_amplitudes, source_locations) pair a propagator takes; the raw pair is the low-level path, and forgetting to apply weights silently changes the beam pattern.

deepwave.transducers.disc_element(centre: Sequence[float] | Tensor, radius: float, normal: Sequence[float] | Tensor, spacing: float = 1.0, window: str = 'rect', dtype: dtype | None = None) → tuple[Tensor, Tensor]

Tessellate a flat circular-disc aperture (3-D).

Points fill a flat disc of the given radius centred on centre and lying in the plane whose normal is normal (an unfocused, planar circular element). The disc is tessellated by concentric rings spaced <= spacing cells, with points spaced <= spacing along each ring, giving an approximately uniform areal density.

Parameters:

• centre – The disc centre, shape [3] (grid cells).

• radius – The disc radius (grid cells).

• normal – The disc-plane normal, shape [3] (need not be normalised).

• spacing – Target point spacing (grid cells).

• window – Window (apodisation) applied as a function of the radial fraction (0 at centre, 1 at the rim); see window_weights() for the available types. Optional, default "rect" (uniform) – note that this differs from window_weights()’s own "hann" default, because an element should not be apodised unless you ask for it.

• dtype – The PyTorch datatype. Optional.

Returns:

(locations, weights) with locations of shape [n, 3]. Prefer passing both to element_source(), which applies the weights (and any delay law) and returns the (source_amplitudes, source_locations) pair a propagator takes; the raw pair is the low-level path, and forgetting to apply weights silently changes the beam pattern.

deepwave.transducers.bowl_element(apex: Sequence[float] | Tensor, radius: float, diameter: float, focus: Sequence[float] | Tensor, spacing: float = 1.0, window: str = 'rect', dtype: dtype | None = None) → tuple[Tensor, Tensor]

Tessellate an axisymmetric focused-bowl aperture (3-D).

Models a spherically focused single-element transducer (a spherical cap), the workhorse geometry of therapeutic and photoacoustic ultrasound. The bowl is the cap of a sphere of radius-of-curvature radius whose centre of curvature is at focus; apex is the pole of the cap on the sphere, and the axis points from apex toward focus. The cap extends out to an aperture of the given diameter. The surface is tessellated by concentric rings of constant polar angle, spaced <= spacing cells along the surface.

This mirrors k-Wave’s makeBowl parameterisation (apex position, radius of curvature, aperture diameter, focus direction).

Parameters:

• apex – The bowl apex/pole, shape [3] (grid cells). This is a point on the bowl surface; the bowl’s centre of curvature is the separate focus argument.

• radius – The radius of curvature (grid cells).

• diameter – The aperture diameter (grid cells); must be <= 2*radius.

• focus – The centre of curvature / geometric focus, shape [3] (grid cells). Its distance from apex should be radius; only its direction from apex is used to define the axis.

• spacing – Target surface point spacing (grid cells).

• window – Window (apodisation) as a function of polar-angle fraction (0 at apex, 1 at rim); see window_weights() for the available types. Optional, default "rect" (uniform) – note that this differs from window_weights()’s own "hann" default, because an element should not be apodised unless you ask for it.

• dtype – The PyTorch datatype. Optional.

Returns:

(locations, weights) with locations of shape [n, 3]. Prefer passing both to element_source(), which applies the weights (and any delay law) and returns the (source_amplitudes, source_locations) pair a propagator takes; the raw pair is the low-level path, and forgetting to apply weights silently changes the beam pattern.

deepwave.transducers.window_weights(n: int, kind: str = 'hann', dtype: dtype | None = None, device: device | None = None) → Tensor

Return per-point window (apodisation) weights for an n-point element.

A window/taper shapes the amplitude across an aperture, most commonly to reduce beam side lobes. The returned window is not normalised (its peak is 1 for the tapered windows, and every value is 1 for "rect"); normalise it yourself if you need the element’s total amplitude to be independent of kind.

kind defaults to "hann" here, while the element builders’ window parameter defaults to "rect". The difference is deliberate: an element should not be apodised unless you ask for it, whereas calling this function at all is asking for a window.

Parameters:

• n – The number of points in the element.

• kind – The window type: "rect" (a.k.a. "uniform", "boxcar" or "none"; all ones), "hann" (a.k.a. "hanning", default), "hamming", or "blackman".

• dtype – The PyTorch datatype. Optional, defaults to PyTorch’s default.

• device – The PyTorch device. Optional, defaults to the CPU.

Returns:

A tensor of shape [n] containing the weights.

deepwave.transducers.delays_to_focus(locations: Tensor, focus: Sequence[float] | Tensor, c: float) → Tensor

Per-point transmit delays to focus/steer a beam at focus.

For each element point, returns the time delay to apply on transmit so that all points’ contributions arrive in phase at the focal point. Points farther from the focus fire earlier; the delays are non-negative and the point nearest the focus has the largest delay. Placing the focus far away in a chosen direction approximates plane-wave steering.

This helper assumes isotropic physical cell spacing. Convert physical speed to cells per unit time with the single common grid spacing; for an anisotropic grid, compute delays from physical coordinates separately.

All quantities are in your units. Positions are in grid cells, c is a speed in cells per unit time, and the returned delays are in that same unit of time – the one to use for dt. If your model is in metres and seconds, c_cells = c_physical / grid_spacing, and the delays come back in seconds.

Parameters:

• locations – Element point locations, shape [n, ndim] (grid cells).

• focus – The focal point, shape [ndim] (grid cells).

• c – The wave speed between the element and the focus, in grid cells per unit time.

Returns:

A tensor of shape [n] of non-negative delays, in the time unit implied by c.

deepwave.transducers.apply_delays(wavelet: Tensor, delays: Tensor, dt: float, nt: int | None = None) → Tensor

Build per-point amplitudes by fractional-delay-shifting a base wavelet.

Each element point emits the same base wavelet shifted in time by its delay. Shifting is done by a spectral (FFT) phase ramp, which applies an arbitrary fractional (sub-sample) delay exactly for a band-limited signal – the same spectral philosophy Deepwave uses for source/receiver time resampling. A positive delay shifts the wavelet later in time.

The operation is differentiable with respect to both wavelet and delays (the phase ramp is a smooth function of the delays), so delay laws can be optimised end-to-end.

delays and dt are in the same unit of time, whichever unit that is – the one delays_to_focus() returns for your c. Nothing here needs to know what a second is.

Note

The FFT shift is circular: energy shifted past the end of the record wraps to the beginning. Make nt long enough (and pad the base wavelet) that the delayed wavelet stays within the record.

Parameters:

• wavelet – The base wavelet, shape [nt0].

• delays – Per-point delays, shape [n], in the same time unit as dt. May be negative.

• dt – The time sample spacing.

• nt – The output number of time samples. Optional; if given and different from len(wavelet) the wavelet is zero-padded or truncated to nt before shifting.

Returns:

A tensor of shape [n, nt] of per-point amplitudes.

deepwave.transducers.element_source(wavelet: Tensor, locations: Tensor, weights: Tensor, *, delays: Tensor | None = None, dt: float | None = None, nt: int | None = None, n_shots: int | None = None) → tuple[Tensor, Tensor]

Assemble an element into the source pair a propagator takes.

Takes the (locations, weights) an element builder returns, a base wavelet, and an optional delay law, and returns (source_amplitudes, source_locations) of shapes [n_shots, n, nt] and [n_shots, n, ndim], ready to pass straight to a propagator. This is the recommended way to drive an element: it applies the apodisation weights, which are easy to forget and whose omission silently changes the beam pattern.

Every step is differentiable, with respect to wavelet, weights, locations and delays alike, so a delay law or an apodisation can itself be optimised.

delays and dt are in the same unit of time, whichever unit that is – the one delays_to_focus() returns for a c in cells per unit time. Nothing here needs to know what a second is.

The shot dimension is added with torch.Tensor.expand(), so a single-shot wavelet costs no extra memory when broadcast to many shots.

Parameters:

• wavelet – The base wavelet, shape [nt0] (the same wavelet in every shot) or [n_shots, nt0] (one per shot).

• locations – Element point locations, shape [n, ndim] (grid cells), from an element builder.

• weights – Per-point apodisation weights, shape [n], from the same builder.

• delays – Per-point transmit delays, shape [n], e.g. from delays_to_focus(). Optional; if given, dt is required.

• dt – The time sample spacing, in the same unit as delays. Required if and only if delays is given.

• nt – The output number of time samples. Optional; if given and different from wavelet.shape[-1] the wavelet is zero-padded or truncated to nt, exactly as in apply_delays().

• n_shots – The number of shots. Optional; defaults to 1 for a 1-D wavelet and to wavelet.shape[0] for a 2-D one. Giving a value that disagrees with a 2-D wavelet’s leading dimension raises rather than broadcasting.

Returns:

(source_amplitudes, source_locations) of shapes [n_shots, n, nt] and [n_shots, n, ndim].

Raises:

DeepwaveValueError – If delays is given without dt; if weights and locations describe different numbers of points; if wavelet is neither 1-D nor 2-D; if delays is not 1-D of length n; or if n_shots disagrees with a 2-D wavelet’s leading dimension.

deepwave.transducers.element_receiver_trace(receiver_amplitudes: Tensor, weights: Tensor) → Tensor

Collapse an element’s per-point receiver traces into one trace.

Returns the discrete (apodisation-)weighted sum of traces at the element’s tessellation points. Builder weights do not include area quadrature or normalisation, so the result and its units depend on point density.

Parameters:

• receiver_amplitudes – Recorded traces with the element’s points on the second-to-last axis, shape [..., n, nt] (e.g. the propagator’s [shot, receiver, time] output with receiver == n).

• weights – Per-point weights, shape [n] (e.g. from an element builder or window_weights()).

Returns:

The weighted-sum element trace, shape [..., nt].

Misfit functions

The deepwave.misfits module provides three differentiable FWI data-misfit (loss) functions: deepwave.misfits.l2() (the conventional sum of squared differences), and two that are less prone to cycle-skipping than it, deepwave.misfits.envelope() (mean squared error between Hilbert envelopes) and deepwave.misfits.normalised_cross_correlation() (one minus the per-trace normalised cross-correlation, which is invariant to amplitude scaling). All three share one shape contract – the time axis is reduced away to one value per trace, and reduction ("mean", "sum" or "none") then says what to do with that per-trace tensor – so swapping one for another really is a one-line change. They do not share a scale, though: deepwave.misfits.l2() sums the squared differences over time while deepwave.misfits.envelope() averages them, and deepwave.misfits.normalised_cross_correlation() is bounded in [0, 2] per trace whatever the amplitudes, so expect to retune the step length after a swap. See Graph Space Optimal Transport for when and why to swap the misfit function during an inversion.

Two of these functions are commonly reached for as a replacement for torch.nn.functional.mse_loss(), and the correspondence is exact in one reduction only. l2(pred, target, reduction="sum") equals mse_loss(pred, target, reduction="sum"), but l2(pred, target, reduction="mean") is nt times mse_loss(pred, target, reduction="mean") – deepwave.misfits.l2() sums over time and averages over traces, whereas mse_loss averages over every element. Nothing raises, so a swap made on the default reduction scales the loss and its gradient by nt (often 1e3 to 1e4); divide the learning rate by nt, or use reduction="sum" on both sides.

deepwave.misfits.l2(pred: Tensor, target: Tensor, *, reduction: Literal['mean', 'sum', 'none'] = 'mean') → Tensor

Squared-error (least-squares) misfit between pred and target.

The conventional FWI data misfit. Per trace, with \(p[n]\) = pred, \(t[n]\) = target, and \(n = 0, \dots, n_t - 1\),

\[L = \sum_{n} \bigl(p[n] - t[n]\bigr)^2,\]

which is then reduced across the leading (non-time) dimensions according to reduction.

The misfit is usually written in the FWI literature with a leading factor of \(1/2\) (which cancels the 2 produced by differentiating the square); omitting it, as here, multiplies the misfit and its gradient by the same constant 2, moving no minimum and only rescaling the step length an optimiser must take.

This is the misfit that is most prone to cycle-skipping: it compares the oscillatory waveforms sample by sample, so a predicted arrival more than half a period from the observed one is pulled towards the wrong cycle. envelope() and normalised_cross_correlation() are drop-in alternatives with the same shape contract, though not with the same scale: this function sums over time where envelope() averages, so a swap usually wants the step length adjusted. The module docstring gives the exact relationship, and the one to torch.nn.functional.mse_loss(), whose "mean" reduction differs from this one by a factor of nt.

Parameters:

• pred – Predicted data, a real floating-point tensor of shape [..., nt] (e.g. [shot, receiver, time]).

• target – Observed/target data. Must have the same shape and dtype as pred.

• reduction – One of "mean" (default), "sum", or "none". With "none", the time axis is still reduced away (it is not meaningful to leave it, since the misfit is inherently a per-trace quantity), returning the per-trace sum of squared differences with shape pred.shape[:-1]; "mean"/"sum" further reduce that per-trace tensor to a scalar.

Returns:

A scalar tensor ("mean"/"sum") or a tensor of shape pred.shape[:-1] ("none").

Raises:

• TypeError – If pred or target is not a floating-point tensor, or their dtypes differ.

• ValueError – If pred and target have different shapes, either has zero dimensions, or reduction is not recognised.

deepwave.misfits.envelope(pred: Tensor, target: Tensor, *, reduction: Literal['mean', 'sum', 'none'] = 'mean') → Tensor

Mean squared error between the Hilbert envelopes of pred and target.

The Hilbert envelope (analytic-signal magnitude) of a real trace \(x[n]\), \(n = 0, \dots, n_t - 1\), is computed from its discrete Fourier transform \(X = \mathrm{FFT}(x)\) by discarding the negative-frequency half of the spectrum and doubling the positive frequencies, following the standard FFT-based construction of the discrete analytic signal (Marple, 1999):

\[\begin{split}A[k] = \begin{cases} 1, & k = 0 \\ 2, & 1 \le k < \lceil n_t / 2 \rceil \\ 1, & n_t \text{ even and } k = n_t / 2 \\ 0, & \text{otherwise,} \end{cases} \qquad z = \mathrm{IFFT}(A \odot X), \qquad e[n] = |z[n]|,\end{split}\]

where \(z\) is the discrete analytic signal and \(e\) its envelope (the \(k = n_t / 2\) “Nyquist” case only arises for even \(n_t\); it is self-conjugate and so kept at weight 1, not doubled). Writing \(e_{\text{pred}}\) and \(e_{\text{target}}\) for the envelopes of pred and target, the per-trace value is their mean squared difference over time,

\[L = \frac{1}{n_t} \sum_{n} \bigl(e_{\text{pred}}[n] - e_{\text{target}}[n]\bigr)^2,\]

which is then reduced across the leading (non-time) dimensions according to reduction. Because it responds to envelope (amplitude-of-arrival) differences rather than the raw, oscillatory waveform, this misfit is markedly less sensitive to the phase/traveltime errors that cause cycle-skipping in the raw squared-error misfit, and is a standard FWI objective for that reason (Bozdağ, Trampert & Tromp, 2011).

Parameters:

• pred – Predicted data, a real floating-point tensor of shape [..., nt] (e.g. [shot, receiver, time]).

• target – Observed/target data. Must have the same shape and dtype as pred.

• reduction – One of "mean" (default), "sum", or "none". With "none", the time axis is still reduced away (it is not meaningful to leave it, since the misfit is inherently a per-trace quantity), returning the per-trace mean squared envelope difference with shape pred.shape[:-1]; "mean"/ "sum" further reduce that per-trace tensor to a scalar.

Returns:

A scalar tensor ("mean"/"sum") or a tensor of shape pred.shape[:-1] ("none").

Raises:

• TypeError – If pred or target is not a floating-point tensor, or their dtypes differ.

• ValueError – If pred and target have different shapes, either has zero dimensions, or reduction is not recognised.

References

S. L. Marple, “Computing the discrete-time analytic signal via FFT,” IEEE Transactions on Signal Processing, vol. 47, no. 9, 1999.

E. Bozdağ, J. Trampert, J. Tromp, “Misfit functions for full waveform inversion based on instantaneous phase and envelope measurements,” Geophysical Journal International, vol. 185, no. 2, 2011.

deepwave.misfits.normalised_cross_correlation(pred: Tensor, target: Tensor, *, reduction: Literal['mean', 'sum', 'none'] = 'mean') → Tensor

One minus the per-trace normalised cross-correlation.

Per trace, with \(p[n]\) = pred, \(t[n]\) = target, \(n = 0, \dots, n_t - 1\), and \(\lVert \cdot \rVert_2\) the L2 norm over the time axis,

\[\begin{split}\mathrm{NCC} = \begin{cases} \dfrac{\sum_n p[n]\, t[n]}{\lVert p \rVert_2 \, \lVert t \rVert_2}, & \lVert p \rVert_2 \, \lVert t \rVert_2 > 0, \\[6pt] 0, & \lVert p \rVert_2 \, \lVert t \rVert_2 = 0, \end{cases} \qquad L = 1 - \mathrm{NCC},\end{split}\]

No mean is subtracted from either trace before normalising, so \(\mathrm{NCC}\) is the cosine similarity of the raw traces, not the statistical (mean centred) correlation coefficient: the misfit stays sensitive to any DC offset, which is physically meaningful for wave data, since a non-zero mean is itself a mismatch worth penalising.

The \(\lVert p \rVert_2 \, \lVert t \rVert_2 = 0\) branch applies only when at least one of pred/target is an exactly all-zero trace (its L2 norm is exactly 0.0, not merely small): then \(\mathrm{NCC} = 0\) and \(L = 1\), a finite, defined result (rather than the 0 / 0 that an unguarded division would give) that reads as “no correlation” for a trace with no signal. It is a branch on the value 0.0, not a small-denominator threshold, so it does not fire at small-but-nonzero amplitudes: traces many orders of magnitude below unit scale are compared exactly like unit-scale data (see Normalised cross-correlation: numerical conditioning for the internal rescaling that makes that so).

\(L\) is exactly 0 whenever pred equals target scaled by any positive constant on that trace (a pure amplitude error), and otherwise ranges over \([0, 2]\). That amplitude-scale invariance, with undiminished sensitivity to phase/traveltime shifts, is why normalised-correlation (“global correlation norm”) objectives are used as an FWI misfit robust to the amplitude mismatches that a squared-error misfit penalises but that need not indicate a poor model (Choi & Alkhalifah, 2012). The per-trace \(L\) is then reduced across the leading (non-time) dimensions according to reduction.

Parameters:

• pred – Predicted data, a real floating-point tensor of shape [..., nt] (e.g. [shot, receiver, time]).

• target – Observed/target data. Must have the same shape and dtype as pred.

• reduction – One of "mean" (default), "sum", or "none". With "none", the time axis is still reduced away (it is not meaningful to leave it, since the misfit is inherently a per-trace quantity), returning the per-trace value with shape pred.shape[:-1]; "mean"/"sum" further reduce that per-trace tensor to a scalar.

Returns:

A scalar tensor ("mean"/"sum") or a tensor of shape pred.shape[:-1] ("none").

Raises:

• TypeError – If pred or target is not a floating-point tensor, or their dtypes differ.

• ValueError – If pred and target have different shapes, either has zero dimensions, or reduction is not recognised.

References

Y. Choi, T. Alkhalifah, “Application of multi-source waveform inversion to marine streamer data using the global correlation norm,” Geophysical Prospecting, vol. 60, no. 4, 2012.

Wavefield-space tools

The deepwave.wavefields module widens Deepwave from a differentiable shot-data operator (point sources in, receiver traces out) to a differentiable wavefield operator. Reach for it when the thing you want to force, reduce, or observe is the field itself rather than a trace: source_field_to_sources() and merge_point_sources() turn a dense forcing field into ordinary point sources, for wavefield-reconstruction inversion, source encoding and adjoint-source experiments; measure() with a RunningDFT accumulates a frequency-domain reduction as a run proceeds, so a frequency-domain objective costs O(n_freqs) storage instead of O(nt) and still carries gradients; and record_boundary() / reinject_boundary() record traces on a rectangular surface and replay them as source amplitudes at its cells. Everything here is pure Python built on the public propagator interface, and works on any propagator. Wavefield-space tools covers each in depth, including the limits of trace replay.

deepwave.wavefields.source_field_to_sources(source_field: Tensor, *, origin: Sequence[int] | None = None, mask: Tensor | None = None) → tuple[Tensor, Tensor]

Convert a volumetric source field into per-cell point sources.

Stability: provisional (see Stability policy).

A volumetric (extended) source is a forcing f(shot, t, *cell) added to the wavefield at every cell of a region on every outer time step. Deepwave injects a point source at an integer grid cell as an exact single-cell forcing (the Hicks windowed-sinc interpolation collapses to a Kronecker delta at an on-node location), so a volumetric field is represented exactly by one integer-located point source per cell, with that cell’s time series as its amplitude. Because each cell appears once, the resulting source locations are automatically unique – the unique-cell constraint that governs hand-placed point sources never bites for field forcing.

The mapping is a pure reshape/gather, so the returned amplitudes are differentiable w.r.t. source_field: a gradient on the receiver data or on a later wavefield flows straight back to the field.

Parameters:

• source_field – The forcing, shape [shot, nt, *region_shape] where region_shape has ndim entries (e.g. [shot, nt, ny, nx] in 2-D). nt is the number of time steps and defines the propagation length.

• origin – The model cell index (length ndim) of the region’s first cell, i.e. region cell (0, ..., 0) sits at model cell origin. Defaults to all zeros (the region starts at the model origin). This is added to the in-region cell indices to give model-relative source locations.

• mask – An optional boolean tensor of shape region_shape selecting which cells are active. If given, only the True cells become sources (the field values at False cells are ignored); this is the region-masked form and keeps the source count (and cost) down for sparse regions.

Returns:

A tuple (source_amplitudes, source_locations) with source_amplitudes of shape [shot, n_cell, nt] and source_locations of shape [shot, n_cell, ndim] (integer, model-relative), ready to pass to any Deepwave propagator.

deepwave.wavefields.merge_point_sources(field_amplitudes: Tensor, field_locations: Tensor, source_amplitudes: Tensor | None, source_locations: Tensor | None) → tuple[Tensor, Tensor]

Concatenates field-derived sources with explicit point sources.

The second half of the volumetric-source recipe. source_field_to_sources() turns an extended forcing into one exact single-cell source per cell; this merges that set with any point sources of your own, so that both can be passed through the propagator’s single source-amplitude/location pair:

amp, loc = deepwave.wavefields.source_field_to_sources(
    field, origin=origin, mask=mask
)
amp, loc = deepwave.wavefields.merge_point_sources(
    amp, loc, source_amplitudes, source_locations
)
out = deepwave.scalar(v, dx, dt, sources={"u": (amp, loc)},
                      receivers={"u": receiver_locations})

The combined location set still has to be unique within each shot, so the explicit source cells must be disjoint from the field’s cells (fold any coincident point source into the field instead).

Stability: provisional (see Stability policy).

Parameters:

• field_amplitudes – The amplitudes returned by source_field_to_sources().

• field_locations – The locations returned by source_field_to_sources().

• source_amplitudes – The explicit point sources’ amplitudes, or None for “no point sources”, in which case the field’s own amplitudes and locations are returned unchanged.

• source_locations – The explicit point sources’ locations, as an integer torch.Tensor. Off-grid (floating-point or deepwave.Hicks) point sources cannot share the field’s exact-node source slot and are rejected.

Returns:

• The concatenated source amplitudes.

• The concatenated source locations.

Return type:

Tuple

Raises:

ValueError – If source_amplitudes is given without source_locations, if the two sets have different numbers of time samples, or if source_locations is not an integer tensor.

class deepwave.wavefields.RunningDFT(frequencies: Sequence[float], on: str | None = None)

Declarative running-DFT measurement, accumulated segment-wise.

Stability: provisional (see Stability policy).

Accumulates X[f] = sum_t x[t] * exp(-2j*pi*f*t*dt) over the propagation, where x is either a receiver-amplitude output or a wavefield of the propagator. The sum is built incrementally as segments complete. The DFT accumulators scale with the number of frequencies; each segment’s trace and propagation history also contribute to peak memory. The result matches an offline DFT of the assembled trace to round-off and is differentiable.

Parameters:

• frequencies – The frequencies at which to evaluate the transform, in the reciprocal of dt’s time unit (Hz when dt is in seconds).

• on – The name of the propagator output to transform. For receiver transforms this is a receiver-amplitude output name (e.g. "receiver_amplitudes" for scalar, "receiver_amplitudes_p" for acoustic); None (default) selects the propagator’s first receiver-amplitude output. It may instead be a wavefield name (e.g. "wavefield_0"), which requires per-step segmentation (see measure()), or one of the keys of the propagator’s auxiliary_state (e.g. "psiy_m1"), which is treated the same way. Those keys are implementation state, not part of the public interface: their names, number and shapes may change in any release (see Stability policy).

deepwave.wavefields.measure(propagator: Callable[[...], Any], model_args: Sequence[Tensor], grid_spacing: float | int | Sequence[float], dt: float, source_amplitudes: Mapping[str, Tensor], measurements: Sequence[RunningDFT], *, n_segments: int = 1, initial_state: PropagatorOutput | Mapping[str, Tensor | AuxiliaryState] | None = None, return_output: bool = False, **kwargs: Any) → MeasureResult

Run a propagator while accumulating declarative reductions on the fly.

Stability: provisional (see Stability policy).

Runs propagator from step 0 to nt in n_segments contiguous segments (threading state through initial_state), accumulating each RunningDFT in measurements as the segments complete. Nothing but the running accumulators (and, if return_output, the final state) is kept across segments. The accumulators cost O(n_freqs) per receiver; peak memory includes a segment’s trace and propagation history. The result is returned inside the autograd graph.

Receiver-amplitude measurements work at any n_segments. A measurement on a wavefield (on names a wavefield, or one of the propagator’s auxiliary_state keys) needs the field at every outer step, so it forces per-step segmentation (n_segments = nt); this makes the call nt propagator invocations. If a different n_segments was explicitly requested, this override is reported via a warning (the per-step segmentation itself is exact but has Python-per-step overhead – prefer receiver measurements when you can).

Requires ``step_ratio == 1`` when the run is segmented (whenever the resolved n_segments > 1, including the per-step wavefield path). Otherwise a ValueError is raised, because splitting a propagation into segments is only bit-exact – and hence the accumulated reduction only equals an offline reduction of the whole trace – when the user dt already satisfies the CFL condition, so Deepwave does no internal sub-stepping (the source amplitudes are otherwise resampled to the internal time step with a global FFT that does not decompose across a segment boundary; see deepwave.checkpointing’s module docstring). A single-segment (n_segments == 1) receiver measurement runs at any step_ratio.

Parameters:

• propagator – A Deepwave propagator function (deepwave.scalar etc.).

• model_args – Model tensors in positional order ((v,) for scalar).

• grid_spacing – Grid spacing, as passed to the propagator.

• dt – The user time step.

• source_amplitudes – Mapping from source-amplitude argument name to tensor, defining nt (e.g. {"source_amplitudes": src}).

• measurements – The reductions to accumulate.

• n_segments – Number of contiguous segments (default 1). Ignored (forced to nt) if any measurement is on a wavefield.

• initial_state – Optional starting state threaded into the first segment.

• return_output – If True, populate the result’s output field with the final PropagatorOutput; it is None otherwise. The shape of the return value does not depend on this: MeasureResult is a tuple subclass, so a result that was sometimes a bare list and sometimes a pair would make for m in measure(...) iterate the wrong thing.

• **kwargs – Remaining propagator keyword arguments forwarded to every segment. storage_mode and deepwave.checkpointing.run_segments()’s checkpointing parameters must not appear.

Returns:

A MeasureResult; unpack it as measurements, output = measure(...). measurements holds one accumulated complex tensor per measurement, in the order of measurements (result.measurements[i] is the reduction of measurements[i]). Receiver transforms have shape [shot, receiver, n_freq]; wavefield transforms have shape [shot, *wavefield_spatial, n_freq]. output is the final PropagatorOutput if return_output is True, and None otherwise.

class deepwave.wavefields.MeasureResult(measurements: tuple[Tensor, ...], output: PropagatorOutput | None)

The result of deepwave.wavefields.measure().

Returned for both values of return_output; unpack it as measurements, output = measure(...), or read measure(...).measurements.

Stability: provisional (see Stability policy).

measurements

One accumulated complex tensor per RunningDFT, in the order they were given.

Type:

tuple[torch.Tensor, …]

output

The final PropagatorOutput, or None unless return_output=True.

Type:

deepwave.PropagatorOutput | None

Create new instance of MeasureResult(measurements, output)

deepwave.wavefields.boundary_surface_locations(box: Sequence[tuple[int, int]]) → Tensor

Integer cell coordinates of the surface (faces) of a rectangular box.

Stability: provisional (see Stability policy).

Enumerates every cell lying on the boundary faces of the axis-aligned box box (each dimension given as a (start, stop) half-open range, in model cell indices). Corner and edge cells appear exactly once. The result is the dense receiver/source set for record_boundary() / reinject_boundary().

Parameters:

box – A sequence of ndim (start, stop) tuples giving the box extent in model cell indices (half-open: cells start .. stop - 1).

Returns:

An integer tensor of shape [n_surface, ndim] of the surface cell coordinates.

deepwave.wavefields.record_boundary(propagator: Callable[[...], Any], model_args: Sequence[Tensor], grid_spacing: float | int | Sequence[float], dt: float, source_amplitudes: Tensor, source_locations: Tensor, box: Sequence[tuple[int, int]], *, channel: str | None = None, source_channel: str | None = None, receiver_kwarg: str | None = None, output_name: str | None = None, **kwargs: Any) → BoundaryRecording

Record the wavefield on a rectangular surface with dense receivers.

Stability: provisional (see Stability policy).

Places a receiver at every cell on the faces of box (via boundary_surface_locations()) and runs the propagator, returning the recorded surface traces and the surface cell coordinates. The source traces supplied by the caller excite the selected source channel.

Use the same explicit pml_freq when comparing this run with a replay run, so their absorbing boundaries have the same configuration.

Parameters:

• propagator – A Deepwave propagator function – any that deepwave.spec() accepts.

• model_args – Model tensors in positional order.

• grid_spacing – Grid spacing.

• dt – Time step.

• source_amplitudes – Source amplitudes [shot, source, nt].

• source_locations – Source locations [shot, source, ndim].

• box – The recording surface, as (start, stop) per dimension.

• channel – The receiver channel to record, e.g. "u" on deepwave.scalar(), "p" on deepwave.elastic(). Both the receiver keyword and the output field are resolved from it through deepwave.spec(), so callers name a channel – which is stable – rather than a keyword spelling, which is not. Defaults to the propagator’s first receiver channel, which is "u" on scalar and reproduces this function’s original behaviour exactly.

• source_channel – The source channel to excite, resolved independently of channel. Defaults to the propagator’s first source channel for the dimension of box. For example, use "volume_rate" for the acoustic Born background source.

• receiver_kwarg – An explicit override of the receiver-locations keyword, for anything the spec cannot express. Passing it together with channel raises.

• output_name – An explicit override of the output field holding the recorded traces. Passing it together with channel raises.

• **kwargs – Remaining propagator keyword arguments.

Returns:

A BoundaryRecording, a named (recording, surface_locations) tuple where recording has shape [shot, n_surface, nt] and surface_locations is [n_surface, ndim] – the two arguments reinject_boundary() takes, in that order.

class deepwave.wavefields.BoundaryRecording(recording: Tensor, surface_locations: Tensor)

The result of deepwave.wavefields.record_boundary().

A named two-tuple, so recording, surface = record_boundary(...) still unpacks and still composes with reinject_boundary().

Stability: provisional (see Stability policy).

recording

The recorded surface traces, [shot, surface_cell, time].

Type:

torch.Tensor

surface_locations

The surface cell coordinates, [surface_cell, ndim].

Type:

torch.Tensor

Create new instance of BoundaryRecording(recording, surface_locations)

deepwave.wavefields.reinject_boundary(propagator: Callable[[...], Any], model_args: Sequence[Tensor], grid_spacing: float | int | Sequence[float], dt: float, recording: Tensor, surface_locations: Tensor, *, channel: str | None = None, amplitude_scale: float = 1.0, **kwargs: Any) → PropagatorOutput

Replay boundary traces as sources at the supplied surface cells.

Stability: provisional (see Stability policy).

This helper multiplies the supplied traces by amplitude_scale and uses them as source amplitudes at the supplied surface cells. It applies no boundary-integral weights, normal derivatives, dipole terms, or automatic conversion from recorded-field units to source units. It does not in general reconstruct the original field or preserve its phase or amplitude.

Use the same explicit pml_freq when comparing the recording and replay runs, so their absorbing boundaries have the same configuration.

Parameters:

• propagator – A Deepwave propagator function.

• model_args – Model tensors for the replay run.

• grid_spacing – Grid spacing.

• dt – Time step.

• recording – Boundary recording [shot, n_surface, nt] from record_boundary().

• surface_locations – Surface cell coordinates [n_surface, ndim].

• channel – The physical channel to inject into, e.g. "u" on deepwave.scalar(), "p" on deepwave.elastic(). The source amplitude and location keywords are resolved from it through deepwave.spec(), so this works on every propagator rather than only the ones that spell their source keywords the way scalar does. Defaults to the propagator’s first source channel, reproducing this function’s original behaviour on scalar.

• amplitude_scale – A scalar multiplying the supplied traces before injection. Default 1.0.

• **kwargs – Remaining propagator keyword arguments.

Returns:

The propagator output of the replay run.