From tensors to GPR waveforms.

Installation

DeepGPR is published on PyPI. Install the package into the environment that already contains the PyTorch build appropriate for your hardware.

pip install DeepGPR

Requirements

Run a differentiable forward model

The model tensors use relative permittivity eps_r and electrical conductivity sigma. A 2D model is represented with a singleton third spatial dimension.

import torch
import DeepGPR

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dx, dt, nt = 0.02, 3e-11, 500

eps_r = torch.full((64, 96, 1), 4.0, device=device, requires_grad=True)
sigma = torch.zeros_like(eps_r)
source_location = torch.tensor([[[16, 12, 0]]], device=device)
receiver_location = torch.tensor([[[48, 12, 0]]], device=device)

source = DeepGPR.wavelet.ricker(
    2e8, nt, dt, 5e-9, device=device
).reshape(1, nt, 1)

result = DeepGPR.compute(
    device=device,
    dx=dx,
    dt=dt,
    source_amplitudes=source,
    source_location=source_location,
    receiver_location=receiver_location,
    eps_r=eps_r,
    sigma=sigma,
)

receiver_data = result[-1]
receiver_data.square().sum().backward()
print(receiver_data.shape, eps_r.grad.shape)

The final tuple element is receiver data with shape (nstep, nt, nrx). Calling backward() returns material gradients with exactly the input model shape, including air. PML is outside that shape.

Device and GPU configuration

Pass either a torch.device or a device string. DeepGPR selects the native backend for that device type.

# Automatic selection
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Explicit choices
cpu_device = torch.device("cpu")
gpu_device = torch.device("cuda:0")

# Optional CUDA memory control for saved model-gradient wavefields
result = DeepGPR.compute(
    device=gpu_device,
    # model, source, and acquisition arguments...
    wavefield_storage_dtype=torch.float16,
    use_async_offload=True,
)

Forward modelling

DeepGPR.compute advances the 2D or 3D Maxwell system with FDTD, applies CPML absorbing boundaries, records the requested electric-field component, and preserves the final solver state for continuation or differentiation.

Gradient modes

mode=2 is the Ez-only material-gradient path for 2D Ez-TM modelling. mode=3 uses Ex, Ey, and Ez contributions and is required for 3D or other electric-field component model gradients.

Main API

This table covers the parameters most often needed for modelling. The complete signature and constraints remain in the repository README.

Physical model and automatic PML

Supply only air and the target region in eps_r, sigma, and optional mu_r. Each call replicates current edge values outward by the configured PML thickness. Acquisition indices start at zero in the input model; the solver shifts them internally. Remove old manual PML padding and its coordinate offsets when migrating.

For input (nx, ny, nz) and face widths [px0, px1, py0, py1, pz0, pz1], the solver uses Nx = nx + px0 + px1, and likewise for Y and Z. In 2D, Nz = 1. Update the physical model with your FWI optimizer; PML is regenerated on the next call. Apply an air mask if air must remain fixed.

Backward holds PML materials and coefficient averages fixed and crops material gradients to the model. It does not sum PML sensitivities into edge cells. For finite-difference checks, freeze the boundary values or the PML explicitly. Native CPU/CUDA libraries must be rebuilt with the deepgpr_supports_external_pml capability.

Open the full interface documentation on GitHub ↗

Outputs

1. E_saved — saved electric-field history for model-gradient work and diagnostics.
2. (Ex, Ey, Ez) — electric-field state at the final time step.
3. (Hx, Hy, Hz) — magnetic-field state at the final time step.
4. PML_Tuple — 24 CPML auxiliary state tensors.
5. receiver_amplitudes — recorded data with shape (nstep, nt, nrx).

Histories (including compressed and saved-file histories) retain the full (Nx, Ny, Nz) solver grid. E/H tensors have shape (nstep, Nx+1, Ny+1, Nz+1), including the Yee halo. For physical history plots, slice each spatial axis from its low PML width for the input model length.

Continuation and checkpointing

checkpoint_initial_field takes the physical model and allocates full-grid E/H/PML states. Pass all returned states back without cropping or padding, using the same geometry, PML settings, and shot ordering. The material model stays unextended on every segment.

The native solver advances supplied states in place. Clone all 30 boundary tensors inside a PyTorch checkpointed function before propagation, and return their final states with the receiver traces. Keep materials fixed throughout the trajectory and backward pass; start the next FWI iteration from the appropriate initial state after optimization.

Open the temporal checkpoint example ↗

Examples and source

The repository contains notebooks for forward modelling, 2D FWI, and 3D FWI, plus numerical verification and runtime benchmark tools.