tide-GPR 0.0.9__py3-none-manylinux_2_28_x86_64.whl

tide/padding.py

@@ -0,0 +1,139 @@
+from typing import Union
+
+import torch
+
+
+def reverse_pad(pad: list[int]) -> list[int]:
+    """Reverse the padding order for use with torch.nn.functional.pad.
+
+    PyTorch's pad function expects padding in reverse order (last dim first).
+    This function converts [y0, y1, x0, x1] to [x0, x1, y0, y1].
+
+    Args:
+        pad: Padding values in [y0, y1, x0, x1] format.
+
+    Returns:
+        Padding values in PyTorch format [x0, x1, y0, y1].
+    """
+    # For 2D: [y0, y1, x0, x1] -> [x0, x1, y0, y1]
+    result = []
+    for i in range(len(pad) // 2 - 1, -1, -1):
+        result.extend([pad[i * 2], pad[i * 2 + 1]])
+    return result
+
+
+def create_or_pad(
+    tensor: torch.Tensor,
+    pad: Union[int, list[int]],
+    device: torch.device,
+    dtype: torch.dtype,
+    size: tuple[int, ...],
+    mode: str = "constant",
+) -> torch.Tensor:
+    """Creates a zero tensor of specified size or pads an existing tensor.
+
+    If the input tensor is empty (numel == 0), a new zero tensor with the
+    given size is created. Otherwise, the tensor is padded according to
+    the specified mode.
+
+    This is a unified padding function that supports:
+    - Zero padding (mode='constant') for wavefields
+    - Replicate padding (mode='replicate') for models
+
+    Args:
+        tensor: The input tensor to be created or padded.
+        pad: The padding to apply. Can be an integer (for uniform padding)
+            or a list of integers [y0, y1, x0, x1] for per-side padding.
+        device: The PyTorch device for the tensor.
+        dtype: The PyTorch data type for the tensor.
+        size: The desired size of the tensor if it needs to be created.
+        mode: Padding mode ('constant', 'replicate', 'reflect', 'circular').
+            Default is 'constant' (zero padding)
+
+    Returns:
+        The created or padded tensor.
+
+    Example:
+        >>> # Create a zero tensor of size [2, 110, 110] (batch=2, with padding)
+        >>> wf = create_or_pad(torch.empty(0), 5, device, dtype, (2, 110, 110))
+        >>>
+        >>> # Pad a wavefield with zeros [2, 100, 100] -> [2, 110, 110]
+        >>> wf_padded = create_or_pad(wf, [5, 5, 5, 5], device, dtype, (2, 110, 110))
+        >>>
+        >>> # Pad a model with replicate mode [100, 100] -> [110, 110]
+        >>> eps_padded = create_or_pad(eps, [5, 5, 5, 5], device, dtype, (110, 110), mode='replicate')
+    """
+    if isinstance(pad, int):
+        # Convert single int to [pad, pad, pad, pad, ...] for each spatial dimension
+        # size includes batch dimension if len > 2, so spatial ndim = len(size) - 1 or len(size)
+        ndim = len(size) - 1 if len(size) > 2 else len(size)
+        pad = [pad] * ndim * 2
+
+    if tensor.numel() == 0:
+        return torch.zeros(size, device=device, dtype=dtype)
+
+    if max(pad) == 0:
+        return tensor.clone()
+
+    # Reverse padding for PyTorch's pad function
+    reversed_pad = reverse_pad(pad)
+
+    # For non-constant padding modes (replicate, reflect, circular),
+    # PyTorch requires:
+    # - 2D spatial padding: 3D or 4D input
+    # - 3D spatial padding: 4D or 5D input
+    original_ndim = tensor.ndim
+    needs_unsqueeze = original_ndim in {2, 3} and mode != "constant"
+
+    if needs_unsqueeze:
+        tensor = tensor.unsqueeze(0)
+
+    result = torch.nn.functional.pad(tensor, reversed_pad, mode=mode)
+
+    if needs_unsqueeze:
+        result = result.squeeze(0)
+
+    # PyTorch's autograd system automatically tracks gradients through operations.
+    # Explicitly calling requires_grad_() is incompatible with torch.func transforms.
+    # Simply return the result; gradient tracking is handled automatically.
+    return result
+
+
+def zero_interior(
+    tensor: torch.Tensor,
+    fd_pad: Union[int, list[int]],
+    pml_width: list[int],
+    dim: int,
+) -> torch.Tensor:
+    """Zero out the interior region of a tensor (keeping only PML regions).
+
+    This is used for CPML auxiliary variables which should only be non-zero
+    in the PML regions. Setting the interior to zero allows the propagator
+    to skip unnecessary PML calculations in those regions.
+
+    Args:
+        tensor: The input tensor with shape [batch, ny, nx].
+        fd_pad: Finite difference padding. Can be an int or list [y0, y1, x0, x1].
+        pml_width: The width of PML regions [top, bottom, left, right].
+        dim: The spatial dimension to zero (0 for y, 1 for x).
+
+    Returns:
+        The tensor with interior region zeroed out.
+    """
+    shape = tensor.shape[1:]  # Spatial dimensions (without batch)
+    ndim = len(shape)
+
+    if isinstance(fd_pad, int):
+        fd_pad = [fd_pad] * 2 * ndim
+
+    # Calculate interior slice for the specified dimension
+    interior_start = fd_pad[dim * 2] + pml_width[dim * 2]
+    interior_end = shape[dim] - pml_width[dim * 2 + 1] - fd_pad[dim * 2 + 1]
+
+    # Zero out the interior
+    if dim == 0:  # y dimension
+        tensor[:, interior_start:interior_end, :].fill_(0)
+    else:  # x dimension
+        tensor[:, :, interior_start:interior_end].fill_(0)
+
+    return tensor

tide/resampling.py

@@ -0,0 +1,246 @@
+"""Signal resampling utilities for CFL handling."""
+
+import math
+
+import torch
+
+
+def cosine_taper_end(signal: torch.Tensor, taper_len: int) -> torch.Tensor:
+    """Apply a cosine taper to the end of the signal in the last dimension.
+
+    Args:
+        signal: Input tensor to taper.
+        taper_len: Number of samples to taper at the end.
+
+    Returns:
+        Tapered signal.
+    """
+    if taper_len <= 0 or signal.shape[-1] <= taper_len:
+        return signal
+
+    # Create taper: 1 -> 0 over taper_len samples
+    taper = 0.5 * (
+        1 + torch.cos(torch.linspace(0, math.pi, taper_len, device=signal.device))
+    )
+    # Apply taper to the last taper_len elements
+    signal = signal.clone()
+    signal[..., -taper_len:] = signal[..., -taper_len:] * taper
+    return signal
+
+
+def zero_last_element_of_final_dimension(signal: torch.Tensor) -> torch.Tensor:
+    """Zero the last element of the final dimension (Nyquist frequency).
+
+    This is used to avoid aliasing when resampling signals in the frequency domain.
+
+    Args:
+        signal: Input tensor.
+
+    Returns:
+        Signal with last element of final dimension set to zero.
+    """
+    signal = signal.clone()
+    signal[..., -1] = 0
+    return signal
+
+
+def upsample(
+    signal: torch.Tensor,
+    step_ratio: int,
+    freq_taper_frac: float = 0.0,
+    time_pad_frac: float = 0.0,
+    time_taper: bool = False,
+) -> torch.Tensor:
+    """Upsample the final dimension of a tensor by a factor.
+
+    Low-pass upsampling is used to produce an upsampled signal without
+    introducing higher frequencies than were present in the input.
+
+    This is typically used when the CFL condition requires a smaller internal
+    time step than the user-provided time step.
+
+    Args:
+        signal: Tensor to upsample (time should be the last dimension).
+        step_ratio: Integer factor by which to upsample.
+        freq_taper_frac: Fraction of frequency spectrum end to taper (0.0-1.0).
+            Helps reduce ringing artifacts.
+        time_pad_frac: Fraction of signal length for zero padding (0.0-1.0).
+            Helps reduce wraparound artifacts.
+        time_taper: Whether to apply a Hann window in time.
+            Useful for correctness tests to ensure signals taper to zero.
+
+    Returns:
+        Upsampled signal.
+
+    Example:
+        >>> # Source with 100 time samples, need 3x internal steps for CFL
+        >>> source = torch.randn(1, 1, 100)
+        >>> source_upsampled = upsample(source, step_ratio=3)
+        >>> print(source_upsampled.shape)  # [1, 1, 300]
+    """
+    if signal.numel() == 0 or step_ratio == 1:
+        return signal
+
+    # Optional zero padding to reduce wraparound artifacts
+    n_time_pad = int(time_pad_frac * signal.shape[-1]) if time_pad_frac > 0.0 else 0
+    if n_time_pad > 0:
+        signal = torch.nn.functional.pad(signal, (0, n_time_pad))
+
+    nt = signal.shape[-1]
+    up_nt = nt * step_ratio
+
+    # Transform to frequency domain
+    signal_f = torch.fft.rfft(signal, norm="ortho") * math.sqrt(step_ratio)
+
+    # Apply frequency taper or zero Nyquist
+    if freq_taper_frac > 0.0:
+        freq_taper_len = int(freq_taper_frac * signal_f.shape[-1])
+        signal_f = cosine_taper_end(signal_f, freq_taper_len)
+    elif signal_f.shape[-1] > 1:
+        signal_f = zero_last_element_of_final_dimension(signal_f)
+
+    # Zero-pad in frequency domain for upsampling
+    pad_len = up_nt // 2 + 1 - signal_f.shape[-1]
+    if pad_len > 0:
+        signal_f = torch.nn.functional.pad(signal_f, (0, pad_len))
+
+    # Back to time domain
+    signal = torch.fft.irfft(signal_f, n=up_nt, norm="ortho")
+
+    # Remove padding
+    if n_time_pad > 0:
+        signal = signal[..., : signal.shape[-1] - n_time_pad * step_ratio]
+
+    # Optional time taper
+    if time_taper:
+        signal = signal * torch.hann_window(
+            signal.shape[-1],
+            periodic=False,
+            device=signal.device,
+        )
+
+    return signal
+
+
+def downsample(
+    signal: torch.Tensor,
+    step_ratio: int,
+    freq_taper_frac: float = 0.0,
+    time_pad_frac: float = 0.0,
+    time_taper: bool = False,
+    shift: float = 0.0,
+) -> torch.Tensor:
+    """Downsample the final dimension of a tensor by a factor.
+
+    Frequencies higher than or equal to the Nyquist frequency of the
+    downsampled signal will be zeroed before downsampling.
+
+    This is typically used when the internal time step is smaller than the
+    user-provided time step due to CFL requirements.
+
+    Args:
+        signal: Tensor to downsample (time should be the last dimension).
+        step_ratio: Integer factor by which to downsample.
+        freq_taper_frac: Fraction of frequency spectrum end to taper (0.0-1.0).
+            Helps reduce ringing artifacts.
+        time_pad_frac: Fraction of signal length for zero padding (0.0-1.0).
+            Helps reduce wraparound artifacts.
+        time_taper: Whether to apply a Hann window in time.
+            Useful for correctness tests.
+        shift: Amount to shift in time before downsampling (in time samples).
+
+    Returns:
+        Downsampled signal.
+
+    Example:
+        >>> # Receiver data at internal rate, downsample to user rate
+        >>> data = torch.randn(300, 1, 1)  # [nt_internal, shot, receiver]
+        >>> data_ds = downsample(data.movedim(0, -1), step_ratio=3).movedim(-1, 0)
+        >>> print(data_ds.shape)  # [100, 1, 1]
+    """
+    if signal.numel() == 0 or (step_ratio == 1 and shift == 0.0):
+        return signal
+
+    # Optional time taper
+    if time_taper:
+        signal = signal * torch.hann_window(
+            signal.shape[-1],
+            periodic=False,
+            device=signal.device,
+        )
+
+    # Optional zero padding
+    n_time_pad = (
+        int(time_pad_frac * (signal.shape[-1] // step_ratio))
+        if time_pad_frac > 0.0
+        else 0
+    )
+    if n_time_pad > 0:
+        signal = torch.nn.functional.pad(signal, (0, n_time_pad * step_ratio))
+
+    nt = signal.shape[-1]
+    down_nt = nt // step_ratio
+
+    # Transform to frequency domain, keeping only frequencies below new Nyquist
+    signal_f = torch.fft.rfft(signal, norm="ortho")[..., : down_nt // 2 + 1]
+
+    # Apply frequency taper or zero Nyquist
+    if freq_taper_frac > 0.0:
+        freq_taper_len = int(freq_taper_frac * signal_f.shape[-1])
+        signal_f = cosine_taper_end(signal_f, freq_taper_len)
+    elif signal_f.shape[-1] > 1:
+        signal_f = zero_last_element_of_final_dimension(signal_f)
+
+    # Apply time shift in frequency domain
+    if shift != 0.0:
+        freqs = torch.fft.rfftfreq(signal.shape[-1], device=signal.device)[
+            : down_nt // 2 + 1
+        ]
+        signal_f = signal_f * torch.exp(-1j * 2 * math.pi * freqs * shift)
+
+    # Back to time domain
+    signal = torch.fft.irfft(signal_f, n=down_nt, norm="ortho") / math.sqrt(step_ratio)
+
+    # Remove padding
+    if n_time_pad > 0:
+        signal = signal[..., : signal.shape[-1] - n_time_pad]
+
+    return signal
+
+
+def downsample_and_movedim(
+    receiver_amplitudes: torch.Tensor,
+    step_ratio: int,
+    freq_taper_frac: float = 0.0,
+    time_pad_frac: float = 0.0,
+    time_taper: bool = False,
+    shift: float = 0.0,
+) -> torch.Tensor:
+    """Downsample receiver data and move time dimension to last axis.
+
+    Convenience function that combines downsampling with moving the time
+    dimension to the expected output format [shot, receiver, time].
+
+    Args:
+        receiver_amplitudes: Receiver data [time, shot, receiver].
+        step_ratio: Integer factor by which to downsample.
+        freq_taper_frac: Fraction of frequency spectrum to taper.
+        time_pad_frac: Fraction for zero padding.
+        time_taper: Whether to apply Hann window.
+        shift: Time shift before downsampling.
+
+    Returns:
+        Processed receiver data [shot, receiver, time].
+    """
+    if receiver_amplitudes.numel() > 0:
+        # Move time to last dimension for processing
+        receiver_amplitudes = torch.movedim(receiver_amplitudes, 0, -1)
+        receiver_amplitudes = downsample(
+            receiver_amplitudes,
+            step_ratio,
+            freq_taper_frac=freq_taper_frac,
+            time_pad_frac=time_pad_frac,
+            time_taper=time_taper,
+            shift=shift,
+        )
+    return receiver_amplitudes