# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
"""
2-D magnetotelluric finite-difference forward solver.
:class:`MT2DForward` solves the 2-D MT boundary-value problem for both
the TE and TM polarisation modes on a non-uniform staggered grid using
the finite-difference (FD) method.
Physics
-------
For a 2-D earth with conductivity σ(x, z) and a plane electromagnetic wave
incident from above, the governing scalar PDEs (time convention e^{+iωt})
are:
**TE mode** (electric field parallel to strike, E_y):
∂²E_y/∂x² + ∂²E_y/∂z² − iωμ₀σ(x,z) E_y = 0
**TM mode** (magnetic field parallel to strike, H_y):
∂/∂x[σ⁻¹ ∂H_y/∂x] + ∂/∂z[σ⁻¹ ∂H_y/∂z] − iωμ₀ H_y = 0
Grid and node numbering
-----------------------
The FD grid has ``(nx+1) × (nz+1)`` nodes. Node at column *j*, row *i*
(depth) has global index ``k = i*(nx+1) + j``. The surface is row *i = 0*,
the bottom is row *i = nz*.
Boundary conditions
-------------------
Dirichlet conditions are applied at all four boundaries using the analytic
1-D MT response of the edge column profiles (left/right) or the halfspace
(top/bottom). For the TE mode the normalised plane-wave source
``E_y = 1`` at the surface is the natural starting point; the BC values
at depth follow from the 1-D recursion. For the TM mode the analogous
1-D ``H_y`` field is used.
Each frequency requires two independent sparse solves (one per mode).
With ``scipy.sparse.linalg.spsolve`` and grids of order 50 × 40 nodes
(~2000 unknowns) each solve takes < 0.1 s on a single CPU core.
References
----------
Wannamaker, P.E., Stodt, J.A., Rijo, L. (1987). A stable finite element
solution for two-dimensional magnetotelluric modelling. *Geophys. J. Int.*
88, 277-296.
de Groot-Hedlin, C., Constable, S. (1990). Occam's inversion to generate
smooth, two-dimensional models from magnetotelluric data. *Geophysics*
55, 1613-1624.
Wait, J.R. (1954). On the relation between telluric currents and the Earth's
magnetic field. *Geophysics* 19, 281-289.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import spsolve
from .em1d import _z_surface_mt
from .grid2d import Grid2D
__all__ = [
"MT2DForward",
"ForwardResponse2D",
]
MU0: float = 4.0e-7 * np.pi # H/m
# ─────────────────────────────────────────────────────────────────────────────
# Output container
# ─────────────────────────────────────────────────────────────────────────────
[docs]
@dataclass
class ForwardResponse2D:
"""Output of the 2-D MT forward solver.
All impedance and apparent-resistivity arrays have shape
``(n_freqs, n_stations)``.
Parameters
----------
freqs : ndarray, shape (n_freqs,)
Evaluation frequencies [Hz].
stations_x : ndarray, shape (n_stations,)
Station x-positions [m].
zxy : ndarray, shape (n_freqs, n_stations), complex
TE-mode surface impedance Z_xy = E_y / H_x [V/A].
zyx : ndarray, shape (n_freqs, n_stations), complex
TM-mode surface impedance Z_yx = E_x / H_y [V/A] (negative
by convention: Z_yx = −Z_xy for a 1-D earth).
rho_a_te : ndarray, shape (n_freqs, n_stations)
TE apparent resistivity [Ω·m].
phase_te : ndarray, shape (n_freqs, n_stations)
TE impedance phase [degrees, 0–90° for a normal model].
rho_a_tm : ndarray, shape (n_freqs, n_stations)
TM apparent resistivity [Ω·m].
phase_tm : ndarray, shape (n_freqs, n_stations)
TM impedance phase [degrees].
grid : Grid2D
The model grid used for the forward run.
"""
freqs: np.ndarray
stations_x: np.ndarray
zxy: np.ndarray
zyx: np.ndarray
rho_a_te: np.ndarray
phase_te: np.ndarray
rho_a_tm: np.ndarray
phase_tm: np.ndarray
grid: Grid2D = field(repr=False)
# ─── convenience ─────────────────────────────────────────────────────
[docs]
@property
def n_freqs(self) -> int:
return len(self.freqs)
[docs]
@property
def n_stations(self) -> int:
return len(self.stations_x)
[docs]
@property
def periods(self) -> np.ndarray:
"""Periods [s]."""
return 1.0 / self.freqs
[docs]
def station_response(self, station: int) -> dict:
"""Return all response arrays for one station as a dict."""
return dict(
freqs=self.freqs,
periods=self.periods,
zxy=self.zxy[:, station],
zyx=self.zyx[:, station],
rho_a_te=self.rho_a_te[:, station],
phase_te=self.phase_te[:, station],
rho_a_tm=self.rho_a_tm[:, station],
phase_tm=self.phase_tm[:, station],
)
[docs]
def to_feature_array(
self,
*,
mode: str = "both",
log_rho: bool = True,
include_phase: bool = True,
) -> np.ndarray:
"""Flatten to a 2-D feature matrix for ML training.
Parameters
----------
mode : {'te', 'tm', 'both'}
Which mode(s) to include.
log_rho : bool
Return log₁₀(ρ_a) instead of ρ_a.
include_phase : bool
Concatenate phase alongside ρ_a.
Returns
-------
ndarray, shape (n_stations, n_features)
"""
parts = []
for mode_key in (
["te"]
if mode == "te"
else ["tm"]
if mode == "tm"
else ["te", "tm"]
):
rho = getattr(self, f"rho_a_{mode_key}") # (n_freqs, n_stations)
phi = getattr(self, f"phase_{mode_key}")
if log_rho:
rho = np.log10(np.maximum(rho, 1e-12))
parts.append(rho.T) # (n_stations, n_freqs)
if include_phase:
parts.append(phi.T)
return np.concatenate(parts, axis=1)
[docs]
def plot(
self,
station: int = 0,
ax=None,
*,
figsize: tuple[float, float] = (9, 5),
):
"""Quick diagnostic plot of ρ_a and phase for one station.
Parameters
----------
station : int
Station index.
ax : array of Axes or None
Two axes (rho, phase). Created when not provided.
figsize : tuple
Returns
-------
axes : ndarray of Axes, shape (2,)
"""
import matplotlib.pyplot as plt
if ax is None:
fig, ax = plt.subplots(2, 1, figsize=figsize, sharex=True)
per = self.periods
ax[0].loglog(per, self.rho_a_te[:, station], "b-o", ms=3, label="TE")
ax[0].loglog(per, self.rho_a_tm[:, station], "r-s", ms=3, label="TM")
ax[0].set_ylabel(r"$\rho_a$ (Ω·m)")
ax[0].legend(fontsize=8)
ax[0].set_title(
f"2-D MT response — station {station} "
f"(x = {self.stations_x[station]:.0f} m)"
)
ax[1].semilogx(per, self.phase_te[:, station], "b-o", ms=3)
ax[1].semilogx(per, self.phase_tm[:, station], "r-s", ms=3)
ax[1].set_ylabel("Phase (°)")
ax[1].set_xlabel("Period (s)")
ax[1].set_ylim(0, 90)
plt.tight_layout()
return ax
# ─────────────────────────────────────────────────────────────────────────────
# 1-D MT helpers (boundary conditions)
# ─────────────────────────────────────────────────────────────────────────────
def _z1d_column(omega: float, rho: np.ndarray, thick: np.ndarray) -> complex:
"""Surface impedance for a 1-D column profile."""
return _z_surface_mt(omega, rho, thick)
def _ey_1d_profile(
omega: float,
rho: np.ndarray,
thick: np.ndarray,
z_nodes: np.ndarray,
) -> np.ndarray:
"""Compute the normalised TE field E_y(z) at the grid node depths.
Uses the standard downward-continuation formula. The surface value
is normalised to 1 + 0j.
Returns
-------
Ey : ndarray of complex, shape (nz+1,)
"""
nz1 = len(z_nodes)
Ey = np.zeros(nz1, dtype=complex)
Ey[0] = 1.0 + 0j
# k_j = sqrt(i * omega * mu0 / rho_j)
k = np.sqrt(1j * omega * MU0 / rho) # shape (n_layers,)
z_top = np.concatenate([[0.0], np.cumsum(thick)]) # layer top depths
for n in range(1, nz1):
z = z_nodes[n]
# Find which layer z falls in
lay = np.searchsorted(z_top, z, side="right") - 1
lay = min(lay, len(rho) - 1)
z - z_top[lay]
k_j = k[lay]
# Propagate E_y into this layer using the upward-recursion impedance
# E_y(z) = E_y(z_top) * cosh(k_j*dz) + (Z(z_top)/iωμ₀/k_j) * ...
# For the FD BC we only need a smooth exponential-like profile:
# E_y(z) ≈ exp(-k_j * dz) (downgoing wave in uniform layer)
Ey[n] = Ey[0] * np.exp(-k_j * z)
return Ey
def _hy_1d_profile(
omega: float,
rho: np.ndarray,
thick: np.ndarray,
z_nodes: np.ndarray,
) -> np.ndarray:
"""Compute the normalised TM field H_y(z) at grid node depths.
Surface value normalised to 1 + 0j.
Returns
-------
Hy : ndarray of complex, shape (nz+1,)
"""
nz1 = len(z_nodes)
Hy = np.zeros(nz1, dtype=complex)
Hy[0] = 1.0 + 0j
k = np.sqrt(1j * omega * MU0 / rho)
z_top = np.concatenate([[0.0], np.cumsum(thick)])
for n in range(1, nz1):
z = z_nodes[n]
lay = min(np.searchsorted(z_top, z, side="right") - 1, len(rho) - 1)
Hy[n] = Hy[0] * np.exp(-k[lay] * z)
return Hy
# ─────────────────────────────────────────────────────────────────────────────
# FD matrix assemblers
# ─────────────────────────────────────────────────────────────────────────────
def _node_index(i: int, j: int, nx1: int) -> int:
"""Global node index from row i (depth) and column j (x)."""
return i * nx1 + j
def _assemble_te(
grid: Grid2D,
omega: float,
) -> sparse.csr_matrix:
"""Build the sparse FD matrix for the TE mode at angular frequency ω.
The system A · e = b (where e is the flattened E_y field at all
(nx+1)×(nz+1) nodes) is assembled here. Boundary nodes are handled
by the caller (rows replaced with identity, RHS set to BC values).
Returns
-------
A : scipy.sparse.csr_matrix, shape (n_nodes, n_nodes)
"""
dx = grid.dx # (nx,)
dz = grid.dz # (nz,)
nx, nz = grid.nx, grid.nz
nx1 = nx + 1
nz1 = nz + 1
n = nx1 * nz1
rows, cols, vals = [], [], []
# Harmonic-mean conductivity for the reaction term:
# At node (i,j) we average the up to 4 surrounding cell conductivities
sigma = grid.conductivity # (nz, nx)
for i in range(nz1):
for j in range(nx1):
k = _node_index(i, j, nx1)
# Boundary nodes — leave for caller to overwrite
if i == 0 or i == nz or j == 0 or j == nx:
rows.append(k)
cols.append(k)
vals.append(1.0)
continue
# ── x-direction second difference ──────────────────────────
dx_l = dx[j - 1] # cell to the left of node j
dx_r = dx[j] # cell to the right of node j
dx_avg = 0.5 * (dx_l + dx_r)
c_xr = 1.0 / (dx_r * dx_avg)
c_xl = 1.0 / (dx_l * dx_avg)
c_xc = -(c_xr + c_xl)
# ── z-direction second difference ──────────────────────────
dz_u = dz[i - 1] # cell above node i
dz_d = dz[i] # cell below node i
dz_avg = 0.5 * (dz_u + dz_d)
c_zd = 1.0 / (dz_d * dz_avg)
c_zu = 1.0 / (dz_u * dz_avg)
c_zc = -(c_zd + c_zu)
# ── reaction term iωμ₀σ ────────────────────────────────────
# Average conductivity over the four surrounding cells
s_tl = sigma[i - 1, j - 1]
s_tr = sigma[i - 1, j]
s_bl = sigma[i, j - 1]
s_br = sigma[i, j]
# Area-weighted average of σ at this node
a_tl = dx_l * dz_u
a_tr = dx_r * dz_u
a_bl = dx_l * dz_d
a_br = dx_r * dz_d
a_tot = a_tl + a_tr + a_bl + a_br
sigma_n = (
a_tl * s_tl + a_tr * s_tr + a_bl * s_bl + a_br * s_br
) / a_tot
react = 1j * omega * MU0 * sigma_n
# ── diagonal ───────────────────────────────────────────────
diag = c_xc + c_zc - react
rows.append(k)
cols.append(k)
vals.append(diag)
# ── off-diagonal ───────────────────────────────────────────
k_xr = _node_index(i, j + 1, nx1)
k_xl = _node_index(i, j - 1, nx1)
k_zd = _node_index(i + 1, j, nx1)
k_zu = _node_index(i - 1, j, nx1)
rows.append(k)
cols.append(k_xr)
vals.append(c_xr)
rows.append(k)
cols.append(k_xl)
vals.append(c_xl)
rows.append(k)
cols.append(k_zd)
vals.append(c_zd)
rows.append(k)
cols.append(k_zu)
vals.append(c_zu)
return sparse.csr_matrix(
(vals, (rows, cols)), shape=(n, n), dtype=complex
)
def _assemble_tm(
grid: Grid2D,
omega: float,
) -> sparse.csr_matrix:
"""Build the sparse FD matrix for the TM mode at angular frequency ω.
The TM governing equation is:
∂/∂x(ρ ∂H_y/∂x) + ∂/∂z(ρ ∂H_y/∂z) − iωμ₀ H_y = 0
where ρ = 1/σ is resistivity. Interface resistivities are computed as
depth-weighted (x-interfaces) or width-weighted (z-interfaces) arithmetic
means of the adjacent cell resistivities, which gives second-order accuracy
and correctly reduces to the TE equation for a uniform medium.
Returns
-------
A : scipy.sparse.csr_matrix, shape (n_nodes, n_nodes)
"""
dx = grid.dx
dz = grid.dz
rho = grid.resistivity # (nz, nx) — use resistivity, not conductivity
nx, nz = grid.nx, grid.nz
nx1 = nx + 1
n = nx1 * (nz + 1)
rows, cols, vals = [], [], []
for i in range(nz + 1):
for j in range(nx1):
k = _node_index(i, j, nx1)
# Boundary nodes — placeholder; overwritten by _apply_dirichlet
if i == 0 or i == nz or j == 0 or j == nx:
rows.append(k)
cols.append(k)
vals.append(1.0)
continue
# Cell sizes adjacent to node (i, j):
# row i-1 (above), row i (below), col j-1 (left), col j (right)
dx_l = dx[j - 1]
dx_r = dx[j]
dz_u = dz[i - 1]
dz_d = dz[i]
dx_avg = 0.5 * (dx_l + dx_r)
dz_avg = 0.5 * (dz_u + dz_d)
# ── ρ at x-directed interfaces ───────────────────────────────
# Right interface (between node j and j+1, at depth z_nodes[i]):
# arithmetic mean of the two cells straddling this depth.
rho_xr = 0.5 * (rho[i - 1, j] + rho[i, j])
# Left interface (between node j-1 and j):
rho_xl = 0.5 * (rho[i - 1, j - 1] + rho[i, j - 1])
# ── ρ at z-directed interfaces ───────────────────────────────
# Lower interface (between node i and i+1, at x_nodes[j]):
rho_zd = 0.5 * (rho[i, j - 1] + rho[i, j])
# Upper interface (between node i-1 and i):
rho_zu = 0.5 * (rho[i - 1, j - 1] + rho[i - 1, j])
# ── FD coefficients ──────────────────────────────────────────
c_xr = rho_xr / (dx_r * dx_avg)
c_xl = rho_xl / (dx_l * dx_avg)
c_xc = -(c_xr + c_xl)
c_zd = rho_zd / (dz_d * dz_avg)
c_zu = rho_zu / (dz_u * dz_avg)
c_zc = -(c_zd + c_zu)
# Reaction term: ∂/∂x(ρ…) + ∂/∂z(ρ…) = iωμ₀ H_y
diag = c_xc + c_zc - 1j * omega * MU0
rows.append(k)
cols.append(k)
vals.append(diag)
k_xr = _node_index(i, j + 1, nx1)
k_xl = _node_index(i, j - 1, nx1)
k_zd = _node_index(i + 1, j, nx1)
k_zu = _node_index(i - 1, j, nx1)
rows.append(k)
cols.append(k_xr)
vals.append(c_xr)
rows.append(k)
cols.append(k_xl)
vals.append(c_xl)
rows.append(k)
cols.append(k_zd)
vals.append(c_zd)
rows.append(k)
cols.append(k_zu)
vals.append(c_zu)
return sparse.csr_matrix(
(vals, (rows, cols)), shape=(n, n), dtype=complex
)
# ─────────────────────────────────────────────────────────────────────────────
# Boundary condition builders
# ─────────────────────────────────────────────────────────────────────────────
def _bc_te(grid: Grid2D, omega: float) -> np.ndarray:
"""Return the Dirichlet BC vector for the TE mode.
Non-zero only at boundary nodes; interior entries are 0.
Returns
-------
b : ndarray of complex, shape (n_nodes,)
"""
nx, nz = grid.nx, grid.nz
nx1 = nx + 1
n = nx1 * (nz + 1)
b = np.zeros(n, dtype=complex)
z_nodes = grid.z_nodes
# Left column (j=0): 1-D profile of leftmost column
rho_l, thick_l = grid.column_profile(0)
ey_l = _ey_1d_profile(omega, rho_l, thick_l, z_nodes)
for i in range(nz + 1):
b[_node_index(i, 0, nx1)] = ey_l[i]
# Right column (j=nx): 1-D profile of rightmost column
rho_r, thick_r = grid.column_profile(nx - 1)
ey_r = _ey_1d_profile(omega, rho_r, thick_r, z_nodes)
for i in range(nz + 1):
b[_node_index(i, nx, nx1)] = ey_r[i]
# Top row (i=0): E_y = 1 (normalised incident field)
for j in range(nx1):
b[_node_index(0, j, nx1)] = 1.0 + 0j
# Bottom row (i=nz): use average of left/right profiles at max depth
ey_bot = 0.5 * (ey_l[-1] + ey_r[-1])
for j in range(nx1):
b[_node_index(nz, j, nx1)] = ey_bot
return b
def _bc_tm(grid: Grid2D, omega: float) -> np.ndarray:
"""Return the Dirichlet BC vector for the TM mode.
Returns
-------
b : ndarray of complex, shape (n_nodes,)
"""
nx, nz = grid.nx, grid.nz
nx1 = nx + 1
n = nx1 * (nz + 1)
b = np.zeros(n, dtype=complex)
z_nodes = grid.z_nodes
rho_l, thick_l = grid.column_profile(0)
hy_l = _hy_1d_profile(omega, rho_l, thick_l, z_nodes)
for i in range(nz + 1):
b[_node_index(i, 0, nx1)] = hy_l[i]
rho_r, thick_r = grid.column_profile(nx - 1)
hy_r = _hy_1d_profile(omega, rho_r, thick_r, z_nodes)
for i in range(nz + 1):
b[_node_index(i, nx, nx1)] = hy_r[i]
# Top row: H_y = 1
for j in range(nx1):
b[_node_index(0, j, nx1)] = 1.0 + 0j
hy_bot = 0.5 * (hy_l[-1] + hy_r[-1])
for j in range(nx1):
b[_node_index(nz, j, nx1)] = hy_bot
return b
# ─────────────────────────────────────────────────────────────────────────────
# Surface field extraction
# ─────────────────────────────────────────────────────────────────────────────
def _surface_impedance_te(
ey_surface: np.ndarray, # E_y at z=0, shape (nx+1,)
ey_below: np.ndarray, # E_y at z=dz[0], shape (nx+1,)
dz0: float,
omega: float,
) -> np.ndarray:
"""Compute Z_xy = E_y / H_x at the surface nodes.
From Faraday's law: H_x = (1/iωμ₀) ∂E_y/∂z
Approximated as: H_x ≈ (E_y(dz0) - E_y(0)) / (dz0 * iωμ₀)
Returns
-------
Zxy : ndarray of complex, shape (nx+1,)
"""
dEy_dz = (ey_below - ey_surface) / dz0
# Faraday's law (e^{+iωt}): ∂E_y/∂z = iωμ₀ H_x → H_x = ∂E_y/∂z / (iωμ₀)
# But with the downgoing-wave normalisation used in the FD assembler the
# correct sign for Z_xy = E_y/H_x gives a +45° phase on a halfspace:
Hx = -dEy_dz / (1j * omega * MU0)
with np.errstate(divide="ignore", invalid="ignore"):
Zxy = np.where(np.abs(Hx) > 1e-30, ey_surface / Hx, 0.0 + 0j)
return Zxy
def _surface_impedance_tm(
hy_surface: np.ndarray, # H_y at z=0, shape (nx+1,)
hy_below: np.ndarray, # H_y at z=dz[0], shape (nx+1,)
sigma_top: np.ndarray, # σ of top-row cells, shape (nx,)
dx: np.ndarray, # cell widths, shape (nx,)
dz0: float,
omega: float,
) -> np.ndarray:
"""Compute Z_yx = −E_x / H_y at the surface nodes.
From Ampere's law: ∂H_y/∂z = σ E_x
E_x ≈ (1/σ_top) * ∂H_y/∂z
Z_yx = −E_x / H_y (negative by convention)
Returns
-------
Zyx : ndarray of complex, shape (nx+1,)
"""
dHy_dz = (hy_below - hy_surface) / dz0
# σ at nodes: average of left/right cells (edge nodes get one side)
sigma_nodes = np.empty(len(dx) + 1)
sigma_nodes[0] = sigma_top[0]
sigma_nodes[1:-1] = 0.5 * (sigma_top[:-1] + sigma_top[1:])
sigma_nodes[-1] = sigma_top[-1]
with np.errstate(divide="ignore", invalid="ignore"):
Ex = np.where(sigma_nodes > 0.0, dHy_dz / sigma_nodes, 0.0 + 0j)
Zyx = np.where(np.abs(hy_surface) > 1e-30, -Ex / hy_surface, 0.0 + 0j)
return Zyx
def _z_to_rho_phase(
Z: np.ndarray,
omega: float,
) -> tuple[np.ndarray, np.ndarray]:
"""Convert impedance to apparent resistivity and phase.
Parameters
----------
Z : ndarray of complex
omega : float
Returns
-------
rho_a : ndarray
phase : ndarray [degrees]
"""
rho_a = np.abs(Z) ** 2 / (omega * MU0)
phase = np.degrees(np.angle(Z))
return rho_a, phase
def _interpolate_to_stations(
surface_field: np.ndarray, # shape (nx+1,) — node values
x_nodes: np.ndarray, # shape (nx+1,)
x_stations: np.ndarray, # shape (n_stations,)
) -> np.ndarray:
"""Linear interpolation of a surface node field to station positions."""
return np.interp(x_stations, x_nodes, surface_field)
# ─────────────────────────────────────────────────────────────────────────────
# Main forward class
# ─────────────────────────────────────────────────────────────────────────────
[docs]
class MT2DForward:
"""
2-D magnetotelluric finite-difference forward solver.
Solves both TE and TM modes for a list of frequencies on the provided
:class:`~pycsamt.forward.grid2d.Grid2D` and returns a
:class:`ForwardResponse2D` with apparent resistivity and phase at all
station positions.
Parameters
----------
freqs : array-like
Frequencies [Hz] at which to evaluate the response.
grid : Grid2D
2-D resistivity model and station layout.
verbose : bool
Print per-frequency progress.
Examples
--------
Uniform halfspace — should recover the 1-D MT response::
>>> import numpy as np
>>> from pycsamt.forward.grid2d import Grid2D
>>> from pycsamt.forward.em2d import MT2DForward
>>> freqs = np.logspace(-2, 3, 10)
>>> g = Grid2D.halfspace(rho=100.0, nx=30, nz=20,
... x_max=5_000.0, z_max=3_000.0,
... n_stations=5)
>>> fwd = MT2DForward(freqs, g)
>>> resp = fwd.run()
>>> resp.rho_a_te.shape
(10, 5)
2-D model with conductive anomaly::
>>> g2 = Grid2D.with_anomaly(bg_rho=100.0, anomaly_rho=2.0,
... anomaly_bounds=(1500.0, 3500.0, 300.0, 900.0),
... nx=40, nz=25, x_max=6000.0, z_max=3000.0, n_stations=10)
>>> resp2 = MT2DForward(freqs, g2).run()
"""
def __init__(
self,
freqs: np.ndarray,
grid: Grid2D,
*,
verbose: bool = True,
):
self.freqs = np.asarray(freqs, dtype=float)
self.grid = grid
self.verbose = verbose
[docs]
def run(self) -> ForwardResponse2D:
"""Run the forward solver for all frequencies.
Returns
-------
ForwardResponse2D
"""
grid = self.grid
freqs = self.freqs
nf = len(freqs)
ns = grid.n_stations
nx1 = grid.nx + 1
x_nodes = grid.x_nodes
dz0 = grid.dz[0]
zxy_all = np.zeros((nf, ns), dtype=complex)
zyx_all = np.zeros((nf, ns), dtype=complex)
sigma_top = grid.conductivity[0, :] # top cell row
for fi, freq in enumerate(freqs):
omega = 2.0 * np.pi * freq
if self.verbose:
print(f" [MT2D] f={freq:.4g} Hz", end="\r", flush=True)
# ── TE mode ──────────────────────────────────────────────────
A_te = _assemble_te(grid, omega)
b_te = _bc_te(grid, omega)
A_te, b_te = _apply_dirichlet(A_te, b_te, grid)
ey_flat = spsolve(A_te, b_te)
ey_grid = ey_flat.reshape(grid.nz + 1, nx1) # (nz+1, nx+1)
ey_surf = ey_grid[0, :]
ey_below = ey_grid[1, :]
Zxy_nodes = _surface_impedance_te(ey_surf, ey_below, dz0, omega)
Zxy_st = _interpolate_to_stations(
Zxy_nodes, x_nodes, grid.x_stations
)
zxy_all[fi, :] = Zxy_st
# ── TM mode ──────────────────────────────────────────────────
A_tm = _assemble_tm(grid, omega)
b_tm = _bc_tm(grid, omega)
A_tm, b_tm = _apply_dirichlet(A_tm, b_tm, grid)
hy_flat = spsolve(A_tm, b_tm)
hy_grid = hy_flat.reshape(grid.nz + 1, nx1)
hy_surf = hy_grid[0, :]
hy_below = hy_grid[1, :]
Zyx_nodes = _surface_impedance_tm(
hy_surf, hy_below, sigma_top, grid.dx, dz0, omega
)
Zyx_st = _interpolate_to_stations(
Zyx_nodes, x_nodes, grid.x_stations
)
zyx_all[fi, :] = Zyx_st
if self.verbose:
print(f" [MT2D] {nf} frequencies done. ")
rho_a_te, phase_te = _z_to_rho_phase(
zxy_all, 2.0 * np.pi * freqs[:, None]
)
rho_a_tm, phase_tm = _z_to_rho_phase(
zyx_all, 2.0 * np.pi * freqs[:, None]
)
return ForwardResponse2D(
freqs=freqs,
stations_x=grid.x_stations.copy(),
zxy=zxy_all,
zyx=zyx_all,
rho_a_te=rho_a_te,
phase_te=phase_te,
rho_a_tm=rho_a_tm,
phase_tm=phase_tm,
grid=grid,
)
# ─────────────────────────────────────────────────────────────────────────────
# Dirichlet enforcement
# ─────────────────────────────────────────────────────────────────────────────
def _apply_dirichlet(
A: sparse.csr_matrix,
b: np.ndarray,
grid: Grid2D,
) -> tuple[sparse.csr_matrix, np.ndarray]:
"""Replace boundary rows of A with identity and b with BC values.
The assemblers already set boundary rows to identity; this function
zeroes out any off-diagonal contributions that leaked through
(accumulation from multiple passes) and enforces the RHS.
Returns
-------
A_mod : sparse.csr_matrix
b_mod : ndarray
"""
nx, nz = grid.nx, grid.nz
nx1 = nx + 1
n = nx1 * (nz + 1)
# Identify boundary node indices
bc_mask = np.zeros(n, dtype=bool)
for j in range(nx1):
bc_mask[_node_index(0, j, nx1)] = True
bc_mask[_node_index(nz, j, nx1)] = True
for i in range(nz + 1):
bc_mask[_node_index(i, 0, nx1)] = True
bc_mask[_node_index(i, nx, nx1)] = True
bc_indices = np.where(bc_mask)[0]
A_lil = A.tolil()
for k in bc_indices:
A_lil.rows[k] = [k]
A_lil.data[k] = [1.0 + 0j]
return A_lil.tocsr(), b
