API reference¶

Generated from the source docstrings. Everything here is also importable from the package root, e.g. from sweep_nn import VelocityINR.

Velocity networks¶

sweep_nn.velocity_inr.VelocityINR ¶

VelocityINR(
    base_velocity: Tensor,
    *,
    vp_mean: float = 0.0,
    vp_std: float = 50.0,
    hidden_features: int = 64,
    hidden_layers: int = 3,
    first_omega0: float = 30.0,
    hidden_omega0: float = 30.0,
    use_bias: bool = False,
    use_hash_encoding: bool = True,
    hash_levels: int = 16,
    hash_features_per_level: int = 2,
    hash_log2_size: int = 15,
    hash_base_resolution: int | list[int] = 4,
    hash_finest_resolution: int | list[int] = 512,
    hash_c2f: bool = False,
    hash_c2f_base_levels: int = 2,
    hash_c2f_ramp: str = "cosine",
    hash_growing: bool = False,
    hash_backend: str = "pytorch",
    use_fourier_encoding: bool = False,
    fourier_levels: int = 6,
    fourier_include_input: bool = True,
    direct_velocity: bool = False,
    coord_min: float = 0.0,
    coord_max: float = 1.0,
    bounds: Tuple[float, float] | None = None,
    water_mask: Tensor | None = None,
    water_vp: float = 1500.0,
    lateral_downsample: Tuple[int, int] | int = 1,
    compile_render: bool = False
)

Bases: torch.nn.modules.module.Module

Hash-encoded SIREN representation of a 2-D or 3-D velocity field.

Parameters:

base_velocity (Tensor) –

Tensor (nz, nx) (2-D) or (nz, ny, nx) (3-D). Treated as a non-trainable buffer — the network learns the perturbation on top (delta mode) or replaces it entirely (direct_velocity=True).
vp_mean (float, default: 0.0 ) –

Added to the scaled network output: velocity = base + (mlp_out * vp_std + vp_mean).
vp_std (float, default: 50.0 ) –

Scales the network output, so it sets the typical perturbation in m/s: the network's natural output range is ~[-1, 1], so vp_std=50 gives ±50 m/s typical updates.
use_hash_encoding (bool, default: True ) –

Wrap coordinates with :class:MultiResHashGrid before the SIREN head. Required for high-frequency velocity detail; without it the SIREN must do all the spatial-frequency work alone (slower, less expressive).
hash_levels (int, default: 16 ) –

Number of hash-grid levels. See :class:MultiResHashGrid.
hash_features_per_level (int, default: 2 ) –

Features stored per level.
hash_log2_size (int, default: 15 ) –

log2 of the hash-table size of each level.
hash_base_resolution (int | list[int], default: 4 ) –

Grid resolution of the coarsest level (one int, or one per axis).
hash_finest_resolution (int | list[int], default: 512 ) –

Grid resolution of the finest level (one int, or one per axis).
hash_c2f (bool, default: False ) –

If True, build a :class:CoarseToFineHashGrid instead: only the hash_c2f_base_levels coarsest levels are open at init, and the training loop drives the unfreeze schedule via self.encoder.set_progress(...). Same hash hyperparameters.
hash_c2f_base_levels (int, default: 2 ) –

Levels open at init when hash_c2f=True.
hash_c2f_ramp (str, default: 'cosine' ) –

Shape of the unfreeze ramp when hash_c2f=True.
hidden_features (int, default: 64 ) –

SIREN width. See :class:SirenMLP.
hidden_layers (int, default: 3 ) –

Number of hidden sine layers after the first one.
first_omega0 (float, default: 30.0 ) –

Sine frequency of the first layer.
hidden_omega0 (float, default: 30.0 ) –

Sine frequency of the hidden layers.
direct_velocity (bool, default: False ) –

If True, ignore the base and return mlp_out * vp_std + vp_mean. Useful for from-scratch reconstruction.
coord_min (float, default: 0.0 ) –

Lower end of the coordinate range. Coordinates span [coord_min, coord_max), which must match the convention the hash encoder expects.
coord_max (float, default: 1.0 ) –

Upper end of the coordinate range (exclusive).
bounds (Tuple[float, float] | None, default: None ) –

Optional (vp_min, vp_max) clamp applied at render time. Set to None (default) to let the caller clamp.
water_mask (Tensor | None, default: None ) –

