Merge pull request #179 from MiraGeoscience/GEOPY-2075_c

GEOPY-2075: c

Author: benk-mira

simpeg_drivers/components/models.py

@@ -29,6 +29,25 @@
     from simpeg_drivers.driver import InversionDriver
 
 
+MODEL_TYPES = [
+    "starting",
+    "reference",
+    "lower_bound",
+    "upper_bound",
+    "conductivity",
+    "alpha_s",
+    "length_scale_x",
+    "length_scale_y",
+    "length_scale_z",
+    "gradient_dip",
+    "gradient_direction",
+    "s_norm",
+    "x_norm",
+    "y_norm",
+    "z_norm",
+]
+
+
 class InversionModelCollection:
     """
     Collection of inversion models.
@@ -39,50 +58,41 @@ class InversionModelCollection:
 
     """
 
-    model_types = [
-        "starting",
-        "reference",
-        "lower_bound",
-        "upper_bound",
-        "conductivity",
-        "alpha_s",
-        "length_scale_x",
-        "length_scale_y",
-        "length_scale_z",
-        "s_norm",
-        "x_norm",
-        "y_norm",
-        "z_norm",
-    ]
-
     def __init__(self, driver: InversionDriver):
         """
         :param driver: Parental InversionDriver class.
         """
         self._active_cells: np.ndarray | None = None
         self._driver = driver
         self.is_sigma = self.driver.params.physical_property == "conductivity"
