Repository: simpeg/simpeg
Pull request #1177 (2023-02-03 19:42:49)

Linear mapping
Add a simple (yet flexible) `LinearMap`.


Patch:
diff --git a/SimPEG/__init__.py b/SimPEG/__init__.py
index 51961e283..3f2f48210 100644
--- a/SimPEG/__init__.py
+++ b/SimPEG/__init__.py
@@ -73,6 +73,7 @@ Mappings
   maps.ComboMap
   maps.ComplexMap
   maps.ExpMap
+  maps.LinearMap
   maps.IdentityMap
   maps.InjectActiveCells
   maps.MuRelative
diff --git a/SimPEG/maps.py b/SimPEG/maps.py
index b8f2c3533..9a9a04e5b 100644
--- a/SimPEG/maps.py
+++ b/SimPEG/maps.py
@@ -534,6 +534,83 @@ class ComboMap(IdentityMap):
         return len(self.maps)
 
 
+class LinearMap(IdentityMap):
+    """A generalized linear mapping.
+
+    A simple map that implements the linear mapping,
+
+    >>> y = A @ x + b
+
+    Parameters
+    ----------
+    A : (M, N) array_like, optional
+        The matrix operator, can be any object that implements `__matmul__`
+        and has a `shape` attribute.
+    b : (M) array_like, optional
+        Additive part of the linear operation.
+    """
+
+    def __init__(self, A, b=None, **kwargs):
+        mesh = kwargs.pop("mesh", None)
+        nP = kwargs.pop("nP", None)
+        super().__init__(**kwargs)
+        self.A = A
+        self.b = b
+
+    @property
+    def A(self):
+        """The linear operator matrix.
+
+        Returns
+        -------
+        LinearOperator
+            Must support matrix multiplication and have a shape attribute.
+        """
+        return self._A
+
+    @A.setter
+    def A(self, value):
+        if not hasattr(value, "__matmul__"):
+            raise TypeError(
+                f"{repr(value)} does not implement the matrix multiplication operator."
+            )
+        if not hasattr(value, "shape"):
+            raise TypeError(f"{repr(value)} does not have a shape attribute.")
+        self._A = value
+        self._nP = value.shape[1]
+        self._shape = value.shape
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def b(self):
+        """Added part of the linear operation.
+
+        Returns
+        -------
+        numpy.ndarray
+        """
+        return self._b
+
+    @b.setter
+    def b(self, value):
+        if value is not None:
+            value = validate_ndarray_with_shape("b", value, shape=(self.shape[0],))
+        self._b = value
+
+    def _transform(self, m):
+        if self.b is None:
+            return self.A @ m
+        return self.A @ m + self.b
+
+    def deriv(self, m, v=None):
+        if v is None:
+            return self.A
+        return self.A @ v
+
+
 class Projection(IdentityMap):
     r"""Projection mapping.
 


Test patch:
diff --git a/tests/base/test_maps.py b/tests/base/test_maps.py
index 199809cfb..699688048 100644
--- a/tests/base/test_maps.py
+++ b/tests/base/test_maps.py
@@ -1,6 +1,9 @@
 import numpy as np
 import unittest
 import discretize
+import pytest
+import scipy.sparse as sp
+
 from SimPEG import maps, models, utils
 from discretize.utils import mesh_builder_xyz, refine_tree_xyz
 import inspect
@@ -23,6 +26,7 @@ MAPS_TO_EXCLUDE_2D = [
     "ActiveCells",
     "InjectActiveCells",
     "LogMap",
+    "LinearMap",
     "ReciprocalMap",
     "PolynomialPetroClusterMap",
     "Surject2Dto3D",
@@ -49,6 +53,7 @@ MAPS_TO_EXCLUDE_3D = [
     "ActiveCells",
     "InjectActiveCells",
     "LogMap",
+    "LinearMap",
     "ReciprocalMap",
     "PolynomialPetroClusterMap",
     "CircleMap",
@@ -575,5 +580,73 @@ class TestSCEMT(unittest.TestCase):
         mapping.test(m=m, dx=0.05 * np.ones(mesh.n_cells), num=3)
 
 
+@pytest.mark.parametrize(
+    "A, b",
+    [
+        (np.random.rand(20, 40), np.random.rand(20)),  # test with np.array
+        (np.random.rand(20, 40), None),  # Test without B
+        (sp.random(20, 40, density=0.1), np.random.rand(20)),  # test with sparse matrix
+        (
+            sp.csr_array(sp.random(20, 40, density=0.1)),
+            np.random.rand(20),
+        ),  # test with sparse array
+        (sp.linalg.aslinearoperator(np.random.rand(20, 40)), np.random.rand(20)),
+    ],
+)
+def test_LinearMap(A, b):
+    map1 = maps.LinearMap(A, b)
+    x = np.random.rand(map1.nP)
+    y1 = A @ x
+    if b is not None:
+        y1 += b
+    y2 = map1 * x
+    np.testing.assert_allclose(y1, y2)
+
+
+@pytest.mark.parametrize(
+    "A, b",
+    [
+        (np.random.rand(20, 40), np.random.rand(20)),  # test with np.array
+        (np.random.rand(20, 40), None),  # Test without B
+        (sp.random(20, 40, density=0.1), np.random.rand(20)),  # test with sparse matrix
+        (
+            sp.csr_array(sp.random(20, 40, density=0.1)),
+            np.random.rand(20),
+        ),  # test with sparse array
+        (sp.linalg.aslinearoperator(np.random.rand(20, 40)), np.random.rand(20)),
+    ],
+)
+def test_LinearMapDerivs(A, b):
+    mapping = maps.LinearMap(A, b)
+    # check passing v=None
+    m = np.random.rand(mapping.nP)
+    v = np.random.rand(mapping.nP)
+    y1 = mapping.deriv(m) @ v
+    y2 = mapping.deriv(m, v=v)
+    np.testing.assert_equal(y1, y2)
+    mapping.test()
+
+
+def test_LinearMap_errors():
+    # object with no matmul operator
+    with pytest.raises(TypeError):
+        maps.LinearMap("No Matmul")
+
+    # object with matmul and no shape
+    class IncompleteLinearOperator:
+        def __matmul__(self, other):
+            return other
+
+    incomplete_A = IncompleteLinearOperator()
+    with pytest.raises(TypeError):
+        maps.LinearMap(incomplete_A)
+
+    # poorly shaped b vector
+    A = np.random.rand(20, 40)
+    b = np.random.rand(40)
+    with pytest.raises(ValueError):
+        maps.LinearMap(A, b=b)
+
+
 if __name__ == "__main__":
     unittest.main()