Optional boolean tensor with the SAME shape as base_velocity. Voxels where True are pinned to water_vp at render time — the SIREN's output for those cells is ignored entirely. This is the right tool for "I know the water column is 1500 m/s, do NOT let SIREN init noise contaminate it" cases: the SIREN at init has ~std=0.08 raw output, so with vp_std=500 the water layer would otherwise sit at 1500±40 m/s of garbage from epoch 0. Because the rendered output doesn't depend on SIREN params at masked voxels, gradients there are exactly zero — equivalent to freezing those cells AND giving SIREN free model capacity to spend on the rest of the model.
water_vp (float, default: 1500.0 ) –

Velocity in m/s used at water-mask voxels. Default 1500.0.

Notes

Multi-stage FWI usage: net = VelocityINR(init_vp_75m, vp_std=50.0, ...) # ... train at stage 0 ... net.update_base_velocity(init_vp_37p5m) # resample base, keep params # ... train at stage 1 ...

update_base_velocity ¶

update_base_velocity(new_base: Tensor, water_mask: Tensor | None = None) -> None

Swap in a new base (possibly different shape); rebuild coords.

When the new base has a different shape than the old, any previously installed water_mask becomes stale. Pass a water_mask matched to new_base.shape to refresh; leave None to clear the existing mask (and warn loudly that a stale mask was dropped).

forward ¶

forward() -> torch.Tensor

Render the velocity model on the current base grid.

render_shape ¶

render_shape(shape: Tuple[int, ...]) -> torch.Tensor

Render at an arbitrary grid shape (bilinear/trilinear-resample the base).

render_window ¶

render_window(*bounds: int) -> torch.Tensor

Render a rectangular window on the base grid.

Bounds layout: 2-D: (z0, z1, x0, x1) 3-D: (z0, z1, y0, y1, x0, x1)

backward_velocity_gradient ¶

backward_velocity_gradient(velocity_grad: Tensor, *, chunk_rows: int = 64) -> None

Back-propagate a full-grid velocity gradient row-by-row.

Use this when the full-grid render() + .backward() would OOM: e.g. a 3-D model with millions of voxels and a deep SIREN. The method renders one slab of rows at a time, calls backward() on that slab with the corresponding slice of velocity_grad, and accumulates gradients onto the trainable parameters without ever holding the full graph in memory.

Parameters:

velocity_grad (Tensor) –

Tensor with the same shape as :attr:base_velocity.
chunk_rows (int, default: 64 ) –

Rows along the slow axis (axis 0) per chunk. Smaller = less peak memory, more overhead. 64 is a sensible default for 2-D; for 3-D models start with 4-8.

sweep_nn.multi_param_inr.MultiParamINR ¶

MultiParamINR(
    base_models: Sequence[Tensor],
    *,
    means: Sequence[float],
    stds: Sequence[float],
    bounds: Sequence[Tuple[float, float] | None],
    water_mask: Tensor | None = None,
    water_values: Sequence[float] | None = None,
    hidden_features: int = 64,
    hidden_layers: int = 3,
    first_omega0: float = 30.0,
    hidden_omega0: float = 30.0,
    use_bias: bool = False,
    use_hash_encoding: bool = True,
    hash_levels: int = 16,
    hash_features_per_level: int = 2,
    hash_log2_size: int = 15,
    hash_base_resolution: int | List[int] = 4,
    hash_finest_resolution: int | List[int] = 512,
    hash_c2f: bool = False,
    hash_c2f_base_levels: int = 2,
    hash_c2f_ramp: str = "cosine",
    hash_growing: bool = False,
    hash_backend: str = "pytorch",
    direct_velocity: bool = False,
    coord_min: float = 0.0,
    coord_max: float = 1.0
)

Bases: torch.nn.modules.module.Module

Shared-trunk INR with one output channel per FWI model parameter.

Parameters:

base_models (Sequence[Tensor]) –

List of n_params initial-model tensors, all the same shape ((nz, nx) or (nz, ny, nx)). base_models[0] is the primary parameter (vp). Each is kept as a (non-persistent) buffer; the channel renders base_i + delta_i.
means (Sequence[float]) –

Per-channel offset in delta_i = raw_i * stds[i] + means[i]. Length n_params.
stds (Sequence[float]) –

Per-channel scale in the same affine, length n_params. It is the per-parameter "scale" (its effective learning rate is lr * stds[i]); pick it per parameter's magnitude (vp ~ 500, impedance z ~ 2).
bounds (Sequence[Tuple[float, float] | None]) –

Per-channel (min, max) clamp, or None to leave a channel unbounded. Length n_params.
water_mask (Tensor | None, default: None ) –

Optional shared boolean water mask (shape = model shape). Masked voxels of channel i are pinned to water_values[i] (gradient there is 0, so the SIREN never spends capacity on the water column).
water_values (Sequence[float] | None, default: None ) –

