Repository: simpeg/simpeg

Issue:
Full Anisotropy
_From @lheagy on June 19, 2015 18:41_

At the moment, MeSigmaIDeriv, MfRhoDeriv, MfRhoDeriv will not take fully anisotropic media - this actually needs to be taken care of inside of getFaceInnerProductDeriv and getEdgeInnerProductDeriv within simpeg, but I am creating the issue here too so it is documented that the fully anisotropic inversion won't be working yet

_Copied from original issue: simpeg/simpegem#22_


Discussion:
jcapriot: > This looks great Joe! I am looking forward to playing with it. My comments are pretty minor, mostly trying to think about clarity. One thing that wasn't obvious to me is when we would encounter a scenario where `v.dim > 2`? I don't think we should encounter this for the forward? but maybe these is a possibility with the adjoint? (or maybe I am just missing the use case!)

`v.dim` > 2 happens when we are doing an adjoint multiplication, with multiple field polarizations, and many `vs` (i.e. constructing the sensitivity matrix for an MT problem). So `v` in that case would be like (`v` n_m by n_d by n_pol)
lheagy: > > This looks great Joe! I am looking forward to playing with it. My comments are pretty minor, mostly trying to think about clarity. One thing that wasn't obvious to me is when we would encounter a scenario where `v.dim > 2`? I don't think we should encounter this for the forward? but maybe these is a possibility with the adjoint? (or maybe I am just missing the use case!)
> 
> `v.dim` > 2 happens when we are doing an adjoint multiplication, with multiple field polarizations, and many `vs` (i.e. constructing the sensitivity matrix for an MT problem). So `v` in that case would be like (`v` n_m by n_d by n_pol)

That makes sense. Thanks for clarifying! 

Patch:
diff --git a/SimPEG/base/pde_simulation.py b/SimPEG/base/pde_simulation.py
index a38ad4a46..a213d922a 100644
--- a/SimPEG/base/pde_simulation.py
+++ b/SimPEG/base/pde_simulation.py
@@ -1,6 +1,6 @@
 import numpy as np
 import scipy.sparse as sp
-from discretize.utils import Zero
+from discretize.utils import Zero, TensorType
 from ..simulation import BaseSimulation
 from .. import props
 from scipy.constants import mu_0
@@ -8,34 +8,70 @@ from scipy.constants import mu_0
 
 def __inner_mat_mul_op(M, u, v=None, adjoint=False):
     u = np.squeeze(u)
-    if v is not None:
-        if v.ndim > 1:
-            v = np.squeeze(v)
-        if u.ndim > 1:
-            # u has multiple fields
-            if v.ndim == 1:
-                v = v[:, None]
-        else:
+    if sp.issparse(M):
+        if v is not None:
             if v.ndim > 1:
-                u = u[:, None]
-        if v.ndim > 2:
-            u = u[:, None, :]
-        if adjoint:
-            if u.ndim > 1 and u.shape[-1] > 1:
-                return M.T * (u * v).sum(axis=-1)
-            return M.T * (u * v)
-        if u.ndim > 1 and u.shape[1] > 1:
-            return np.squeeze(u[:, None, :] * (M * v)[:, :, None])
-        return u * (M * v)
-    else:
+                v = np.squeeze(v)
+            if u.ndim > 1:
+                # u has multiple fields
+                if v.ndim == 1:
+                    v = v[:, None]
+            else:
+                if v.ndim > 1:
+                    u = u[:, None]
+            if v.ndim > 2:
+                u = u[:, None, :]
+            if adjoint:
+                if u.ndim > 1 and u.shape[-1] > 1:
+                    return M.T * (u * v).sum(axis=-1)
+                return M.T * (u * v)
+            if u.ndim > 1 and u.shape[1] > 1:
+                return np.squeeze(u[:, None, :] * (M * v)[:, :, None])
+            return u * (M * v)
+        else:
+            if u.ndim > 1:
+                UM = sp.vstack([sp.diags(u[:, i]) @ M for i in range(u.shape[1])])
+            else:
+                U = sp.diags(u, format="csr")
+                UM = U @ M
+            if adjoint:
+                return UM.T
+            return UM
+    elif isinstance(M, tuple):
+        # assume it was a tuple of M_func, prop_deriv
+        M_deriv_func, prop_deriv = M
         if u.ndim > 1:
