Store tiles on driver

Author: domfournier

simpeg_drivers/components/factories/misfit_factory.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from dask.distributed import wait
+from dask.distributed import Client, wait
 from simpeg import objective_function
 from simpeg.dask import objective_function as dask_objective_function
 from simpeg.objective_function import ComboObjectiveFunction
@@ -30,32 +30,38 @@
 
 
 class MisfitFactory(SimPEGFactory):
-    """Build SimPEG global misfit function."""
-
-    def __init__(self, params, client, simulation, workers):
-        """
-        :param params: Options object containing SimPEG object parameters.
-        """
+    """
+    Build SimPEG global misfit function.
+
+    :param params: Options object containing SimPEG object parameters.
+    :param simulation: SimPEG simulation object.
+    :param tiles: Dictionary of nested lists with arrays of indices for the tiles.
+    :param client: Dask client or boolean to indicate whether to use dask.
+    :param workers: List of worker addresses to use for dask computations.
+    """
+
+    def __init__(
+        self,
+        params,
+        simulation,
+        tiles: dict[list[np.ndarray]],
+        client: Client | bool,
+        workers: list[tuple[str]],
+    ):
         super().__init__(params)
 
         self.simpeg_object = self.concrete_object()
-        self.factory_type = self.params.inversion_type
         self.simulation = simulation
+        self.tiles = tiles
         self.client = client
         self.workers = workers
 
     def concrete_object(self):
         return objective_function.ComboObjectiveFunction
 
     def assemble_arguments(  # pylint: disable=arguments-differ
-        self, tiles
+        self,
     ):
-        # Base slice over frequencies
-        if self.factory_type in ["magnetotellurics", "tipper", "fdem"]:
-            channels = self.simulation.survey.frequencies
-        else:
-            channels = [None]
-
         use_futures = self.client
 
         # Pickle the simulation to the temporary file
@@ -66,8 +72,9 @@ def assemble_arguments(  # pylint: disable=arguments-differ
 
         misfits = []
         tile_count = 0
-        for channel in channels:
+        for channel, tiles in self.tiles.items():
             for local_indices in tiles:
+                # Split again but use the same mesh extent based on tile vertices
                 for sub_ind in local_indices:
                     if len(sub_ind) == 0:
                         continue
@@ -117,10 +124,10 @@ def assemble_arguments(  # pylint: disable=arguments-differ
     def assemble_keyword_arguments(self, **_):
         """Implementation of abstract method from SimPEGFactory."""
 
-    def build(self, tiles, **_):
+    def build(self, **_):
         """To be over-ridden in factory implementations."""
 
-        misfits = self.assemble_arguments(tiles)
+        misfits = self.assemble_arguments()
 
         if self.client:
             return dask_objective_function.DistributedComboMisfits(

simpeg_drivers/driver.py

@@ -103,7 +103,7 @@ def __init__(
                 "Disk storage of sensitivities is not compatible with distributed processing."
             )
 
-        self._workers: list[tuple[str]] | None = self.validate_workers(workers)
+        self._workers: list[tuple[str]] = self.validate_workers(workers)
 
     @property
     def out_group(self) -> SimPEGGroup:
@@ -142,7 +142,7 @@ def validate_out_group(self, out_group: SimPEGGroup | None) -> SimPEGGroup:
         return out_group
 
     @property
-    def client(self) -> Client | bool | None:
+    def client(self) -> Client | bool:
         """
         Dask client or False if not using Dask.distributed.
         """
@@ -279,6 +279,7 @@ def __init__(
         self._ordering: list[np.ndarray] | None = None
         self._mappings: list[maps.IdentityMap] | None = None
         self._window = None
+        self.tiles: dict[list[np.ndarray]]
 
     def split_list(self, tiles: list[np.ndarray]) -> list[np.ndarray]:
         """
@@ -309,7 +310,9 @@ def split_list(self, tiles: list[np.ndarray]) -> list[np.ndarray]:
 
         flat_tile_list = []
         for tile, split in zip(tiles, split_list):
-            flat_tile_list.append(np.array_split(tile, split))
+            flat_tile_list.append(
+                sub for sub in np.array_split(tile, split) if len(sub) > 0
+            )
         return flat_tile_list
 
     @property
@@ -318,16 +321,19 @@ def data_misfit(self):
         if getattr(self, "_data_misfit", None) is None:
             with fetch_active_workspace(self.workspace, mode="r+"):
                 # Tile locations
-                tiles = self.get_tiles()
-
-                if self.logger:
-                    self.logger.write(f"Setting up {len(tiles)} tile(s) . . .\n")
+                if self.logger and self.params.compute.tile_spatial > 1:
+                    self.logger.write(
+                        f"Setting up {self.params.compute.tile_spatial} tiles . . .\n"
+                    )
 
+                self.tiles = self.get_tiles()
                 self._data_misfit = MisfitFactory(
-                    self.params, self.client, self.simulation, self.workers
-                ).build(
-                    self.split_list(tiles),
-                )
+                    self.params,
+                    self.simulation,
+                    self.tiles,
+                    client=self.client,
+                    workers=self.workers,
+                ).build()
 
         return self._data_misfit
 
@@ -776,13 +782,24 @@ def get_tiles(self):
 
             return np.array_split(indices, n_chunks)
 
-        return tile_locations(
+        tiles = tile_locations(
             self.inversion_data.locations,
             self.params.compute.tile_spatial,
             labels=self.inversion_data.parts,
             sorting=self.simulation.survey.sorting,
         )
 
+        self.split_list(tiles)
+
+        # Base slice over frequencies
+        if self.params.inversion_type in ["magnetotellurics", "tipper", "fdem"]:
+            channels = self.simulation.survey.frequencies
+        else:
+            channels = [None]
+
+        # Duplicate tiles for each channel
+        return {channel: tiles for channel in channels}
+
     @classmethod
     def start(cls, filepath: str | Path | InputFile, **kwargs) -> Self:
         """