Per-channel value at the masked voxels, length n_params (e.g. [1500.0, 1.5] for vp and Gardner-water impedance z).

Notes

The remaining arguments mirror :class:VelocityINR (shared SIREN + hash encoder hyperparameters).

base_velocity `property` ¶

base_velocity

Primary-parameter (channel-0) base — mirrors VelocityINR.base_velocity.

update_base_models ¶

update_base_models(
    new_bases: Sequence[Tensor],
    water_mask: Tensor | None = None,
    water_values: Sequence[float] | None = None,
) -> None

Swap in new per-channel bases (possibly a different grid) + rebuild coords — multi-parameter analogue of VelocityINR.update_base_velocity.

The shared encoder + SIREN parameters are KEPT (the whole point of carrying the net across multiscale stages); only base_stack / coords / water mask are refreshed. new_bases is one tensor per channel (same order as construction), each resampled to the new grid. Pass a water_mask matched to the new shape to refresh the water pin (a stale mask on a shape change is dropped with a warning).

render_all ¶

render_all(chunk_rows: int | None = None) -> torch.Tensor

Render every parameter -> (n_params, *shape) (channel 0 = vp).

chunk_rows bounds peak memory by rendering z-slabs (bit-identical, pointwise). Under no_grad the caller controls the graph.

render_param ¶

render_param(i: int, chunk_rows: int | None = None) -> torch.Tensor

Render parameter i alone -> *shape.

render ¶

render(chunk_rows: int | None = None) -> torch.Tensor

Render the PRIMARY parameter (vp, channel 0) — VelocityINR-compatible.

render_window ¶

render_window(*bounds: int) -> torch.Tensor

Render a rectangular window over ALL channels -> (n_params, *win).

Bounds: 2-D (z0, z1, x0, x1), 3-D (z0, z1, y0, y1, x0, x1) — the N-channel analogue of VelocityINR.render_window, used by the DD tile render. coords/base_stack may live on CPU for huge global grids; only this window's slices move to the compute device.

backward_gradients ¶

backward_gradients(grads: Sequence[Tensor], *, chunk_rows: int = 64) -> None

Back-propagate a per-parameter gradient list onto the shared trunk.

grads[i] is dL/d(field_i) (same shape as the model). Renders one z-slab of ALL channels at a time and backwards them together, so the shared trunk accumulates every parameter's contribution while peak memory stays O(chunk_rows) — the multi-parameter analogue of VelocityINR.backward_velocity_gradient.

Encoders and SIREN¶

sweep_nn.hash_encoding.MultiResHashGrid ¶

MultiResHashGrid(
    dim: int,
    *,
    n_levels: int = 16,
    n_features_per_level: int = 2,
    log2_hashmap_size: int = 15,
    base_resolution: int | List[int] = 2,
    finest_resolution: int | List[int] = 16,
    dtype: dtype = torch.float32,
    backend: str = "pytorch"
)

Bases: torch.nn.modules.module.Module

Anisotropic multi-resolution hash-grid encoder.

forward ¶

forward(pos: Tensor) -> torch.Tensor

Encode normalized coords (..., dim) -> features (..., L*F).

sweep_nn.siren.SirenMLP ¶

SirenMLP(
    in_features: int,
    out_features: int = 1,
    *,
    hidden_features: int = 64,
    hidden_layers: int = 3,
    first_omega0: float = 30.0,
    hidden_omega0: float = 30.0,
    bias: bool = False
)

Bases: torch.nn.modules.module.Module

Generic SIREN MLP: (..., in_features) -> (..., out_features).

Unlike :class:SIREN (which caches a coord grid and ties to an output shape), SirenMLP is a pure functional MLP. It is the building block used by :class:~sweep_nn.velocity_inr.VelocityINR (where the input comes from a hash encoder) and :class:~sweep_nn.wavelet.SirenWavelet.

First layer uses first_omega0; hidden/output layers use hidden_omega0. The final layer is a plain nn.Linear (no sin), SIREN-style.

sweep_nn.siren.SIREN ¶

SIREN(
    out_shape: Tuple[int, ...],
    *,
    hidden_features: int = 128,
    hidden_layers: int = 4,
    w0: float = 30.0,
    vp_min: float = 1500.0,
    vp_max: float = 4500.0,
    squash: str = "tanh"
)

Bases: sweep_nn.reparam.Reparameterizer

SIREN reparameterizer.

Parameters:

out_shape (Tuple[int, ...]) –

(nz, nx) 2-D or (nz, ny, nx) 3-D output grid.
hidden_features (int, default: 128 ) –

Width of the MLP.
hidden_layers (int, default: 4 ) –