-            UM = sp.vstack([sp.diags(u[:, i]) @ M for i in range(u.shape[1])])
+            Mu = [M_deriv_func(u[:, i]) for i in range(u.shape[1])]
+            if v is None:
+                Mu = sp.vstack([M @ prop_deriv for M in Mu])
+                if adjoint:
+                    return Mu.T
+                return Mu
+            elif v.ndim > 1:
+                v = np.squeeze(v)
+            if adjoint:
+                return sum(
+                    [prop_deriv.T @ (Mu[i].T @ v[..., i]) for i in range(u.shape[1])]
+                )
+            pv = prop_deriv @ v
+            return np.stack([M @ pv for M in Mu], axis=-1)
         else:
-            U = sp.diags(u, format="csr")
-            UM = U @ M
-        if adjoint:
-            return UM.T
-        return UM
+            Mu = M_deriv_func(u)
+            if v is None:
+                Mu = Mu @ prop_deriv
+                if adjoint:
+                    return Mu.T
+                return Mu
+            elif v.ndim > 1:
+                v = np.squeeze(v)
+            if adjoint:
+                return prop_deriv.T @ (Mu.T @ v)
+            return Mu @ (prop_deriv @ v)
+    else:
+        raise TypeError(
+            "The stashed property derivative is an unexpected type. Expected either a `tuple` or a "
+            f"sparse matrix. Received a {type(M)}."
+        )
 
 
 def with_property_mass_matrices(property_name):
@@ -232,10 +268,22 @@ def with_property_mass_matrices(property_name):
                 return Zero()
             stash_name = f"_Mf_{arg}_deriv"
             if getattr(self, stash_name, None) is None:
-                M_prop_deriv = self.mesh.get_face_inner_product_deriv(
-                    np.ones(self.mesh.n_cells)
-                )(np.ones(self.mesh.n_faces)) * getattr(self, f"{arg.lower()}Deriv")
-                setattr(self, stash_name, M_prop_deriv)
+                prop = getattr(self, arg.lower())
+                t_type = TensorType(self.mesh, prop)
+
+                M_deriv_func = self.mesh.get_face_inner_product_deriv(model=prop)
+                prop_deriv = getattr(self, f"{arg.lower()}Deriv")
+                # t_type == 3 for full tensor model, t_type < 3 for scalar, isotropic, or axis-aligned anisotropy.
+                if t_type < 3 and self.mesh._meshType.lower() in (
+                    "cyl",
+                    "tensor",
+                    "tree",
+                ):
+                    M_prop_deriv = M_deriv_func(np.ones(self.mesh.n_faces)) @ prop_deriv
+                    setattr(self, stash_name, M_prop_deriv)
+                else:
+                    setattr(self, stash_name, (M_deriv_func, prop_deriv))
+
             return __inner_mat_mul_op(
                 getattr(self, stash_name), u, v=v, adjoint=adjoint
             )
@@ -252,10 +300,21 @@ def with_property_mass_matrices(property_name):
                 return Zero()
             stash_name = f"_Me_{arg}_deriv"
             if getattr(self, stash_name, None) is None:
-                M_prop_deriv = self.mesh.get_edge_inner_product_deriv(
-                    np.ones(self.mesh.n_cells)
-                )(np.ones(self.mesh.n_edges)) * getattr(self, f"{arg.lower()}Deriv")
-                setattr(self, stash_name, M_prop_deriv)
+                prop = getattr(self, arg.lower())
+                t_type = TensorType(self.mesh, prop)
+
+                M_deriv_func = self.mesh.get_edge_inner_product_deriv(model=prop)
+                prop_deriv = getattr(self, f"{arg.lower()}Deriv")
+                # t_type == 3 for full tensor model, t_type < 3 for scalar, isotropic, or axis-aligned anisotropy.
+                if t_type < 3 and self.mesh._meshType.lower() in (
+                    "cyl",
+                    "tensor",
+                    "tree",
+                ):
+                    M_prop_deriv = M_deriv_func(np.ones(self.mesh.n_edges)) @ prop_deriv
+                    setattr(self, stash_name, M_prop_deriv)
+                else:
+                    setattr(self, stash_name, (M_deriv_func, prop_deriv))
             return __inner_mat_mul_op(
                 getattr(self, stash_name), u, v=v, adjoint=adjoint
             )


Test patch:
diff --git a/tests/base/test_mass_matrices.py b/tests/base/test_mass_matrices.py
index 396a75cd2..7dd1d82eb 100644
--- a/tests/base/test_mass_matrices.py
+++ b/tests/base/test_mass_matrices.py
@@ -6,6 +6,8 @@ import numpy as np
 from scipy.constants import mu_0
 from discretize.tests import check_derivative
 from discretize.utils import Zero
