Bring back getJtJdiag for mag to deal with spherical

Author: domfournier

simpeg/dask/potential_fields/magnetics/simulation.py

@@ -1,5 +1,40 @@
+import numpy as np
 from ....potential_fields.magnetics import Simulation3DIntegral as Sim
-from ...simulation import getJtJdiag
+from ....utils import sdiag, mkvc
+
+
+def dask_getJtJdiag(self, m, W=None, f=None):
+    """
+    Return the diagonal of JtJ
+    """
+
+    self.model = m
+
+    self.model = m
+    if W is None:
+        W = np.ones(self.Jmatrix.shape[0])
+    else:
+        W = W.diagonal()
+
+    if getattr(self, "_gtg_diagonal", None) is None:
+        if not self.is_amplitude_data:
+            diag = np.asarray(np.einsum("i,ij,ij->j", W**2, self.Jmatrix, self.Jmatrix))
+        else:
+            ampDeriv = self.ampDeriv
+            J = (
+                ampDeriv[0, :, None] * self.Jmatrix[::3]
+                + ampDeriv[1, :, None] * self.Jmatrix[1::3]
+                + ampDeriv[2, :, None] * self.Jmatrix[2::3]
+            )
+            diag = ((W[:, None] * J) ** 2).sum(axis=0).compute()
+        self._gtg_diagonal = diag
+    else:
+        diag = self._gtg_diagonal
+
+    return mkvc((sdiag(np.sqrt(diag)) @ self.chiDeriv).power(2).sum(axis=0))
+
+
+Sim.getJtJdiag = dask_getJtJdiag
 
 
 @property
@@ -14,5 +49,4 @@ def G(self):
 
 
 Sim._delete_on_model_update = []
-Sim.getJtJdiag = getJtJdiag
 Sim.G = G