Depth of the MLP.
w0 (float, default: 30.0 ) –

Initial-layer frequency. The paper recommends 30 for image-scale tasks; higher values fit higher spatial frequencies at the cost of slower convergence.

sweep_nn.siren.SineLayer ¶

SineLayer(
    in_features: int,
    out_features: int,
    *,
    w0: float = 30.0,
    is_first: bool = False,
    bias: bool = True
)

Bases: torch.nn.modules.module.Module

Linear + sin(w0 * x), with paper-faithful initialization.

First-layer init uses U(-1/in, 1/in); hidden/output layers use U(-sqrt(6/in)/w0, +sqrt(6/in)/w0) (Sitzmann et al. 2020 sec. 3.2).

Priors¶

sweep_nn.priors.TVPrior ¶

TVPrior(
    *,
    order: str = "first",
    x_weight: float = 1.0,
    z_weight: float = 1.0,
    y_weight: float = 1.0,
    velocity_scale_m_s: float = 1000.0
)

Bases: torch.nn.modules.module.Module

Total-variation-style smoothness penalty on a velocity volume.

Computes a derivative-L2 (Sobolev) penalty along the lateral (x and, in 3-D, y) and depth (z) axes. Supports first-order, second-order, or both. Used as a soft regularizer in the FWI loss::

L_total = L_data + weight * TVPrior(...)(velocity)

The velocity tensor is internally normalized by velocity_scale_m_s so the weight stays interpretable across surveys with different vp magnitudes (m/s).

In 2-D the model is (nz, nx); in 3-D it is (nz, ny, nx). The depth axis is always axis 0.

Parameters:

order (str, default: 'first' ) –

"first" (default) — sum of squared first differences. "second" — sum of squared second differences (curvature). "both" / "mixed" — both first and second.
x_weight (float, default: 1.0 ) –

Multiplier for the lateral (x) differences.
z_weight (float, default: 1.0 ) –

Multiplier for the depth (z) differences.
y_weight (float, default: 1.0 ) –

Multiplier for the y differences; ignored on 2-D inputs.
velocity_scale_m_s (float, default: 1000.0 ) –

Normalization for vp before differentiation (m/s). Default 1000.0.

sweep_nn.priors.SeabedFreezeMask ¶

SeabedFreezeMask(
    seabed_depth_m: Tensor | ndarray, *, dz_m: float, buffer_cells: int = 0
)

Mask that zeros gradients above the seabed (water column).

For marine surveys the water-column vp is essentially constant at ~1500 m/s and well known a priori — letting FWI update those cells introduces noise that doesn't help the inversion. Multiplying the gradient by mask = (depth_idx >= floor(seabed_idx)) before the optimizer step keeps the water column frozen at its initial value while letting everything below the seabed update normally.

Parameters:

seabed_depth_m (Tensor | ndarray) –

(ny, nx) (3-D) or (nx,) (2-D) per-trace seabed depths in meters from the model's top edge (z = 0).
dz_m (float) –

Grid spacing along the depth axis (m).
buffer_cells (int, default: 0 ) –

Optional integer pad — keep the first buffer_cells rows below the seabed also frozen. Useful when the wavelet has a non-zero rise time and the seabed reflection straddles a few cells.

Use

Build the mask once, store on the optimizer's device, and call apply_to(grad) (in-place) after backward and before optim.step. The mask broadcasts: a 2-D (nx,) seabed against a 2-D (nz, nx) vp; a 2-D (ny, nx) seabed against a 3-D (nz, ny, nx) vp.

build_mask ¶

build_mask(vp_shape: tuple[int, ...], *, device, dtype) -> torch.Tensor

Return a vp_shape boolean mask broadcast against the seabed map.

apply_to ¶

apply_to(grad: Tensor) -> None

In-place multiply grad by the seabed mask.

sweep_nn.diffusion.prior.DiffusionVelocityPrior ¶

DiffusionVelocityPrior(
    ckpt_path: str,
    device: str = "cuda",
    mode: str = "patch",
    strength: float = 0.3,
    ddim_steps: int = 10,
    patch: int = 64,
    stride: int = 32,
    vmin: float | None = None,
    vmax: float | None = None,
    use_ema: bool = True,
)

red_loss ¶

red_loss(vp: Tensor) -> torch.Tensor

RED term computed in NORMALISED [-1,1] space so its magnitude is O(1) and comparable to a normalised data misfit.

L = 0.5 * mean((x_norm(vp) - D(vp).detach())^2), with x_norm = to_norm(vp) differentiable in vp (linear in-range), so grad_vp carries the correct 2/(vmax-vmin) chain-rule factor automatically.

