MOve different tiling strategies to sub classes

Author: domfournier

simpeg_drivers/driver.py

@@ -806,35 +806,13 @@ def get_tiles(self) -> dict[str, list[np.ndarray]]:
 
         :return: Dictionary with channels as keys and list of tiles as values.
         """
-        n_data = self.inversion_data.mask.sum()
-        indices = np.arange(n_data)
-
-        # Split tiles based on inversion type
-        if "1d" in self.params.inversion_type:
-            # Heuristic to avoid too many chunks
-            n_chunks = n_data // self.params.compute.max_chunk_size
-
-            if self.workers:
-                n_chunks /= len(self.workers)
-                n_chunks = int(n_chunks) * len(self.workers)
-
-            n_chunks = np.max([n_chunks, 1, len(self.workers)])
-            tiles = [[tile] for tile in np.array_split(indices, n_chunks)]
-        # Split per line for 2D inversions
-        elif "2d" in self.params.inversion_type:
-            tiles = [
-                [np.where(self.params.line_selection.line_object.values == line_id)[0]]
-                for line_id in np.unique(self.params.line_selection.line_object.values)
-            ]
-        # Kmeans split with subsequent splitting to optimize load
-        else:
-            tiles = tile_locations(
-                self.inversion_data.locations,
-                self.params.compute.tile_spatial,
-                labels=self.inversion_data.parts,
-                sorting=self.simulation.survey.sorting,
-            )
-            tiles = self.split_list(tiles)
+        tiles = tile_locations(
+            self.inversion_data.locations,
+            self.params.compute.tile_spatial,
+            labels=self.inversion_data.parts,
+            sorting=self.simulation.survey.sorting,
+        )
+        tiles = self.split_list(tiles)
 
         # Base slice over frequencies
         if self.params.inversion_type in [

simpeg_drivers/electricals/base_2d.py

@@ -13,6 +13,7 @@
 
 from logging import getLogger
 
+import numpy as np
 from geoh5py.objects import DrapeModel, Octree, PotentialElectrode
 from geoh5py.ui_json.ui_json import fetch_active_workspace
 from pydantic import field_validator, model_validator
@@ -24,42 +25,15 @@
     DrapeModelOptions,
     LineSelectionOptions,
 )
-from simpeg_drivers.utils.surveys import create_mesh_by_line_id
+from simpeg_drivers.utils.surveys import (
+    create_mesh_by_line_id,
+    get_parts_from_electrodes,
+)
 
 
 logger = getLogger(__name__)
 
 
-class Base2DDriver(InversionDriver):
-    """
-    Base class for 2D DC and IP forward and inversion drivers.
-
-    Survey lines are inverted independently and internally stacked as a single
-    long survey. The inversion mesh is created as a drape mesh over the survey lines.
-    """
-
-    @property
-    def inversion_mesh(self) -> InversionMesh:
-        """Inversion mesh"""
-        if getattr(self, "_inversion_mesh", None) is None:
-            with fetch_active_workspace(self.workspace, mode="r+"):
-                entity = None
-                if self.params.mesh is None:
-                    entity = create_mesh_by_line_id(
-                        self.workspace,
-                        self.params.line_selection.line_object,
-                        self.params.drape_model,
-                        parent=self.out_group,
-                    )
-                    self.params.mesh = entity
-
-                self._inversion_mesh = InversionMesh(
-                    self.workspace, self.params, entity=entity
-                )
-
-        return self._inversion_mesh
-
-
 class Base2DOptions(CoreOptions):
     """
     Base options for the Direct Current 2D forward and inverse driver.
@@ -72,10 +46,13 @@ class Base2DOptions(CoreOptions):
     """
 
     data_object: PotentialElectrode
-    line_selection: LineSelectionOptions = LineSelectionOptions()
+    line_selection: LineSelectionOptions | None = None
     mesh: DrapeModel | Octree | None = None
     drape_model: DrapeModelOptions = DrapeModelOptions()
 
+    _line_parts: np.ndarray | None = None
+    _selected_parts: list[int] | None = None
+
     @field_validator("mesh", mode="before")
     @classmethod
     def mesh_cannot_be_octree(cls, value: Octree | DrapeModel):
@@ -105,3 +82,84 @@ def deprecated_pseudo(cls, data: dict):
             data["line_selection"] = line_selection
 
         return data
+
+    @property
+    def line_parts(self) -> np.ndarray:
+        """
+        Generate monotonic line parts from line identifier or inferred from graph of potentials.
+        """
+        if self._line_parts is None:
+            if (
+                self.line_selection is not None
+                and self.line_selection.property is not None
+            ):
+                _, self._line_parts = np.unique(
+                    self.line_selection.property.values, return_inverse=True
+                )
+            else:
+                self._line_parts = get_parts_from_electrodes(self.data_object)
+
+        return self._line_parts
+
+    @property
+    def selected_parts(self) -> list[int]:
+        """
+        Translate line section ids to monotonic parts.
+        """
+        if self._selected_parts is None:
+            parts = []
+            if (
+                self.line_selection is not None
+                and self.line_selection.property is not None
+            ):
+                for count, val in enumerate(
+                    np.unique(self.line_selection.property.values)
+                ):
+                    if val in self.line_selection.value:
+                        parts.append(count)
+            else:
+                parts = np.arange(len(np.unique(self.line_parts)))
+
+            self._selected_parts = parts
+
+        return self._selected_parts
+
+
+class Base2DDriver(InversionDriver):
+    """
+    Base class for 2D DC and IP forward and inversion drivers.
+
+    Survey lines are inverted independently and internally stacked as a single
+    long survey. The inversion mesh is created as a drape mesh over the survey lines.
+    """
+
+    @property
+    def inversion_mesh(self) -> InversionMesh:
+        """Inversion mesh"""
+        if getattr(self, "_inversion_mesh", None) is None:
+            with fetch_active_workspace(self.workspace, mode="r+"):
+                entity = None
+                if self.params.mesh is None:
+                    entity = create_mesh_by_line_id(
+                        self.workspace,
+                        self.params.line_ids,
+                        self.params.drape_model,
+                        parent=self.out_group,
+                    )
+                    self.params.mesh = entity
+
+                self._inversion_mesh = InversionMesh(
+                    self.workspace, self.params, entity=entity
+                )
+
+        return self._inversion_mesh
+
+    def get_tiles(self) -> dict[str, list[np.ndarray]]:
+        """
+        Generate tiles from survey parts.
+        """
+        tiles = [
+            [np.where(self.params.line_parts == part)[0]]
+            for part in self.params.selected_parts
+        ]
+        return {None: tiles}

simpeg_drivers/electromagnetics/base_1d_driver.py

@@ -86,6 +86,20 @@ def get_1d_mesh(self) -> TensorMesh:
         )
         return layers_mesh
 
+    def get_tiles(self) -> dict[str, list[np.ndarray]]:
+        n_data = self.inversion_data.mask.sum()
+        indices = np.arange(n_data)
+
+        # Heuristic to avoid too many chunks
+        n_chunks = n_data // self.params.compute.max_chunk_size
+
+        if self.workers:
+            n_chunks /= len(self.workers)
+            n_chunks = int(n_chunks) * len(self.workers)
+
+        n_chunks = np.max([n_chunks, 1, len(self.workers)])
+        return {None: [[tile] for tile in np.array_split(indices, n_chunks)]}
+
     @property
     def simulation(self):
         """