Add delayed methods for updateing simualtion within Future

Author: domfournier

simpeg/directives/_vector_models.py

@@ -1,5 +1,5 @@
 import numpy as np
-
+from dask.distributed import Future
 from . import (
     BaseSaveGeoH5,
     InversionDirective,
@@ -56,8 +56,11 @@ def update(self):
         self.invProb.model = m
         self.opt.xc = self.invProb.model
 
-        for misfit in self.dmisfit.objfcts:
-            misfit.simulation.model = m
+        if isinstance(self.dmisfit.objfcts[0], Future):
+            self.dmisfit.model = m
+        else:
+            for misfit in self.dmisfit.objfcts:
+                misfit.simulation.model = m
 
     def _reproject(self, m):
         """
@@ -71,6 +74,15 @@ def _reproject(self, m):
         return m
 
 
+def update_map(misfit):
+    if isinstance(misfit.simulation, MetaSimulation):
+        misfit.simulation.simulations[0].chiMap = (
+            SphericalSystem() * misfit.simulation.simulations[0].chiMap
+        )
+    else:
+        misfit.simulation.chiMap = SphericalSystem() * misfit.simulation.chiMap
+
+
 class VectorInversion(InversionDirective):
     """
     Control a vector inversion from Cartesian to spherical coordinates.
@@ -90,19 +102,15 @@ def __init__(
         self, simulations: list, regularizations: ComboObjectiveFunction, **kwargs
     ):
         self.reference_angles = (False, False, False)
-        self.simulations = simulations
+        self.misfits = simulations
         self.regularizations = regularizations
 
         set_kwargs(self, **kwargs)
 
     @property
     def target(self):
         if getattr(self, "_target", None) is None:
-            nD = 0
-            for survey in self.survey:
-                nD += survey.nD
-
-            self._target = nD * self.chifact_target
+            self._target = np.hstack(self.invProb.dpred).shape[0] * self.chifact_target
 
         return self._target
 
@@ -116,10 +124,6 @@ def initialize(self):
 
         self.reference_model = reg.reference_model
 
-        for dmisfit in self.dmisfit.objfcts:
-            if getattr(dmisfit.simulation, "coordinate_system", None) is not None:
-                dmisfit.simulation.coordinate_system = self.mode
-
     def endIter(self):
 
         model = self.invProb.model.copy()
@@ -217,20 +221,12 @@ def endIter(self):
             self.opt.upper[indices[nC:]] = np.inf
 
             updates = {}
-            for simulation in self.simulations:
-                if isinstance(simulation, MetaSimulation):
-
-                    if hasattr(self.dmisfit, "client"):
-                        updates[simulation] = (
-                            "chiMap",
-                            SphericalSystem() * simulation.simulations[0].chiMap,
-                        )
-                    else:
-                        simulation.simulations[0].chiMap = (
-                            SphericalSystem() * simulation.simulations[0].chiMap
-                        )
+            for misfit in self.misfits:
+
+                if isinstance(misfit, Future):
+                    updates[misfit] = ("", update_map)
                 else:
-                    simulation.chiMap = SphericalSystem() * simulation.chiMap
+                    update_map(misfit)
 
             if hasattr(self.dmisfit, "client"):
                 self.dmisfit.broadcast_updates(updates)