Revert "Store tiles on driver"

This reverts commit 834bb6a.

Revert "Re-assign computed list of dask tasks"

This reverts commit 07fc674.

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 Client, wait
+from dask.distributed import wait
 from simpeg import objective_function
 from simpeg.dask import objective_function as dask_objective_function
 from simpeg.objective_function import ComboObjectiveFunction
@@ -30,38 +30,32 @@
 
 
 class MisfitFactory(SimPEGFactory):
-    """
-    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]],
-    ):
+    """Build SimPEG global misfit function."""
+
+    def __init__(self, params, client, simulation, workers):
+        """
+        :param params: Options object containing SimPEG object parameters.
+        """
         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,
+        self, tiles
     ):
+        # 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
@@ -72,9 +66,8 @@ def assemble_arguments(  # pylint: disable=arguments-differ
 
         misfits = []
         tile_count = 0
-        for channel, tiles in self.tiles.items():
+        for channel in channels:
             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
@@ -124,10 +117,10 @@ def assemble_arguments(  # pylint: disable=arguments-differ
     def assemble_keyword_arguments(self, **_):
         """Implementation of abstract method from SimPEGFactory."""
 
-    def build(self, **_):
+    def build(self, tiles, **_):
         """To be over-ridden in factory implementations."""
 
-        misfits = self.assemble_arguments()
+        misfits = self.assemble_arguments(tiles)
 
         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]] = self.validate_workers(workers)
+        self._workers: list[tuple[str]] | None = 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:
+    def client(self) -> Client | bool | None:
         """
         Dask client or False if not using Dask.distributed.
         """
@@ -279,7 +279,6 @@ 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]:
         """
@@ -310,9 +309,7 @@ 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(
-                sub for sub in np.array_split(tile, split) if len(sub) > 0
-            )
+            flat_tile_list.append(np.array_split(tile, split))
         return flat_tile_list
 
     @property
@@ -321,19 +318,16 @@ def data_misfit(self):
         if getattr(self, "_data_misfit", None) is None:
             with fetch_active_workspace(self.workspace, mode="r+"):
                 # Tile locations
-                if self.logger and self.params.compute.tile_spatial > 1:
-                    self.logger.write(
-                        f"Setting up {self.params.compute.tile_spatial} tiles . . .\n"
-                    )
+                tiles = self.get_tiles()
+
+                if self.logger:
+                    self.logger.write(f"Setting up {len(tiles)} tile(s) . . .\n")
 
-                self.tiles = self.get_tiles()
                 self._data_misfit = MisfitFactory(
-                    self.params,
-                    self.simulation,
-                    self.tiles,
-                    client=self.client,
-                    workers=self.workers,
-                ).build()
+                    self.params, self.client, self.simulation, self.workers
+                ).build(
+                    self.split_list(tiles),
+                )
 
         return self._data_misfit
 
@@ -782,24 +776,13 @@ def get_tiles(self):
 
             return np.array_split(indices, n_chunks)
 
-        tiles = tile_locations(
+        return 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:
         """

simpeg_drivers/plate_simulation/match/driver.py

@@ -223,7 +223,7 @@ def run(self):
             # Display progress bar
             if isinstance(tasks[0], Future):
                 progress(tasks)
-                tasks = self.client.gather(tasks)
+                self.client.gather(tasks)
 
             scores = np.hstack(tasks)
             ranked = np.argsort(scores)[::-1]