Fix square in getJtJdiag

Author: fourndo

SimPEG/dask/electromagnetics/static/induced_polarization/simulation.py

@@ -110,7 +110,7 @@ def dask_getJtJdiag(self, m, W=None):
         if W is None:
             W = self._scale * np.ones(self.nD)
         else:
-            W = self._scale * W.diagonal()
+            W = (self._scale * W.diagonal())**2.0
 
         diag = da.einsum('i,ij,ij->j', W, self.Jmatrix, self.Jmatrix)
 

SimPEG/dask/electromagnetics/static/resistivity/simulation.py

@@ -52,7 +52,7 @@ def dask_getJtJdiag(self, m, W=None):
         if W is None:
             W = np.ones(self.nD)
         else:
-            W = W.diagonal()
+            W = W.diagonal()**2.
 
         diag = da.einsum('i,ij,ij->j', W, self.Jmatrix, self.Jmatrix)
 

SimPEG/dask/electromagnetics/time_domain/simulation.py

@@ -161,7 +161,9 @@ def compute_J(self, f=None, Ainv=None):
     # ATinv_df_duT_v = {}
     f = dask.delayed(f)
     field_derivs_t = {}
+
     for tInd, dt in tqdm(zip(reversed(range(self.nT)), reversed(self.time_steps))):
+
         AdiagTinv = Ainv[dt]
         Asubdiag = self.getAsubdiag(tInd)
         d_count = 0

SimPEG/dask/potential_fields/gravity/simulation.py

@@ -72,7 +72,7 @@ def dask_getJtJdiag(self, m, W=None):
     if W is None:
         W = np.ones(self.nD)
     else:
-        W = W.diagonal()
+        W = W.diagonal()**2.
 
     if getattr(self, "_gtg_diagonal", None) is None:
         diag = array.einsum('i,ij,ij->j', W, self.G, self.G)

SimPEG/dask/potential_fields/magnetics/simulation.py

@@ -75,7 +75,7 @@ def dask_getJtJdiag(self, m, W=None):
         self.Jmatrix  # Wait to finish
     if getattr(self, "_gtg_diagonal", None) is None:
         if not self.is_amplitude_data:
-            diag = array.einsum('i,ij,ij->j', W, self.G, self.G)
+            diag = array.einsum('i,ij,ij->j', W**2., self.G, self.G)
         else:  # self.modelType is amplitude
             fieldDeriv = self.fieldDeriv
             J = (

SimPEG/dask/simulation.py

@@ -73,7 +73,7 @@ def dask_getJtJdiag(self, m, W=None):
         if W is None:
             W = np.ones(self.nD)
         else:
-            W = W.diagonal()
+            W = W.diagonal()**2.
 
         diag = array.einsum('i,ij,ij->j', W, self.Jmatrix, self.Jmatrix)