residual ¶

residual(vp: Tensor) -> torch.Tensor

vp - D(vp) in physical units (the raw prior descent direction).

sds_loss ¶

sds_loss(vp: Tensor, t_lo: float = 0.02, t_hi: float = 0.5) -> torch.Tensor

Surrogate whose grad w.r.t. vp is the SDS prior gradient (score-matching distillation of the diffusion model into whatever produces vp — pixels or an INR). Computed in normalised space; magnitude O(1).

Other¶

sweep_nn.wavelet.SirenWavelet ¶

SirenWavelet(
    nt: int,
    *,
    hidden_features: int = 64,
    hidden_layers: int = 3,
    first_omega0: float = 30.0,
    hidden_omega0: float = 30.0,
    bias: bool = True
)

Bases: torch.nn.modules.module.Module

1-D SIREN parameterization of a time-domain source wavelet.

Parameters:

nt (int) –

Number of time samples in the output wavelet.
hidden_features (int, default: 64 ) –

SIREN width. See :class:SirenMLP.
hidden_layers (int, default: 3 ) –

Number of hidden sine layers after the first one.
first_omega0 (float, default: 30.0 ) –

Sine frequency of the first layer.
hidden_omega0 (float, default: 30.0 ) –

Sine frequency of the hidden layers.
bias (bool, default: True ) –

Per the docstring: keep this True. False is exposed only for completeness (e.g. fitting an explicitly odd wavelet).

sweep_nn.dip.DIPReparam ¶

DIPReparam(
    out_shape: Tuple[int, int],
    *,
    latent_channels: int = 8,
    base_ch: int = 32,
    depth: int = 4,
    vp_min: float = 1500.0,
    vp_max: float = 4500.0,
    squash: str = "tanh"
)

Bases: sweep_nn.reparam.Reparameterizer

Deep Image Prior reparameterizer (2-D).

Parameters:

out_shape (Tuple[int, int]) –

(nz, nx).
latent_channels (int, default: 8 ) –

Number of channels in the fixed input noise.
base_ch (int, default: 32 ) –

U-Net width: channels of the first level.
depth (int, default: 4 ) –

U-Net depth: number of levels.

sweep_nn.reparam.Reparameterizer ¶

Reparameterizer(
    out_shape: Tuple[int, ...],
    *,
    vp_min: float = 1500.0,
    vp_max: float = 4500.0,
    squash: str = "tanh"
)

Bases: torch.nn.modules.module.Module

Abstract base: net() -> vp_tensor in physical units.

Subclasses implement :meth:_raw_forward, which must return a tensor with shape out_shape whose values lie roughly in [-1, 1] (or any range — see :meth:_to_physical).

Parameters:

out_shape (Tuple[int, ...]) –

Spatial shape of the velocity model, e.g. (nz, nx) for 2-D.
vp_min (float, default: 1500.0 ) –

Lower end of the physical-unit range.
vp_max (float, default: 4500.0 ) –

Upper end. The raw output is rescaled and (by default, when squash="tanh") tanh-squashed to live in [vp_min, vp_max].
squash (str, default: 'tanh' ) –

How the raw output is mapped to [vp_min, vp_max]: - "tanh" — vp_min + (vp_max-vp_min) * 0.5*(tanh(raw)+1) - "sigmoid" — vp_min + (vp_max-vp_min) * sigmoid(raw) - "linear" — no squashing; raw values used directly (loss must enforce bounds).

API reference¶

Velocity networks¶

sweep_nn.velocity_inr.VelocityINR ¶

update_base_velocity ¶

forward ¶

render_shape ¶

render_window ¶

backward_velocity_gradient ¶

sweep_nn.multi_param_inr.MultiParamINR ¶

base_velocity property ¶

update_base_models ¶

render_all ¶

render_param ¶

render ¶

render_window ¶

backward_gradients ¶

Encoders and SIREN¶

sweep_nn.hash_encoding.MultiResHashGrid ¶

forward ¶

sweep_nn.siren.SirenMLP ¶

sweep_nn.siren.SIREN ¶

sweep_nn.siren.SineLayer ¶

Priors¶

sweep_nn.priors.TVPrior ¶

sweep_nn.priors.SeabedFreezeMask ¶

build_mask ¶

apply_to ¶

sweep_nn.diffusion.prior.DiffusionVelocityPrior ¶

red_loss ¶

residual ¶

sds_loss ¶

Other¶

sweep_nn.wavelet.SirenWavelet ¶

sweep_nn.dip.DIPReparam ¶

sweep_nn.reparam.Reparameterizer ¶

base_velocity `property` ¶