+import scipy.sparse as sp
+import pytest
 
 
 # define a very simple class...
@@ -41,9 +43,24 @@ class TestSim(unittest.TestCase):
         self.mesh = discretize.TensorMesh([5, 6, 7])
 
         self.sim = SimpleSim(self.mesh, sigmaMap=maps.ExpMap())
-        self.start_mod = np.log(1e-2 * np.ones(self.mesh.n_cells)) + np.random.randn(
-            self.mesh.n_cells
-        )
+        n_cells = self.mesh.n_cells
+        self.start_mod = np.log(np.full(n_cells, 1e-2)) + np.random.randn(n_cells)
+        self.start_diag_mod = np.r_[
+            np.log(np.full(n_cells, 1e-2)),
+            np.log(np.full(n_cells, 2e-2)),
+            np.log(np.full(n_cells, 3e-2)),
+        ] + np.random.randn(3 * n_cells)
+
+        self.sim_full_aniso = SimpleSim(self.mesh, sigmaMap=maps.IdentityMap())
+
+        self.start_full_mod = np.r_[
+            np.full(n_cells, 1),
+            np.full(n_cells, 2),
+            np.full(n_cells, 3),
+            np.full(n_cells, -1),
+            np.full(n_cells, 1),
+            np.full(n_cells, -2),
+        ]
 
     def test_zero_returns(self):
         n_c = self.mesh.n_cells
@@ -180,6 +197,66 @@ class TestSim(unittest.TestCase):
             UM @ v, sim.MfSigmaDeriv(u2, v).reshape(-1, order="F")
         )
 
+    def test_forward_anis_expected_shapes(self):
+        sim = self.sim
+        sim.model = self.start_full_mod
+
+        n_f = self.mesh.n_faces
+        n_p = sim.model.size
+        # if U.shape (*, )
+        u = np.random.rand(n_f)
+        v = np.random.randn(n_p)
+        u2 = np.random.rand(n_f, 2)
+        v2 = np.random.randn(n_p, 4)
+
+        # These cases should all return an array of shape (n_f, )
+        # if V.shape (*, )
+        out = sim.MfSigmaDeriv(u, v)
+        assert out.shape == (n_f,)
+        out = sim.MfSigmaDeriv(u, v[:, None])
+        assert out.shape == (n_f,)
+        out = sim.MfSigmaDeriv(u[:, None], v)
+        assert out.shape == (n_f,)
+        out = sim.MfSigmaDeriv(u[:, None], v[:, None])
+        assert out.shape == (n_f,)
+
+        # now check passing multiple V's
+        out = sim.MfSigmaDeriv(u, v2)
+        assert out.shape == (n_f, 4)
+        out = sim.MfSigmaDeriv(u[:, None], v2)
+        assert out.shape == (n_f, 4)
+
+        # also ensure it properly broadcasted the operation....
+        out_2 = np.empty_like(out)
+        for i in range(v2.shape[1]):
+            out_2[:, i] = sim.MfSigmaDeriv(u[:, None], v2[:, i])
+        np.testing.assert_equal(out, out_2)
+
+        # now check for multiple source polarizations
+        out = sim.MfSigmaDeriv(u2, v)
+        assert out.shape == (n_f, 2)
+        out = sim.MfSigmaDeriv(u2, v[:, None])
+        assert out.shape == (n_f, 2)
+
+        # and with multiple RHS
+        out = sim.MfSigmaDeriv(u2, v2)
+        assert out.shape == (n_f, v2.shape[1], 2)
+
+        # and test broadcasting here...
+        out_2 = np.empty_like(out)
+        for i in range(v2.shape[1]):
+            out_2[:, i, :] = sim.MfSigmaDeriv(u2, v2[:, i])
+        np.testing.assert_equal(out, out_2)
+
+        # test None as v
+        UM = sim.MfSigmaDeriv(u)
+        np.testing.assert_allclose(UM @ v, sim.MfSigmaDeriv(u, v))
+
+        UM = sim.MfSigmaDeriv(u2)
+        np.testing.assert_allclose(
+            UM @ v, sim.MfSigmaDeriv(u2, v).reshape(-1, order="F")
+        )
+
     def test_adjoint_expected_shapes(self):
         sim = self.sim
         sim.model = self.start_mod
@@ -242,7 +319,69 @@ class TestSim(unittest.TestCase):
             UMT @ v2_2.reshape(-1, order="F"), sim.MfSigmaDeriv(u2, v2_2, adjoint=True)
         )
 
