Update total gradient loop

Author: domfournier

SimPEG/regularization/base.py

@@ -30,8 +30,6 @@ class BaseRegularization(BaseObjectiveFunction):
     :param weights: Weight multipliers to customize the least-squares function.
     """
 
-    _model = None
-
     def __init__(
         self,
         mesh: RegularizationMesh | BaseMesh,
@@ -51,6 +49,8 @@ def __init__(
                 f"Value of type {type(mesh)} provided."
             )
 
+        self._model = None
+        self._parent = None
         self._regularization_mesh = mesh
         self._weights = {}
 
@@ -136,6 +136,23 @@ def mapping(self, mapping: maps.IdentityMap):
             )
         self._mapping = mapping
 
+    @property
+    def parent(self):
+        """
+        The parent objective function
+        """
+        return self._parent
+
+    @parent.setter
+    def parent(self, parent):
+        combo_class = ComboObjectiveFunction
+        if not isinstance(parent, combo_class):
+            raise TypeError(
+                f"Invalid parent of type '{parent.__class__.__name__}'. "
+                f"Parent must be a {combo_class.__name__}."
+            )
+        self._parent = parent
+
     @property
     def units(self) -> str | None:
         """Specify the model units. Special care given to 'radian' values"""
@@ -555,10 +572,7 @@ def _cell_distances(self):
         """
         Distances between cell centers for the cell center difference.
         """
-        if self.__cell_distances is None:
-            self.__cell_distances = 1.0 / np.max(self.cell_gradient, axis=1).data
-
-        return self.__cell_distances
+        return getattr(self.regularization_mesh, f"cell_distances_{self.orientation}")
 
     @property
     def orientation(self):
@@ -784,6 +798,10 @@ def __init__(
         else:
             objfcts = kwargs.pop("objfcts")
         super().__init__(objfcts=objfcts, **kwargs)
+
+        for fun in objfcts:
+            fun.parent = self
+
         self.mapping = mapping
         self.reference_model = reference_model
         self.reference_model_in_smooth = reference_model_in_smooth

SimPEG/regularization/sparse.py

@@ -140,30 +140,12 @@ def update_weights(self, m):
         """
         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