From 57a2f232046454d0f194072e73439928b3894b56 Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 10 Apr 2026 10:34:11 -0700 Subject: [PATCH 1/8] Set RHS to same as A matrix --- simpeg/potential_fields/magnetics/simulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py index 733c92be87..75fa4be23a 100644 --- a/simpeg/potential_fields/magnetics/simulation.py +++ b/simpeg/potential_fields/magnetics/simulation.py @@ -1862,7 +1862,7 @@ def _getRHS(self, m): ).diagonal() ) - return rhs + return rhs.astype(self.solver_dtype) def _getA(self): A = self._Div * self.MfMuiI * self._DivT From 6eecc1c30ededea5425e050072a0149e428ef41c Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 10 Apr 2026 10:34:42 -0700 Subject: [PATCH 2/8] Repair instantiation of Wires with tuple of int --- simpeg/maps/_base.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/simpeg/maps/_base.py b/simpeg/maps/_base.py index e1ae8109d2..cb186ed54f 100644 --- a/simpeg/maps/_base.py +++ b/simpeg/maps/_base.py @@ -1150,6 +1150,10 @@ def __init__(self, *args): for arg in args: if isinstance(arg[1], (int, np.integer)): + + if not getattr(self, "_nP", None): + self._nP = int(np.sum([w[1] for w in args])) + wire = Projection(self.nP, slice(start, start + arg[1])) start += arg[1] else: From 89c6ac149adb4739df53a786773c7af8f3d565aa Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 10 Apr 2026 13:12:26 -0700 Subject: [PATCH 3/8] Add dask method for getjtjdiag for mvi pde --- simpeg/dask/__init__.py | 1 + .../magnetics/simulation_pde.py | 74 +++++++++++++++++++ simpeg/directives/_vector_models.py | 4 +- 3 files changed, 77 insertions(+), 2 deletions(-) create mode 100644 simpeg/dask/potential_fields/magnetics/simulation_pde.py diff --git a/simpeg/dask/__init__.py b/simpeg/dask/__init__.py index 9a58d91b7f..3bfeb9c8ea 100644 --- a/simpeg/dask/__init__.py +++ b/simpeg/dask/__init__.py @@ -11,6 +11,7 @@ import simpeg.dask.potential_fields.base import simpeg.dask.potential_fields.gravity.simulation import simpeg.dask.potential_fields.magnetics.simulation + import simpeg.dask.potential_fields.magnetics.simulation_pde import simpeg.dask.simulation import simpeg.dask.inverse_problem import simpeg.dask.objective_function diff --git a/simpeg/dask/potential_fields/magnetics/simulation_pde.py b/simpeg/dask/potential_fields/magnetics/simulation_pde.py new file mode 100644 index 0000000000..fe1df07d13 --- /dev/null +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -0,0 +1,74 @@ +import numpy as np +from ....potential_fields.magnetics import Simulation3DDifferential as Sim +from ....utils import sdiag, mkvc + + +def distance_weights(locations, cell_centers, exponent=3, threshold=1e-2): + distance_weights = np.zeros(len(cell_centers)) + for ind, loc in enumerate(locations): + distance = np.linalg.norm(cell_centers - loc, axis=1) + distance_weights += (distance + threshold) ** (-2 * exponent) + + return distance_weights + +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() + + client, worker = self._get_client_worker() + + n_threads = self.n_threads(client=client, worker=worker) + + chunks = np.array_split(self.survey.receiver_locations, n_threads) + cell_centers = self.mesh.cell_centers.copy() + + if client: + cell_centers = client.scatter(cell_centers, workers=worker) + else: + delayed_distance_weights = delayed(distance_weights) + + futures = [] + for block in chunks: + if client: + futures.append( + client.submit( + distance_weights, + block, + cell_centers, + workers=worker, + ) + ) + else: + futures.append( + array.from_delayed( + delayed_compute_rows( + block, + cell_centers, + ), + dtype=np.float32, + shape=( + len(block), + self.active_cells.sum(), + ), + ) + ) + + if client: + diag = client.gather(futures) + else: + diag = compute(futures) + + diag = np.tile(np.vstack(diag).sum(axis=0) * self.mesh.cell_volumes**2.,3)**0.5 + return mkvc((sdiag(np.sqrt(diag)) @ self.remDeriv).power(2).sum(axis=0)) + + +Sim.getJtJdiag = dask_getJtJdiag \ No newline at end of file diff --git a/simpeg/directives/_vector_models.py b/simpeg/directives/_vector_models.py index b2b58c7dd3..64706d433e 100644 --- a/simpeg/directives/_vector_models.py +++ b/simpeg/directives/_vector_models.py @@ -91,7 +91,7 @@ class VectorInversion(InversionDirective): chifact_target = 1.0 reference_model = None mode = "cartesian" - inversion_type = "mvis" + inversion_type = "magnetic vector" norms = [] alphas = [] cartesian_model = None @@ -162,7 +162,7 @@ def endIter(self): if ( self.invProb.phi_d < self.target - ) and self.mode == "cartesian": # and self.inversion_type == 'mvis': + ) and self.mode == "cartesian" and self.inversion_type == "magnetic vector": print("Switching MVI to spherical coordinates") self.mode = "spherical" self.cartesian_model = model From 092f610c7557e6e185a304fe02dc2604254269a6 Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 10 Apr 2026 13:46:26 -0700 Subject: [PATCH 4/8] Review distance weights comps --- .../magnetics/simulation_pde.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/simpeg/dask/potential_fields/magnetics/simulation_pde.py b/simpeg/dask/potential_fields/magnetics/simulation_pde.py index fe1df07d13..9d6e0d16cd 100644 --- a/simpeg/dask/potential_fields/magnetics/simulation_pde.py +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -1,16 +1,20 @@ +from dask import array, compute, delayed import numpy as np from ....potential_fields.magnetics import Simulation3DDifferential as Sim from ....utils import sdiag, mkvc -def distance_weights(locations, cell_centers, exponent=3, threshold=1e-2): +def distance_weights(locations, cell_centers, cell_volumes, exponent=3, threshold=1e-2): distance_weights = np.zeros(len(cell_centers)) - for ind, loc in enumerate(locations): + for loc in locations: distance = np.linalg.norm(cell_centers - loc, axis=1) - distance_weights += (distance + threshold) ** (-2 * exponent) + distance_weights += cell_volumes**2.0 * (distance + threshold) ** ( + -2 * exponent + ) return distance_weights + def dask_getJtJdiag(self, m, W=None, f=None): """ Return the diagonal of JtJ @@ -30,9 +34,11 @@ def dask_getJtJdiag(self, m, W=None, f=None): chunks = np.array_split(self.survey.receiver_locations, n_threads) cell_centers = self.mesh.cell_centers.copy() + cell_volumes = self.mesh.cell_volumes.copy() if client: cell_centers = client.scatter(cell_centers, workers=worker) + cell_volumes = client.scatter(cell_volumes, workers=worker) else: delayed_distance_weights = delayed(distance_weights) @@ -44,15 +50,17 @@ def dask_getJtJdiag(self, m, W=None, f=None): distance_weights, block, cell_centers, + cell_volumes, workers=worker, ) ) else: futures.append( array.from_delayed( - delayed_compute_rows( + delayed_distance_weights( block, cell_centers, + cell_volumes, ), dtype=np.float32, shape=( @@ -67,8 +75,8 @@ def dask_getJtJdiag(self, m, W=None, f=None): else: diag = compute(futures) - diag = np.tile(np.vstack(diag).sum(axis=0) * self.mesh.cell_volumes**2.,3)**0.5 + diag = np.tile(np.vstack(diag).sum(axis=0), 3) return mkvc((sdiag(np.sqrt(diag)) @ self.remDeriv).power(2).sum(axis=0)) -Sim.getJtJdiag = dask_getJtJdiag \ No newline at end of file +Sim.getJtJdiag = dask_getJtJdiag From 2a2e8dde0e2c84ce839c65e48491afaa03b9a8a3 Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 10 Apr 2026 14:23:10 -0700 Subject: [PATCH 5/8] Fix dtype everywhere Ainv is called --- .../potential_fields/magnetics/simulation.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py index 75fa4be23a..9f3cb44c36 100644 --- a/simpeg/potential_fields/magnetics/simulation.py +++ b/simpeg/potential_fields/magnetics/simulation.py @@ -1991,13 +1991,23 @@ def _Jtvec(self, m, v, f): if v is None: v = np.eye(Q.shape[0]) divt_solve_q = ( - self._DivT * (self._Ainv * ((Q * self.MfMuiI * -self._DivT).T * v)) + self._DivT + * ( + self._Ainv + * ((Q * self.MfMuiI * -self._DivT).T * v).astype(self.solver_dtype) + ) + Q.T * v ) del v else: divt_solve_q = ( - self._DivT * (self._Ainv * ((-self._Div * (self.MfMuiI.T * (Q.T * v))))) + self._DivT + * ( + self._Ainv + * ((-self._Div * (self.MfMuiI.T * (Q.T * v)))).astype( + self.solver_dtype + ) + ) + Q.T * v ) @@ -2071,7 +2081,9 @@ def _Jvec(self, m, v, f): self.MfMuiI * Mf_r_mui_deriv * v ) - Ainv_Ddm = self._Ainv * (self._Div * (-dCmu_dm + db_dm)) + Ainv_Ddm = self._Ainv * (self._Div * (-dCmu_dm + db_dm)).astype( + self.solver_dtype + ) Jv = Q * (C * Ainv_Ddm + (-dCmu_dm + db_dm)) From 19624197c4e51c685be15ccfce33919cc834c405 Mon Sep 17 00:00:00 2001 From: domfournier Date: Mon, 13 Apr 2026 07:56:04 -0700 Subject: [PATCH 6/8] Minor clean ups --- simpeg/dask/potential_fields/magnetics/simulation_pde.py | 2 +- simpeg/directives/_save_geoh5.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/simpeg/dask/potential_fields/magnetics/simulation_pde.py b/simpeg/dask/potential_fields/magnetics/simulation_pde.py index 9d6e0d16cd..f23bfbca68 100644 --- a/simpeg/dask/potential_fields/magnetics/simulation_pde.py +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -65,7 +65,7 @@ def dask_getJtJdiag(self, m, W=None, f=None): dtype=np.float32, shape=( len(block), - self.active_cells.sum(), + len(cell_centers), ), ) ) diff --git a/simpeg/directives/_save_geoh5.py b/simpeg/directives/_save_geoh5.py index 8647c96993..72712b31a2 100644 --- a/simpeg/directives/_save_geoh5.py +++ b/simpeg/directives/_save_geoh5.py @@ -486,7 +486,7 @@ def write(self, iteration: int, **_): if (channel_name in child.name and isinstance(child, FloatData)) ] - if children[0] is not None: + if children: properties += children if len(properties) == 0: From f1a6f6eb83aca975080c4d9466ccfb0a1335d50b Mon Sep 17 00:00:00 2001 From: domfournier Date: Tue, 14 Apr 2026 08:46:15 -0700 Subject: [PATCH 7/8] Update simpeg/dask/potential_fields/magnetics/simulation_pde.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- simpeg/dask/potential_fields/magnetics/simulation_pde.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/simpeg/dask/potential_fields/magnetics/simulation_pde.py b/simpeg/dask/potential_fields/magnetics/simulation_pde.py index f23bfbca68..922ab811bd 100644 --- a/simpeg/dask/potential_fields/magnetics/simulation_pde.py +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -5,14 +5,14 @@ def distance_weights(locations, cell_centers, cell_volumes, exponent=3, threshold=1e-2): - distance_weights = np.zeros(len(cell_centers)) + weights = np.zeros(len(cell_centers)) for loc in locations: distance = np.linalg.norm(cell_centers - loc, axis=1) - distance_weights += cell_volumes**2.0 * (distance + threshold) ** ( + weights += cell_volumes**2.0 * (distance + threshold) ** ( -2 * exponent ) - return distance_weights + return weights def dask_getJtJdiag(self, m, W=None, f=None): From 465bc8dc352a226d931c113eb27c7bd3fc771451 Mon Sep 17 00:00:00 2001 From: domfournier Date: Tue, 14 Apr 2026 10:37:19 -0700 Subject: [PATCH 8/8] Update simpeg/dask/potential_fields/magnetics/simulation_pde.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- simpeg/dask/potential_fields/magnetics/simulation_pde.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/simpeg/dask/potential_fields/magnetics/simulation_pde.py b/simpeg/dask/potential_fields/magnetics/simulation_pde.py index 922ab811bd..61252b1e4a 100644 --- a/simpeg/dask/potential_fields/magnetics/simulation_pde.py +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -73,7 +73,7 @@ def dask_getJtJdiag(self, m, W=None, f=None): if client: diag = client.gather(futures) else: - diag = compute(futures) + diag = compute(futures)[0] diag = np.tile(np.vstack(diag).sum(axis=0), 3) return mkvc((sdiag(np.sqrt(diag)) @ self.remDeriv).power(2).sum(axis=0))