Merge branch 'develop' into GEOPY-789

Author: domfournier

recipe.yaml

@@ -2,7 +2,7 @@ schema_version: 1
 
 context:
   name: "mira-simpeg"
-  version: "0.23.0.1a4"
+  version: "0.23.0.2a1"
   python_min: "3.10"
 
 package:
@@ -27,7 +27,7 @@ requirements:
   run:
     - python >=${{ python_min }}
     # Mira packages
-    - geoh5py >=0.11.0a, <0.12.dev
+    - geoh5py >=0.12.0a, <0.13.dev
     # direct dependencies
     - discretize >=0.11
     - geoana >=0.7.0

simpeg/dask/electromagnetics/time_domain/simulation.py

@@ -11,6 +11,7 @@
 from dask import array, delayed
 from dask.distributed import get_client
 
+from time import time
 from simpeg.dask.utils import get_parallel_blocks
 from simpeg.utils import mkvc
 
@@ -405,10 +406,6 @@ def get_field_deriv_block(
         # Cut out early data
         time_check = np.kron(time_mask, np.ones(shape, dtype=bool))[rx_ind]
         local_ind = np.arange(rx_ind.shape[0])[time_check]
-
-        if len(local_ind) < 1:
-            continue
-
         indices.append(
             (np.arange(count, count + len(local_ind)), local_ind),
         )
@@ -425,14 +422,14 @@ def get_field_deriv_block(
 
         stacked_blocks.append(stacked_block)
 
-    if len(stacked_blocks) > 0:
-        blocks = np.hstack(stacked_blocks)
-
+    blocks = np.hstack(stacked_blocks)
+    if blocks.ndim == 2 and blocks.shape[1] > 0:
         solve = (AdiagTinv * blocks).reshape(blocks.shape)
     else:
         solve = None
 
     updated_ATinv_df_duT_v = []
+
     for (_, arrays), field_deriv, ATinv_chunk, (columns, local_ind) in zip(
         block, field_derivs, ATinv_df_duT_v, indices, strict=True
     ):
@@ -444,10 +441,9 @@ def get_field_deriv_block(
             )
             ATinv_chunk = np.zeros(shape, dtype=np.float32)
 
-        if solve is None:
-            continue
+        if solve is not None:
+            ATinv_chunk[:, local_ind] = solve[:, columns]
 
-        ATinv_chunk[:, local_ind] = solve[:, columns]
         updated_ATinv_df_duT_v.append(ATinv_chunk)
 
     return updated_ATinv_df_duT_v

simpeg/dask/objective_function.py

@@ -72,14 +72,15 @@ def _validate_type_or_future_of_type(
         property_name, objects, obj_type, ensure_unique=True
     )
     workload = [[]]
+    lookup = {}
     count = 0
     for obj in objects:
         if count == len(workers):
             count = 0
             workload.append([])
         obj.simulation.simulations[0].worker = workers[count]
         future = client.scatter([obj], workers=workers[count])[0]
-
+        lookup[obj] = (future, workers[count])
         if hasattr(obj, "name"):
             future.name = obj.name
 
@@ -100,7 +101,7 @@ def _validate_type_or_future_of_type(
         raise TypeError(f"{property_name} futures must be an instance of {obj_type}")
 
     if return_workers:
-        return workload, workers
+        return workload, workers, lookup
     else:
         return workload
 
@@ -390,7 +391,7 @@ def objfcts(self):
     def objfcts(self, objfcts):
         client = self.client
 
-        futures, workers = _validate_type_or_future_of_type(
+        futures, workers, lookup = _validate_type_or_future_of_type(
             "objfcts",
             objfcts,
             L2DataMisfit,
@@ -404,8 +405,8 @@ def objfcts(self, objfcts):
         self._workers = workers
 
         self._lookup = {
-            obj.simulation: (future, worker)
-            for future, worker, obj in zip(futures[0], workers, objfcts)
+            misfit.simulation: (future, worker)
+            for misfit, (future, worker) in lookup.items()
         }
 
     def residuals(self, m, f=None):
@@ -443,12 +444,12 @@ def broadcast_updates(self, updates: dict):
             if fun not in self._lookup:
                 continue
 
-            objfct, worker = self._lookup[fun]
+            future, worker = self._lookup[fun]
 
             stores.append(
                 client.submit(
                     _setter_broadcast,
-                    objfct,
+                    future,
                     key,
                     value,
                     workers=worker,

simpeg/directives/_regularization.py

@@ -229,19 +229,20 @@ def adjust_cooling_schedule(self):
         """
         Adjust the cooling schedule based on the misfit.
         """
-        ratio = self.invProb.phi_d / self.misfit_from_chi_factor(self.chifact_target)
+        if self.metrics.start_irls_iter is None:
+            return
 
-        if (
-            np.abs(1.0 - ratio) > self.misfit_tolerance
-            and self.metrics.start_irls_iter is not None
-        ):
+        ratio = self.invProb.phi_d / self.misfit_from_chi_factor(self.chifact_target)
+        if np.abs(1.0 - ratio) > self.misfit_tolerance:
 
             if ratio > 1:
-                ratio = np.mean([2.0, ratio])
+                update_ratio = 1 / np.mean([0.75, 1 / ratio])
             else:
-                ratio = np.mean([0.75, ratio])
+                update_ratio = 1 / np.mean([2.0, 1 / ratio])
 
-            self.cooling_factor = ratio
+            self.cooling_factor = update_ratio
+        else:
+            self.cooling_factor = 1.0
 
     def initialize(self):
         """