Clip out rows of zero gradients. Simplify get_regularization

Author: domfournier

simpeg_drivers/driver.py

@@ -480,38 +480,31 @@ def get_regularization(self):
                 weight = mapping * getattr(self.models, weight_name)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
 
-                if isinstance(fun, SparseSmoothness):
-                    if is_rotated:
-                        if forward_mesh is None:
-                            fun = set_rotated_operators(
-                                fun,
-                                neighbors,
-                                comp,
-                                self.models.gradient_dip,
-                                self.models.gradient_direction,
-                            )
-
-                    else:
-                        weight = (
-                            getattr(
-                                reg_func.regularization_mesh,
-                                f"aveCC2F{fun.orientation}",
-                            )
-                            * weight
-                        )
-                        norm = (
-                            getattr(
-                                reg_func.regularization_mesh,
-                                f"aveCC2F{fun.orientation}",
-                            )
-                            * norm
+                if not isinstance(fun, SparseSmoothness):
+                    fun.set_weights(**{comp: weight})
+                    fun.norm = norm
+                    functions.append(fun)
+                    continue
+
+                if is_rotated:
+                    if forward_mesh is None:
+                        fun = set_rotated_operators(
+                            fun,
+                            neighbors,
+                            comp,
+                            self.models.gradient_dip,
+                            self.models.gradient_direction,
                         )
 
-                fun.set_weights(**{comp: weight})
-                fun.norm = norm
+                average_op = getattr(
+                    reg_func.regularization_mesh,
+                    f"aveCC2F{fun.orientation}",
+                )
+                fun.set_weights(**{comp: average_op @ weight})
+                fun.norm = average_op @ norm
                 functions.append(fun)
 
-                if isinstance(fun, SparseSmoothness) and is_rotated:
+                if is_rotated:
                     fun.gradient_type = "components"
                     backward_fun = deepcopy(fun)
                     setattr(backward_fun, "_regularization_mesh", backward_mesh)
@@ -526,6 +519,12 @@ def get_regularization(self):
                             self.models.gradient_direction,
                             forward=False,
                         )
+                    average_op = getattr(
+                        backward_fun.regularization_mesh,
+                        f"aveCC2F{fun.orientation}",
+                    )
+                    backward_fun.set_weights(**{comp: average_op @ weight})
+                    backward_fun.norm = average_op @ norm
                     functions.append(backward_fun)
 
             # Will avoid recomputing operators if the regularization mesh is the same

simpeg_drivers/utils/regularization.py

@@ -415,16 +415,20 @@ def set_rotated_operators(
     grad_op = rotated_gradient(
         function.regularization_mesh.mesh, neighbors, axis, dip, direction, forward
     )
+    grad_op_active = function.regularization_mesh.Pac.T @ (
+        grad_op @ function.regularization_mesh.Pac
+    )
+    active_faces = grad_op_active.max(axis=1).toarray().ravel() > 0
+
     setattr(
         function.regularization_mesh,
         f"_cell_gradient_{function.orientation}",
-        function.regularization_mesh.Pac.T
-        @ (grad_op @ function.regularization_mesh.Pac),
+        grad_op_active[active_faces, :],
     )
     setattr(
         function.regularization_mesh,
         f"_aveCC2F{function.orientation}",
-        sdiag(np.ones(function.regularization_mesh.n_cells)),
+        sdiag(np.ones(function.regularization_mesh.n_cells))[active_faces, :],
     )
 
     return function