Merge pull request #132 from MiraGeoscience/GEOPY-2735

Geopy 2735

Author: domfournier

simpeg/dask/__init__.py

@@ -1,6 +1,7 @@
 try:
     import simpeg.dask.simulation
     import simpeg.dask.electromagnetics.frequency_domain.simulation
+    import simpeg.dask.electromagnetics.frequency_domain.simulation_1d
     import simpeg.dask.electromagnetics.static.resistivity.simulation
     import simpeg.dask.electromagnetics.static.resistivity.simulation_2d
     import simpeg.dask.electromagnetics.static.induced_polarization.simulation

simpeg/dask/electromagnetics/frequency_domain/simulation_1d.py

@@ -0,0 +1,26 @@
+from ....electromagnetics.frequency_domain.simulation_1d import (
+    Simulation1DLayered as Sim,
+)
+
+from ...simulation import getJtJdiag, Jvec, Jtvec
+
+Sim._delete_on_model_update = ["_Jmatrix", "_jtjdiag", "_J"]
+
+
+@property
+def Jmatrix(self):
+    """
+    Sensitivity matrix stored on disk
+    Return the diagonal of JtJ
+    """
+    if getattr(self, "_Jmatrix", None) is None:
+        Jmat = self._getJ(self.model)
+        self._Jmatrix = Jmat["ds"] * self.sigmaDeriv
+
+    return self._Jmatrix
+
+
+Sim.getJtJdiag = getJtJdiag
+Sim.Jvec = Jvec
+Sim.Jtvec = Jtvec
+Sim.Jmatrix = Jmatrix

simpeg/dask/electromagnetics/time_domain/simulation_1d.py

@@ -12,8 +12,8 @@ def Jmatrix(self):
     Return the diagonal of JtJ
     """
     if getattr(self, "_Jmatrix", None) is None:
-        Jmat = self.getJ(self.model)
-        self._Jmatrix = Jmat * self.sigmaDeriv
+        Jmat = self._getJ(self.model)
+        self._Jmatrix = Jmat["ds"] * self.sigmaDeriv
 
     return self._Jmatrix