__all__ = ["Spread"]
import logging
from typing import Callable, Optional
import numpy as np
from pylops import LinearOperator
from pylops.utils import deps
from pylops.utils.decorators import reshaped
from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray
jit_message = deps.numba_import("the spread module")
if jit_message is None:
from numba import jit
from ._spread_numba import (
_matvec_numba_onthefly,
_matvec_numba_table,
_rmatvec_numba_onthefly,
_rmatvec_numba_table,
)
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING)
[docs]class Spread(LinearOperator):
r"""Spread operator.
Spread values from the input model vector arranged as a 2-dimensional
array of size :math:`[n_{x_0} \times n_{t_0}]` into the data vector of size
:math:`[n_x \times n_t]`. Note that the value at each single pair
:math:`(x_0, t_0)` in the input is spread over the entire :math:`x` axis
in the output.
Spreading is performed along parametric curves provided as look-up table
of pre-computed indices (``table``) or computed on-the-fly using a
function handle (``fh``).
In adjont mode, values from the data vector are instead stacked
along the same parametric curves.
Parameters
----------
dims : :obj:`tuple`
Dimensions of model vector (vector will be reshaped internally into
a two-dimensional array of size :math:`[n_{x_0} \times n_{t_0}]`,
where the first dimension is the spreading direction)
dimsd : :obj:`tuple`
Dimensions of data vector (vector will be reshaped internal into
a two-dimensional array of size :math:`[n_x \times n_t]`,
where the first dimension is the stacking direction)
table : :obj:`np.ndarray`, optional
Look-up table of indices of size
:math:`[n_{x_0} \times n_{t_0} \times n_x]` (if ``None`` use function
handle ``fh``). When ``dtable`` is not provided, the ``data`` will be created
as follows
.. code-block:: python
data[ix, table[ix0, it0, ix]] += model[ix0, it0]
.. note::
When using ``table`` without ``dtable``, its elements must be
between 0 and :math:`n_{t_0} - 1` (or ``numpy.nan``).
dtable : :obj:`np.ndarray`, optional
Look-up table of decimals remainders for linear interpolation of size
:math:`[n_{x_0} \times n_{t_0} \times n_x]` (if ``None`` use function
handle ``fh``). When provided, the ``data`` will be created as follows
.. code-block:: python
data[ix, table[ix0, it0, ix]] += (1 - dtable[ix0, it0, ix]) * model[ix0, it0]
data[ix, table[ix0, it0, ix] + 1] += dtable[ix0, it0, ix] * model[ix0, it0]
.. note::
When using ``table`` and ``dtable``, the elements of ``table`` indices must be
between 0 and :math:`n_{t_0} - 2` (or ``numpy.nan``).
fh : :obj:`callable`, optional
If ``None`` will use look-up table ``table``. When provided, should be a
function which takes indices ``ix0`` and ``it0`` and returns
an array of size :math:`n_x` containing each respective time index.
Alternatively, if linear interpolation is required, it should output in
addition to the time indices, a weight for interpolation with linear
interpolation, to be used as follows
.. code-block:: python
data[ix, index] += (1 - dindices[ix]) * model[ix0, it0]
data[ix, index + 1] += dindices[ix] * model[ix0, it0]
where ``index`` refers to a time index in the first array returned by ``fh``
and ``dindices`` refers to the weight in the second array returned by ``fh``.
.. note::
When using ``fh`` with one output (time indices), the time indices must be
between 0 and :math:`n_{t_0} - 1` (or ``numpy.nan``). When using ``fh`` with two outputs
(time indices and weights), they must be within the between 0 and
:math:`n_{t_0} - 2` (or ``numpy.nan``).
interp : :obj:`bool`, optional
Use only if engine ``engine='numba'``. Apply linear interpolation (``True``)
or nearest interpolation (``False``) during stacking/spreading along
parametric curve.
When using ``engine="numpy"``, it will be inferred directly from ``fh`` or
the presence of ``dtable``.
engine : :obj:`str`, optional
Engine used for spread computation (``numpy`` or ``numba``). Note that
``numba`` can only be used when providing a look-up table
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
------
KeyError
If ``engine`` is neither ``numpy`` nor ``numba``
NotImplementedError
If both ``table`` and ``fh`` are not provided
ValueError
If ``table`` has shape different from
:math:`[n_{x_0} \times n_{t_0} \times n_x]`
Notes
-----
The Spread operator applies the following linear transform in forward mode
to the model vector after reshaping it into a 2-dimensional array of size
:math:`[n_x \times n_t]`:
.. math::
m(x_0, t_0) \rightarrow d(x, t=f(x_0, x, t_0)) \quad \forall x
where :math:`f(x_0, x, t)` is a mapping function that returns a value :math:`t`
given values :math:`x_0`, :math:`x`, and :math:`t_0`. Note that for each
:math:`(x_0, t_0)` pair, spreading is done over the entire :math:`x` axis
in the data domain.
In adjoint mode, the model is reconstructed by means of the following
stacking operation:
.. math::
m(x_0, t_0) = \int{d(x, t=f(x_0, x, t_0))} \,\mathrm{d}x
Note that ``table`` (or ``fh``) must return integer numbers
representing indices in the axis :math:`t`. However it also possible to
perform linear interpolation as part of the spreading/stacking process by
providing the decimal part of the mapping function (:math:`t - \lfloor
t \rfloor`) either in ``dtable`` input parameter or as second value in
the return of ``fh`` function.
"""
def __init__(
self,
dims: InputDimsLike,
dimsd: InputDimsLike,
table: Optional[NDArray] = None,
dtable: Optional[NDArray] = None,
fh: Optional[Callable] = None,
interp: Optional[bool] = None,
engine: str = "numpy",
dtype: DTypeLike = "float64",
name: str = "S",
) -> None:
super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dimsd, name=name)
if engine not in ["numpy", "numba"]:
raise KeyError("engine must be numpy or numba")
if engine == "numba" and jit_message is None:
self.engine = "numba"
else:
if engine == "numba" and jit is not None:
logging.warning(jit_message)
self.engine = "numpy"
# axes
self.nx0, self.nt0 = self.dims[0], self.dims[1]
self.nx, self.nt = self.dimsd[0], self.dimsd[1]
self.table = table
self.dtable = dtable
self.fh = fh
# find out if mapping is in table of function handle
if self.table is None and fh is None:
raise NotImplementedError("provide either table or fh.")
elif self.table is not None:
if fh is not None:
raise ValueError("provide only one of table or fh.")
if self.table.shape != (self.nx0, self.nt0, self.nx):
raise ValueError("table must have shape [nx0 x nt0 x nx]")
self.usetable = True
if np.any(self.table > self.nt):
raise ValueError("values in table must be smaller than nt")
else:
self.usetable = False
# find out if linear interpolation has to be carried out
self.interp = False
if self.usetable:
if self.dtable is not None:
if self.dtable.shape != (self.nx0, self.nt0, self.nx):
raise ValueError("dtable must have shape [nx0 x nt x nx]")
self.interp = True
else:
if self.engine == "numba":
self.interp = interp
else:
if len(fh(0, 0)) == 2:
self.interp = True
if interp is not None and self.interp != interp:
logging.warning("interp has been overridden to %r.", self.interp)
def _matvec_numpy(self, x: NDArray) -> NDArray:
y = np.zeros(self.dimsd, dtype=self.dtype)
for it in range(self.dims[1]):
for ix0 in range(self.dims[0]):
if self.usetable:
indices = self.table[ix0, it]
if self.interp:
dindices = self.dtable[ix0, it]
else:
if self.interp:
indices, dindices = self.fh(ix0, it)
else:
indices = self.fh(ix0, it)
mask = np.argwhere(~np.isnan(indices))
if mask.size > 0:
indices = (indices[mask]).astype(int)
if not self.interp:
y[mask, indices] += x[ix0, it]
else:
y[mask, indices] += (1 - dindices[mask]) * x[ix0, it]
y[mask, indices + 1] += dindices[mask] * x[ix0, it]
return y
def _rmatvec_numpy(self, x: NDArray) -> NDArray:
y = np.zeros(self.dims, dtype=self.dtype)
for it in range(self.dims[1]):
for ix0 in range(self.dims[0]):
if self.usetable:
indices = self.table[ix0, it]
if self.interp:
dindices = self.dtable[ix0, it]
else:
if self.interp:
indices, dindices = self.fh(ix0, it)
else:
indices = self.fh(ix0, it)
mask = np.argwhere(~np.isnan(indices))
if mask.size > 0:
indices = (indices[mask]).astype(int)
if not self.interp:
y[ix0, it] = np.sum(x[mask, indices])
else:
y[ix0, it] = np.sum(
x[mask, indices] * (1 - dindices[mask])
) + np.sum(x[mask, indices + 1] * dindices[mask])
return y
@reshaped
def _matvec(self, x: NDArray) -> NDArray:
if self.engine == "numba":
y = np.zeros(self.dimsd, dtype=self.dtype)
if self.usetable:
y = _matvec_numba_table(
x,
y,
self.dims,
self.interp,
self.table,
self.table if self.dtable is None else self.dtable,
)
else:
y = _matvec_numba_onthefly(x, y, self.dims, self.interp, self.fh)
else:
y = self._matvec_numpy(x)
return y
@reshaped
def _rmatvec(self, x: NDArray) -> NDArray:
if self.engine == "numba":
y = np.zeros(self.dims, dtype=self.dtype)
if self.usetable:
y = _rmatvec_numba_table(
x,
y,
self.dims,
self.dimsd,
self.interp,
self.table,
self.table if self.dtable is None else self.dtable,
)
else:
y = _rmatvec_numba_onthefly(
x, y, self.dims, self.dimsd, self.interp, self.fh
)
else:
y = self._rmatvec_numpy(x)
return y
