Source code for pylops.signalprocessing.shift

__all__ = ["Shift"]

from typing import Tuple, Union

import numpy as np
import numpy.typing as npt
from numpy.core.multiarray import normalize_axis_index

from pylops.basicoperators import Diagonal
from pylops.signalprocessing import FFT
from pylops.utils._internal import _value_or_sized_to_array
from pylops.utils.typing import DTypeLike


[docs]def Shift( dims: Tuple, shift: Union[float, npt.ArrayLike], axis: int = -1, nfft: int = None, sampling: float = 1.0, real: bool = False, engine: str = "numpy", dtype: DTypeLike = "complex128", name: str = "S", **kwargs_fftw ): r"""Shift operator Apply fractional shift in the frequency domain along an ``axis`` of a multi-dimensional array of size ``dims``. Parameters ---------- dims : :obj:`tuple` Number of samples for each dimension shift : :obj:`float` or :obj:`numpy.ndarray` Fractional shift to apply in the same unit as ``sampling``. For multi-dimensional inputs, this can be a scalar to apply to every trace along the chosen axis or an array of shifts to be applied to each trace. axis : :obj:`int`, optional .. versionadded:: 2.0.0 Axis along which shift is applied nfft : :obj:`int`, optional Number of samples in Fourier Transform (same as input if ``nfft=None``) sampling : :obj:`float`, optional Sampling step :math:`\Delta t`. real : :obj:`bool`, optional Model to which fft is applied has real numbers (``True``) or not (``False``). Used to enforce that the output of adjoint of a real model is real. engine : :obj:`str`, optional Engine used for fft computation (``numpy``, ``scipy``, or ``fftw``). Choose ``numpy`` when working with CuPy arrays. dtype : :obj:`str`, optional Type of elements in input array. name : :obj:`str`, optional .. versionadded:: 2.0.0 Name of operator (to be used by :func:`pylops.utils.describe.describe`) **kwargs_fftw Arbitrary keyword arguments for :py:class:`pyfftw.FTTW` Attributes ---------- shape : :obj:`tuple` Operator shape explicit : :obj:`bool` Operator contains a matrix that can be solved explicitly (``True``) or not (``False``) Raises ------ ValueError If ``dims`` is provided and ``axis`` is bigger than ``len(dims)`` NotImplementedError If ``engine`` is neither ``numpy``, ``scipy``, nor ``fftw`` Notes ----- The Shift operator applies the forward Fourier transform, an element-wise complex scaling, and inverse fourier transform .. math:: \mathbf{y}= \mathbf{F}^{-1} \mathbf{S} \mathbf{F} \mathbf{x} Here :math:`\mathbf{S}` is a diagonal operator that scales the Fourier transformed input by :math:`e^{-j2\pi f t_S}`, where :math:`t_S` is the chosen ``shift``. """ Fop = FFT( dims, axis=axis, nfft=nfft, sampling=sampling, real=real, engine=engine, dtype=dtype, **kwargs_fftw ) if isinstance(dims, int): dimsdiag = None else: dimsdiag = list(dims) dimsdiag[axis] = len(Fop.f) shift = _value_or_sized_to_array(shift) if shift.size == 1: shift = np.exp(-1j * 2 * np.pi * Fop.f * shift).astype(Fop.cdtype) Sop = Diagonal(shift, dims=dimsdiag, axis=axis, dtype=Fop.cdtype) else: # add dimensions to shift to match dimensions of model and data axis = normalize_axis_index(axis, len(dims)) fdims = np.ones(shift.ndim + 1, dtype=int) fdims[axis] = Fop.f.size f = Fop.f.reshape(fdims) sdims = np.ones(shift.ndim + 1, dtype=int) sdims[:axis] = shift.shape[:axis] sdims[axis + 1 :] = shift.shape[axis:] shift = np.exp(-1j * 2 * np.pi * f * shift.reshape(sdims)).astype(Fop.cdtype) Sop = Diagonal(shift, dtype=Fop.cdtype) Op = Fop.H * Sop * Fop Op.dims = Op.dimsd = Fop.dims # force dtype to that of input (FFT always upcasts it to complex) Op.dtype = dtype Op.name = name return Op