Repository: simpeg/simpeg
Pull request #1103 (2022-09-15 20:54:41)

Amplitude regularization
New implementation for MVI sparse inversion to keep the problem (mostly) linear.
 - Implementation of `VectorAmplitude` regularization

$\phi_s = \|| \mathbf{W} \mathbf{P} \mathbf{a}(m - m_{ref}) \||_2^2$

 - Implementation of `Group` map as extension of `Wire` with derivatives of wires returned as a list.

Solution with amplitude regularizer

![Sparse_MVI](https://user-images.githubusercontent.com/55204635/190506462-a1fc11fb-b449-40b5-935e-4d65bba8e464.png)






Patch:
diff --git a/.gitignore b/.gitignore
index 0b177b79a..9418df3e9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,6 +17,7 @@ develop-eggs
 .installed.cfg
 lib
 lib64
+.idea
 __pycache__
 
 # Installer logs
diff --git a/SimPEG/directives/directives.py b/SimPEG/directives/directives.py
index 405039ff4..e9ceb80e1 100644
--- a/SimPEG/directives/directives.py
+++ b/SimPEG/directives/directives.py
@@ -2082,12 +2082,17 @@ class Update_IRLS(InversionDirective):
         if self.mode == 1:
             self.norms = []
             for reg in self.reg.objfcts:
+                if not isinstance(reg, Sparse):
+                    continue
                 self.norms.append(reg.norms)
                 reg.norms = [2.0 for obj in reg.objfcts]
                 reg.model = self.invProb.model
 
         # Update the model used by the regularization
         for reg in self.reg.objfcts:
+            if not isinstance(reg, Sparse):
+                continue
+
             reg.model = self.invProb.model
 
         if self.sphericalDomain:
@@ -2140,6 +2145,9 @@ class Update_IRLS(InversionDirective):
 
             # Print to screen
             for reg in self.reg.objfcts:
+                if not isinstance(reg, Sparse):
+                    continue
+
                 for obj in reg.objfcts:
                     if isinstance(reg, (Sparse, BaseSparse)):
                         obj.irls_threshold = obj.irls_threshold / self.coolEpsFact
@@ -2149,8 +2157,10 @@ class Update_IRLS(InversionDirective):
             # Reset the regularization matrices so that it is
             # recalculated for current model. Do it to all levels of comboObj
             for reg in self.reg.objfcts:
-                if isinstance(reg, (Sparse, BaseSparse)):
-                    reg.update_weights(reg.model)
+                if not isinstance(reg, Sparse):
+                    continue
+
+                reg.update_weights(reg.model)
 
             self.update_beta = True
             self.invProb.phi_m_last = self.reg(self.invProb.model)
@@ -2181,7 +2191,6 @@ class Update_IRLS(InversionDirective):
                 threshold = np.percentile(
                     np.abs(obj.mapping * obj._delta_m(self.invProb.model)), self.prctile
                 )
-
                 if isinstance(obj, SmoothnessFirstOrder):
                     threshold /= reg.regularization_mesh.base_length
 
@@ -2189,6 +2198,8 @@ class Update_IRLS(InversionDirective):
 
         # Re-assign the norms supplied by user l2 -> lp
         for reg, norms in zip(self.reg.objfcts, self.norms):
+            if not isinstance(reg, Sparse):
+                continue
             reg.norms = norms
 
         # Save l2-model
@@ -2196,6 +2207,8 @@ class Update_IRLS(InversionDirective):
 
         # Print to screen
         for reg in self.reg.objfcts:
+            if not isinstance(reg, Sparse):
+                continue
             if not self.silent:
                 print("irls_threshold " + str(reg.objfcts[0].irls_threshold))
 
@@ -2704,18 +2717,20 @@ class UpdateSensitivityWeights(InversionDirective):
         # Compute and sum root-mean squared sensitivities for all objective functions
         wr = np.zeros_like(self.invProb.model)
         for reg in self.reg.objfcts:
-            if not isinstance(reg, BaseSimilarityMeasure):
-                mesh = reg.regularization_mesh
-                n_cells = mesh.nC
+            if isinstance(reg, BaseSimilarityMeasure):
+                continue
+
+            mesh = reg.regularization_mesh
+            n_cells = mesh.nC
+            mapped_jtj_diag = reg.mapping * jtj_diag
+            # reshape the mapped, so you can divide by volume
+            # (let's say it was a vector or anisotropic model)
+            mapped_jtj_diag = mapped_jtj_diag.reshape((n_cells, -1), order="F")
+            wr_temp = mapped_jtj_diag / reg.regularization_mesh.vol[:, None] ** 2.0
+            wr_temp = wr_temp.reshape(-1, order="F")
 
-                mapped_jtj_diag = reg.mapping * jtj_diag
-                # reshape the mapped, so you can divide by volume
-                # (let's say it was a vector or anisotropic model)
-                mapped_jtj_diag = mapped_jtj_diag.reshape((n_cells, -1), order="F")
-                wr_temp = mapped_jtj_diag / reg.regularization_mesh.vol[:, None] ** 2.0
-                wr_temp = wr_temp.reshape(-1, order="F")
+            wr += reg.mapping.deriv(self.invProb.model).T * wr_temp
 
-                wr += reg.mapping.deriv(self.invProb.model).T * wr_temp
         wr **= 0.5
 
         # Apply thresholding
diff --git a/SimPEG/maps.py b/SimPEG/maps.py
index 9dd4787ba..443774b24 100644
--- a/SimPEG/maps.py
+++ b/SimPEG/maps.py
@@ -3240,6 +3240,8 @@ class InjectActiveCells(IdentityMap):
         return int(self.indActive.sum())
 
     def _transform(self, m):
+        if m.ndim > 1:
+            return self.P * m + self.valInactive[:, None]
         return self.P * m + self.valInactive
 
     def inverse(self, u):
diff --git a/SimPEG/objective_function.py b/SimPEG/objective_function.py
index b63d1bfb7..2ac618541 100644
--- a/SimPEG/objective_function.py
+++ b/SimPEG/objective_function.py
@@ -27,7 +27,7 @@ class BaseObjectiveFunction(BaseSimPEG):
     mapPair = IdentityMap  #: Base class of expected maps
     _mapping = None  #: An IdentityMap instance.
     _has_fields = False  #: should we have the option to store fields
-
+    _parent = None  #: parent objective function
     _nP = None  #: number of parameters
 
     def __init__(self, nP=None, **kwargs):
@@ -265,6 +265,8 @@ class ComboObjectiveFunction(BaseObjectiveFunction):
                     else:
                         self._nP = fct.nP
 
+                fct.parent = self
+
         validate_list(objfcts, multipliers)
 
         self.objfcts = objfcts
diff --git a/SimPEG/regularization/__init__.py b/SimPEG/regularization/__init__.py
index 40a9e3eda..998e1ca73 100644
--- a/SimPEG/regularization/__init__.py
+++ b/SimPEG/regularization/__init__.py
@@ -206,6 +206,11 @@ from .pgi import PGIsmallness, PGI
 from .cross_gradient import CrossGradient
 from .correspondence import LinearCorrespondence
 from .jtv import JointTotalVariation
+from .vector import (
+    BaseVectorRegularization,
+    CrossReferenceRegularization,
+    VectorAmplitude,
+)
 
 from .vector import BaseVectorRegularization, CrossReferenceRegularization
 
diff --git a/SimPEG/regularization/base.py b/SimPEG/regularization/base.py
index 0979c8dcd..834bca1b0 100644
--- a/SimPEG/regularization/base.py
+++ b/SimPEG/regularization/base.py
@@ -32,6 +32,8 @@ class BaseRegularization(BaseObjectiveFunction):
     """
 
     _model = None
+    _parent = None
+    _W = None
 
     def __init__(
         self,
@@ -137,6 +139,19 @@ class BaseRegularization(BaseObjectiveFunction):
             )
         self._mapping = mapping
 
+    @property
+    def parent(self):
+        """
+        The parent objective function
+        """
+        return self._parent
+
+    @parent.setter
+    def parent(self, parent):
+        if not isinstance(parent, ComboObjectiveFunction):
+            raise TypeError("Parent must be a ComboObjectiveFunction")
+        self._parent = parent
+
     @property
     def units(self) -> str | None:
         """Specify the model units. Special care given to 'radian' values"""
diff --git a/SimPEG/regularization/sparse.py b/SimPEG/regularization/sparse.py
index 071d0822d..a10cd45e4 100644
--- a/SimPEG/regularization/sparse.py
+++ b/SimPEG/regularization/sparse.py
@@ -140,30 +140,12 @@ class SparseSmoothness(BaseSparse, SmoothnessFirstOrder):
         """
         Compute and store the irls weights.
         """
-        if self.gradient_type == "total":
-            delta_m = self.mapping * self._delta_m(m)
-            f_m = np.zeros_like(delta_m)
-            for ii, comp in enumerate("xyz"):
-                if self.regularization_mesh.dim > ii:
-                    dm = (
-                        getattr(self.regularization_mesh, f"cell_gradient_{comp}")
-                        * delta_m
-                    )
-
-                    if self.units is not None and self.units.lower() == "radian":
-                        Ave = getattr(self.regularization_mesh, f"aveCC2F{comp}")
-                        length_scales = Ave * (
-                            self.regularization_mesh.Pac.T
-                            * self.regularization_mesh.mesh.h_gridded[:, ii]
-                        )
-                        dm = (
-                            utils.mat_utils.coterminal(dm * length_scales)
-                            / length_scales
-                        )
-
-                    f_m += np.abs(
-                        getattr(self.regularization_mesh, f"aveF{comp}2CC") * dm
-                    )
+        if self.gradient_type == "total" and self.parent is not None:
+            f_m = np.zeros(self.regularization_mesh.nC)
+            for obj in self.parent.objfcts:
+                if isinstance(obj, SparseSmoothness):
+                    avg = getattr(self.regularization_mesh, f"aveF{obj.orientation}2CC")
+                    f_m += np.abs(avg * obj.f_m(m))
 
             f_m = getattr(self.regularization_mesh, f"aveCC2F{self.orientation}") * f_m
 
@@ -224,6 +206,7 @@ class Sparse(WeightedLeastSquares):
         gradient_type="total",
         irls_scaled=True,
         irls_threshold=1e-8,
+        objfcts=None,
         **kwargs,
     ):
         if not isinstance(mesh, RegularizationMesh):
@@ -238,20 +221,21 @@ class Sparse(WeightedLeastSquares):
         if active_cells is not None:
             self._regularization_mesh.active_cells = active_cells
 
-        objfcts = [
-            SparseSmallness(mesh=self.regularization_mesh),
-            SparseSmoothness(mesh=self.regularization_mesh, orientation="x"),
-        ]
+        if objfcts is None:
+            objfcts = [
+                SparseSmallness(mesh=self.regularization_mesh),
+                SparseSmoothness(mesh=self.regularization_mesh, orientation="x"),
+            ]
 
-        if mesh.dim > 1:
-            objfcts.append(
-                SparseSmoothness(mesh=self.regularization_mesh, orientation="y")
-            )
+            if mesh.dim > 1:
+                objfcts.append(
+                    SparseSmoothness(mesh=self.regularization_mesh, orientation="y")
+                )
 
-        if mesh.dim > 2:
-            objfcts.append(
-                SparseSmoothness(mesh=self.regularization_mesh, orientation="z")
-            )
+            if mesh.dim > 2:
+                objfcts.append(
+                    SparseSmoothness(mesh=self.regularization_mesh, orientation="z")
+                )
 
         gradientType = kwargs.pop("gradientType", None)
         super().__init__(
diff --git a/SimPEG/regularization/vector.py b/SimPEG/regularization/vector.py
index 531633e0c..0ae8c1705 100644
--- a/SimPEG/regularization/vector.py
+++ b/SimPEG/regularization/vector.py
@@ -1,10 +1,15 @@
 from __future__ import annotations
-
-# Regularizations for vector models.
+from typing import TYPE_CHECKING
 
 import scipy.sparse as sp
 import numpy as np
-from .base import BaseRegularization, Smallness
+from .base import Smallness
+from discretize.base import BaseMesh
+from .base import RegularizationMesh, BaseRegularization
+from .sparse import Sparse, SparseSmallness, SparseSmoothness
+
+if TYPE_CHECKING:
+    from scipy.sparse import csr_matrix
 
 
 class BaseVectorRegularization(BaseRegularization):
@@ -14,6 +19,13 @@ class BaseVectorRegularization(BaseRegularization):
     your model is made up of vector values in each cell or it is an anisotropic model.
     """
 
+    @property
+    def n_comp(self):
+        """Number of components in the model."""
+        if self.mapping.shape[0] == "*":
+            return self.regularization_mesh.dim
+        return int(self.mapping.shape[0] / self.regularization_mesh.nC)
+
     @property
     def _weights_shapes(self) -> list[tuple[int]]:
         """Acceptable lengths for the weights
@@ -24,7 +36,8 @@ class BaseVectorRegularization(BaseRegularization):
             Each tuple represents accetable shapes for the weights
         """
         mesh = self.regularization_mesh
-        return [(mesh.nC,), (mesh.dim * mesh.nC,), (mesh.nC, mesh.dim)]
+
+        return [(mesh.nC,), (self.n_comp * mesh.nC,), (mesh.nC, self.n_comp)]
 
 
 class CrossReferenceRegularization(Smallness, BaseVectorRegularization):
@@ -151,3 +164,217 @@ class CrossReferenceRegularization(Smallness, BaseVectorRegularization):
                 diag = np.r_[weights, weights, weights]
             self._W = sp.diags(diag, format="csr")
         return self._W
+
+
+class BaseAmplitude(BaseVectorRegularization):
+    """
+    Base vector amplitude class.
+    Requires a mesh and a :obj:`SimPEG.maps.Wires` mapping.
+    """
+
+    def amplitude(self, m):
+        return np.linalg.norm(
+            (self.mapping * self._delta_m(m)).reshape(
+                (self.regularization_mesh.nC, self.n_comp), order="F"
+            ),
+            axis=1,
+        )
+
+    def deriv(self, m) -> np.ndarray:
+        """ """
+        d_m = self._delta_m(m)
+
+        return self.f_m_deriv(m).T * (
+            self.W.T
+            @ self.W
+            @ (self.f_m_deriv(m) @ d_m).reshape((-1, self.n_comp), order="F")
+        ).flatten(order="F")
+
+    def deriv2(self, m, v=None) -> csr_matrix:
+        """ """
+        f_m_deriv = self.f_m_deriv(m)
+
+        if v is None:
+            return f_m_deriv.T * (
+                sp.block_diag([self.W.T * self.W] * self.n_comp) * f_m_deriv
+            )
+
+        return f_m_deriv.T * (
+            self.W.T @ self.W @ (f_m_deriv * v).reshape((-1, self.n_comp), order="F")
+        ).flatten(order="F")
+
+
+class AmplitudeSmallness(SparseSmallness, BaseAmplitude):
+    """
+    Sparse smallness regularization on vector amplitude.
+    """
+
+    def f_m(self, m):
+        """
+        Compute the amplitude of a vector model.
+        """
+
+        return self.amplitude(m)
+
+    @property
+    def W(self):
+        if getattr(self, "_W", None) is None:
+            mesh = self.regularization_mesh
+            nC = mesh.nC
+
+            weights = np.ones(
+                nC,
+            )
+            for value in self._weights.values():
+                if value.shape == (nC,):
+                    weights *= value
+                elif value.shape == (self.n_comp * nC,):
+                    weights *= np.linalg.norm(
+                        value.reshape((nC, self.n_comp), order="F"), axis=1
+                    )
+
+            self._W = sp.diags(np.sqrt(weights), format="csr")
+
+        return self._W
+
+
+class AmplitudeSmoothnessFirstOrder(SparseSmoothness, BaseAmplitude):
+    """
+    Sparse first spatial derivatives of amplitude.
+    """
+
+    @property
+    def _weights_shapes(self) -> list[tuple[int]]:
+        """Acceptable lengths for the weights
+
+        Returns
+        -------
+        list of tuple
+            Each tuple represents accetable shapes for the weights
+        """
+        nC = self.regularization_mesh.nC
+        nF = getattr(
+            self.regularization_mesh, "aveCC2F{}".format(self.orientation)
+        ).shape[0]
+        return [
+            (nF,),
+            (self.n_comp * nF,),
+            (nF, self.n_comp),
+            (nC,),
+            (self.n_comp * nC,),
+            (nC, self.n_comp),
+        ]
+
+    def f_m(self, m):
+        a = self.amplitude(m)
+
+        return self.cell_gradient @ a
+
+    def f_m_deriv(self, m) -> csr_matrix:
+        """"""
+        return sp.block_diag([self.cell_gradient] * self.n_comp) @ self.mapping.deriv(
+            self._delta_m(m)
+        )
+
+    @property
+    def W(self):
+        """
+        Weighting matrix that takes the volumes, free weights, fixed weights and
+        length scales of the difference operator (normalized optional).
+        """
+        if getattr(self, "_W", None) is None:
+            average_cell_2_face = getattr(
+                self.regularization_mesh, "aveCC2F{}".format(self.orientation)
+            )
+            nC = self.regularization_mesh.nC
+            nF = average_cell_2_face.shape[0]
+            weights = 1.0
+            for values in self._weights.values():
+                if values.shape[0] == nC:
+                    values = average_cell_2_face * values
+                elif not values.shape == (nF,):
+                    values = np.linalg.norm(
+                        values.reshape((-1, self.n_comp), order="F"), axis=1
+                    )
+                    if values.size == nC:
+                        values = average_cell_2_face * values
+
+                weights *= values
+
+            self._W = sp.diags(np.sqrt(weights), format="csr")
+
+        return self._W
+
+
+class VectorAmplitude(Sparse):
+    r"""
+    The regularization is:
+
+    The function defined here approximates:
+
+    .. math::
+
+        \phi_m(\mathbf{m}) = \alpha_s \| \mathbf{W}_s \; \mathbf{a}(\mathbf{m} - \mathbf{m_{ref}) \|_p
+        + \alpha_x \| \mathbf{W}_x \; \frac{\partial}{\partial x} \mathbf{a}(\mathbf{m} - \mathbf{m_{ref}) \|_p
+        + \alpha_y \| \mathbf{W}_y \; \frac{\partial}{\partial y} \mathbf{a}(\mathbf{m} - \mathbf{m_{ref}) \|_p
+        + \alpha_z \| \mathbf{W}_z \; \frac{\partial}{\partial z} \mathbf{a}(\mathbf{m} - \mathbf{m_{ref}) \|_p
+
+    where $\mathbf{a}(\mathbf{m} - \mathbf{m_{ref})$ is the vector amplitude of the difference between
+    the model and the reference model.
+
+    .. math::
+
+        \mathbf{a}(\mathbf{m} - \mathbf{m_{ref}) = [\sum_{i}^{N}(\mathbf{P}_i\;(\mathbf{m} - \mathbf{m_{ref}}))^{2}]^{1/2}
+
+    where :math:`\mathbf{P}_i` is the projection of i-th component of the vector model with N-dimensions.
+
+    """
+
+    def __init__(
+        self,
+        mesh,
+        mapping=None,
+        active_cells=None,
+        **kwargs,
+    ):
+        if not isinstance(mesh, (BaseMesh, RegularizationMesh)):
+            raise TypeError(
+                f"'regularization_mesh' must be of type {RegularizationMesh} or {BaseMesh}. "
+                f"Value of type {type(mesh)} provided."
+            )
+
+        if not isinstance(mesh, RegularizationMesh):
+            mesh = RegularizationMesh(mesh)
+
+        self._regularization_mesh = mesh
+
+        if active_cells is not None:
+            self._regularization_mesh.active_cells = active_cells
+
+        objfcts = [
+            AmplitudeSmallness(mesh=self.regularization_mesh, mapping=mapping),
+            AmplitudeSmoothnessFirstOrder(
+                mesh=self.regularization_mesh, orientation="x", mapping=mapping
+            ),
+        ]
+
+        if mesh.dim > 1:
+            objfcts.append(
+                AmplitudeSmoothnessFirstOrder(
+                    mesh=self.regularization_mesh, orientation="y", mapping=mapping
+                )
+            )
+
+        if mesh.dim > 2:
+            objfcts.append(
+                AmplitudeSmoothnessFirstOrder(
+                    mesh=self.regularization_mesh, orientation="z", mapping=mapping
+                )
+            )
+
+        super().__init__(
+            self.regularization_mesh,
+            objfcts=objfcts,
+            mapping=mapping,
+            **kwargs,
+        )
diff --git a/examples/03-magnetics/plot_inv_mag_MVI_Sparse_TreeMesh.py b/examples/03-magnetics/plot_inv_mag_MVI_Sparse_TreeMesh.py
index cc9176afd..c24b4f6b8 100644
--- a/examples/03-magnetics/plot_inv_mag_MVI_Sparse_TreeMesh.py
+++ b/examples/03-magnetics/plot_inv_mag_MVI_Sparse_TreeMesh.py
@@ -16,7 +16,6 @@ that
 
 """
 
-from discretize import TreeMesh
 from SimPEG import (
     data,
     data_misfit,
@@ -117,100 +116,6 @@ mesh = refine_tree_xyz(
 actv = active_from_xyz(mesh, topo)
 nC = int(actv.sum())
 
-###########################################################################
-# A simple function to plot vectors in TreeMesh
-#
-# Should eventually end up on discretize
-#
-
-
-def plotVectorSectionsOctree(
-    mesh,
-    m,
-    normal="X",
-    ind=0,
-    vmin=None,
-    vmax=None,
-    scale=1.0,
-    vec="k",
-    axs=None,
-    actvMap=None,
-    fill=True,
-):
-    """
-    Plot section through a 3D tensor model
-    """
-    # plot recovered model
-    normalInd = {"X": 0, "Y": 1, "Z": 2}[normal]
-    antiNormalInd = {"X": [1, 2], "Y": [0, 2], "Z": [0, 1]}[normal]
-
-    h2d = (mesh.h[antiNormalInd[0]], mesh.h[antiNormalInd[1]])
-    x2d = (mesh.x0[antiNormalInd[0]], mesh.x0[antiNormalInd[1]])
-
-    #: Size of the sliced dimension
-    szSliceDim = len(mesh.h[normalInd])
-    if ind is None:
-        ind = int(szSliceDim // 2)
-
-    cc_tensor = [None, None, None]
-    for i in range(3):
-        cc_tensor[i] = np.cumsum(np.r_[mesh.x0[i], mesh.h[i]])
-        cc_tensor[i] = (cc_tensor[i][1:] + cc_tensor[i][:-1]) * 0.5
-    slice_loc = cc_tensor[normalInd][ind]
-
-    # Create a temporary TreeMesh with the slice through
-    temp_mesh = TreeMesh(h2d, x2d)
-    level_diff = mesh.max_level - temp_mesh.max_level
-
-    XS = [None, None, None]
-    XS[antiNormalInd[0]], XS[antiNormalInd[1]] = np.meshgrid(
-        cc_tensor[antiNormalInd[0]], cc_tensor[antiNormalInd[1]]
-    )
-    XS[normalInd] = np.ones_like(XS[antiNormalInd[0]]) * slice_loc
-    loc_grid = np.c_[XS[0].reshape(-1), XS[1].reshape(-1), XS[2].reshape(-1)]
-    inds = np.unique(mesh._get_containing_cell_indexes(loc_grid))
-
-    grid2d = mesh.gridCC[inds][:, antiNormalInd]
-    levels = mesh._cell_levels_by_indexes(inds) - level_diff
-    temp_mesh.insert_cells(grid2d, levels)
-    tm_gridboost = np.empty((temp_mesh.nC, 3))
-    tm_gridboost[:, antiNormalInd] = temp_mesh.gridCC
-    tm_gridboost[:, normalInd] = slice_loc
-
-    # Interpolate values to mesh.gridCC if not 'CC'
-    mx = actvMap * m[:, 0]
-    my = actvMap * m[:, 1]
-    mz = actvMap * m[:, 2]
-
-    m = np.c_[mx, my, mz]
-
-    # Interpolate values from mesh.gridCC to grid2d
-    ind_3d_to_2d = mesh._get_containing_cell_indexes(tm_gridboost)
-    v2d = m[ind_3d_to_2d, :]
-    amp = np.sum(v2d**2.0, axis=1) ** 0.5
-
-    if axs is None:
-        axs = plt.subplot(111)
-
-    if fill:
-        temp_mesh.plot_image(amp, ax=axs, clim=[vmin, vmax], grid=True)
-
-    axs.quiver(
-        temp_mesh.gridCC[:, 0],
-        temp_mesh.gridCC[:, 1],
-        v2d[:, antiNormalInd[0]],
-        v2d[:, antiNormalInd[1]],
-        pivot="mid",
-        scale_units="inches",
-        scale=scale,
-        linewidths=(1,),
-        edgecolors=(vec),
-        headaxislength=0.1,
-        headwidth=10,
-        headlength=30,
-    )
-
-
 ###########################################################################
 # Forward modeling data
 # ---------------------
@@ -271,16 +176,19 @@ plt.gca().set_aspect("equal", adjustable="box")
 
 # Plot the vector model
 ax = plt.subplot(2, 1, 2)
-plotVectorSectionsOctree(
-    mesh,
-    model,
-    axs=ax,
+mesh.plot_slice(
+    actv_plot * model.reshape((-1, 3), order="F"),
+    v_type="CCv",
+    view="vec",
+    ax=ax,
     normal="Y",
     ind=66,
-    actvMap=actv_plot,
-    scale=0.5,
-    vmin=0.0,
-    vmax=0.025,
+    grid=True,
+    quiver_opts={
+        "pivot": "mid",
+        "scale": 5 * np.abs(model).max(),
+        "scale_units": "inches",
+    },
 )
 ax.set_xlim([-200, 200])
 ax.set_ylim([-100, 75])
@@ -380,26 +288,35 @@ wires = maps.Wires(("amp", nC), ("theta", nC), ("phi", nC))
 # Create a Combo Regularization
 # Regularize the amplitude of the vectors
 reg_a = regularization.Sparse(
-    mesh, gradient_type="components", active_cells=actv, mapping=wires.amp
+    mesh,
+    gradient_type="total",
+    active_cells=actv,
+    mapping=wires.amp,
+    norms=[0.0, 1.0, 1.0, 1.0],  # Only norm on gradients used,
+    reference_model=np.zeros(3 * nC),
 )
-reg_a.norms = [0.0, 0.0, 0.0, 0.0]  # Sparse on the model and its gradients
-reg_a.reference_model = np.zeros(3 * nC)
 
 # Regularize the vertical angle of the vectors
 reg_t = regularization.Sparse(
-    mesh, gradient_type="components", active_cells=actv, mapping=wires.theta
+    mesh,
+    gradient_type="total",
+    active_cells=actv,
+    mapping=wires.theta,
+    alpha_s=0.0,  # No reference angle,
+    norms=[0.0, 1.0, 1.0, 1.0],  # Only norm on gradients used,
 )
-reg_t.alpha_s = 0.0  # No reference angle
 reg_t.units = "radian"
-reg_t.norms = [0.0, 0.0, 0.0, 0.0]  # Only norm on gradients used
 
 # Regularize the horizontal angle of the vectors
 reg_p = regularization.Sparse(
-    mesh, gradient_type="components", active_cells=actv, mapping=wires.phi
+    mesh,
+    gradient_type="total",
+    active_cells=actv,
+    mapping=wires.phi,
+    alpha_s=0.0,  # No reference angle,
+    norms=[0.0, 1.0, 1.0, 1.0],  # Only norm on gradients used,
 )
-reg_p.alpha_s = 0.0  # No reference angle
 reg_p.units = "radian"
-reg_p.norms = [0.0, 0.0, 0.0, 0.0]  # Only norm on gradients used
 
 reg = reg_a + reg_t + reg_p
 reg.reference_model = np.zeros(3 * nC)
@@ -456,16 +373,19 @@ mrec_MVI_S = inv.run(m_start)
 
 plt.figure(figsize=(8, 8))
 ax = plt.subplot(2, 1, 1)
-plotVectorSectionsOctree(
-    mesh,
-    mrec_MVIC.reshape((nC, 3), order="F"),
-    axs=ax,
+mesh.plot_slice(
+    actv_plot * mrec_MVIC.reshape((nC, 3), order="F"),
+    v_type="CCv",
+    view="vec",
+    ax=ax,
     normal="Y",
-    ind=65,
-    actvMap=actv_plot,
-    scale=0.05,
-    vmin=0.0,
-    vmax=0.005,
+    ind=66,
+    grid=True,
+    quiver_opts={
+        "pivot": "mid",
+        "scale": 5 * np.abs(mrec_MVIC).max(),
+        "scale_units": "inches",
+    },
 )
 ax.set_xlim([-200, 200])
 ax.set_ylim([-100, 75])
@@ -476,23 +396,26 @@ plt.gca().set_aspect("equal", adjustable="box")
 
 ax = plt.subplot(2, 1, 2)
 vec_xyz = utils.mat_utils.spherical2cartesian(
-    invProb.model.reshape((nC, 3), order="F")
+    mrec_MVI_S.reshape((nC, 3), order="F")
 ).reshape((nC, 3), order="F")
 
-plotVectorSectionsOctree(
-    mesh,
-    vec_xyz,
-    axs=ax,
+mesh.plot_slice(
+    actv_plot * vec_xyz,
+    v_type="CCv",
+    view="vec",
+    ax=ax,
     normal="Y",
-    ind=65,
-    actvMap=actv_plot,
-    scale=0.4,
-    vmin=0.0,
-    vmax=0.025,
+    ind=66,
+    grid=True,
+    quiver_opts={
+        "pivot": "mid",
+        "scale": 5 * np.abs(vec_xyz).max(),
+        "scale_units": "inches",
+    },
 )
 ax.set_xlim([-200, 200])
 ax.set_ylim([-100, 75])
-ax.set_title("Sparse model (Spherical)")
+ax.set_title("Sparse model (L0L2)")
 ax.set_xlabel("x")
 ax.set_ylabel("y")
 plt.gca().set_aspect("equal", adjustable="box")
diff --git a/examples/03-magnetics/plot_inv_mag_MVI_VectorAmplitude.py b/examples/03-magnetics/plot_inv_mag_MVI_VectorAmplitude.py
new file mode 100644
index 000000000..1479020d9
--- /dev/null
+++ b/examples/03-magnetics/plot_inv_mag_MVI_VectorAmplitude.py
@@ -0,0 +1,248 @@
+"""
+Magnetic inversion on a TreeMesh
+================================
+
+In this example, we demonstrate the use of a Magnetic Vector Inverison
+on 3D TreeMesh for the inversion of magnetic data.
+
+The inverse problem uses the :class:'SimPEG.regularization.VectorAmplitude'
+regularization borrowed from ...
+
+"""
+
+from SimPEG import (
+    data,
+    data_misfit,
+    directives,
+    maps,
+    inverse_problem,
+    optimization,
+    inversion,
+    regularization,
+)
+
+from SimPEG import utils
+from SimPEG.utils import mkvc, sdiag
+
+from discretize.utils import mesh_builder_xyz, refine_tree_xyz, active_from_xyz
+from SimPEG.potential_fields import magnetics
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+# sphinx_gallery_thumbnail_number = 3
+
+###############################################################################
+# Setup
+# -----
+#
+# Define the survey and model parameters
+#
+# First we need to define the direction of the inducing field
+# As a simple case, we pick a vertical inducing field of magnitude 50,000 nT.
+#
+#
+np.random.seed(1)
+# We will assume a vertical inducing field
+H0 = (50000.0, 90.0, 0.0)
+
+# Create grid of points for topography
+# Lets create a simple Gaussian topo and set the active cells
+[xx, yy] = np.meshgrid(np.linspace(-200, 200, 50), np.linspace(-200, 200, 50))
+b = 100
+A = 50
+zz = A * np.exp(-0.5 * ((xx / b) ** 2.0 + (yy / b) ** 2.0))
+topo = np.c_[utils.mkvc(xx), utils.mkvc(yy), utils.mkvc(zz)]
+
+# Create an array of observation points
+xr = np.linspace(-100.0, 100.0, 20)
+yr = np.linspace(-100.0, 100.0, 20)
+X, Y = np.meshgrid(xr, yr)
+Z = A * np.exp(-0.5 * ((X / b) ** 2.0 + (Y / b) ** 2.0)) + 5
+
+# Create a MAGsurvey
+xyzLoc = np.c_[mkvc(X.T), mkvc(Y.T), mkvc(Z.T)]
+rxLoc = magnetics.receivers.Point(xyzLoc)
+srcField = magnetics.sources.SourceField(receiver_list=[rxLoc], parameters=H0)
+survey = magnetics.survey.Survey(srcField)
+
+###############################################################################
+# Inversion Mesh
+# --------------
+#
+# Here, we create a TreeMesh with base cell size of 5 m.
+#
+
+# Create a mesh
+h = [5, 5, 5]
+padDist = np.ones((3, 2)) * 100
+
+mesh = mesh_builder_xyz(
+    xyzLoc, h, padding_distance=padDist, depth_core=100, mesh_type="tree"
+)
+mesh = refine_tree_xyz(
+    mesh, topo, method="surface", octree_levels=[2, 6], finalize=True
+)
+
+
+# Define an active cells from topo
+actv = active_from_xyz(mesh, topo)
+nC = int(actv.sum())
+
+###########################################################################
+# Forward modeling data
+# ---------------------
+#
+# We can now create a magnetization model and generate data.
+#
+model_azm_dip = np.zeros((mesh.nC, 2))
+model_amp = np.ones(mesh.nC) * 1e-8
+ind = utils.model_builder.getIndicesBlock(
+    np.r_[-30, -20, -10],
+    np.r_[30, 20, 25],
+    mesh.gridCC,
+)[0]
+model_amp[ind] = 0.05
+model_azm_dip[ind, 0] = 45.0
+model_azm_dip[ind, 1] = 90.0
+
+# Remove air cells
+model_azm_dip = model_azm_dip[actv, :]
+model_amp = model_amp[actv]
+model = sdiag(model_amp) * utils.mat_utils.dip_azimuth2cartesian(
+    model_azm_dip[:, 0], model_azm_dip[:, 1]
+)
+
+# Create reduced identity map
+idenMap = maps.IdentityMap(nP=nC * 3)
+
+# Create the simulation
+simulation = magnetics.simulation.Simulation3DIntegral(
+    survey=survey, mesh=mesh, chiMap=idenMap, ind_active=actv, model_type="vector"
+)
+
+# Compute some data and add some random noise
+d = simulation.dpred(mkvc(model))
+std = 10  # nT
+synthetic_data = d + np.random.randn(len(d)) * std
+wd = np.ones(len(d)) * std
+
+# Assign data and uncertainties to the survey
+data_object = data.Data(survey, dobs=synthetic_data, standard_deviation=wd)
+
+# Create a projection matrix for plotting later
+actv_plot = maps.InjectActiveCells(mesh, actv, np.nan)
+
+
+######################################################################
+# Inversion
+# ---------
+#
+# We can now attempt the inverse calculations.
+#
+
+# Create sensitivity weights from our linear forward operator
+rxLoc = survey.source_field.receiver_list[0].locations
+
+# This Mapping connects the regularizations for the three-component
+# vector model
+wires = maps.Wires(("p", nC), ("s", nC), ("t", nC))
+m0 = np.ones(3 * nC) * 1e-4  # Starting model
+
+# Create the regularization on the amplitude of magnetization
+reg = regularization.VectorAmplitude(
+    mesh,
+    mapping=idenMap,
+    active_cells=actv,
+    reference_model_in_smooth=True,
+    norms=[0.0, 2.0, 2.0, 2.0],
+    gradient_type="total",
+)
+
+# Data misfit function
+dmis = data_misfit.L2DataMisfit(simulation=simulation, data=data_object)
+dmis.W = 1.0 / data_object.standard_deviation
+
+# The optimization scheme
+opt = optimization.ProjectedGNCG(
+    maxIter=20, lower=-10, upper=10.0, maxIterLS=20, maxIterCG=20, tolCG=1e-4
+)
+
+# The inverse problem
+invProb = inverse_problem.BaseInvProblem(dmis, reg, opt)
+
+# Estimate the initial beta factor
+betaest = directives.BetaEstimate_ByEig(beta0_ratio=1e1)
+
+# Add sensitivity weights
+sensitivity_weights = directives.UpdateSensitivityWeights()
+
+# Here is where the norms are applied
+IRLS = directives.Update_IRLS(f_min_change=1e-3, max_irls_iterations=10, beta_tol=5e-1)
+
+# Pre-conditioner
+update_Jacobi = directives.UpdatePreconditioner()
+
+
+inv = inversion.BaseInversion(
+    invProb, directiveList=[sensitivity_weights, IRLS, update_Jacobi, betaest]
+)
+
+# Run the inversion
+mrec = inv.run(m0)
+
+
+#############################################################
+# Final Plot
+# ----------
+#
+# Let's compare the smooth and compact model
+#
+#
+#
+
+plt.figure(figsize=(12, 6))
+ax = plt.subplot(2, 2, 1)
+im = utils.plot_utils.plot2Ddata(xyzLoc, synthetic_data, ax=ax)
+plt.colorbar(im[0])
+ax.set_title("Predicted data.")
+plt.gca().set_aspect("equal", adjustable="box")
+
+for ii, (title, mvec) in enumerate(
+    [("True model", model), ("Smooth model", invProb.l2model), ("Sparse model", mrec)]
+):
+    ax = plt.subplot(2, 2, ii + 2)
+    mesh.plot_slice(
+        actv_plot * mvec.reshape((-1, 3), order="F"),
+        v_type="CCv",
+        view="vec",
+        ax=ax,
+        normal="Y",
+        grid=True,
+        quiver_opts={
+            "pivot": "mid",
+            "scale": 8 * np.abs(mvec).max(),
+            "scale_units": "inches",
+        },
+    )
+    ax.set_xlim([-200, 200])
+    ax.set_ylim([-100, 75])
+    ax.set_title(title)
+    ax.set_xlabel("x")
+    ax.set_ylabel("z")
+    plt.gca().set_aspect("equal", adjustable="box")
+
+plt.show()
+
+print("END")
+# Plot the final predicted data and the residual
+# plt.figure()
+# ax = plt.subplot(1, 2, 1)
+# utils.plot_utils.plot2Ddata(xyzLoc, invProb.dpred, ax=ax)
+# ax.set_title("Predicted data.")
+# plt.gca().set_aspect("equal", adjustable="box")
+#
+# ax = plt.subplot(1, 2, 2)
+# utils.plot_utils.plot2Ddata(xyzLoc, synthetic_data - invProb.dpred, ax=ax)
+# ax.set_title("Data residual.")
+# plt.gca().set_aspect("equal", adjustable="box")


Test patch:
diff --git a/tests/base/test_regularization.py b/tests/base/test_regularization.py
index 4d48f2278..497d48678 100644
--- a/tests/base/test_regularization.py
+++ b/tests/base/test_regularization.py
@@ -28,6 +28,7 @@ IGNORE_ME = [
     "CrossGradient",
     "LinearCorrespondence",
     "JointTotalVariation",
+    "VectorAmplitude",
     "CrossReferenceRegularization",
 ]
 
@@ -551,6 +552,25 @@ class RegularizationTests(unittest.TestCase):
 
             assert reg.gradient_type == "total"  # Check default
 
+    def test_vector_amplitude(self):
+        n_comp = 4
+        mesh = discretize.TensorMesh([8, 7])
+        model = np.random.randn(mesh.nC, n_comp)
+
+        with pytest.raises(TypeError, match="'regularization_mesh' must be of type"):
+            regularization.VectorAmplitude("abc")
+
+        reg = regularization.VectorAmplitude(
+            mesh, maps.IdentityMap(nP=n_comp * mesh.nC)
+        )
+
+        with pytest.raises(ValueError, match=f"'weights' must be one of"):  # noqa: W605
+            reg.set_weights(abc=(1.0, 1.0))
+
+        np.testing.assert_almost_equal(
+            reg.objfcts[0].f_m(model.flatten(order="F")), np.linalg.norm(model, axis=1)
+        )
+
 
 def test_WeightedLeastSquares():
     mesh = discretize.TensorMesh([3, 4, 5])
diff --git a/tests/pf/test_mag_vector_amplitude.py b/tests/pf/test_mag_vector_amplitude.py
new file mode 100644
index 000000000..bfd94afc1
--- /dev/null
+++ b/tests/pf/test_mag_vector_amplitude.py
@@ -0,0 +1,184 @@
+import unittest
+from SimPEG import (
+    directives,
+    maps,
+    inverse_problem,
+    optimization,
+    data_misfit,
+    inversion,
+    utils,
+    regularization,
+)
+
+
+from discretize.utils import mesh_builder_xyz, refine_tree_xyz, active_from_xyz
+import numpy as np
+from SimPEG.potential_fields import magnetics as mag
+import shutil
+
+
+class MVIProblemTest(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(0)
+        H0 = (50000.0, 90.0, 0.0)
+
+        # The magnetization is set along a different
+        # direction (induced + remanence)
+        M = np.array([45.0, 90.0])
+
+        # Create grid of points for topography
+        # Lets create a simple Gaussian topo
+        # and set the active cells
+        [xx, yy] = np.meshgrid(np.linspace(-200, 200, 50), np.linspace(-200, 200, 50))
+        b = 100
+        A = 50
+        zz = A * np.exp(-0.5 * ((xx / b) ** 2.0 + (yy / b) ** 2.0))
+
+        # We would usually load a topofile
+        topo = np.c_[utils.mkvc(xx), utils.mkvc(yy), utils.mkvc(zz)]
+
+        # Create and array of observation points
+        xr = np.linspace(-100.0, 100.0, 20)
+        yr = np.linspace(-100.0, 100.0, 20)
+        X, Y = np.meshgrid(xr, yr)
+        Z = A * np.exp(-0.5 * ((X / b) ** 2.0 + (Y / b) ** 2.0)) + 5
+
+        # Create a MAGsurvey
+        xyzLoc = np.c_[utils.mkvc(X.T), utils.mkvc(Y.T), utils.mkvc(Z.T)]
+        rxLoc = mag.Point(xyzLoc)
+        srcField = mag.SourceField([rxLoc], parameters=H0)
+        survey = mag.Survey(srcField)
+
+        # Create a mesh
+        h = [5, 5, 5]
+        padDist = np.ones((3, 2)) * 100
+
+        mesh = mesh_builder_xyz(
+            xyzLoc, h, padding_distance=padDist, depth_core=100, mesh_type="tree"
+        )
+        mesh = refine_tree_xyz(
+            mesh, topo, method="surface", octree_levels=[4, 4], finalize=True
+        )
+        self.mesh = mesh
+        # Define an active cells from topo
+        actv = active_from_xyz(mesh, topo)
+        nC = int(actv.sum())
+
+        model = np.zeros((mesh.nC, 3))
+
+        # Convert the inclination declination to vector in Cartesian
+        M_xyz = utils.mat_utils.dip_azimuth2cartesian(M[0], M[1])
+
+        # Get the indicies of the magnetized block
+        ind = utils.model_builder.getIndicesBlock(
+            np.r_[-20, -20, -10],
+            np.r_[20, 20, 25],
+            mesh.gridCC,
+        )[0]
+
+        # Assign magnetization values
+        model[ind, :] = np.kron(np.ones((ind.shape[0], 1)), M_xyz * 0.05)
+
+        # Remove air cells
+        self.model = model[actv, :]
+
+        # Create active map to go from reduce set to full
+        self.actvMap = maps.InjectActiveCells(mesh, actv, np.nan)
+
+        # Creat reduced identity map
+        idenMap = maps.IdentityMap(nP=nC * 3)
+
+        # Create the forward model operator
+        sim = mag.Simulation3DIntegral(
+            self.mesh,
+            survey=survey,
+            model_type="vector",
+            chiMap=idenMap,
+            ind_active=actv,
+            store_sensitivities="disk",
+        )
+        self.sim = sim
+
+        # Compute some data and add some random noise
+        data = sim.make_synthetic_data(
+            utils.mkvc(self.model), relative_error=0.0, noise_floor=5.0, add_noise=True
+        )
+
+        reg = regularization.VectorAmplitude(
+            mesh,
+            mapping=idenMap,
+            active_cells=actv,
+            reference_model_in_smooth=True,
+            norms=[0.0, 0.0, 0.0, 0.0],
+            gradient_type="components",
+        )
+
+        # Data misfit function
+        dmis = data_misfit.L2DataMisfit(simulation=sim, data=data)
+        # dmis.W = 1./survey.std
+
+        # Add directives to the inversion
+        opt = optimization.ProjectedGNCG(
+            maxIter=10, lower=-10, upper=10.0, maxIterLS=5, maxIterCG=5, tolCG=1e-4
+        )
+
+        invProb = inverse_problem.BaseInvProblem(dmis, reg, opt)
+
+        # A list of directive to control the inverson
+        betaest = directives.BetaEstimate_ByEig(beta0_ratio=1e1)
+
+        # Here is where the norms are applied
+        # Use pick a treshold parameter empirically based on the distribution of
+        #  model parameters
+        IRLS = directives.Update_IRLS(
+            f_min_change=1e-3, max_irls_iterations=10, beta_tol=5e-1
+        )
+
+        # Pre-conditioner
+        update_Jacobi = directives.UpdatePreconditioner()
+        sensitivity_weights = directives.UpdateSensitivityWeights(everyIter=False)
+        self.inv = inversion.BaseInversion(
+            invProb, directiveList=[sensitivity_weights, IRLS, update_Jacobi, betaest]
+        )
+
+        # Run the inversion
+        self.mstart = np.ones(3 * nC) * 1e-4  # Starting model
+
+    def test_vector_amplitude_inverse(self):
+        # Run the inversion
+        mrec = self.inv.run(self.mstart)
+
+        nC = int(mrec.shape[0] / 3)
+        vec_xyz = mrec.reshape((nC, 3), order="F")
+        residual = np.linalg.norm(vec_xyz - self.model) / np.linalg.norm(self.model)
+
+        # import matplotlib.pyplot as plt
+        # ax = plt.subplot()
+        # self.mesh.plot_slice(
+        #     self.actvMap * mrec.reshape((-1, 3), order="F"),
+        #     v_type="CCv",
+        #     view="vec",
+        #     ax=ax,
+        #     normal="Y",
+        #     grid=True,
+        #     quiver_opts={
+        #         "pivot": "mid",
+        #         "scale": 8 * np.abs(mrec).max(),
+        #         "scale_units": "inches",
+        #     },
+        # )
+        # plt.gca().set_aspect("equal", adjustable="box")
+        #
+        # plt.show()
+
+        self.assertLess(residual, 1)
+        # self.assertTrue(residual < 0.05)
+
+    def tearDown(self):
+        # Clean up the working directory
+        if self.sim.store_sensitivities == "disk":
+            shutil.rmtree(self.sim.sensitivity_path)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/pf/test_pf_quadtree_inversion_linear.py b/tests/pf/test_pf_quadtree_inversion_linear.py
index 028966c81..88f76433e 100644
--- a/tests/pf/test_pf_quadtree_inversion_linear.py
+++ b/tests/pf/test_pf_quadtree_inversion_linear.py
@@ -189,9 +189,6 @@ class QuadTreeLinProblemTest(unittest.TestCase):
             grav_srcField = gravity.SourceField([grav_rxLoc])
             grav_survey = gravity.Survey(grav_srcField)
 
-            # Set only non-zero cells as active
-            self.active_cells = ~(self.model == 0.0)
-
             # Create the gravity forward model operator
             self.grav_sim_active = gravity.SimulationEquivalentSourceLayer(
                 self.mesh,
@@ -268,7 +265,7 @@ class QuadTreeLinProblemTest(unittest.TestCase):
                 lower=-1.0,
                 upper=1.0,
                 maxIterLS=5,
-                maxIterCG=10,
+                maxIterCG=20,
                 tolCG=1e-4,
             )
 
@@ -323,9 +320,17 @@ class QuadTreeLinProblemTest(unittest.TestCase):
             1.0,
         )
 
+        self.active_cells = utils.model_builder.addBlock(
+            self.mesh.cell_centers,
+            np.zeros(self.mesh.nC, dtype=bool),
+            np.r_[-40, -40],
+            np.r_[40, 40],
+            True,
+        )
+
         # Set only non-zero cells as active. Some tests use all cells
         # (by not using `self.active_cells`), and others use the active cells
-        self.active_cells = ~(self.model == 0.0)
+        # self.active_cells = ~(self.model == 0.0)
 
         # Create reduced identity maps. Two versions: for the all-active
         # and the active-subset models
@@ -495,7 +500,7 @@ class QuadTreeLinProblemTest(unittest.TestCase):
         # Wide difference in results run locally (0.04) versus the pipeline
         # (0.21), so seems to need unusually large tolerance.
         print("MODEL RESIDUAL: {}".format(model_residual))
-        self.assertAlmostEqual(model_residual, 0.14, delta=0.1)
+        self.assertAlmostEqual(model_residual, 0.1, delta=0.1)
 
         # Check data converged to less than 10% of target misfit
         data_misfit = 2.0 * self.grav_inv_active.invProb.dmisfit(
@@ -509,6 +514,15 @@ class QuadTreeLinProblemTest(unittest.TestCase):
         # Run the inversion from a zero starting model
         mrec = self.mag_inv_active.run(np.zeros(int(self.active_cells.sum())))
 
+        # import matplotlib.pyplot as plt
+        #
+        # fig = plt.figure()
+        # ax = plt.subplot()
+        # m_out = np.zeros(self.mesh.nC) * np.nan
+        # m_out[self.active_cells] = mrec
+        # self.mesh.plot_image(m_out, ax=ax)
+        # fig.savefig("mrec.png")
+
         # Compute predicted data
         dpred = self.mag_sim_active.dpred(self.mag_model[self.active_cells])
 
@@ -517,7 +531,7 @@ class QuadTreeLinProblemTest(unittest.TestCase):
             mrec - self.mag_model[self.active_cells]
         ) / np.linalg.norm(self.mag_model[self.active_cells])
         print("MODEL RESIDUAL: {}".format(model_residual))
-        self.assertAlmostEqual(model_residual, 0.11, delta=0.05)
+        self.assertAlmostEqual(model_residual, 0.01, delta=0.05)
 
         # Check data converged to less than 10% of target misfit
         data_misfit = 2.0 * self.mag_inv_active.invProb.dmisfit(