-    def test_adjoint_opp_shapes(self):
+    def test_adjoint_anis_expected_shapes(self):
+        sim = self.sim
+        sim.model = self.start_full_mod
+
+        n_f = self.mesh.n_faces
+        n_p = sim.model.size
+
+        u = np.random.rand(n_f)
+        v = np.random.randn(n_f)
+        v2 = np.random.randn(n_f, 4)
+        u2 = np.random.rand(n_f, 2)
+        v2_2 = np.random.randn(n_f, 2)
+        v3 = np.random.rand(n_f, 4, 2)
+
+        # These cases should all return an array of shape (n_c, )
+        # if V.shape (n_f, )
+        out = sim.MfSigmaDeriv(u, v, adjoint=True)
+        assert out.shape == (n_p,)
+        out = sim.MfSigmaDeriv(u, v[:, None], adjoint=True)
+        assert out.shape == (n_p,)
+        out = sim.MfSigmaDeriv(u[:, None], v, adjoint=True)
+        assert out.shape == (n_p,)
+        out = sim.MfSigmaDeriv(u[:, None], v[:, None], adjoint=True)
+        assert out.shape == (n_p,)
+
+        # now check passing multiple V's
+        out = sim.MfSigmaDeriv(u, v2, adjoint=True)
+        assert out.shape == (n_p, 4)
+        out = sim.MfSigmaDeriv(u[:, None], v2, adjoint=True)
+        assert out.shape == (n_p, 4)
+
+        # also ensure it properly broadcasted the operation....
+        out_2 = np.empty_like(out)
+        for i in range(v2.shape[1]):
+            out_2[:, i] = sim.MfSigmaDeriv(u, v2[:, i], adjoint=True)
+        np.testing.assert_equal(out, out_2)
+
+        # now check for multiple source polarizations
+        out = sim.MfSigmaDeriv(u2, v2_2, adjoint=True)
+        assert out.shape == (n_p,)
+        out = sim.MfSigmaDeriv(u2, v2_2, adjoint=True)
+        assert out.shape == (n_p,)
+
+        # and with multiple RHS
+        out = sim.MfSigmaDeriv(u2, v3, adjoint=True)
+        assert out.shape == (n_p, v3.shape[1])
+
+        # and test broadcasting here...
+        out_2 = np.empty_like(out)
+        for i in range(v2.shape[1]):
+            out_2[:, i] = sim.MfSigmaDeriv(u2, v3[:, i, :], adjoint=True)
+        np.testing.assert_equal(out, out_2)
+
+        # test None as v
+        UMT = sim.MfSigmaDeriv(u, adjoint=True)
+        np.testing.assert_allclose(UMT @ v, sim.MfSigmaDeriv(u, v, adjoint=True))
+
+        UMT = sim.MfSigmaDeriv(u2, adjoint=True)
+        np.testing.assert_allclose(
+            UMT @ v2_2.reshape(-1, order="F"), sim.MfSigmaDeriv(u2, v2_2, adjoint=True)
+        )
+
+    def test_adjoint_opp(self):
         sim = self.sim
         sim.model = self.start_mod
 
@@ -301,6 +440,44 @@ class TestSim(unittest.TestCase):
         yJtv = np.sum(y2 * sim.MfSigmaIDeriv(u2, v3, adjoint=True))
         np.testing.assert_allclose(vJy, yJtv)
 
+    def test_anis_adjoint_opp(self):
+        sim = self.sim
+        sim.model = self.start_full_mod
+
+        n_f = self.mesh.n_faces
+        n_p = sim.model.size
+
+        u = np.random.rand(n_f)
+        u2 = np.random.rand(n_f, 2)
+
+        y = np.random.rand(n_p)
+        y2 = np.random.rand(n_p, 4)
+
+        v = np.random.randn(n_f)
+        v2 = np.random.randn(n_f, 4)
+        v2_2 = np.random.randn(n_f, 2)
+        v3 = np.random.rand(n_f, 4, 2)
+
+        # u1, y1 -> v1
+        vJy = v @ sim.MfSigmaDeriv(u, y)
+        yJtv = y @ sim.MfSigmaDeriv(u, v, adjoint=True)
+        np.testing.assert_allclose(vJy, yJtv)
+
+        # u1, y2 -> v2
+        vJy = np.sum(v2 * sim.MfSigmaDeriv(u, y2))
+        yJtv = np.sum(y2 * sim.MfSigmaDeriv(u, v2, adjoint=True))
+        np.testing.assert_allclose(vJy, yJtv)
+
+        # u2, y1 -> v2_2
+        vJy = np.sum(v2_2 * sim.MfSigmaDeriv(u2, y))
+        yJtv = np.sum(y * sim.MfSigmaDeriv(u2, v2_2, adjoint=True))
+        np.testing.assert_allclose(vJy, yJtv)
+
+        # u2, y2 -> v3
+        vJy = np.sum(v3 * sim.MfSigmaDeriv(u2, y2))
+        yJtv = np.sum(y2 * sim.MfSigmaDeriv(u2, v3, adjoint=True))
+        np.testing.assert_allclose(vJy, yJtv)
+
     def test_Mcc_deriv(self):
         u = np.random.randn(self.mesh.n_cells)
         sim = self.sim
