Wavefield-space tools

Deepwave is, at its core, a differentiable shot-data operator: you inject point sources and read receiver traces, and gradients flow back to the model, the source, and the initial state. The deepwave.wavefields module widens that surface toward a differentiable wavefield operator, for research workflows that need to force, reduce, or observe the field itself. Everything here is additive, pure Python, built on the public propagator interface (and, for the segment drivers, on deepwave.checkpointing.run_segments()); no compiled-backend or ABI change is involved, and the existing propagators are unaffected.

There are three capabilities: volumetric source fields, on-the-fly reductions, and boundary record/reinject. The first is wired into deepwave.scalar() as the source_field keyword; the others are drivers you call on any propagator.

Volumetric / extended source fields

A point source injects forcing at a single grid cell. A volumetric (extended) source injects a forcing field f(shot, t, *cell) at every cell of a region on every time step – the building block of wavefield-reconstruction inversion (WRI), source encoding, and adjoint-source experiments. Convert the field (shape [shot, nt, *spatial]) into ordinary point sources with source_field_to_sources(), optionally merge it with point sources of your own with merge_point_sources(), and pass the result to any propagator:

import torch
import deepwave
from deepwave.wavefields import merge_point_sources, source_field_to_sources

# A dense forcing over the whole model, differentiable w.r.t. the field.
source_field = torch.zeros(n_shots, nt, ny, nx, requires_grad=True)
amplitudes, locations = source_field_to_sources(source_field)
# Optional: add explicit point sources on cells the field does not use.
amplitudes, locations = merge_point_sources(
    amplitudes, locations, source_amplitudes, source_locations
)
out = deepwave.scalar(
    v, dx, dt,
    sources={"u": (amplitudes, locations)},
    receivers={"u": receiver_locations},
    pml_freq=freq,
)
out.receiver_amplitudes.pow(2).sum().backward()   # d/d source_field

The field is implemented by mapping each (active) cell to one point source at that integer grid node. On a regular grid an on-node source is an exact single-cell forcing – the Hicks windowed-sinc interpolation collapses to a Kronecker delta at an integer location – so the mapping introduces no smoothing and is bit-exact against building the per-cell sources by hand. Because each cell is used exactly once, the field’s sources are automatically unique: the unique-cell constraint that governs hand-placed point sources never applies to field forcing. If you also pass explicit source_amplitudes/source_locations, the two sets are concatenated and their cells must be disjoint (fold any coincident point source into the field).

Two options restrict the field to a sub-region, keeping the source count (and cost) down:

  • origin – the model cell index of the field’s first cell, so a field of shape [shot, nt, *region_shape] covers a rectangular sub-box.

  • mask – a boolean tensor of shape region_shape selecting which cells are active (a non-rectangular, region-masked source).

Both helpers are propagator-agnostic, so acoustic and elastic users build the same per-cell sources and pass them on whichever channel they want.

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.

Differentiable on-the-fly reductions

Frequency-domain FWI and related objectives need only a handful of Fourier components of the receiver data, not the whole time series. deepwave.wavefields.measure() runs a propagator segment-by-segment and accumulates a declarative reduction – a RunningDFT at chosen frequencies – as the segments complete, returning the result inside the autograd graph. The full trace is never retained across segments. DFT accumulators cost O(n_freqs) per receiver; peak memory also includes a segment’s receiver trace and propagation history. The accumulated transform equals an offline FFT/DFT of the assembled trace to floating-point round-off, in both value and gradient:

from deepwave.wavefields import measure, RunningDFT

result = measure(
    deepwave.scalar,
    (v,), dx, dt,
    {"source_amplitudes": source_amplitudes},
    [RunningDFT(frequencies=[5.0, 9.0, 13.0], on="u")],
    n_segments=8,
    source_locations=source_locations,
    receiver_locations=receiver_locations,
    max_vel=max_vel, pml_freq=freq,
)
dft = result.measurements[0]     # [shot, receiver, n_freq], complex, in-graph
(dft.abs() ** 2).sum().backward()

measure returns a MeasureResult: a named (measurements, output) pair, so result.measurements and measurements, output = measure(...) both read naturally. measurements holds one complex tensor per measurement, in the order of the measurements argument. output is None unless you ask for it with return_output=True, which populates it with the run’s final PropagatorOutput – the receiver amplitudes and final state a direct call would have returned, alongside the reduction:

result = measure(
    deepwave.scalar,
    (v,), dx, dt,
    {"source_amplitudes": source_amplitudes},
    [RunningDFT(frequencies=[5.0, 9.0, 13.0], on="u")],
    n_segments=8,
    return_output=True,
    source_locations=source_locations,
    receiver_locations=receiver_locations,
    max_vel=max_vel, pml_freq=freq,
)
dft = result.measurements[0]
print(result.output.info.step_ratio)          # the run's own metadata
final_wavefield = result.output.wavefield_0   # in the autograd graph

The shape of the return value does not depend on return_output: the pair is always a pair, so unpacking code does not have to change when you turn the output on. Keeping the output does cost the final state’s memory, which is why it is off by default.

on is required, and names either a channel of the propagator ("u" on deepwave.scalar(), "p" on deepwave.elastic()) or a full output field name ("receiver_amplitudes_p"). Naming a channel is the better habit: it is stable, and it is unambiguous on a Born propagator, where “the first receiver output” is the background data.

Receiver-amplitude measurements work at any n_segments (the segment-wise partial DFT sums are exact regardless of where the boundaries fall). A measurement whose on names a wavefield (e.g. RunningDFT([...], on="wavefield_0")) needs the field at every outer step, so it forces per-step segmentation (n_segments = nt, reported by a warning): it is exact but has Python-per-step overhead, so prefer receiver measurements where you can.

on may also name one of the keys of the propagator’s auxiliary_state (e.g. RunningDFT([...], on="psiy_m1")), which is treated the same way as a wavefield target and forces the same per-step segmentation. Those keys are the internal auxiliary-field names Internal callback names lists, and they keep working (see Stability policy), but what they hold is the absorbing boundary’s own state rather than physics: a measurement built this way is a debugging tool rather than a stable part of an objective function, because the numbers in it follow the absorbing-boundary implementation.

Note

Exactness requires step_ratio == 1 whenever the run is segmented (the resolved n_segments > 1, including the per-step wavefield path). measure raises a ValueError otherwise (reduce dt): its running DFT is accumulated on user-rate samples, so, unlike deepwave.checkpointing.run_segments(), it cannot simply run its segments at the internal time step when the user dt needs internal sub-stepping. A single-segment (n_segments == 1) receiver measurement runs at any step_ratio.

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)

Boundary recording and trace replay

Use dense receivers to record traces at the cells of a rectangular surface, then replay those traces as source amplitudes at the same cells:

from deepwave.wavefields import record_boundary, reinject_boundary

box = [(30, 55), (30, 55)]
recording, surface = record_boundary(
    deepwave.scalar, (v,), dx, dt,
    source_amplitudes, source_locations, box,
    pml_freq=freq,
)
replay = reinject_boundary(
    deepwave.scalar, (v,), dx, dt,
    recording, surface,
    receiver_locations=interior_locations, pml_freq=freq,
)

record_boundary selects its receiver quantity with channel= and its excitation independently with source_channel=. reinject_boundary uses channel= to select the source channel. Names resolve through deepwave.spec():

recording, surface = record_boundary(
    deepwave.elastic, (lamb, mu, buoyancy), dx, dt,
    source_amplitudes, source_locations, box,
    source_channel="pressure_rate", channel="p", pml_freq=freq,
)
replay = reinject_boundary(
    deepwave.elastic, (lamb, mu, buoyancy), dx, dt,
    recording, surface, channel="pressure_rate", pml_freq=freq,
)

For recording, channel=None selects the first receiver channel and source_channel=None selects the first source channel for the box dimension. For replay, channel=None selects the first source channel for the surface dimension. A Born source can therefore be selected separately from a scattered receiver, for example source_channel="volume_rate", channel="p" on acoustic_born (where channel="p" records scattered pressure).

record_boundary() returns a BoundaryRecording, a named (recording, surface_locations) pair, so the unpacking above and rec.recording/rec.surface_locations are equivalent.

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.

The default amplitude_scale=1.0 injects the traces unchanged. A different value scales the injected amplitudes linearly; it is not a general field reconstruction calibration. Use the same explicit pml_freq in recording and replay runs when comparing them, so the absorbing boundaries are configured consistently.

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.