Merge pull request #78 from MiraGeoscience/GEOPY-2045

GEOPY-2045: Phi-d at Iter 0 shows as nan in the *.out file

Author: domfournier

.pre-commit-config.yaml

@@ -21,6 +21,7 @@ repos:
           - flake8-mutable==1.2.0
           - flake8-rst-docstrings==0.3.0
           - flake8-docstrings==1.7.0
+          - flake8-pyproject==1.2.3
   - repo: https://github.com/MiraGeoscience/pre-commit-hooks
     rev: v1.1.0
     hooks:

pyproject.toml

@@ -136,7 +136,7 @@ exclude_also = [
 directory = "coverage_html_report"
 
 [tool.black]
-required-version = '24.3.0'
+required-version = '25.1.0'
 target-version = ['py38', 'py39', 'py310', 'py311']
 
 [tool.flake8]

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

simpeg/directives/__init__.py

@@ -118,15 +118,24 @@
     ScalingMultipleDataMisfits_ByEig,
     JointScalingSchedule,
     UpdateSensitivityWeights,
-    VectorInversion,
+    Update_IRLS,
+    ScaleMisfitMultipliers,
+)
+
+from ._save_geoh5 import (
+    BaseSaveGeoH5,
     SaveDataGeoH5,
     SaveLogFilesGeoH5,
     SaveModelGeoH5,
     SavePropertyGroup,
     SaveSensitivityGeoH5,
-    Update_IRLS,
+)
+
+from ._regularization import UpdateIRLS, SphericalUnitsWeights
+
+from ._vector_models import (
+    VectorInversion,
     ProjectSphericalBounds,
-    ScaleMisfitMultipliers,
 )
 
 from .pgi_directives import (
@@ -135,8 +144,6 @@
     PGI_AddMrefInSmooth,
 )
 
-from ._regularization import UpdateIRLS, SphericalUnitsWeights
-
 from .sim_directives import (
     SimilarityMeasureInversionDirective,
     SimilarityMeasureSaveOutputEveryIteration,