Skip to content

PropTorch

The PyTorch propagator wrapper that drives any WaveEquation from sweep.equations through a time loop. It builds a backend (the eager or the compiled impl='c' propagator, both PropBase subclasses) and forwards the shared keyword arguments to it — those are the ones the user mostly tunes.

sweep.propagator.torch.PropTorch 

PropTorch(
    *args,
    backend=None,
    impl=None,
    backend_options=None,
    eager_options=None,
    cuda_options=None,
    memory=None,
    **kwargs
)

Bases: torch.nn.modules.module.Module

PyTorch propagator. impl='eager' is pure torch (CPU or GPU); impl='c' is the prebuilt CUDA core (sweep/lib/cu<major>/libsweep_core.so) driven through the pure-Python ctypes layer sweep.backend.c -- CUDA GPU only, nothing compiles after pip install. impl=None means 'auto': 'c' when sweep.is_torch_binding_available(), the equation declares C_NAME and the propagator's device is CUDA (device=, else the equation's device), else 'eager'; an explicit impl='c' that cannot be honoured falls back to eager with a UserWarning. backend=None inherits the equation's backend. cuda_options= applies to impl='c' only, eager_options= to impl='eager'; memory= (Full() / BoundarySaving() / Ckpt()) is impl-agnostic (default: boundary saving on 'c' -- full storage for equations without compiled boundary saving -- and chunked ckpt on eager).

Shared keywords (PropBase)

PropTorch is not itself a PropBase subclass: it forwards every shared keyword argument (shape, dh, dt, abcn, pml_type, free_surface, use_ckpt, checkpointing options, …) to a PropBase-derived backend, whose constructor is documented below. memory=, impl=, backend= and the backend_options= / eager_options= / cuda_options= blocks are PropTorch's own (see above).

sweep.propagator.base.PropBase 

PropBase(
    equation,
    shape,
    source_type=None,
    receiver_type=None,
    abcn=50,
    free_surface=False,
    topography=None,
    topo_method="auto",
    dh=10.0,
    dt=0.002,
    dev=None,
    device=None,
    use_ckpt=None,
    ckpt_chunks=100,
    ckpt_mode="chunk",
    ckpt_num=0,
    ckpt_storage="gpu",
    ckpt_pinned_memory=None,
    pml_type=None,
    nt=-1,
    B=1,
    allow_growth=True,
    boundary_saving_config=None,
    boundary_buffer=None,
    **kwargs
)

Base class for the Propagator

Parameters:

  • equation (class) –

    The wave equation class from sweep.equations

  • shape (tupel or list) –

    The shape of the model

  • source_type (list, default: None ) –

    List of strings for the source type. Defaults to [].

  • receiver_type (list, default: None ) –

    List of strings for the receiver type. Defaults to [].

  • abcn (int, default: 50 ) –

    The number of layers of absorbing boundary conditions. Defaults to 50.

  • free_surface (bool, default: False ) –

    If the model has a free surface. Defaults to False.

  • topography (array_like, default: None ) –

    Irregular free-surface topography — a 1-D integer array of length nx_phys giving the per-column surface row index in the physical grid (0 = top of physical domain). When given, a free surface is implicit — you do NOT need to set free_surface=True. topo_method selects the discretisation (image method vs APM); the propagator's PML layout is auto-set to match. Defaults to None.

  • topo_method (str, default: 'auto' ) –

    Which surface scheme to use when topography is given. One of:

    • 'auto' (default) — pick 'apm' if the equation declares supports_apm=True (currently :class:Elastic/:class:ElasticAPM), else 'image' (vacuum / Robertsson staircase).
    • 'image' — staircase image method. Acoustic uses Mittet 2002 vacuum cells; Elastic uses Robertsson 1996 odd-parity stress mirror. Sets free_surface=True internally (top PML suppressed).
    • 'apm' — Cao & Chen 2018 parameter-modified method (elastic only). Sets free_surface=False internally (full PML, including top). Best long-time stability on rough staircase topography.

    Ignored when topography is None. Legacy: free_surface=True + topography (without topo_method) still selects image method, with a DeprecationWarning.

  • dh (float or sequence, default: 10.0 ) –

    Grid spacing in model-axis order. For 2D use (dz, dx) and for 3D use (dz, dy, dx). Defaults to 10..

  • dt (float, default: 0.002 ) –

    Time step (seconds). Defaults to 0.002.

  • dev (str, default: None ) –

    Deprecated alias for device. Defaults to None.

  • device (str | device, default: None ) –

    The device to run the simulation on. When None, the equation's device is used. Preferred over dev.

  • use_ckpt (bool | None, default: None ) –

    Legacy request for / exclusion of the checkpointing mode. The gradient-memory mode is a three-way choice (full / boundary / ckpt) resolved by options.resolve_memory_strategy; prefer memory=Full() / BoundarySaving(...) / Ckpt(...). None (default) picks the backend default: 'boundary' for impl='c', 'ckpt' for eager/jax.

  • ckpt_chunks (int, default: 100 ) –

    The number of time steps to chunk for checkpointing. Defaults to 100.

  • ckpt_mode (str, default: 'chunk' ) –

    Checkpointing mode. "chunk" stores periodic checkpoints and replays each chunk, while "recursive" stores a fixed number of checkpoints and recursively recomputes intermediate states. Defaults to "chunk".

  • ckpt_num (int, default: 0 ) –

    Number of persistent checkpoints to save when ckpt_mode="recursive". Defaults to 0.

  • ckpt_storage (str, default: 'gpu' ) –

    Store CUDA checkpoints on "gpu" or "cpu". CPU storage uses host memory to reduce device-memory pressure.

  • ckpt_pinned_memory (bool, default: None ) –

    Use pinned host memory when ckpt_storage="cpu". Defaults to True for CPU checkpoint storage.

  • pml_type (str, default: None ) –

    The type of PML to use. You almost never need to set this — leave it None and the propagator falls back to equation.default_pml_type, which is the only CPML formulation each equation ships. The kwarg exists for advanced experiments (e.g. Acoustic1st accepts 'spml' in addition to its default 'cpmls'). A value outside equation.supported_pml raises ValueError. Possible string values across the codebase: 'cpmlr', 'cpmls', 'spml'. Defaults to None.

  • nt (int, default: -1 ) –

    The number of time steps. Defaults to -1, which means it will be determined by the length of the source time function.

  • B (int, default: 1 ) –

    The batch size for the simulation. Defaults to 1.

  • allow_growth (bool, default: True ) –

    Whether to allow GPU memory growth. Defaults to True.

  • boundary_saving_config (dict, default: None ) –

    Legacy dict spelling of memory=BoundarySaving(...), which is what to pass to PropTorch (it warns on this dict). The keys mirror the BoundarySaving fields -- enabled, storage ('gpu', 'cpu' or 'disk'), transfer_interval, pinned_memory, ring_buffers, disk_dir, disk_async_read, storage_dtype, tail_steps -- with the same per-storage defaults (storage='cpu': transfer_interval=64, pinned_memory=True). Defaults to None: the backend's default strategy, which on impl='c' is boundary saving with a GPU ring (full storage for equations without compiled boundary saving).

  • boundary_buffer (int, default: None ) –

    Number of sigma=0 cells between the physical box and the PML ramp on every absorbing face. None (default) gives the equation the buffer its boundary-saving reconstruction needs (none for pointwise imaging), under every memory strategy, so all of them solve the same grid. An explicit value is honoured, but boundary saving refuses one below the need, and a non-zero buffer cannot be combined with topography= yet.