Merge pull request #725 from mrava87/doc-literal

Doc: use Literal type hint

Author: mrava87

pylops/avo/avo.py

@@ -12,14 +12,13 @@
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
-import numpy.typing as npt
 from numpy import cos, sin, tan
 
 from pylops import LinearOperator
 from pylops.utils._internal import _value_or_sized_to_tuple
 from pylops.utils.backend import get_array_module
 from pylops.utils.decorators import reshaped
-from pylops.utils.typing import DTypeLike, NDArray
+from pylops.utils.typing import DTypeLike, NDArray, Tavolinearization
 
 
 def zoeppritz_scattering(
@@ -29,7 +28,7 @@ def zoeppritz_scattering(
     vp0: float,
     vs0: float,
     rho0: float,
-    theta1: Union[float, npt.ArrayLike],
+    theta1: Union[float, NDArray],
 ) -> NDArray:
     r"""Zoeppritz solution.
 
@@ -250,12 +249,12 @@ def zoeppritz_pp(
 
 
 def approx_zoeppritz_pp(
-    vp1: Union[List, Tuple, npt.ArrayLike],
-    vs1: Union[List, Tuple, npt.ArrayLike],
-    rho1: Union[List, Tuple, npt.ArrayLike],
-    vp0: Union[List, Tuple, npt.ArrayLike],
-    vs0: Union[List, Tuple, npt.ArrayLike],
-    rho0: Union[List, Tuple, npt.ArrayLike],
+    vp1: Union[List, Tuple, NDArray],
+    vs1: Union[List, Tuple, NDArray],
+    rho1: Union[List, Tuple, NDArray],
+    vp0: Union[List, Tuple, NDArray],
+    vs0: Union[List, Tuple, NDArray],
+    rho0: Union[List, Tuple, NDArray],
     theta1: Union[float, NDArray],
 ) -> NDArray:
     """PP reflection coefficient from the approximate Zoeppritz equation.
@@ -335,7 +334,7 @@ def approx_zoeppritz_pp(
 
 
 def akirichards(
-    theta: npt.ArrayLike,
+    theta: NDArray,
     vsvp: Union[float, NDArray],
     n: int = 1,
 ) -> Tuple[NDArray, NDArray, NDArray]:
@@ -409,7 +408,7 @@ def akirichards(
 
 
 def fatti(
-    theta: npt.ArrayLike,
+    theta: NDArray,
     vsvp: Union[float, NDArray],
     n: int = 1,
 ) -> Tuple[NDArray, NDArray, NDArray]:
@@ -485,7 +484,7 @@ def fatti(
 
 
 def ps(
-    theta: npt.ArrayLike,
+    theta: NDArray,
     vsvp: Union[float, NDArray],
     n: int = 1,
 ) -> Tuple[NDArray, NDArray, NDArray]:
@@ -650,7 +649,7 @@ def __init__(
         vsvp: Union[float, NDArray] = 0.5,
         nt0: int = 1,
         spatdims: Optional[Union[int, Tuple[int]]] = None,
-        linearization: str = "akirich",
+        linearization: Tavolinearization = "akirich",
         dtype: DTypeLike = "float64",
         name: str = "A",
     ) -> None:

pylops/avo/poststack.py

@@ -3,10 +3,9 @@
     "PoststackInversion",
 ]
 
-from typing import Optional, Tuple, Union
+from typing import Literal, Optional, Tuple, Union
 
 import numpy as np
-import numpy.typing as npt
 from scipy.sparse.linalg import lsqr
 
 from pylops import (
@@ -136,12 +135,12 @@ def _PoststackLinearModelling(
 
 
 def PoststackLinearModelling(
-    wav: npt.ArrayLike,
+    wav: NDArray,
     nt0: int,
     spatdims: Optional[Union[int, ShapeLike]] = None,
     explicit: bool = False,
     sparse: bool = False,
-    kind: str = "centered",
+    kind: Literal["centered", "forward"] = "centered",
     name: Optional[str] = None,
 ) -> LinearOperator:
     r"""Post-stack linearized seismic modelling operator.
@@ -173,7 +172,7 @@ def PoststackLinearModelling(
         Create a sparse matrix (``True``) or dense  (``False``) when
         ``explicit=True``
     kind : :obj:`str`, optional
-        Derivative kind (``forward`` or ``centered``).
+        Derivative kind (``centered`` or ``forward``).
     name : :obj:`str`, optional
         .. versionadded:: 2.0.0
 
@@ -224,7 +223,7 @@ def PoststackLinearModelling(
 
 def PoststackInversion(
     data: NDArray,
-    wav: npt.ArrayLike,
+    wav: NDArray,
     m0: Optional[NDArray] = None,
     explicit: bool = False,
     simultaneous: bool = False,

pylops/avo/prestack.py

@@ -4,7 +4,7 @@
     "PrestackInversion",
 ]
 
-from typing import Optional, Tuple, Union
+from typing import Callable, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 from scipy.sparse.linalg import lsqr
@@ -33,7 +33,7 @@
     inplace_set,
 )
 from pylops.utils.signalprocessing import convmtx
-from pylops.utils.typing import NDArray, ShapeLike
+from pylops.utils.typing import NDArray, ShapeLike, Tavolinearization
 
 _linearizations = {"akirich": 3, "fatti": 3, "ps": 3}
 
@@ -44,9 +44,9 @@ def PrestackLinearModelling(
     vsvp: Union[float, NDArray] = 0.5,
     nt0: int = 1,
     spatdims: Optional[Union[int, ShapeLike]] = None,
-    linearization: str = "akirich",
+    linearization: Tavolinearization = "akirich",
     explicit: bool = False,
-    kind: str = "centered",
+    kind: Literal["centered", "forward"] = "centered",
     name: Optional[str] = None,
 ) -> LinearOperator:
     r"""Pre-stack linearized seismic modelling operator.
@@ -90,7 +90,7 @@ def PrestackLinearModelling(
         or a ``MatrixMult`` linear operator with dense matrix
         (``True``, preferred for small data)
     kind : :obj:`str`, optional
-        Derivative kind (``forward`` or ``centered``).
+        Derivative kind (``centered`` or ``forward``).
     name : :obj:`str`, optional
         .. versionadded:: 2.0.0
 
@@ -228,7 +228,7 @@ def PrestackWaveletModelling(
     nwav: int,
     wavc: Optional[int] = None,
     vsvp: Union[float, NDArray] = 0.5,
-    linearization: str = "akirich",
+    linearization: Union[Tavolinearization, Callable] = "akirich",
     name: Optional[str] = None,
 ) -> LinearOperator:
     r"""Pre-stack linearized seismic modelling operator for wavelet.
@@ -357,15 +357,15 @@ def PrestackInversion(
     theta: NDArray,
     wav: NDArray,
     m0: Optional[NDArray] = None,
-    linearization: str = "akirich",
+    linearization: Union[Tavolinearization, List[Tavolinearization]] = "akirich",
     explicit: bool = False,
     simultaneous: bool = False,
     epsI: Optional[float] = None,
     epsR: Optional[float] = None,
     dottest: bool = False,
     returnres: bool = False,
     epsRL1: Optional[float] = None,
-    kind: str = "centered",
+    kind: Literal["centered", "forward"] = "centered",
     vsvp: Union[float, NDArray] = 0.5,
     **kwargs_solver
 ) -> Union[NDArray, Tuple[NDArray, NDArray]]:
@@ -421,7 +421,7 @@ def PrestackInversion(
     epsRL1 : :obj:`float`, optional
         Damping factor for additional blockiness regularization term
     kind : :obj:`str`, optional
-        Derivative kind (``forward`` or ``centered``).
+        Derivative kind (``centered`` or ``forward``).
     vsvp : :obj:`float` or :obj:`numpy.ndarray`
         :math:`V_S/V_P` ratio (constant or time/depth variant)
     **kwargs_solver

pylops/basicoperators/block.py

@@ -4,7 +4,7 @@
 
 from pylops import LinearOperator
 from pylops.basicoperators import HStack, VStack
-from pylops.utils.typing import DTypeLike, NDArray
+from pylops.utils.typing import DTypeLike, NDArray, Tparallel_kind
 
 
 class _Block(LinearOperator):
@@ -17,7 +17,7 @@ class _Block(LinearOperator):
     def __init__(
         self,
         ops: Iterable[Iterable[LinearOperator]],
-        forceflat: bool = None,
+        forceflat: Optional[bool] = None,
         dtype: Optional[DTypeLike] = None,
         _HStack=HStack,
         _VStack=VStack,
@@ -153,8 +153,8 @@ def __init__(
         self,
         ops: Iterable[Iterable[LinearOperator]],
         nproc: int = 1,
-        forceflat: bool = None,
-        parallel_kind: str = "multiproc",
+        forceflat: Optional[bool] = None,
+        parallel_kind: Tparallel_kind = "multiproc",
         dtype: Optional[DTypeLike] = None,
     ):
         super().__init__(

pylops/basicoperators/blockdiag.py

@@ -23,7 +23,7 @@
 from pylops import LinearOperator
 from pylops.basicoperators import MatrixMult
 from pylops.utils.backend import get_array_module, get_module, inplace_set
-from pylops.utils.typing import DTypeLike, NDArray
+from pylops.utils.typing import DTypeLike, NDArray, Tinoutengine, Tparallel_kind
 
 
 def _matvec_rmatvec_map(op, x: NDArray) -> NDArray:
@@ -142,9 +142,9 @@ def __init__(
         self,
         ops: Sequence[LinearOperator],
         nproc: int = 1,
-        forceflat: bool = None,
-        inoutengine: Optional[tuple] = None,
-        parallel_kind: str = "multiproc",
+        forceflat: Optional[bool] = None,
+        inoutengine: Optional[Tinoutengine] = None,
+        parallel_kind: Tparallel_kind = "multiproc",
         dtype: Optional[DTypeLike] = None,
     ) -> None:
         if parallel_kind not in ["multiproc", "multithread"]:

pylops/basicoperators/causalintegration.py

@@ -1,6 +1,6 @@
 __all__ = ["CausalIntegration"]
 
-from typing import Union
+from typing import Literal, Union
 
 import numpy as np
 
@@ -100,8 +100,8 @@ def __init__(
         self,
         dims: Union[int, InputDimsLike],
         axis: int = -1,
-        sampling: float = 1,
-        kind: str = "full",
+        sampling: float = 1.0,
+        kind: Literal["full", "half", "trapezoidal"] = "full",
         removefirst: bool = False,
         dtype: DTypeLike = "float64",
         name: str = "C",

pylops/basicoperators/directionalderivative.py

@@ -3,9 +3,11 @@
     "SecondDirectionalDerivative",
 ]
 
+from typing import Union
+
 from pylops import LinearOperator
 from pylops.basicoperators import Diagonal, Gradient, Sum
-from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray
+from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray, Tderivkind
 
 
 class FirstDirectionalDerivative(LinearOperator):
@@ -25,8 +27,9 @@ class FirstDirectionalDerivative(LinearOperator):
     v : :obj:`numpy.ndarray`, optional
         Single direction (array of size :math:`n_\text{dims}`) or group of directions
         (array of size :math:`[n_\text{dims} \times n_{d_0} \times ... \times n_{d_{n_\text{dims}}}]`)
-    sampling : :obj:`tuple`, optional
-        Sampling steps for each direction.
+    sampling : :obj:`tuple` or :obj:`float`, optional
+        Sampling steps for each direction. If a single float
+        is provided, it is used for all directions.
     edge : :obj:`bool`, optional
         Use reduced order derivative at edges (``True``) or
         ignore them (``False``).
@@ -77,9 +80,9 @@ def __init__(
         self,
         dims: InputDimsLike,
         v: NDArray,
-        sampling: int = 1,
+        sampling: Union[float, InputDimsLike] = 1.0,
         edge: bool = False,
-        kind: str = "centered",
+        kind: Tderivkind = "centered",
         dtype: DTypeLike = "float64",
         name: str = "F",
     ):
@@ -88,7 +91,12 @@ def __init__(
         self.kind = kind
         self.v = v
         Op = self._calc_first_ddop(
-            dims=dims, sampling=sampling, edge=edge, kind=kind, dtype=dtype, v=v
+            dims=dims,
+            v=v,
+            sampling=sampling,
+            edge=edge,
+            kind=kind,
+            dtype=dtype,
         )
         super().__init__(Op=Op, name=name)
 
@@ -102,9 +110,9 @@ def _rmatvec(self, x: NDArray) -> NDArray:
     def _calc_first_ddop(
         dims: InputDimsLike,
         v: NDArray,
-        sampling: int,
+        sampling: Union[float, InputDimsLike],
         edge: bool,
-        kind: str,
+        kind: Tderivkind,
         dtype: DTypeLike,
     ):
         Gop = Gradient(dims, sampling=sampling, edge=edge, kind=kind, dtype=dtype)
@@ -133,11 +141,16 @@ class SecondDirectionalDerivative(LinearOperator):
     v : :obj:`numpy.ndarray`, optional
         Single direction (array of size :math:`n_\text{dims}`) or group of directions
         (array of size :math:`[n_\text{dims} \times n_{d_0} \times ... \times n_{d_{n_\text{dims}}}]`)
-    sampling : :obj:`tuple`, optional
-        Sampling steps for each direction.
+    sampling : :obj:`tuple` or :obj:`float`, optional
+        Sampling steps for each direction. If a single float
+        is provided, it is used for all directions.
     edge : :obj:`bool`, optional
         Use reduced order derivative at edges (``True``) or
         ignore them (``False``).
+    kind : :obj:`str`, optional
+        .. versionadded:: 2.7.0
+
+        Derivative kind (``forward``, ``centered``, or ``backward``).
     dtype : :obj:`str`, optional
         Type of elements in input array.
 
@@ -175,17 +188,18 @@ def __init__(
         self,
         dims: InputDimsLike,
         v: NDArray,
-        sampling: int = 1,
+        sampling: Union[float, InputDimsLike] = 1.0,
         edge: bool = False,
+        kind: Tderivkind = "centered",
         dtype: DTypeLike = "float64",
         name: str = "S",
     ):
-        self.dims = dims
-        self.v = v
         self.sampling = sampling
         self.edge = edge
+        self.kind = kind
+        self.v = v
         Op = self._calc_second_ddop(
-            dims=dims, v=v, sampling=sampling, edge=edge, dtype=dtype
+            dims=dims, v=v, sampling=sampling, edge=edge, kind=kind, dtype=dtype
         )
         super().__init__(Op=Op, name=name)
 
@@ -197,10 +211,15 @@ def _rmatvec(self, x: NDArray) -> NDArray:
 
     @staticmethod
     def _calc_second_ddop(
-        dims: InputDimsLike, v: NDArray, sampling: int, edge: bool, dtype: DTypeLike
+        dims: InputDimsLike,
+        v: NDArray,
+        sampling: int,
+        edge: bool,
+        kind: Tderivkind,
+        dtype: DTypeLike,
     ):
         Dop = FirstDirectionalDerivative(
-            dims=dims, v=v, sampling=sampling, edge=edge, dtype=dtype
+            dims=dims, v=v, sampling=sampling, edge=edge, kind=kind, dtype=dtype
         )
         ddop = -Dop.H * Dop
         return ddop