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..f23bfbca68 --- /dev/null +++ b/simpeg/dask/potential_fields/magnetics/simulation_pde.py @@ -0,0 +1,82 @@ +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, cell_volumes, exponent=3, threshold=1e-2): + distance_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) ** ( + -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() + 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) + + futures = [] + for block in chunks: + if client: + futures.append( + client.submit( + distance_weights, + block, + cell_centers, + cell_volumes, + workers=worker, + ) + ) + else: + futures.append( + array.from_delayed( + delayed_distance_weights( + block, + cell_centers, + cell_volumes, + ), + dtype=np.float32, + shape=( + len(block), + len(cell_centers), + ), + ) + ) + + if client: + diag = client.gather(futures) + else: + diag = compute(futures) + + 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 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: 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 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: diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py index 733c92be87..9f3cb44c36 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 @@ -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))