From ec5d00cc06c83af585e0f3782f8a2e5cbeb3401c Mon Sep 17 00:00:00 2001 From: domfournier Date: Wed, 9 Oct 2024 07:48:42 -0700 Subject: [PATCH 01/13] Move sparse conversion one level --- simpeg/electromagnetics/natural_source/receivers.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/simpeg/electromagnetics/natural_source/receivers.py b/simpeg/electromagnetics/natural_source/receivers.py index 7c6eed5207..a7e1d34b91 100644 --- a/simpeg/electromagnetics/natural_source/receivers.py +++ b/simpeg/electromagnetics/natural_source/receivers.py @@ -555,11 +555,12 @@ def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): else: ghx_v += gh_v - gh_v = ( - Phx.T @ sp.csr_matrix(ghx_v) - + Phy.T @ sp.csr_matrix(ghy_v) - + Phz.T @ sp.csr_matrix(ghz_v) + gh_v = sp.csr_matrix( + Phx.T @ ghx_v + + Phy.T @ ghy_v + + Phz.T @ ghz_v ) + return f._hDeriv(src, None, gh_v, adjoint=True) dh_v = f._hDeriv(src, du_dm_v, v, adjoint=False) From d592933cafd0ce3be67cb8df11ba43df01732503 Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 08:23:54 -0700 Subject: [PATCH 02/13] Add Point3DTipper to dask module with sparse deriv comps --- simpeg/dask/__init__.py | 1 + .../frequency_domain/simulation.py | 7 +- .../natural_source/receivers.py | 82 +++++++++++++++++++ 3 files changed, 88 insertions(+), 2 deletions(-) create mode 100644 simpeg/dask/electromagnetics/natural_source/receivers.py diff --git a/simpeg/dask/__init__.py b/simpeg/dask/__init__.py index f5a00b7334..6de861ddb1 100644 --- a/simpeg/dask/__init__.py +++ b/simpeg/dask/__init__.py @@ -1,6 +1,7 @@ try: import simpeg.dask.simulation import simpeg.dask.electromagnetics.frequency_domain.simulation + import simpeg.dask.electromagnetics.natural_source.receivers import simpeg.dask.electromagnetics.static.resistivity.simulation import simpeg.dask.electromagnetics.static.resistivity.simulation_2d import simpeg.dask.electromagnetics.static.induced_polarization.simulation diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index eb8527ed8a..0013d3281d 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -4,6 +4,7 @@ import scipy.sparse as sp from multiprocessing import cpu_count from dask import array, compute, delayed +from simpeg.utils import sdiag # from dask.distributed import get_client, Client, performance_report from simpeg.dask.simulation import dask_Jvec, dask_Jtvec, dask_getJtJdiag @@ -222,8 +223,9 @@ def compute_J(self, f=None): # with performance_report(filename="dask-report.html"): # Dask process for all derivatives + tc = time() blocks_receiver_derivs = compute(blocks_receiver_derivs)[0] - + print(f"Time to compute all derivatives: {time() - tc}") for block_derivs_chunks, addresses_chunks in tqdm( zip(blocks_receiver_derivs, blocks), ncols=len(blocks_receiver_derivs), @@ -298,7 +300,8 @@ def receiver_derivs(survey, mesh, fields, blocks): receiver = source.receiver_list[address[0][1]] if isinstance(source, PlanewaveXYPrimary): - v = np.eye(receiver.nD, dtype=float) + v = sdiag(np.ones(receiver.nD)) + # v = np.eye(receiver.nD, dtype=float) else: v = sp.csr_matrix(np.ones(receiver.nD), dtype=float) diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py new file mode 100644 index 0000000000..e902a55fdf --- /dev/null +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -0,0 +1,82 @@ +import scipy.sparse as sp +import numpy as np +from simpeg.electromagnetics.natural_source.receivers import Point3DTipper +from simpeg.utils import sdiag + + +def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): + # will grab both primary and secondary and sum them! + + if not isinstance(f, np.ndarray): + h = f[src, "h"] + else: + h = f + + Phx = self.getP(mesh, "Fx", "h") + Phy = self.getP(mesh, "Fy", "h") + Phz = self.getP(mesh, "Fz", "h") + hx = Phx @ h + hy = Phy @ h + hz = Phz @ h + + if self.orientation[1] == "x": + h = -hy + else: + h = hx + + top = h[:, 0] * hz[:, 1] - h[:, 1] * hz[:, 0] + bot = hx[:, 0] * hy[:, 1] - hx[:, 1] * hy[:, 0] + imp = top / bot + + if adjoint: + + gtop_v = v @ sdiag(bot**-1) + gbot_v = sdiag(-imp) @ v @ sdiag(bot**-1) + + diag_blocks = gbot_v @ np.c_[hy[:, 1], -hy[:, 0]] + ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gbot_v @ np.c_[-hx[:, 1], hx[:, 0]] + ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gtop_v @ np.c_[-h[:, 1], h[:, 0]] + ghz_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gtop_v @ np.c_[hz[:, 1], -hz[:, 0]] + gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + if self.orientation[1] == "x": + ghy_v -= gh_v + else: + ghx_v += gh_v + + gh_v = Phx.T @ ghx_v + Phy.T @ ghy_v + Phz.T @ ghz_v + + return f._hDeriv(src, None, gh_v, adjoint=True) + + dh_v = f._hDeriv(src, du_dm_v, v, adjoint=False) + dhx_v = Phx @ dh_v + dhy_v = Phy @ dh_v + dhz_v = Phz @ dh_v + if self.orientation[1] == "x": + dh_v = -dhy_v + else: + dh_v = dhx_v + + dtop_v = ( + h[:, 0] * dhz_v[:, 1] + + dh_v[:, 0] * hz[:, 1] + - h[:, 1] * dhz_v[:, 0] + - dh_v[:, 1] * hz[:, 0] + ) + dbot_v = ( + hx[:, 0] * dhy_v[:, 1] + + dhx_v[:, 0] * hy[:, 1] + - hx[:, 1] * dhy_v[:, 0] + - dhx_v[:, 1] * hy[:, 0] + ) + + return (bot * dtop_v - top * dbot_v) / (bot * bot) + + +Point3DTipper._eval_tipper_deriv = _eval_tipper_deriv From af1a770b8ffc297395eaa7409d1d4806e82d9973 Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 08:43:38 -0700 Subject: [PATCH 03/13] Remove time testers --- simpeg/dask/electromagnetics/frequency_domain/simulation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 0013d3281d..721f60c8f1 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -223,9 +223,8 @@ def compute_J(self, f=None): # with performance_report(filename="dask-report.html"): # Dask process for all derivatives - tc = time() blocks_receiver_derivs = compute(blocks_receiver_derivs)[0] - print(f"Time to compute all derivatives: {time() - tc}") + for block_derivs_chunks, addresses_chunks in tqdm( zip(blocks_receiver_derivs, blocks), ncols=len(blocks_receiver_derivs), @@ -301,7 +300,6 @@ def receiver_derivs(survey, mesh, fields, blocks): if isinstance(source, PlanewaveXYPrimary): v = sdiag(np.ones(receiver.nD)) - # v = np.eye(receiver.nD, dtype=float) else: v = sp.csr_matrix(np.ones(receiver.nD), dtype=float) From ca38e5021f67e0c16c42f0e96848d1c38dacbac6 Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 09:24:45 -0700 Subject: [PATCH 04/13] Implement sparse impedence receiver deriv. Fix shape issue --- .../natural_source/receivers.py | 158 +++++++++++++++++- 1 file changed, 155 insertions(+), 3 deletions(-) diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py index e902a55fdf..b145fb43a3 100644 --- a/simpeg/dask/electromagnetics/natural_source/receivers.py +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -1,9 +1,161 @@ import scipy.sparse as sp import numpy as np -from simpeg.electromagnetics.natural_source.receivers import Point3DTipper +from simpeg.electromagnetics.natural_source.receivers import ( + PointNaturalSource, + Point3DTipper, +) from simpeg.utils import sdiag +def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): + if mesh.dim < 3 and self.orientation in ["xx", "yy"]: + if adjoint: + return 0 * v + else: + return 0 * du_dm_v + e = f[src, "e"] + h = f[src, "h"] + if mesh.dim == 3: + if self.orientation[0] == "x": + Pe = self.getP(mesh, "Ex", "e") + e = Pe @ e + else: + Pe = self.getP(mesh, "Ey", "e") + e = Pe @ e + + Phx = self.getP(mesh, "Fx", "h") + Phy = self.getP(mesh, "Fy", "h") + hx = Phx @ h + hy = Phy @ h + if self.orientation[1] == "x": + h = hy + else: + h = -hx + + top = e[:, 0] * h[:, 1] - e[:, 1] * h[:, 0] + bot = hx[:, 0] * hy[:, 1] - hx[:, 1] * hy[:, 0] + imp = top / bot + else: + if mesh.dim == 1: + e_loc = f.aliasFields["e"][1] + h_loc = f.aliasFields["h"][1] + PE = self.getP(mesh, e_loc) + PH = self.getP(mesh, h_loc) + elif mesh.dim == 2: + if self.orientation == "xy": + PE = self.getP(mesh, "Ex") + PH = self.getP(mesh, "CC") + elif self.orientation == "yx": + PE = self.getP(mesh, "CC") + PH = self.getP(mesh, "Ex") + + top = PE @ e[:, 0] + bot = PH @ h[:, 0] + + if mesh.dim == 1 and self.orientation != f.field_directions: + bot = -bot + + imp = top / bot + + if adjoint: + if self.component == "phase": + # gradient of arctan2(y, x) is (-y/(x**2 + y**2), x/(x**2 + y**2)) + v = 180 / np.pi * imp / (imp.real**2 + imp.imag**2) * v + # switch real and imaginary, and negate real part of output + v = -v.imag - 1j * v.real + # imaginary part gets extra (-) due to conjugate transpose + elif self.component == "apparent_resistivity": + v = 2 * _alpha(src) * imp * v + v = v.real - 1j * v.imag + elif self.component == "imag": + v = -1j * v + + # Work backwards! + gtop_v = sdiag(bot**-1) @ v + gbot_v = sdiag(-imp / bot) @ v + + if mesh.dim == 3: + diag_blocks = gbot_v @ np.c_[hy[:, 1], -hy[:, 0]] + ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gbot_v @ np.c_[-hx[:, 1], hx[:, 0]] + ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gtop_v @ np.c_[h[:, 1], -h[:, 0]] + ge_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + diag_blocks = gtop_v @ np.c_[-e[:, 1], e[:, 0]] + gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") + + if self.orientation[1] == "x": + ghy_v += gh_v + else: + ghx_v -= gh_v + + gh_v = Phx.T @ ghx_v + Phy.T @ ghy_v + ge_v = Pe.T @ ge_v + else: + if mesh.dim == 1 and self.orientation != f.field_directions: + gbot_v = -gbot_v + + gh_v = PH.T @ gbot_v + ge_v = PE.T @ gtop_v + + gfu_h_v, gfm_h_v = f._hDeriv(src, None, gh_v, adjoint=True) + gfu_e_v, gfm_e_v = f._eDeriv(src, None, ge_v, adjoint=True) + + return gfu_h_v + gfu_e_v, gfm_h_v + gfm_e_v + + if mesh.dim == 3: + de_v = Pe @ f._eDeriv(src, du_dm_v, v, adjoint=False) + dh_v = f._hDeriv(src, du_dm_v, v, adjoint=False) + dhx_v = Phx @ dh_v + dhy_v = Phy @ dh_v + if self.orientation[1] == "x": + dh_dm_v = dhy_v + else: + dh_dm_v = -dhx_v + + dtop_v = ( + e[:, 0] * dh_dm_v[:, 1] + + de_v[:, 0] * h[:, 1] + - e[:, 1] * dh_dm_v[:, 0] + - de_v[:, 1] * h[:, 0] + ) + dbot_v = ( + hx[:, 0] * dhy_v[:, 1] + + dhx_v[:, 0] * hy[:, 1] + - hx[:, 1] * dhy_v[:, 0] + - dhx_v[:, 1] * hy[:, 0] + ) + imp_deriv = (bot * dtop_v - top * dbot_v) / (bot * bot) + else: + de_v = PE @ f._eDeriv(src, du_dm_v, v, adjoint=False) + dh_v = PH @ f._hDeriv(src, du_dm_v, v, adjoint=False) + + if mesh.dim == 1 and self.orientation != f.field_directions: + dh_v = -dh_v + + imp_deriv = (de_v - imp * dh_v) / bot + + if self.component == "apparent_resistivity": + rx_deriv = ( + 2 * _alpha(src) * (imp.real * imp_deriv.real + imp.imag * imp_deriv.imag) + ) + elif self.component == "phase": + amp2 = imp.imag**2 + imp.real**2 + deriv_re = -imp.imag / amp2 * imp_deriv.real + deriv_im = imp.real / amp2 * imp_deriv.imag + + rx_deriv = (180 / np.pi) * (deriv_re + deriv_im) + else: + rx_deriv = getattr(imp_deriv, self.component) + return rx_deriv + + +PointNaturalSource._eval_impedance_deriv = _eval_impedance_deriv + + def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): # will grab both primary and secondary and sum them! @@ -30,8 +182,8 @@ def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): if adjoint: - gtop_v = v @ sdiag(bot**-1) - gbot_v = sdiag(-imp) @ v @ sdiag(bot**-1) + gtop_v = sdiag(bot**-1) @ v + gbot_v = sdiag(-imp / bot) @ v diag_blocks = gbot_v @ np.c_[hy[:, 1], -hy[:, 0]] ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") From ba2a51b02496043166502919d49b1978e685910b Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 09:30:58 -0700 Subject: [PATCH 05/13] Fix transpose operation --- .../electromagnetics/natural_source/receivers.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py index b145fb43a3..c571a8fd65 100644 --- a/simpeg/dask/electromagnetics/natural_source/receivers.py +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -75,16 +75,16 @@ def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=Fals gbot_v = sdiag(-imp / bot) @ v if mesh.dim == 3: - diag_blocks = gbot_v @ np.c_[hy[:, 1], -hy[:, 0]] + diag_blocks = gbot_v.T @ np.c_[hy[:, 1], -hy[:, 0]] ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gbot_v @ np.c_[-hx[:, 1], hx[:, 0]] + diag_blocks = gbot_v.T @ np.c_[-hx[:, 1], hx[:, 0]] ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gtop_v @ np.c_[h[:, 1], -h[:, 0]] + diag_blocks = gbot_v.T @ np.c_[h[:, 1], -h[:, 0]] ge_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gtop_v @ np.c_[-e[:, 1], e[:, 0]] + diag_blocks = gbot_v.T @ np.c_[-e[:, 1], e[:, 0]] gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") if self.orientation[1] == "x": @@ -185,16 +185,16 @@ def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): gtop_v = sdiag(bot**-1) @ v gbot_v = sdiag(-imp / bot) @ v - diag_blocks = gbot_v @ np.c_[hy[:, 1], -hy[:, 0]] + diag_blocks = gbot_v.T @ np.c_[hy[:, 1], -hy[:, 0]] ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gbot_v @ np.c_[-hx[:, 1], hx[:, 0]] + diag_blocks = gbot_v.T @ np.c_[-hx[:, 1], hx[:, 0]] ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gtop_v @ np.c_[-h[:, 1], h[:, 0]] + diag_blocks = gtop_v.T @ np.c_[-h[:, 1], h[:, 0]] ghz_v = sp.block_diag(diag_blocks.tolist(), format="csr") - diag_blocks = gtop_v @ np.c_[hz[:, 1], -hz[:, 0]] + diag_blocks = gtop_v.T @ np.c_[hz[:, 1], -hz[:, 0]] gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") if self.orientation[1] == "x": From c841df0bfc8e359ab971a96e883ae5d6640bed1d Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 14:18:24 -0700 Subject: [PATCH 06/13] Re-work sparse sub blocks --- .../frequency_domain/simulation.py | 8 +--- .../natural_source/receivers.py | 40 +++++++++++-------- 2 files changed, 25 insertions(+), 23 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 721f60c8f1..4bbbea39b6 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -297,13 +297,7 @@ def receiver_derivs(survey, mesh, fields, blocks): for address in blocks: source = survey.source_list[address[0][0]] receiver = source.receiver_list[address[0][1]] - - if isinstance(source, PlanewaveXYPrimary): - v = sdiag(np.ones(receiver.nD)) - else: - v = sp.csr_matrix(np.ones(receiver.nD), dtype=float) - - # Assume the derivatives in terms of model are Zero (seems to always be case) + v = sdiag(np.ones(receiver.nD)) dfduT, _ = receiver.evalDeriv( source, mesh, fields, v=v[:, address[1][0]], adjoint=True ) diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py index c571a8fd65..a00ce46514 100644 --- a/simpeg/dask/electromagnetics/natural_source/receivers.py +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -75,17 +75,21 @@ def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=Fals gbot_v = sdiag(-imp / bot) @ v if mesh.dim == 3: - diag_blocks = gbot_v.T @ np.c_[hy[:, 1], -hy[:, 0]] - ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(hy[:, 1]) @ gbot_v, [1, 0]) + block_b = sp.kron(sdiag(-hy[:, 0]) @ gbot_v, [0, 1]) + ghx_v = block_a + block_b - diag_blocks = gbot_v.T @ np.c_[-hx[:, 1], hx[:, 0]] - ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(-hx[:, 1]) @ gbot_v, [1, 0]) + block_b = sp.kron(sdiag(hx[:, 0]) @ gbot_v, [0, 1]) + ghy_v = block_a + block_b - diag_blocks = gbot_v.T @ np.c_[h[:, 1], -h[:, 0]] - ge_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(h[:, 1]) @ gtop_v, [1, 0]) + block_b = sp.kron(sdiag(-h[:, 0]) @ gtop_v, [0, 1]) + ge_v = block_a + block_b - diag_blocks = gbot_v.T @ np.c_[-e[:, 1], e[:, 0]] - gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(-e[:, 1]) @ gtop_v, [1, 0]) + block_b = sp.kron(sdiag(e[:, 0]) @ gtop_v, [0, 1]) + gh_v = block_a + block_b if self.orientation[1] == "x": ghy_v += gh_v @@ -185,17 +189,21 @@ def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): gtop_v = sdiag(bot**-1) @ v gbot_v = sdiag(-imp / bot) @ v - diag_blocks = gbot_v.T @ np.c_[hy[:, 1], -hy[:, 0]] - ghx_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(hy[:, 1]) @ gbot_v, [1, 0]) + block_b = sp.kron(sdiag(-hy[:, 0]) @ gbot_v, [0, 1]) + ghx_v = block_a + block_b - diag_blocks = gbot_v.T @ np.c_[-hx[:, 1], hx[:, 0]] - ghy_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(-hx[:, 1]) @ gbot_v, [1, 0]) + block_b = sp.kron(sdiag(hx[:, 0]) @ gbot_v, [0, 1]) + ghy_v = block_a + block_b - diag_blocks = gtop_v.T @ np.c_[-h[:, 1], h[:, 0]] - ghz_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(-h[:, 1]) @ gtop_v, [1, 0]) + block_b = sp.kron(sdiag(h[:, 0]) @ gtop_v, [0, 1]) + ghz_v = block_a + block_b - diag_blocks = gtop_v.T @ np.c_[hz[:, 1], -hz[:, 0]] - gh_v = sp.block_diag(diag_blocks.tolist(), format="csr") + block_a = sp.kron(sdiag(hz[:, 1]) @ gtop_v, [1, 0]) + block_b = sp.kron(sdiag(-hz[:, 0]) @ gtop_v, [0, 1]) + gh_v = block_a + block_b if self.orientation[1] == "x": ghy_v -= gh_v From cd80581add94fe926d01a13b1acdb097b9510848 Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 14:40:22 -0700 Subject: [PATCH 07/13] Try with block optimization --- simpeg/dask/electromagnetics/frequency_domain/simulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 4bbbea39b6..8b6d33600e 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -201,7 +201,7 @@ def compute_J(self, f=None): compute_row_size = np.ceil(self.max_chunk_size / (A_i.A.shape[0] * 32.0 * 1e-6)) blocks = get_parallel_blocks( - self.survey.source_list, compute_row_size, optimize=False + self.survey.source_list, compute_row_size, optimize=True ) fields_array = delayed(f[:, self._solutionType]) fields = delayed(f) From 26235ffac3ed2c783ec852e7264e1386ccf83894 Mon Sep 17 00:00:00 2001 From: domfournier Date: Thu, 10 Oct 2024 17:07:18 -0700 Subject: [PATCH 08/13] Try looping outside the delayed --- .../frequency_domain/simulation.py | 39 ++++++++++--------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 8b6d33600e..37e031c07f 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -210,14 +210,20 @@ def compute_J(self, f=None): blocks_receiver_derivs = [] for block in blocks: - blocks_receiver_derivs.append( - receiver_derivs( - survey, - mesh, - fields, - block, + if len(block) == 0: + continue + + sub_blocks = [] + for address in block: + sub_blocks.append( + receiver_derivs( + survey, + mesh, + fields, + address, + ) ) - ) + blocks_receiver_derivs.append(sub_blocks) # with Client(processes=False) as client: # with performance_report(filename="dask-report.html"): @@ -292,18 +298,15 @@ def parallel_block_compute( @delayed -def receiver_derivs(survey, mesh, fields, blocks): - field_derivatives = [] - for address in blocks: - source = survey.source_list[address[0][0]] - receiver = source.receiver_list[address[0][1]] - v = sdiag(np.ones(receiver.nD)) - dfduT, _ = receiver.evalDeriv( - source, mesh, fields, v=v[:, address[1][0]], adjoint=True - ) - field_derivatives.append(dfduT) +def receiver_derivs(survey, mesh, fields, address): + source = survey.source_list[address[0][0]] + receiver = source.receiver_list[address[0][1]] + v = sdiag(np.ones(receiver.nD)) + dfduT, _ = receiver.evalDeriv( + source, mesh, fields, v=v[:, address[1][0]], adjoint=True + ) - return field_derivatives + return dfduT @delayed From ed6f172802e0051e41e568ebebc22c9fc447b42b Mon Sep 17 00:00:00 2001 From: domfournier Date: Fri, 11 Oct 2024 15:46:17 -0700 Subject: [PATCH 09/13] Run all source/receiver derivs as flat list --- .../frequency_domain/simulation.py | 58 +++++++++++-------- simpeg/dask/utils.py | 51 ++++++++++++++++ 2 files changed, 85 insertions(+), 24 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 37e031c07f..99b672b63f 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -5,10 +5,11 @@ from multiprocessing import cpu_count from dask import array, compute, delayed from simpeg.utils import sdiag +from time import time # from dask.distributed import get_client, Client, performance_report from simpeg.dask.simulation import dask_Jvec, dask_Jtvec, dask_getJtJdiag -from simpeg.dask.utils import get_parallel_blocks +from simpeg.dask.utils import get_block_survey from simpeg.electromagnetics.natural_source.sources import PlanewaveXYPrimary import zarr from tqdm import tqdm @@ -200,40 +201,52 @@ def compute_J(self, f=None): Jmatrix = np.zeros((self.survey.nD, m_size), dtype=np.float32) compute_row_size = np.ceil(self.max_chunk_size / (A_i.A.shape[0] * 32.0 * 1e-6)) - blocks = get_parallel_blocks( - self.survey.source_list, compute_row_size, optimize=True - ) + + if getattr(self, "source_list_blocks", None) is None: + self.source_list_blocks, self.addresses = get_block_survey( + self.survey.source_list, compute_row_size, optimize=True + ) + fields_array = delayed(f[:, self._solutionType]) fields = delayed(f) - survey = delayed(self.survey) + mesh = delayed(self.mesh) - blocks_receiver_derivs = [] - for block in blocks: - if len(block) == 0: - continue + parallel_blocks = [] - sub_blocks = [] - for address in block: - sub_blocks.append( + for block in self.source_list_blocks: + deriv_blocks = [] + for source in block: + deriv_blocks.append( receiver_derivs( - survey, + source, + source.receiver_list[0], # Always a list of one mesh, fields, - address, + # address, ) ) - blocks_receiver_derivs.append(sub_blocks) + + parallel_blocks.append(deriv_blocks) # with Client(processes=False) as client: # with performance_report(filename="dask-report.html"): # Dask process for all derivatives - blocks_receiver_derivs = compute(blocks_receiver_derivs)[0] + # blocks_receiver_derivs = [] + tc = time() + parallel_blocks = compute(parallel_blocks)[0] + print("Time to compute receiver derivatives", time() - tc) + # for address in blocks: + # sub_blocks = [] + # for block in address: + # sub_blocks.append(source_blocks[block[0]][block[1]][:, block[2]]) + # blocks_receiver_derivs.append(sub_blocks) + # for block_derivs_chunks, addresses_chunks in tqdm( - zip(blocks_receiver_derivs, blocks), - ncols=len(blocks_receiver_derivs), + zip(parallel_blocks, self.addresses), + ncols=len(parallel_blocks), desc=f"Sensitivities at {list(self.Ainv)[0]} Hz", ): Jmatrix = parallel_block_compute( @@ -298,13 +311,10 @@ def parallel_block_compute( @delayed -def receiver_derivs(survey, mesh, fields, address): - source = survey.source_list[address[0][0]] - receiver = source.receiver_list[address[0][1]] +def receiver_derivs(source, receiver, mesh, fields): + v = sdiag(np.ones(receiver.nD)) - dfduT, _ = receiver.evalDeriv( - source, mesh, fields, v=v[:, address[1][0]], adjoint=True - ) + dfduT, _ = receiver.evalDeriv(source, mesh, fields, v=v, adjoint=True) return dfduT diff --git a/simpeg/dask/utils.py b/simpeg/dask/utils.py index ad292dc9a2..8a76d5fa4d 100644 --- a/simpeg/dask/utils.py +++ b/simpeg/dask/utils.py @@ -1,4 +1,5 @@ import numpy as np +from copy import deepcopy from dask.distributed import get_client from multiprocessing import cpu_count @@ -40,6 +41,56 @@ def compute(self, job): return job.compute() +def get_block_survey( + source_list: list, data_block_size, optimize=True +) -> tuple[list, list]: + row_count = 0 + row_index = 0 + block_count = 0 + blocks = [[]] + parallel_blocks_list = [[]] + for s_id, src in enumerate(source_list): + for r_id, rx in enumerate(src.receiver_list): + indices = np.arange(rx.nD).astype(int) + chunks = np.array_split( + indices, int(np.ceil(len(indices) / data_block_size)) + ) + + for ind, chunk in enumerate(chunks): + chunk_size = len(chunk) + new_receiver = type(rx)( + locations=rx.locations[chunk, :], + orientation=rx.orientation, + component=rx.component, + storeProjections=rx.storeProjections, + ) + new_source = deepcopy(src) + new_source.receiver_list = [new_receiver] + + # Condition to start a new block + if (row_count + chunk_size) > (data_block_size * cpu_count()): + row_count = 0 + block_count += 1 + blocks.append([]) + parallel_blocks_list.append([]) + + parallel_blocks_list[block_count].append(new_source) + blocks[block_count].append( + ( + (s_id, r_id, ind), + ( + chunk, + np.arange(row_index, row_index + chunk_size).astype(int), + rx.locations.shape[0], + ), + ) + ) + row_index += chunk_size + row_count += chunk_size + + return parallel_blocks_list, blocks + + def get_parallel_blocks(source_list: list, data_block_size, optimize=True) -> list: """ Get the blocks of sources and receivers to be computed in parallel. From 4dea7cb5ac01da4cee62cd1ea0b4edb205a1f00b Mon Sep 17 00:00:00 2001 From: domfournier Date: Mon, 14 Oct 2024 08:30:14 -0700 Subject: [PATCH 10/13] Remove fields from delayed call --- .../frequency_domain/simulation.py | 33 +++++++++++-------- .../natural_source/receivers.py | 26 ++++++++------- 2 files changed, 34 insertions(+), 25 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index 99b672b63f..f33a414655 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -11,6 +11,10 @@ from simpeg.dask.simulation import dask_Jvec, dask_Jtvec, dask_getJtJdiag from simpeg.dask.utils import get_block_survey from simpeg.electromagnetics.natural_source.sources import PlanewaveXYPrimary +from simpeg.electromagnetics.natural_source.receivers import ( + PointNaturalSource, + Point3DTipper, +) import zarr from tqdm import tqdm @@ -208,8 +212,10 @@ def compute_J(self, f=None): ) fields_array = delayed(f[:, self._solutionType]) - fields = delayed(f) + # fields = delayed(f) + e = delayed(f[self.survey.source_list[0], "e"]) + h = delayed(f[self.survey.source_list[0], "h"]) mesh = delayed(self.mesh) parallel_blocks = [] @@ -222,8 +228,9 @@ def compute_J(self, f=None): source, source.receiver_list[0], # Always a list of one mesh, - fields, - # address, + e, + h, + self, ) ) @@ -235,14 +242,9 @@ def compute_J(self, f=None): # Dask process for all derivatives # blocks_receiver_derivs = [] tc = time() + print("Computing receiver derivatives") parallel_blocks = compute(parallel_blocks)[0] - print("Time to compute receiver derivatives", time() - tc) - # for address in blocks: - # sub_blocks = [] - # for block in address: - # sub_blocks.append(source_blocks[block[0]][block[1]][:, block[2]]) - # blocks_receiver_derivs.append(sub_blocks) - # + print("Runtime:", time() - tc) for block_derivs_chunks, addresses_chunks in tqdm( zip(parallel_blocks, self.addresses), @@ -311,10 +313,15 @@ def parallel_block_compute( @delayed -def receiver_derivs(source, receiver, mesh, fields): - +def receiver_derivs(source, receiver, mesh, e, h, simulation): v = sdiag(np.ones(receiver.nD)) - dfduT, _ = receiver.evalDeriv(source, mesh, fields, v=v, adjoint=True) + + if isinstance(receiver, Point3DTipper): + if receiver.component == "imag": + v = -1j * v + dfduT = receiver._eval_tipper_deriv(source, mesh, h, simulation, v=v, adjoint=True) + elif isinstance(receiver, PointNaturalSource): + dfduT, _ = receiver._eval_impedance_deriv(source, mesh, e, h, simulation, v=v, adjoint=True) return dfduT diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py index a00ce46514..f9484c9618 100644 --- a/simpeg/dask/electromagnetics/natural_source/receivers.py +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -7,14 +7,14 @@ from simpeg.utils import sdiag -def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): +def _eval_impedance_deriv(self, src, mesh, e, h, simulation, du_dm_v=None, v=None, adjoint=False): if mesh.dim < 3 and self.orientation in ["xx", "yy"]: if adjoint: return 0 * v else: return 0 * du_dm_v - e = f[src, "e"] - h = f[src, "h"] + # e = f[src, "e"] + # h = f[src, "h"] if mesh.dim == 3: if self.orientation[0] == "x": Pe = self.getP(mesh, "Ex", "e") @@ -105,10 +105,9 @@ def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=Fals gh_v = PH.T @ gbot_v ge_v = PE.T @ gtop_v - gfu_h_v, gfm_h_v = f._hDeriv(src, None, gh_v, adjoint=True) - gfu_e_v, gfm_e_v = f._eDeriv(src, None, ge_v, adjoint=True) + gfu_h_v = -1.0 / (1j * 2 * np.pi * src.frequency) * (mesh.edge_curl.T * (simulation.MfMui.T * (simulation.MfI.T * gh_v))) - return gfu_h_v + gfu_e_v, gfm_h_v + gfm_e_v + return gfu_h_v + ge_v, None if mesh.dim == 3: de_v = Pe @ f._eDeriv(src, du_dm_v, v, adjoint=False) @@ -160,13 +159,13 @@ def _eval_impedance_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=Fals PointNaturalSource._eval_impedance_deriv = _eval_impedance_deriv -def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): +def _eval_tipper_deriv(self, src, mesh, h, simulation, du_dm_v=None, v=None, adjoint=False): # will grab both primary and secondary and sum them! - if not isinstance(f, np.ndarray): - h = f[src, "h"] - else: - h = f + # if not isinstance(f, np.ndarray): + # h = f[src, "h"] + # else: + # h = f Phx = self.getP(mesh, "Fx", "h") Phy = self.getP(mesh, "Fy", "h") @@ -212,7 +211,10 @@ def _eval_tipper_deriv(self, src, mesh, f, du_dm_v=None, v=None, adjoint=False): gh_v = Phx.T @ ghx_v + Phy.T @ ghy_v + Phz.T @ ghz_v - return f._hDeriv(src, None, gh_v, adjoint=True) + gfu_h_v = -1.0 / (1j * 2 * np.pi * src.frequency) * ( + mesh.edge_curl.T * (simulation.MfMui.T * (simulation.MfI.T * gh_v))) + + return gfu_h_v dh_v = f._hDeriv(src, du_dm_v, v, adjoint=False) dhx_v = Phx @ dh_v From d5068e6cd0af38264baaffadd0c7c6205f78fd2d Mon Sep 17 00:00:00 2001 From: domfournier Date: Mon, 14 Oct 2024 08:52:35 -0700 Subject: [PATCH 11/13] Remove source from delayed --- .../electromagnetics/frequency_domain/simulation.py | 10 +++++----- .../dask/electromagnetics/natural_source/receivers.py | 8 ++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index f33a414655..c283bfc114 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -212,7 +212,7 @@ def compute_J(self, f=None): ) fields_array = delayed(f[:, self._solutionType]) - # fields = delayed(f) + fields = delayed(f) e = delayed(f[self.survey.source_list[0], "e"]) h = delayed(f[self.survey.source_list[0], "h"]) @@ -225,7 +225,7 @@ def compute_J(self, f=None): for source in block: deriv_blocks.append( receiver_derivs( - source, + source.frequency, source.receiver_list[0], # Always a list of one mesh, e, @@ -313,15 +313,15 @@ def parallel_block_compute( @delayed -def receiver_derivs(source, receiver, mesh, e, h, simulation): +def receiver_derivs(frequency, receiver, mesh, e, h, simulation): v = sdiag(np.ones(receiver.nD)) if isinstance(receiver, Point3DTipper): if receiver.component == "imag": v = -1j * v - dfduT = receiver._eval_tipper_deriv(source, mesh, h, simulation, v=v, adjoint=True) + dfduT = receiver._eval_tipper_deriv(frequency, mesh, h, simulation, v=v, adjoint=True) elif isinstance(receiver, PointNaturalSource): - dfduT, _ = receiver._eval_impedance_deriv(source, mesh, e, h, simulation, v=v, adjoint=True) + dfduT, _ = receiver._eval_impedance_deriv(frequency, mesh, e, h, simulation, v=v, adjoint=True) return dfduT diff --git a/simpeg/dask/electromagnetics/natural_source/receivers.py b/simpeg/dask/electromagnetics/natural_source/receivers.py index f9484c9618..16fcaf1c16 100644 --- a/simpeg/dask/electromagnetics/natural_source/receivers.py +++ b/simpeg/dask/electromagnetics/natural_source/receivers.py @@ -7,7 +7,7 @@ from simpeg.utils import sdiag -def _eval_impedance_deriv(self, src, mesh, e, h, simulation, du_dm_v=None, v=None, adjoint=False): +def _eval_impedance_deriv(self, frequency, mesh, e, h, simulation, du_dm_v=None, v=None, adjoint=False): if mesh.dim < 3 and self.orientation in ["xx", "yy"]: if adjoint: return 0 * v @@ -105,7 +105,7 @@ def _eval_impedance_deriv(self, src, mesh, e, h, simulation, du_dm_v=None, v=Non gh_v = PH.T @ gbot_v ge_v = PE.T @ gtop_v - gfu_h_v = -1.0 / (1j * 2 * np.pi * src.frequency) * (mesh.edge_curl.T * (simulation.MfMui.T * (simulation.MfI.T * gh_v))) + gfu_h_v = -1.0 / (1j * 2 * np.pi * frequency) * (mesh.edge_curl.T * (simulation.MfMui.T * (simulation.MfI.T * gh_v))) return gfu_h_v + ge_v, None @@ -159,7 +159,7 @@ def _eval_impedance_deriv(self, src, mesh, e, h, simulation, du_dm_v=None, v=Non PointNaturalSource._eval_impedance_deriv = _eval_impedance_deriv -def _eval_tipper_deriv(self, src, mesh, h, simulation, du_dm_v=None, v=None, adjoint=False): +def _eval_tipper_deriv(self, frequency, mesh, h, simulation, du_dm_v=None, v=None, adjoint=False): # will grab both primary and secondary and sum them! # if not isinstance(f, np.ndarray): @@ -211,7 +211,7 @@ def _eval_tipper_deriv(self, src, mesh, h, simulation, du_dm_v=None, v=None, adj gh_v = Phx.T @ ghx_v + Phy.T @ ghy_v + Phz.T @ ghz_v - gfu_h_v = -1.0 / (1j * 2 * np.pi * src.frequency) * ( + gfu_h_v = -1.0 / (1j * 2 * np.pi * frequency) * ( mesh.edge_curl.T * (simulation.MfMui.T * (simulation.MfI.T * gh_v))) return gfu_h_v From 40a17de8f9c97fbce4742f7937e982503d5abc8d Mon Sep 17 00:00:00 2001 From: domfournier Date: Mon, 14 Oct 2024 10:18:58 -0700 Subject: [PATCH 12/13] Only compute when requested --- .../frequency_domain/simulation.py | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/simpeg/dask/electromagnetics/frequency_domain/simulation.py b/simpeg/dask/electromagnetics/frequency_domain/simulation.py index c283bfc114..538f2ffaa4 100644 --- a/simpeg/dask/electromagnetics/frequency_domain/simulation.py +++ b/simpeg/dask/electromagnetics/frequency_domain/simulation.py @@ -224,13 +224,17 @@ def compute_J(self, f=None): deriv_blocks = [] for source in block: deriv_blocks.append( - receiver_derivs( - source.frequency, - source.receiver_list[0], # Always a list of one - mesh, - e, - h, - self, + array.from_delayed( + receiver_derivs( + source.frequency, + source.receiver_list[0], # Always a list of one + mesh, + e, + h, + self, + ), + dtype=np.complex128, + shape=(A_i.A.shape[0], source.receiver_list[0].nD * 2), ) ) @@ -241,10 +245,10 @@ def compute_J(self, f=None): # Dask process for all derivatives # blocks_receiver_derivs = [] - tc = time() - print("Computing receiver derivatives") - parallel_blocks = compute(parallel_blocks)[0] - print("Runtime:", time() - tc) + # tc = time() + # print("Computing receiver derivatives") + # parallel_blocks = compute(parallel_blocks)[0] + # print("Runtime:", time() - tc) for block_derivs_chunks, addresses_chunks in tqdm( zip(parallel_blocks, self.addresses), @@ -272,7 +276,10 @@ def parallel_block_compute( self, Jmatrix, blocks_receiver_derivs, A_i, fields_array, addresses ): m_size = self.model.size - block_stack = sp.hstack(blocks_receiver_derivs).toarray() + print("Computing block") + tc = time() + block_stack = array.hstack(blocks_receiver_derivs).compute().toarray() + print("Runtime:", time() - tc) ATinvdf_duT = delayed(A_i * block_stack) count = 0 rows = [] From 6bd21a06d9662a196cd467a2a1f0a1f42491abb3 Mon Sep 17 00:00:00 2001 From: domfournier Date: Mon, 14 Oct 2024 10:46:12 -0700 Subject: [PATCH 13/13] Fix indexing bug --- simpeg/electromagnetics/natural_source/receivers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/simpeg/electromagnetics/natural_source/receivers.py b/simpeg/electromagnetics/natural_source/receivers.py index a7e1d34b91..2683ee53ce 100644 --- a/simpeg/electromagnetics/natural_source/receivers.py +++ b/simpeg/electromagnetics/natural_source/receivers.py @@ -175,7 +175,7 @@ def getP(self, mesh, projected_grid, field="e"): if mesh.dim < 3: return super().getP(mesh, projected_grid) - if (mesh, projected_grid) in self._Ps: + if (mesh, projected_grid, field) in self._Ps: return self._Ps[(mesh, projected_grid, field)] if field == "e":