Added Discrete Cosine Transform and Tests

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
           - --line-length=88
 
   - repo: https://github.com/pycqa/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
       - id: isort
         name: isort (python)

pylops/signalprocessing/__init__.py

@@ -29,6 +29,7 @@
     Patch2D                         2D Patching transform operator.
     Patch3D                         3D Patching transform operator.
     Fredholm1                       Fredholm integral of first kind.
+    DCT                             Discrete Cosine Transform
 
 """
 
@@ -56,6 +57,7 @@
 from .dwt import *
 from .dwt2d import *
 from .seislet import *
+from .dct import *
 
 __all__ = [
     "FFT",
@@ -82,4 +84,5 @@
     "DWT",
     "DWT2D",
     "Seislet",
+    "DCT",
 ]

pylops/signalprocessing/dct.py

@@ -0,0 +1,137 @@
+__all__ = ["DCT"]
+
+from typing import List, Optional, Union
+
+import numpy as np
+from scipy import fft
+
+from pylops import LinearOperator
+from pylops.utils._internal import _value_or_sized_to_tuple
+from pylops.utils.decorators import reshaped
+from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray
+
+
+class DCT(LinearOperator):
+    r"""Discreet Cosine Transform
+
+    Performs discreet cosine transform on the given multi-dimensional
+    array along the given axis.
+
+    Parameters
+    ----------
+    dims : :obj:`list` or :obj:`int`
+        Number of samples for each dimension
+    type : :obj:`int`, optional
+        Type of the DCT (see Notes). Default type is 2.
+    axes : :obj:`list` or :obj:`int`, optional
+        Axes over which the DCT is computed. If not given, the last len(dims) axes are used,
+        or all axes if dims is also not specified.
+    dtype : :obj:`str`, optional
+        Type of elements in input array.
+    workers :obj:`int`, optional
+        Maximum number of workers to use for parallel computation. If negative, the value wraps around from os.cpu_count().
+    name : :obj:`str`, optional
+        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``)
+
+    Notes
+    -----
+    The DiscreetCosine is implemented in normalization mode = "ortho" to make the scaling symmetrical.
+    There are 4 types of DCT available in pylops.dct. 'The' DCT generally refers to `type=2` dct and
+    'the' inverse DCT refers to `type=3`.
+
+    **Type 1**
+    There are several definitions of the DCT-I; we use the following
+    (for ``norm="backward"``)
+
+    .. math::
+
+       y_k = x_0 + (-1)^k x_{N-1} + 2 \sum_{n=1}^{N-2} x_n \cos\left(
+       \frac{\pi k n}{N-1} \right)
+
+    If ``orthogonalize=True``, ``x[0]`` and ``x[N-1]`` are multiplied by a
+    scaling factor of :math:`\sqrt{2}`, and ``y[0]`` and ``y[N-1]`` are divided
+    by :math:`\sqrt{2}`. When combined with ``norm="ortho"``, this makes the
+    corresponding matrix of coefficients orthonormal (``O @ O.T = np.eye(N)``).
+
+    (The DCT-I is only supported for input size > 1.)
+
+    **Type 2**
+     There are several definitions of the DCT-II; we use the following
+    (for ``norm="backward"``)
+
+    .. math::
+
+       y_k = 2 \sum_{n=0}^{N-1} x_n \cos\left(\frac{\pi k(2n+1)}{2N} \right)
+
+    If ``orthogonalize=True``, ``y[0]`` is divided by :math:`\sqrt{2}` which,
+    when combined with ``norm="ortho"``, makes the corresponding matrix of
+    coefficients orthonormal (``O @ O.T = np.eye(N)``).
+
+    **Type 3**
+    There are several definitions, we use the following (for
+    ``norm="backward"``)
+
+    .. math::
+
+       y_k = x_0 + 2 \sum_{n=1}^{N-1} x_n \cos\left(\frac{\pi(2k+1)n}{2N}\right)
+
+    If ``orthogonalize=True``, ``x[0]`` terms are multiplied by
+    :math:`\sqrt{2}` which, when combined with ``norm="ortho"``, makes the
+    corresponding matrix of coefficients orthonormal (``O @ O.T = np.eye(N)``).
+
+    The (unnormalized) DCT-III is the inverse of the (unnormalized) DCT-II, up
+    to a factor `2N`. The orthonormalized DCT-III is exactly the inverse of
+    the orthonormalized DCT-II.
+
+    **Type 4**
+    There are several definitions of the DCT-IV; we use the following
+    (for ``norm="backward"``)
+
+    .. math::
+
+       y_k = 2 \sum_{n=0}^{N-1} x_n \cos\left(\frac{\pi(2k+1)(2n+1)}{4N} \right)
+
+    ``orthogonalize`` has no effect here, as the DCT-IV matrix is already
+    orthogonal up to a scale factor of ``2N``.
+
+    """
+
+    def __init__(
+        self,
+        dims: Union[int, InputDimsLike],
+        type: int = 2,
+        axes: Union[int, List[int]] = None,
+        dtype: DTypeLike = "float64",
+        workers: Optional[int] = None,
+        name: str = "C",
+    ) -> None:
+
+        if type > 4 or type < 1:
+            raise ValueError("wrong value of type it can only be 1, 2, 3 or 4")
+        self.type = type
+        self.axes = axes
+        self.workers = workers
+        self.dims = _value_or_sized_to_tuple(dims)
+        super().__init__(
+            dtype=np.dtype(dtype), dims=self.dims, dimsd=self.dims, name=name
+        )
+
+    @reshaped
+    def _matvec(self, x: NDArray) -> NDArray:
+        return fft.dctn(
+            x, axes=self.axes, type=self.type, norm="ortho", workers=self.workers
+        )
+
+    @reshaped
+    def _rmatvec(self, x: NDArray) -> NDArray:
+        return fft.idctn(
+            x, axes=self.axes, type=self.type, norm="ortho", workers=self.workers
+        )

pytests/test_sparse_transforms.py

@@ -0,0 +1,128 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_array_almost_equal
+
+from pylops.signalprocessing import DCT
+from pylops.utils import dottest
+
+par1 = {"ny": 11, "nx": 11, "imag": 0, "dtype": "float64"}
+par2 = {"ny": 11, "nx": 21, "imag": 0, "dtype": "float64"}
+par3 = {"ny": 21, "nx": 21, "imag": 0, "dtype": "float64"}
+
+
+@pytest.mark.parametrize("par", [(par1), (par3)])
+def test_DCT1D(par):
+    """Dot test for Discrete Cosine Transform Operator 1D"""
+
+    t = np.arange(par["ny"])
+    # testing for various types of dct
+    Dct = DCT(dims=(par["ny"],), type=1, dtype=par["dtype"])
+
+    assert dottest(Dct, par["ny"], par["ny"], rtol=1e-6, complexflag=0, verb=True)
+
+    Dct = DCT(dims=(par["ny"],), type=2, dtype=par["dtype"])
+
+    assert dottest(Dct, par["ny"], par["ny"], rtol=1e-6, complexflag=0, verb=True)
+
+    Dct = DCT(dims=(par["ny"],), type=3, dtype=par["dtype"])
+
+    assert dottest(Dct, par["ny"], par["ny"], rtol=1e-6, complexflag=0, verb=True)
+
+    Dct = DCT(dims=(par["ny"],), type=4, dtype=par["dtype"])
+
+    assert dottest(Dct, par["ny"], par["ny"], rtol=1e-6, complexflag=0, verb=True)
+
+    y = Dct.H * (Dct * t)
+
+    assert_array_almost_equal(t, y, decimal=3)
+
+
+@pytest.mark.parametrize("par", [(par1), (par2), (par3)])
+def test_DCT2D(par):
+    """Dot test for Discrete Cosine Transform Operator 2D"""
+
+    t = np.outer(np.arange(par["ny"]), np.arange(par["nx"]))
+
+    Dct = DCT(dims=t.shape, dtype=par["dtype"])
+
+    assert dottest(
+        Dct,
+        par["nx"] * par["ny"],
+        par["nx"] * par["ny"],
+        rtol=1e-6,
+        complexflag=0,
+        verb=True,
+    )
+
+    Dct = DCT(dims=t.shape, dtype=par["dtype"], axes=1)
+
+    assert dottest(
+        Dct,
+        par["nx"] * par["ny"],
+        par["nx"] * par["ny"],
+        rtol=1e-6,
+        complexflag=0,
+        verb=True,
+    )
+
+    y = Dct.H * (Dct * t)
+
+    assert_array_almost_equal(t, y, decimal=3)
+
+
+@pytest.mark.parametrize("par", [(par1), (par2), (par3)])
+def test_DCT3D(par):
+    """Dot test for Discrete Cosine Transform Operator 2D"""
+
+    t = np.random.rand(par["nx"], par["nx"], par["nx"])
+
+    Dct = DCT(dims=t.shape, dtype=par["dtype"])
+
+    assert dottest(
+        Dct,
+        par["nx"] * par["nx"] * par["nx"],
+        par["nx"] * par["nx"] * par["nx"],
+        rtol=1e-6,
+        complexflag=0,
+        verb=True,
+    )
+
+    Dct = DCT(dims=t.shape, dtype=par["dtype"], axes=1)
+
+    assert dottest(
+        Dct,
+        par["nx"] * par["nx"] * par["nx"],
+        par["nx"] * par["nx"] * par["nx"],
+        rtol=1e-6,
+        complexflag=0,
+        verb=True,
+    )
+
+    Dct = DCT(dims=t.shape, dtype=par["dtype"], axes=2)
+
+    assert dottest(
+        Dct,
+        par["nx"] * par["nx"] * par["nx"],
+        par["nx"] * par["nx"] * par["nx"],
+        rtol=1e-6,
+        complexflag=0,
+        verb=True,
+    )
+
+    y = Dct.H * (Dct * t)
+
+    assert_array_almost_equal(t, y, decimal=3)
+
+
+@pytest.mark.parametrize("par", [(par1), (par3)])
+def test_DCT_workers(par):
+    """Dot test for Discrete Cosine Transform Operator with workers"""
+    t = np.arange(par["ny"])
+
+    Dct = DCT(dims=(par["ny"],), type=1, dtype=par["dtype"], workers=2)
+
+    assert dottest(Dct, par["ny"], par["ny"], rtol=1e-6, complexflag=0, verb=True)
+
+    y = Dct.H * (Dct * t)
+
+    assert_array_almost_equal(t, y, decimal=3)