Source code for pylops.waveeqprocessing.kirchhoff

__all__ = ["Kirchhoff"]


import logging
import os
import warnings
from typing import Optional, Tuple, Union

import numpy as np

from pylops import LinearOperator
from pylops.signalprocessing import Convolve1D
from pylops.utils import deps
from pylops.utils._internal import _value_or_sized_to_array
from pylops.utils.decorators import reshaped
from pylops.utils.tapers import taper
from pylops.utils.typing import DTypeLike, NDArray

skfmm_message = deps.skfmm_import("the kirchhoff module")
jit_message = deps.numba_import("the kirchhoff module")

if skfmm_message is None:
    import skfmm

if jit_message is None:
    from numba import jit, prange

    # detect whether to use parallel or not
    numba_threads = int(os.getenv("NUMBA_NUM_THREADS", "1"))
    parallel = True if numba_threads != 1 else False
else:
    prange = range

logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING)


[docs]class Kirchhoff(LinearOperator): r"""Kirchhoff Demigration operator. Kirchhoff-based demigration/migration operator. Uses a high-frequency approximation of Green's function propagators based on ``trav``. Parameters ---------- z : :obj:`numpy.ndarray` Depth axis x : :obj:`numpy.ndarray` Spatial axis t : :obj:`numpy.ndarray` Time axis for data srcs : :obj:`numpy.ndarray` Sources in array of size :math:`\lbrack 2 (3) \times n_s \rbrack` The first axis should be ordered as (``y``,) ``x``, ``z``. recs : :obj:`numpy.ndarray` Receivers in array of size :math:`\lbrack 2 (3) \times n_r \rbrack` The first axis should be ordered as (``y``,) ``x``, ``z``. vel : :obj:`numpy.ndarray` or :obj:`float` Velocity model of size :math:`\lbrack (n_y\,\times)\; n_x \times n_z \rbrack` (or constant) wav : :obj:`numpy.ndarray` Wavelet. wavcenter : :obj:`int` Index of wavelet center y : :obj:`numpy.ndarray` Additional spatial axis (for 3-dimensional problems) mode : :obj:`str`, optional Computation mode (``analytic``, ``eikonal`` or ``byot``, see Notes for more details) wavfilter : :obj:`bool`, optional .. versionadded:: 2.0.0 Apply wavelet filter (``True``) or not (``False``) dynamic : :obj:`bool`, optional .. versionadded:: 2.0.0 Include dynamic weights in computations (``True``) or not (``False``). Note that when ``mode=byot``, the user is required to provide such weights in ``amp``. trav : :obj:`numpy.ndarray` or :obj:`tuple`, optional Traveltime table of size :math:`\lbrack (n_y) n_x n_z \times n_s n_r \rbrack` or pair of traveltime tables of size :math:`\lbrack (n_y) n_x n_z \times n_s \rbrack` and :math:`\lbrack (n_y) n_x n_z \times n_r \rbrack` (to be provided if ``mode='byot'``). Note that the latter approach is recommended as less memory demanding than the former. amp : :obj:`numpy.ndarray`, optional .. versionadded:: 2.0.0 Amplitude table of size :math:`\lbrack (n_y) n_x n_z \times n_s n_r \rbrack` or pair of amplitude tables of size :math:`\lbrack (n_y) n_x n_z \times n_s \rbrack` and :math:`\lbrack (n_y) n_x n_z \times n_r \rbrack` (to be provided if ``mode='byot'``). Note that the latter approach is recommended as less memory demanding than the former. aperture : :obj:`float` or :obj:`tuple`, optional .. versionadded:: 2.0.0 Maximum allowed aperture expressed as the ratio of offset over depth. If ``None``, no aperture limitations are introduced. If scalar, a taper from 80% to 100% of aperture is applied. If tuple, apertures below the first value are accepted and those after the second value are rejected. A tapering is implemented for those between such values. angleaperture : :obj:`float` or :obj:`tuple`, optional .. versionadded:: 2.0.0 Maximum allowed angle (either source or receiver side) in degrees. If ``None``, angle aperture limitations are not introduced. See ``aperture`` for implementation details regarding scalar and tuple cases. anglerefl : :obj:`np.ndarray`, optional .. versionadded:: 2.0.0 Angle between the normal of the reflectors and the vertical axis in degrees snell : :obj:`float` or :obj:`tuple`, optional .. versionadded:: 2.0.0 Threshold on Snell's law evaluation. If larger, the source-receiver-image point is discarded. If ``None``, no check on the validity of the Snell's law is performed. See ``aperture`` for implementation details regarding scalar and tuple cases. engine : :obj:`str`, optional Engine used for computations (``numpy`` or ``numba``). 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`) Attributes ---------- shape : :obj:`tuple` Operator shape explicit : :obj:`bool` Operator contains a matrix that can be solved explicitly (``True``) or not (``False``) Raises ------ NotImplementedError If ``mode`` is neither ``analytic``, ``eikonal``, or ``byot`` Notes ----- The Kirchhoff demigration operator synthesizes seismic data given a propagation velocity model :math:`v` and a reflectivity model :math:`m`. In forward mode [1]_, [2]_: .. math:: d(\mathbf{x_r}, \mathbf{x_s}, t) = \widetilde{w}(t) * \int_V G(\mathbf{x_r}, \mathbf{x}, t) m(\mathbf{x}) G(\mathbf{x}, \mathbf{x_s}, t)\,\mathrm{d}\mathbf{x} where :math:`m(\mathbf{x})` represents the reflectivity at every location in the subsurface, :math:`G(\mathbf{x}, \mathbf{x_s}, t)` and :math:`G(\mathbf{x_r}, \mathbf{x}, t)` are the Green's functions from source-to-subsurface-to-receiver and finally :math:`\widetilde{w}(t)` is a filtered version of the wavelet :math:`w(t)` [3]_ (or the wavelet itself when ``wavfilter=False``). In our implementation, the following high-frequency approximation of the Green's functions is adopted: .. math:: G(\mathbf{x_r}, \mathbf{x}, \omega) = a(\mathbf{x_r}, \mathbf{x}) e^{j \omega t(\mathbf{x_r}, \mathbf{x})} where :math:`a(\mathbf{x_r}, \mathbf{x})` is the amplitude and :math:`t(\mathbf{x_r}, \mathbf{x})` is the traveltime. When ``dynamic=False`` the amplitude is disregarded leading to a kinematic-only Kirchhoff operator. .. math:: d(\mathbf{x_r}, \mathbf{x_s}, t) = \tilde{w}(t) * \int_V e^{j \omega (t(\mathbf{x_r}, \mathbf{x}) + t(\mathbf{x}, \mathbf{x_s}))} m(\mathbf{x}) \,\mathrm{d}\mathbf{x} On the other hand, when ``dynamic=True``, the amplitude scaling is defined as :math:`a(\mathbf{x}, \mathbf{y})=\frac{1}{\|\mathbf{x} - \mathbf{y}\|}`, that is, the reciprocal of the distance between the two points, approximating the geometrical spreading of the wavefront. Moreover an angle scaling is included in the modelling operator added as follows: .. math:: d(\mathbf{x_r}, \mathbf{x_s}, t) = \tilde{w}(t) * \int_V a(\mathbf{x}, \mathbf{x_s}) a(\mathbf{x}, \mathbf{x_r}) \frac{|cos \theta_s + cos \theta_r|} {v(\mathbf{x})} e^{j \omega (t(\mathbf{x_r}, \mathbf{x}) + t(\mathbf{x}, \mathbf{x_s}))} m(\mathbf{x}) \,\mathrm{d}\mathbf{x} where :math:`\theta_s` and :math:`\theta_r` are the angles between the source-side and receiver-side rays and the normal to the reflector at the image point (or the vertical axis at the image point when ``reflslope=None``), respectively. Depending on the choice of ``mode`` the traveltime and amplitude of the Green's function will be also computed differently: * ``mode=analytic`` or ``mode=eikonal``: traveltimes, geometrical spreading, and angles are computed for every source-image point-receiver triplets and the Green's functions are implemented from traveltime look-up tables, placing scaled reflectivity values at corresponding source-to-receiver time in the data. * ``byot``: bring your own tables. Traveltime table are provided directly by user using ``trav`` input parameter. Similarly, in this case one can provide their own amplitude scaling ``amp`` (which should include the angle scaling too). Three aperture limitations have been also implemented as defined by: * ``aperture``: the maximum allowed aperture is expressed as the ratio of offset over depth. This aperture limitation avoid including grazing angles whose contributions can introduce aliasing effects. A taper is added at the edges of the aperture; * ``angleaperture``: the maximum allowed angle aperture is expressed as the difference between the incident or emerging angle at every image point and the vertical axis (or the normal to the reflector if ``anglerefl`` is provided. This aperture limitation also avoid including grazing angles whose contributions can introduce aliasing effects. Note that for a homogenous medium and slowly varying heterogenous medium the offset and angle aperture limits may work in the same way; * ``snell``: the maximum allowed snell's angle is expressed as the absolute value of the sum between incident and emerging angles defined as in the ``angleaperture`` case. This aperture limitation is introduced to turn a scattering-based Kirchhoff engine into a reflection-based Kirchhoff engine where each image point is not considered as scatter but as a local horizontal (or straight) reflector. Finally, the adjoint of the demigration operator is a *migration* operator which projects data in the model domain creating an image of the subsurface reflectivity. .. [1] Bleistein, N., Cohen, J.K., and Stockwell, J.W.. "Mathematics of Multidimensional Seismic Imaging, Migration and Inversion", 2000. .. [2] Santos, L.T., Schleicher, J., Tygel, M., and Hubral, P. "Seismic modeling by demigration", Geophysics, 65(4), pp. 1281-1289, 2000. .. [3] Safron, L. "Multicomponent least-squares Kirchhoff depth migration", MSc Thesis, 2018. """ def __init__( self, z: NDArray, x: NDArray, t: NDArray, srcs: NDArray, recs: NDArray, vel: NDArray, wav: NDArray, wavcenter: int, y: Optional[NDArray] = None, mode: str = "eikonal", wavfilter: bool = False, dynamic: bool = False, trav: Optional[NDArray] = None, amp: Optional[NDArray] = None, aperture: Optional[Tuple[float, float]] = None, angleaperture: Union[float, Tuple[float, float]] = 90.0, anglerefl: Optional[NDArray] = None, snell: Optional[Tuple[float, float]] = None, engine: str = "numpy", dtype: DTypeLike = "float64", name: str = "K", ) -> None: warnings.warn( "A new implementation of Kirchhoff is provided in v2.1.0. " "This currently affects only the inner working of the " "operator, end-users can continue using the operator in " "the same way. Nevertheless, it is now recommended to provide" "the variables trav (and amp) as a tuples containing the " "traveltime (and amplitude) tables for sources and receivers " "separately. This behaviour will eventually become default in " "version v3.0.0.", FutureWarning, ) # identify geometry ( self.ndims, _, dims, self.ny, self.nx, self.nz, ns, nr, _, _, _, _, _, ) = Kirchhoff._identify_geometry(z, x, srcs, recs, y=y) self.dt = t[1] - t[0] self.nt = len(t) # store ix-iy locations of sources and receivers dx = x[1] - x[0] if self.ndims == 2: self.six = np.tile((srcs[0] - x[0]) // dx, (nr, 1)).T.astype(int).ravel() self.rix = np.tile((recs[0] - x[0]) // dx, (ns, 1)).astype(int).ravel() elif self.ndims == 3: # TODO: 3D normalized distances pass # compute traveltime self.dynamic = dynamic self.travsrcrec = True # use separate tables for src and rec traveltimes if mode in ["analytic", "eikonal", "byot"]: if mode in ["analytic", "eikonal"]: # compute traveltime table ( self.trav_srcs, self.trav_recs, self.dist_srcs, self.dist_recs, trav_srcs_grad, trav_recs_grad, ) = Kirchhoff._traveltime_table(z, x, srcs, recs, vel, y=y, mode=mode) if self.dynamic: # need to add a scalar in the denominator of amplitude term to avoid # division by 0, currently set to 1/100 of max distance to avoid having # to large scaling around the source. This number may change in future # or left to the user to define epsdist = 1e-2 self.maxdist = epsdist * ( np.max(self.dist_srcs) + np.max(self.dist_recs) ) # compute angles if self.ndims == 2: # 2d with vertical if anglerefl is None: self.angle_srcs = np.arctan2( trav_srcs_grad[0], trav_srcs_grad[1] ).reshape(np.prod(dims), ns) self.angle_recs = np.arctan2( trav_recs_grad[0], trav_recs_grad[1] ).reshape(np.prod(dims), nr) self.cosangle_srcs = np.cos(self.angle_srcs) self.cosangle_recs = np.cos(self.angle_recs) else: # TODO: 2D with normal raise NotImplementedError( "angle scaling with anglerefl currently not available" ) else: # TODO: 3D raise NotImplementedError( "dynamic=True currently not available in 3D" ) else: if isinstance(trav, tuple): self.trav_srcs, self.trav_recs = trav else: self.travsrcrec = False self.trav = trav if self.dynamic: if isinstance(amp, tuple): self.amp_srcs, self.amp_recs = amp else: self.amp = amp # in byot mode, angleaperture and snell checks are not performed self.angle_srcs = np.ones( (self.ny * self.nx * self.nz, ns), dtype=dtype ) self.angle_recs = np.ones( (self.ny * self.nx * self.nz, nr), dtype=dtype ) else: raise NotImplementedError("method must be analytic, eikonal or byot") # pre-compute traveltime indices if total traveltime is used if not self.travsrcrec: self.itrav = (self.trav / self.dt).astype("int32") self.travd = self.trav / self.dt - self.itrav # create wavelet operator if wavfilter: self.wav = self._wavelet_reshaping( wav, self.dt, srcs.shape[0], recs.shape[0], self.ndims ) else: self.wav = wav self.cop = Convolve1D( (ns * nr, self.nt), h=self.wav, offset=wavcenter, axis=1, dtype=dtype ) # create fixed-size aperture taper for all apertures self.aperturetap = taper(41, 20, "hanning")[20:] # define aperture if aperture is not None: warnings.warn( "Aperture is currently defined as ratio of offset over depth, " "and may be not ideal for highly heterogenous media" ) self.aperture = ( (2 * self.nx / self.nz,) if aperture is None else _value_or_sized_to_array(aperture) ) if len(self.aperture) == 1: self.aperture = np.array([0.8 * self.aperture[0], self.aperture[0]]) # define angle aperture and snell law angleaperture = [0.0, 1000.0] if angleaperture is None else angleaperture self.angleaperture = np.deg2rad(_value_or_sized_to_array(angleaperture)) if len(self.angleaperture) == 1: self.angleaperture = np.array( [0.8 * self.angleaperture[0], self.angleaperture[0]] ) self.snell = ( (np.pi,) if snell is None else np.deg2rad(_value_or_sized_to_array(snell)) ) if len(self.snell) == 1: self.snell = np.array([0.8 * self.snell[0], self.snell[0]]) # dimensions self.ns, self.nr = ns, nr self.nsnr = ns * nr self.ni = np.prod(dims) dims = tuple(dims) if self.ndims == 2 else (dims[0] * dims[1], dims[2]) dimsd = (ns, nr, self.nt) super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dimsd, name=name) # save velocity if using dynamic to compute amplitudes if self.dynamic: self.vel = ( vel.flatten() if not isinstance(vel, (float, int)) else vel * np.ones(np.prod(dims)) ) self._register_multiplications(engine) @staticmethod def _identify_geometry( z: NDArray, x: NDArray, srcs: NDArray, recs: NDArray, y: Optional[NDArray] = None, ) -> Tuple[ int, int, NDArray, int, int, int, int, int, float, float, float, NDArray, NDArray, ]: """Identify geometry and acquisition size and sampling""" ns, nr = srcs.shape[1], recs.shape[1] nz, nx = len(z), len(x) dz = np.abs(z[1] - z[0]) dx = np.abs(x[1] - x[0]) if y is None: ndims = 2 shiftdim = 0 ny = 1 dy = None dims = np.array([nx, nz]) dsamp = np.array([dx, dz]) origin = np.array([x[0], z[0]]) else: ndims = 3 shiftdim = 1 ny = len(y) dy = np.abs(y[1] - y[0]) dims = np.array([ny, nx, nz]) dsamp = np.array([dy, dx, dz]) origin = np.array([y[0], x[0], z[0]]) return ndims, shiftdim, dims, ny, nx, nz, ns, nr, dy, dx, dz, dsamp, origin @staticmethod def _traveltime_table( z: NDArray, x: NDArray, srcs: NDArray, recs: NDArray, vel: Union[float, NDArray], y: Optional[NDArray] = None, mode: str = "eikonal", ) -> Tuple[NDArray, NDArray, NDArray, NDArray, NDArray, NDArray]: r"""Traveltime table Compute traveltimes along the source-subsurface-receivers triplet in 2- or 3-dimensional media given a constant or depth- and space variable velocity. Parameters ---------- z : :obj:`numpy.ndarray` Depth axis x : :obj:`numpy.ndarray` Spatial axis srcs : :obj:`numpy.ndarray` Sources in array of size :math:`\lbrack 2 (3) \times n_s \rbrack` recs : :obj:`numpy.ndarray` Receivers in array of size :math:`\lbrack 2 (3) \times n_r \rbrack` vel : :obj:`numpy.ndarray` or :obj:`float` Velocity model of size :math:`\lbrack (n_y \times)\, n_x \times n_z \rbrack` (or constant) y : :obj:`numpy.ndarray` Additional spatial axis (for 3-dimensional problems) mode : :obj:`numpy.ndarray`, optional Computation mode (``eikonal``, ``analytic`` - only for constant velocity) Returns ------- trav_srcs : :obj:`numpy.ndarray` Source-to-subsurface traveltime table of size :math:`\lbrack (n_y*) n_x n_z \times n_s \rbrack` trav_recs : :obj:`numpy.ndarray` Receiver-to-subsurface traveltime table of size :math:`\lbrack (n_y) n_x n_z \times n_r \rbrack` dist_srcs : :obj:`numpy.ndarray` Source-to-subsurface distance table of size :math:`\lbrack (n_y*) n_x n_z \times n_s \rbrack` dist_recs : :obj:`numpy.ndarray` Receiver-to-subsurface distance table of size :math:`\lbrack (n_y) n_x n_z \times n_r \rbrack` trav_srcs_gradient : :obj:`numpy.ndarray` Source-to-subsurface traveltime gradient table of size :math:`\lbrack (n_y*) n_x n_z \times n_s \rbrack` trav_recs_gradient : :obj:`numpy.ndarray` Receiver-to-subsurface traveltime gradient table of size :math:`\lbrack (n_y) n_x n_z \times n_r \rbrack` """ # define geometry ( ndims, shiftdim, dims, ny, nx, nz, ns, nr, _, _, _, dsamp, origin, ) = Kirchhoff._identify_geometry(z, x, srcs, recs, y=y) # compute traveltimes if mode == "analytic": if not isinstance(vel, (float, int)): raise ValueError("vel must be scalar for mode=analytical") # compute grid if ndims == 2: X, Z = np.meshgrid(x, z, indexing="ij") X, Z = X.ravel(), Z.ravel() else: Y, X, Z = np.meshgrid(y, x, z, indexing="ij") Y, X, Z = Y.ravel(), X.ravel(), Z.ravel() dist_srcs2 = np.zeros((ny * nx * nz, ns)) dist_recs2 = np.zeros((ny * nx * nz, nr)) for isrc, src in enumerate(srcs.T): dist_srcs2[:, isrc] = (X - src[0 + shiftdim]) ** 2 + ( Z - src[1 + shiftdim] ) ** 2 if ndims == 3: dist_srcs2[:, isrc] += (Y - src[0]) ** 2 for irec, rec in enumerate(recs.T): dist_recs2[:, irec] = (X - rec[0 + shiftdim]) ** 2 + ( Z - rec[1 + shiftdim] ) ** 2 if ndims == 3: dist_recs2[:, irec] += (Y - rec[0]) ** 2 trav_srcs = np.sqrt(dist_srcs2) / vel trav_recs = np.sqrt(dist_recs2) / vel dist_srcs = trav_srcs * vel dist_recs = trav_recs * vel elif mode == "eikonal": if skfmm_message is None: dist_srcs = np.zeros((ny * nx * nz, ns)) dist_recs = np.zeros((ny * nx * nz, nr)) trav_srcs = np.zeros((ny * nx * nz, ns)) trav_recs = np.zeros((ny * nx * nz, nr)) for isrc, src in enumerate(srcs.T): src = np.round((src - origin) / dsamp).astype(np.int32) phi = np.ones_like(vel) if ndims == 2: phi[src[0], src[1]] = -1 else: phi[src[0], src[1], src[2]] = -1 dist_srcs[:, isrc] = (skfmm.distance(phi=phi, dx=dsamp)).ravel() trav_srcs[:, isrc] = ( skfmm.travel_time(phi=phi, speed=vel, dx=dsamp) ).ravel() for irec, rec in enumerate(recs.T): rec = np.round((rec - origin) / dsamp).astype(np.int32) phi = np.ones_like(vel) if ndims == 2: phi[rec[0], rec[1]] = -1 else: phi[rec[0], rec[1], rec[2]] = -1 dist_recs[:, irec] = (skfmm.distance(phi=phi, dx=dsamp)).ravel() trav_recs[:, irec] = ( skfmm.travel_time(phi=phi, speed=vel, dx=dsamp) ).ravel() else: raise NotImplementedError(skfmm_message) else: raise NotImplementedError("method must be analytic or eikonal") # compute traveltime gradients at image points trav_srcs_grad = np.gradient( trav_srcs.reshape(*dims, ns), axis=np.arange(ndims) ) trav_recs_grad = np.gradient( trav_recs.reshape(*dims, nr), axis=np.arange(ndims) ) return ( trav_srcs, trav_recs, dist_srcs, dist_recs, trav_srcs_grad, trav_recs_grad, ) def _wavelet_reshaping( self, wav: NDArray, dt: float, dimsrc: int, dimrec: int, dimv: int, ) -> NDArray: """Apply wavelet reshaping as from theory in [1]_""" f = np.fft.rfftfreq(len(wav), dt) W = np.fft.rfft(wav, n=len(wav)) if dimsrc == 2 and dimv == 2: # 2D Wfilt = W * (2 * np.pi * f) elif (dimsrc == 2 or dimrec == 2) and dimv == 3: # 2.5D raise NotImplementedError("2.D wavelet currently not available") elif dimsrc == 3 and dimrec == 3 and dimv == 3: # 3D Wfilt = W * (-1j * 2 * np.pi * f) wavfilt = np.fft.irfft(Wfilt, n=len(wav)) return wavfilt @staticmethod def _trav_kirch_matvec( x: NDArray, y: NDArray, nsnr: int, nt: int, ni: int, itrav: NDArray, travd: NDArray, ) -> NDArray: for isrcrec in prange(nsnr): itravisrcrec = itrav[:, isrcrec] travdisrcrec = travd[:, isrcrec] for ii in range(ni): itravii = itravisrcrec[ii] travdii = travdisrcrec[ii] if 0 <= itravii < nt - 1: y[isrcrec, itravii] += x[ii] * (1 - travdii) y[isrcrec, itravii + 1] += x[ii] * travdii return y @staticmethod def _trav_kirch_rmatvec( x: NDArray, y: NDArray, nsnr: int, nt: int, ni: int, itrav: NDArray, travd: NDArray, ) -> NDArray: for ii in prange(ni): itravii = itrav[ii] travdii = travd[ii] for isrcrec in range(nsnr): itravisrcrecii = itravii[isrcrec] travdisrcrecii = travdii[isrcrec] if 0 <= itravisrcrecii < nt - 1: y[ii] += ( x[isrcrec, itravisrcrecii] * (1 - travdisrcrecii) + x[isrcrec, itravisrcrecii + 1] * travdisrcrecii ) return y @staticmethod def _travsrcrec_kirch_matvec( x: NDArray, y: NDArray, ns: int, nr: int, nt: int, ni: int, dt: float, trav_srcs: NDArray, trav_recs: NDArray, ) -> NDArray: for isrc in prange(ns): travisrc = trav_srcs[:, isrc] for irec in range(nr): travirec = trav_recs[:, irec] trav = travisrc + travirec itrav = (trav / dt).astype("int32") travd = trav / dt - itrav for ii in range(ni): itravii = itrav[ii] travdii = travd[ii] if 0 <= itravii < nt - 1: y[isrc * nr + irec, itravii] += x[ii] * (1 - travdii) y[isrc * nr + irec, itravii + 1] += x[ii] * travdii return y @staticmethod def _travsrcrec_kirch_rmatvec( x: NDArray, y: NDArray, ns: int, nr: int, nt: int, ni: int, dt: float, trav_srcs: NDArray, trav_recs: NDArray, ) -> NDArray: for ii in prange(ni): trav_srcsii = trav_srcs[ii] trav_recsii = trav_recs[ii] for isrc in prange(ns): trav_srcii = trav_srcsii[isrc] for irec in range(nr): trav_recii = trav_recsii[irec] travii = trav_srcii + trav_recii itravii = int(travii / dt) travdii = travii / dt - itravii if 0 <= itravii < nt - 1: y[ii] += ( x[isrc * nr + irec, itravii] * (1 - travdii) + x[isrc * nr + irec, itravii + 1] * travdii ) return y @staticmethod def _amp_kirch_matvec( x: NDArray, y: NDArray, nsnr: int, nt: int, ni: int, itrav: NDArray, travd: NDArray, amp: NDArray, aperturemin: float, aperturemax: float, aperturetap: NDArray, nz: int, six: NDArray, rix: NDArray, angleaperturemin: float, angleaperturemax: float, angles_srcs: NDArray, angles_recs: NDArray, snellmin: float, snellmax: float, ) -> NDArray: nr = angles_recs.shape[-1] daperture = aperturemax - aperturemin dangleaperture = angleaperturemax - angleaperturemin dsnell = snellmax - snellmin for isrcrec in prange(nsnr): # extract traveltime, amplitude, src/rec coordinates at given src/pair itravisrcrec = itrav[:, isrcrec] travdisrcrec = travd[:, isrcrec] ampisrcrec = amp[:, isrcrec] sixisrcrec = six[isrcrec] rixisrcrec = rix[isrcrec] # extract source and receiver angles angles_src = angles_srcs[:, isrcrec // nr] angles_rec = angles_recs[:, isrcrec % nr] for ii in range(ni): # extract traveltime, amplitude at given image point itravii = itravisrcrec[ii] travdii = travdisrcrec[ii] damp = ampisrcrec[ii] # extract source and receiver angle angle_src = angles_src[ii] angle_rec = angles_rec[ii] abs_angle_src = abs(angle_src) abs_angle_rec = abs(angle_rec) abs_angle_src_rec = abs(angle_src + angle_rec) aptap = 1.0 # angle apertures checks if ( abs_angle_src < angleaperturemax and abs_angle_rec < angleaperturemax and abs_angle_src_rec < snellmax ): if abs_angle_src >= angleaperturemin: # extract source angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_src - angleaperturemin) // dangleaperture ) ] ) if abs_angle_rec >= angleaperturemin: # extract receiver angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_rec - angleaperturemin) // dangleaperture ) ] ) if abs_angle_src_rec >= snellmin: # extract snell taper value aptap = ( aptap * aperturetap[ int(20 * (abs_angle_src_rec - snellmin) // dsnell) ] ) # identify x-index of image point iz = ii % nz # aperture check aperture = abs(sixisrcrec - rixisrcrec) / iz if aperture < aperturemax: if aperture >= aperturemin: # extract aperture taper value aptap = ( aptap * aperturetap[ int(20 * ((aperture - aperturemin) // daperture)) ] ) # time limit check if 0 <= itravii < nt - 1: # assign values y[isrcrec, itravii] += x[ii] * (1 - travdii) * damp * aptap y[isrcrec, itravii + 1] += x[ii] * travdii * damp * aptap return y @staticmethod def _amp_kirch_rmatvec( x: NDArray, y: NDArray, nsnr: int, nt: int, ni: int, itrav: NDArray, travd: NDArray, amp: NDArray, aperturemin: float, aperturemax: float, aperturetap: NDArray, nz: int, six: NDArray, rix: NDArray, angleaperturemin: float, angleaperturemax: float, angles_srcs: NDArray, angles_recs: NDArray, snellmin: float, snellmax: float, ) -> NDArray: nr = angles_recs.shape[-1] daperture = aperturemax - aperturemin dangleaperture = angleaperturemax - angleaperturemin dsnell = snellmax - snellmin for ii in prange(ni): itravii = itrav[ii] travdii = travd[ii] ampii = amp[ii] # extract source and receiver angles angle_srcs = angles_srcs[ii] angle_recs = angles_recs[ii] # identify x-index of image point iz = ii % nz for isrcrec in range(nsnr): itravisrcrecii = itravii[isrcrec] travdisrcrecii = travdii[isrcrec] sixisrcrec = six[isrcrec] rixisrcrec = rix[isrcrec] # extract source and receiver angle angle_src = angle_srcs[isrcrec // nr] angle_rec = angle_recs[isrcrec % nr] abs_angle_src = abs(angle_src) abs_angle_rec = abs(angle_rec) abs_angle_src_rec = abs(angle_src + angle_rec) aptap = 1.0 # angle apertures checks if ( abs_angle_src < angleaperturemax and abs_angle_rec < angleaperturemax and abs_angle_src_rec < snellmax ): if abs_angle_src >= angleaperturemin: # extract source angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_src - angleaperturemin) // dangleaperture ) ] ) if abs_angle_rec >= angleaperturemin: # extract receiver angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_rec - angleaperturemin) // dangleaperture ) ] ) if abs_angle_src_rec >= snellmin: # extract snell taper value aptap = ( aptap * aperturetap[ int(20 * (abs_angle_src_rec - snellmin) // dsnell) ] ) # aperture check aperture = abs(sixisrcrec - rixisrcrec) / iz if aperture < aperturemax: if aperture >= aperturemin: # extract aperture taper value aptap = ( aptap * aperturetap[ int(20 * ((aperture - aperturemin) // daperture)) ] ) # time limit check if 0 <= itravisrcrecii < nt - 1: # assign values y[ii] += ( ( x[isrcrec, itravisrcrecii] * (1 - travdisrcrecii) + x[isrcrec, itravisrcrecii + 1] * travdisrcrecii ) * ampii[isrcrec] * aptap ) return y @staticmethod def _ampsrcrec_kirch_matvec( x: NDArray, y: NDArray, ns: int, nr: int, nt: int, ni: int, dt: float, vel: NDArray, trav_srcs: NDArray, trav_recs: NDArray, dist_srcs: NDArray, dist_recs: NDArray, aperturemin: float, aperturemax: float, aperturetap: NDArray, nz: int, six: NDArray, rix: NDArray, angleaperturemin: float, angleaperturemax: float, angles_srcs: NDArray, angles_recs: NDArray, snellmin: float, snellmax: float, maxdist: float, ) -> NDArray: daperture = aperturemax - aperturemin dangleaperture = angleaperturemax - angleaperturemin dsnell = snellmax - snellmin for isrc in prange(ns): travisrc = trav_srcs[:, isrc] distisrc = dist_srcs[:, isrc] angleisrc = angles_srcs[:, isrc] for irec in range(nr): travirec = trav_recs[:, irec] trav = travisrc + travirec itrav = (trav / dt).astype("int32") travd = trav / dt - itrav distirec = dist_recs[:, irec] angleirec = angles_recs[:, irec] dist = distisrc + distirec amp = np.abs(np.cos(angleisrc) + np.cos(angleirec)) / (dist + maxdist) sixisrcrec = six[isrc * nr + irec] rixisrcrec = rix[isrc * nr + irec] for ii in range(ni): itravii = itrav[ii] travdii = travd[ii] damp = amp[ii] / vel[ii] # extract source and receiver angle at given image point angle_src = angleisrc[ii] angle_rec = angleirec[ii] abs_angle_src = abs(angle_src) abs_angle_rec = abs(angle_rec) abs_angle_src_rec = abs(angle_src + angle_rec) aptap = 1.0 # angle apertures checks if ( abs_angle_src < angleaperturemax and abs_angle_rec < angleaperturemax and abs_angle_src_rec < snellmax ): if abs_angle_src >= angleaperturemin: # extract source angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_src - angleaperturemin) // dangleaperture ) ] ) if abs_angle_rec >= angleaperturemin: # extract receiver angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_rec - angleaperturemin) // dangleaperture ) ] ) if abs_angle_src_rec >= snellmin: # extract snell taper value aptap = ( aptap * aperturetap[ int(20 * (abs_angle_src_rec - snellmin) // dsnell) ] ) # identify x-index of image point iz = ii % nz # aperture check aperture = abs(sixisrcrec - rixisrcrec) / iz if aperture < aperturemax: if aperture >= aperturemin: # extract aperture taper value aptap = ( aptap * aperturetap[ int( 20 * ((aperture - aperturemin) // daperture) ) ] ) # time limit check if 0 <= itravii < nt - 1: y[isrc * nr + irec, itravii] += ( x[ii] * (1 - travdii) * damp * aptap ) y[isrc * nr + irec, itravii + 1] += ( x[ii] * travdii * damp * aptap ) return y @staticmethod def _ampsrcrec_kirch_rmatvec( x: NDArray, y: NDArray, ns: int, nr: int, nt: int, ni: int, dt: float, vel: NDArray, trav_srcs: NDArray, trav_recs: NDArray, dist_srcs: NDArray, dist_recs: NDArray, aperturemin: float, aperturemax: float, aperturetap: NDArray, nz: int, six: NDArray, rix: NDArray, angleaperturemin: float, angleaperturemax: float, angles_srcs: NDArray, angles_recs: NDArray, snellmin: float, snellmax: float, maxdist: float, ) -> NDArray: daperture = aperturemax - aperturemin dangleaperture = angleaperturemax - angleaperturemin dsnell = snellmax - snellmin for ii in prange(ni): trav_srcsii = trav_srcs[ii] trav_recsii = trav_recs[ii] dist_srcsii = dist_srcs[ii] dist_recsii = dist_recs[ii] velii = vel[ii] angle_srcsii = angles_srcs[ii] angle_recsii = angles_recs[ii] # identify x-index of image point iz = ii % nz for isrc in range(ns): trav_srcii = trav_srcsii[isrc] dist_srcii = dist_srcsii[isrc] angle_src = angle_srcsii[isrc] for irec in range(nr): trav_recii = trav_recsii[irec] travii = trav_srcii + trav_recii itravii = int(travii / dt) travdii = travii / dt - itravii dist_recii = dist_recsii[irec] angle_rec = angle_recsii[irec] dist = dist_srcii + dist_recii ampii = np.abs(np.cos(angle_src) + np.cos(angle_rec)) / ( dist + maxdist ) damp = ampii / velii sixisrcrec = six[isrc * nr + irec] rixisrcrec = rix[isrc * nr + irec] abs_angle_src = abs(angle_src) abs_angle_rec = abs(angle_rec) abs_angle_src_rec = abs(angle_src + angle_rec) aptap = 1.0 # angle apertures checks if ( abs_angle_src < angleaperturemax and abs_angle_rec < angleaperturemax and abs_angle_src_rec < snellmax ): if abs_angle_src >= angleaperturemin: # extract source angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_src - angleaperturemin) // dangleaperture ) ] ) if abs_angle_rec >= angleaperturemin: # extract receiver angle aperture taper value aptap = ( aptap * aperturetap[ int( 20 * (abs_angle_rec - angleaperturemin) // dangleaperture ) ] ) if abs_angle_src_rec >= snellmin: # extract snell taper value aptap = ( aptap * aperturetap[ int(20 * (abs_angle_src_rec - snellmin) // dsnell) ] ) # aperture check aperture = abs(sixisrcrec - rixisrcrec) / iz if aperture < aperturemax: if aperture >= aperturemin: # extract aperture taper value aptap = ( aptap * aperturetap[ int( 20 * ((aperture - aperturemin) // daperture) ) ] ) # time limit check if 0 <= itravii < nt - 1: # assign values y[ii] += ( ( x[isrc * nr + irec, itravii] * (1 - travdii) + x[isrc * nr + irec, itravii + 1] * travdii ) * damp * aptap ) return y def _register_multiplications(self, engine: str) -> None: if engine not in ["numpy", "numba"]: raise KeyError("engine must be numpy or numba") if engine == "numba" and jit_message is None: # numba numba_opts = dict( nopython=True, nogil=True, parallel=parallel ) # fastmath=True, if self.dynamic and self.travsrcrec: self._kirch_matvec = jit(**numba_opts)(self._ampsrcrec_kirch_matvec) self._kirch_rmatvec = jit(**numba_opts)(self._ampsrcrec_kirch_rmatvec) elif self.dynamic and not self.travsrcrec: self._kirch_matvec = jit(**numba_opts)(self._amp_kirch_matvec) self._kirch_rmatvec = jit(**numba_opts)(self._amp_kirch_rmatvec) elif self.travsrcrec: self._kirch_matvec = jit(**numba_opts)(self._travsrcrec_kirch_matvec) self._kirch_rmatvec = jit(**numba_opts)(self._travsrcrec_kirch_rmatvec) elif not self.travsrcrec: self._kirch_matvec = jit(**numba_opts)(self._trav_kirch_matvec) self._kirch_rmatvec = jit(**numba_opts)(self._trav_kirch_rmatvec) else: if engine == "numba" and jit_message is not None: logging.warning(jit_message) if self.dynamic and self.travsrcrec: self._kirch_matvec = self._ampsrcrec_kirch_matvec self._kirch_rmatvec = self._ampsrcrec_kirch_rmatvec elif self.dynamic and not self.travsrcrec: self._kirch_matvec = self._amp_kirch_matvec self._kirch_rmatvec = self._amp_kirch_rmatvec elif self.travsrcrec: self._kirch_matvec = self._travsrcrec_kirch_matvec self._kirch_rmatvec = self._travsrcrec_kirch_rmatvec elif not self.travsrcrec: self._kirch_matvec = self._trav_kirch_matvec self._kirch_rmatvec = self._trav_kirch_rmatvec @reshaped def _matvec(self, x: NDArray) -> NDArray: y = np.zeros((self.nsnr, self.nt), dtype=self.dtype) if self.dynamic and self.travsrcrec: inputs = ( x.ravel(), y, self.ns, self.nr, self.nt, self.ni, self.dt, self.vel, self.trav_srcs, self.trav_recs, self.dist_srcs, self.dist_recs, self.aperture[0], self.aperture[1], self.aperturetap, self.nz, self.six, self.rix, self.angleaperture[0], self.angleaperture[1], self.angle_srcs, self.angle_recs, self.snell[0], self.snell[1], self.maxdist, ) elif self.dynamic and not self.travsrcrec: inputs = ( x.ravel(), y, self.nsnr, self.nt, self.ni, self.itrav, self.travd, self.amp, self.aperture[0], self.aperture[1], self.aperturetap, self.nz, self.six, self.rix, self.angleaperture[0], self.angleaperture[1], self.angle_srcs, self.angle_recs, self.snell[0], self.snell[1], ) elif not self.dynamic and self.travsrcrec: inputs = ( x.ravel(), y, self.ns, self.nr, self.nt, self.ni, self.dt, self.trav_srcs, self.trav_recs, ) elif not self.dynamic and not self.travsrcrec: inputs = (x.ravel(), y, self.nsnr, self.nt, self.ni, self.itrav, self.travd) y = self._kirch_matvec(*inputs) y = self.cop._matvec(y.ravel()) return y @reshaped def _rmatvec(self, x: NDArray) -> NDArray: x = self.cop._rmatvec(x.ravel()) x = x.reshape(self.nsnr, self.nt) y = np.zeros(self.ni, dtype=self.dtype) if self.dynamic and self.travsrcrec: inputs = ( x, y, self.ns, self.nr, self.nt, self.ni, self.dt, self.vel, self.trav_srcs, self.trav_recs, self.dist_srcs, self.dist_recs, self.aperture[0], self.aperture[1], self.aperturetap, self.nz, self.six, self.rix, self.angleaperture[0], self.angleaperture[1], self.angle_srcs, self.angle_recs, self.snell[0], self.snell[1], self.maxdist, ) elif self.dynamic and not self.travsrcrec: inputs = ( x, y, self.nsnr, self.nt, self.ni, self.itrav, self.travd, self.amp, self.aperture[0], self.aperture[1], self.aperturetap, self.nz, self.six, self.rix, self.angleaperture[0], self.angleaperture[1], self.angle_srcs, self.angle_recs, self.snell[0], self.snell[1], ) elif not self.dynamic and self.travsrcrec: inputs = ( x, y, self.ns, self.nr, self.nt, self.ni, self.dt, self.trav_srcs, self.trav_recs, ) elif not self.dynamic and not self.travsrcrec: inputs = (x, y, self.nsnr, self.nt, self.ni, self.itrav, self.travd) y = self._kirch_rmatvec(*inputs) return y