Skip to content

PropTorch

The PyTorch propagator wrapper that drives any WaveEquation from sweep.equations through a time loop. Accepts an array of keyword arguments inherited from the base class — 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

Base class (full keyword reference)

PropTorch forwards every shared keyword argument to PropBase. The complete list of solver knobs (shape, dh, dt, abcn, pml_type, free_surface, use_ckpt, checkpointing options, boundary-saving options, …) is documented on the base class.

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=True,
    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,
    full_mode="full",
    boundary_saving_config=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, default: True ) –

    Use checkpointing to save memory. Defaults to True.

  • 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'). 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 ) –

    Configuration for boundary saving. Defaults to None, which means boundary saving is disabled. If provided, it should be a dictionary with the following keys: - enabled (bool): Whether to enable boundary saving. If True, the boundary wavefields will be saved and transferred to CPU for checkpointing. Defaults to False. - storage (str): Where to store the boundary wavefields. Options are 'gpu' and 'cpu'. If 'gpu', the boundary wavefields will be stored in GPU memory. If 'cpu', the boundary wavefields will be transferred to CPU memory. Defaults to 'gpu'. - transfer_interval (int): The interval (in time steps) at which to transfer the boundary wavefields to CPU memory if storage is 'cpu'. For example, if transfer_interval is 10, then every 10 time steps the boundary wavefields will be transferred to CPU memory. Defaults to 1. - pinned_memory (bool): Whether to use pinned memory for the boundary wavefields when storage is 'cpu'. Using pinned memory can speed up the transfer between GPU and CPU. Defaults to False. - disk_async_read (bool): Whether to read disk boundary chunks asynchronously during backward when storage is 'disk'. Defaults to False.