-        self.is_vector = (
+        is_vector = (
             True if self.driver.params.inversion_type == "magnetic vector" else False
         )
-        self.n_blocks = (
-            3 if self.driver.params.inversion_type == "magnetic vector" else 1
+        self._starting = InversionModel(driver, "starting", is_vector=is_vector)
+        self._reference = InversionModel(driver, "reference", is_vector=is_vector)
+        self._lower_bound = InversionModel(driver, "lower_bound", is_vector=is_vector)
+        self._upper_bound = InversionModel(driver, "upper_bound", is_vector=is_vector)
+        self._conductivity = InversionModel(driver, "conductivity", is_vector=is_vector)
+        self._alpha_s = InversionModel(driver, "alpha_s", is_vector=is_vector)
+        self._length_scale_x = InversionModel(
+            driver, "length_scale_x", is_vector=is_vector
+        )
+        self._length_scale_y = InversionModel(
+            driver, "length_scale_y", is_vector=is_vector
+        )
+        self._length_scale_z = InversionModel(
+            driver, "length_scale_z", is_vector=is_vector
+        )
+        self._gradient_dip = InversionModel(
+            driver, "gradient_dip", trim_active_cells=False
         )
-        self._starting = InversionModel(driver, "starting")
-        self._reference = InversionModel(driver, "reference")
-        self._lower_bound = InversionModel(driver, "lower_bound")
-        self._upper_bound = InversionModel(driver, "upper_bound")
-        self._conductivity = InversionModel(driver, "conductivity")
-        self._alpha_s = InversionModel(driver, "alpha_s")
-        self._length_scale_x = InversionModel(driver, "length_scale_x")
-        self._length_scale_y = InversionModel(driver, "length_scale_y")
-        self._length_scale_z = InversionModel(driver, "length_scale_z")
-        self._gradient_dip = InversionModel(driver, "gradient_dip")
-        self._gradient_direction = InversionModel(driver, "gradient_direction")
-        self._s_norm = InversionModel(driver, "s_norm")
-        self._x_norm = InversionModel(driver, "x_norm")
-        self._y_norm = InversionModel(driver, "y_norm")
-        self._z_norm = InversionModel(driver, "z_norm")
+        self._gradient_direction = InversionModel(
+            driver, "gradient_direction", trim_active_cells=False
+        )
+        self._s_norm = InversionModel(driver, "s_norm", is_vector=is_vector)
+        self._x_norm = InversionModel(driver, "x_norm", is_vector=is_vector)
+        self._y_norm = InversionModel(driver, "y_norm", is_vector=is_vector)
+        self._z_norm = InversionModel(driver, "z_norm", is_vector=is_vector)
 
     @property
     def n_active(self) -> int:
@@ -307,7 +317,7 @@ def z_norm(self) -> np.ndarray | None:
     def _model_method_wrapper(self, method, name=None, **kwargs):
         """wraps individual model's specific method and applies in loop over model types."""
         returned_items = {}
-        for mtype in self.model_types:
+        for mtype in MODEL_TYPES:
             model = getattr(self, f"_{mtype}")
             if model.model is not None:
                 f = getattr(model, method)
@@ -364,43 +374,24 @@ class InversionModel:
     remove_air: Use active cells vector to remove air cells from model.
     """
 
-    model_types = [
-        "starting",
-        "reference",
-        "lower_bound",
-        "upper_bound",
-        "conductivity",
-        "alpha_s",
-        "length_scale_x",
-        "length_scale_y",
-        "length_scale_z",
-        "gradient_dip",
-        "gradient_direction",
-        "s_norm",
-        "x_norm",
-        "y_norm",
-        "z_norm",
-    ]
-
     def __init__(
         self,
         driver: InversionDriver,
         model_type: str,
+        is_vector: bool = False,
+        trim_active_cells: bool = True,
     ):
         """
         :param driver: InversionDriver object.
-        :param model_type: Type of inversion model, can be any of "starting", "reference",
-            "lower_bound", "upper_bound".
+        :param model_type: Type of inversion model, can be any of MODEL_TYPES.
+        :param is_vector: If True, model is a vector.
+        :param trim_active_cells: If True, remove air cells from model.
         """
         self.driver = driver
         self.model_type = model_type
         self.model: np.ndarray | None = None
-        self.is_vector = (
-            True if self.driver.params.inversion_type == "magnetic vector" else False
-        )
-        self.n_blocks = (
-            3 if self.driver.params.inversion_type == "magnetic vector" else 1
-        )
+        self.is_vector = is_vector
+        self.trim_active_cells = trim_active_cells
         self._initialize()
 
     def _initialize(self):
@@ -452,7 +443,7 @@ def _initialize(self):
                 and self.is_vector
                 and model.shape[0] == self.driver.inversion_mesh.n_cells
             ):
-                model = np.tile(model, self.n_blocks)
+                model = np.tile(model, 3 if self.is_vector else 1)
 
         if model is not None:
             self.model = mkvc(model)
@@ -461,8 +452,8 @@ def _initialize(self):
     def remove_air(self, active_cells):
         """Use active cells vector to remove air cells from model"""
 
-        if self.model is not None:
-            self.model = self.model[np.tile(active_cells, self.n_blocks)]
+        if self.model is not None and self.trim_active_cells:
+            self.model = self.model[np.tile(active_cells, 3 if self.is_vector else 1)]
 
     def permute_2_octree(self) -> np.ndarray | None:
         """
@@ -604,7 +595,7 @@ def model_type(self):
 
     @model_type.setter
     def model_type(self, v):
-        if v not in self.model_types:
-            msg = f"Invalid model_type: {v}. Must be one of {(*self.model_types,)}."
+        if v not in MODEL_TYPES:
+            msg = f"Invalid model_type: {v}. Must be one of {(*MODEL_TYPES,)}."
             raise ValueError(msg)
         self._model_type = v

simpeg_drivers/driver.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import multiprocessing
-
+from copy import deepcopy
 import sys
 from datetime import datetime, timedelta
 import logging
@@ -42,8 +42,13 @@
     objective_function,
     optimization,
 )
-from simpeg.utils import sdiag
-from simpeg.regularization import BaseRegularization, Sparse
+
+from simpeg.regularization import (
+    BaseRegularization,
+    RegularizationMesh,
+    Sparse,
+    SparseSmoothness,
+)
 
 from simpeg_drivers import DRIVER_MAP
 from simpeg_drivers.components import (
@@ -60,7 +65,7 @@
 )
 from simpeg_drivers.joint.params import BaseJointOptions
 from simpeg_drivers.utils.utils import tile_locations
-from simpeg_drivers.utils.regularization import cell_neighbors, rotated_gradient
+from simpeg_drivers.utils.regularization import cell_neighbors, set_rotated_operators
 
 mlogger = logging.getLogger("distributed")
 mlogger.setLevel(logging.WARNING)
@@ -442,91 +447,101 @@ def get_regularization(self):
             return BaseRegularization(mesh=self.inversion_mesh.mesh)
 
         reg_funcs = []
-        total_reg_funcs = []
         is_rotated = self.params.gradient_rotation is not None
+        neighbors = None
+        backward_mesh = None
+        forward_mesh = None
         for mapping in self.mapping:
-            reg_funcs.append(
-                Sparse(
-                    self.inversion_mesh.mesh,
-                    active_cells=self.models.active_cells,
-                    mapping=mapping,
-                    reference_model=self.models.reference,
-                )
+            reg_func = Sparse(
+                forward_mesh or self.inversion_mesh.mesh,
+                active_cells=self.models.active_cells if forward_mesh is None else None,
+                mapping=mapping,
+                reference_model=self.models.reference,
             )
 
-            if is_rotated:
-                reg_funcs.append(
-                    Sparse(
-                        self.inversion_mesh.mesh,
-                        active_cells=self.models.active_cells,
-                        mapping=mapping,
-                        reference_model=self.models.reference,
-                    )
+            if is_rotated and neighbors is None:
+                backward_mesh = RegularizationMesh(
+                    self.inversion_mesh.mesh, active_cells=self.models.active_cells
                 )
-                neighbors = cell_neighbors(reg_funcs[0].regularization_mesh.mesh)
+                neighbors = cell_neighbors(reg_func.regularization_mesh.mesh)
 
             # Adjustment for 2D versus 3D problems
-            comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
-            avg_comps = "sxy" if "2d" in self.params.inversion_type else "sxyz"
-            weights = ["alpha_s"] + [f"length_scale_{k}" for k in comps[1:]]
-
-            for comp, avg_comp, objfct, weight in zip(
-                comps, avg_comps, reg_funcs[0].objfcts, weights
+            components = "sxz" if "2d" in self.params.inversion_type else "sxyz"
+            weight_names = ["alpha_s"] + [f"length_scale_{k}" for k in components[1:]]
+            functions = []
+            for comp, weight_name, fun in zip(
+                components, weight_names, reg_func.objfcts
             ):
-                if getattr(self.models, weight) is None:
-                    setattr(reg_funcs[0], weight, 0.0)
+                if getattr(self.models, weight_name) is None:
+                    setattr(reg_funcs, weight_name, 0.0)
+                    functions.append(fun)
                     continue
 
-                weight = mapping * getattr(self.models, weight)
+                weight = mapping * getattr(self.models, weight_name)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
-                if comp in "xyz":
+
+                if isinstance(fun, SparseSmoothness):
                     if is_rotated:
-                        for reg, forward in zip(reg_funcs, [True, False]):
-                            grad_op = rotated_gradient(
-                                mesh=reg.regularization_mesh.mesh,
-                                neighbors=neighbors,
-                                axis=comp,
-                                dip=self.models.gradient_dip,
-                                direction=self.models.gradient_direction,
-                                forward=forward,
-                            )
-                            setattr(
-                                reg.regularization_mesh,
-                                f"_cell_gradient_{comp}",
-                                reg.regularization_mesh.Pac.T
-                                @ (grad_op @ reg.regularization_mesh.Pac),
-                            )
-                            setattr(
-                                reg.regularization_mesh,
-                                f"_aveCC2F{avg_comp}",
-                                sdiag(np.ones(reg.regularization_mesh.n_cells)),
+                        if forward_mesh is None:
+                            fun = set_rotated_operators(
+                                fun,
+                                neighbors,
+                                comp,
+                                self.models.gradient_dip,
+                                self.models.gradient_direction,
                             )
+
                     else:
                         weight = (
                             getattr(
-                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                                reg_func.regularization_mesh,
+                                f"aveCC2F{fun.orientation}",
                             )
                             * weight
                         )
                         norm = (
                             getattr(
-                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                                reg_func.regularization_mesh,
+                                f"aveCC2F{fun.orientation}",
                             )
                             * norm
                         )
 
-                objfct.set_weights(**{comp: weight})
-                objfct.norm = norm
-
-            if getattr(self.params, "gradient_type") is not None:
-                for reg in reg_funcs:
-                    setattr(
-                        reg,
-                        "gradient_type",
-                        getattr(self.params, "gradient_type"),
-                    )
-
-            total_reg_funcs += reg_funcs
+                fun.set_weights(**{comp: weight})
+                fun.norm = norm
+                functions.append(fun)
+
+                if isinstance(fun, SparseSmoothness) and is_rotated:
+                    fun.gradient_type = "components"
+                    backward_fun = deepcopy(fun)
+                    setattr(backward_fun, "_regularization_mesh", backward_mesh)
+
+                    # Only do it once for MVI
+                    if not forward_mesh:
+                        backward_fun = set_rotated_operators(
+                            backward_fun,
+                            neighbors,
+                            comp,
+                            self.models.gradient_dip,
+                            self.models.gradient_direction,
+                            forward=False,
+                        )
+                    functions.append(backward_fun)
+
+            # Will avoid recomputing operators if the regularization mesh is the same
+            forward_mesh = reg_func.regularization_mesh
+            reg_func.objfcts = functions
+            reg_func.norms = [fun.norm for fun in functions]
+            reg_funcs.append(reg_func)
+
+        # TODO - To be deprcated on GEOPY-2109
+        if getattr(self.params, "gradient_type") is not None:
+            for reg in reg_funcs:
+                setattr(
+                    reg,
+                    "gradient_type",
+                    getattr(self.params, "gradient_type"),
+                )
 
         return objective_function.ComboObjectiveFunction(objfcts=reg_funcs)