@@ -352,6 +529,40 @@ class TestSim(unittest.TestCase):
 
         assert check_derivative(f, x0=x0, num=3, plotIt=False)
 
+    def test_Me_diagonal_anisotropy_deriv(self):
+        u = np.random.randn(self.mesh.n_edges)
+        sim = self.sim
+        x0 = self.start_diag_mod
+
+        def f(x):
+            sim.model = x
+            d = sim.MeSigma @ u
+
+            def Jvec(v):
+                sim.model = x0
+                return sim.MeSigmaDeriv(u, v)
+
+            return d, Jvec
+
+        assert check_derivative(f, x0=x0, num=3, plotIt=False)
+
+    def test_Me_full_anisotropy_deriv(self):
+        u = np.random.randn(self.mesh.n_edges)
+        sim = self.sim_full_aniso
+        x0 = self.start_full_mod
+
+        def f(x):
+            sim.model = x
+            d = sim.MeSigma @ u
+
+            def Jvec(v):
+                sim.model = x0
+                return sim.MeSigmaDeriv(u, v)
+
+            return d, Jvec
+
+        assert check_derivative(f, x0=x0, num=3, plotIt=False)
+
     def test_Mf_deriv(self):
         u = np.random.randn(self.mesh.n_faces)
         sim = self.sim
@@ -369,6 +580,40 @@ class TestSim(unittest.TestCase):
 
         assert check_derivative(f, x0=x0, num=3, plotIt=False)
 
+    def test_Mf_diagonal_anisotropy_deriv(self):
+        u = np.random.randn(self.mesh.n_faces)
+        sim = self.sim
+        x0 = self.start_diag_mod
+
+        def f(x):
+            sim.model = x
+            d = sim.MfSigma @ u
+
+            def Jvec(v):
+                sim.model = x0
+                return sim.MfSigmaDeriv(u, v)
+
+            return d, Jvec
+
+        assert check_derivative(f, x0=x0, num=3, plotIt=False)
+
+    def test_Mf_full_anisotropy_deriv(self):
+        u = np.random.randn(self.mesh.n_faces)
+        sim = self.sim_full_aniso
+        x0 = self.start_full_mod
+
+        def f(x):
+            sim.model = x
+            d = sim.MfSigma @ u
+
+            def Jvec(v):
+                sim.model = x0
+                return sim.MfSigmaDeriv(u, v)
+
+            return d, Jvec
+
+        assert check_derivative(f, x0=x0, num=3, plotIt=False)
+
     def test_MccI_deriv(self):
         u = np.random.randn(self.mesh.n_cells)
         sim = self.sim
@@ -540,3 +785,24 @@ class TestSim(unittest.TestCase):
         yJv = y @ sim.MfSigmaIDeriv(u, v)
         vJty = v @ sim.MfSigmaIDeriv(u, y, adjoint=True)
         np.testing.assert_allclose(yJv, vJty)
+
+
+def test_bad_derivative_stash():
+    mesh = discretize.TensorMesh([5, 6, 7])
+    sim = SimpleSim(mesh, sigmaMap=maps.ExpMap())
+    sim.model = np.random.rand(mesh.n_cells)
+
+    u = np.random.rand(mesh.n_edges)
+    v = np.random.rand(mesh.n_cells)
+
+    # This should work
+    sim.MeSigmaDeriv(u, v)
+    # stashed derivative operation is a sparse matrix
+    assert sp.issparse(sim._Me_Sigma_deriv)
+
+    # Let's set the stashed item as a bad value which would error
+    # The user shouldn't cause this to happen, but a developer might.
+    sim._Me_Sigma_deriv = [40, 10, 30]
+
+    with pytest.raises(TypeError):
+        sim.MeSigmaDeriv(u, v)