witwin-maxwell 0.3.0__py3-none-win_amd64.whl

witwin/maxwell/fdfd/ams.py

@@ -0,0 +1,300 @@
+"""Auxiliary-space Maxwell (Hiptmair-Xu) preconditioner for the FDFD system.
+
+Structure (all GPU, CuPy sparse):
+
+- Gradient auxiliary space: the PML-stretched discrete gradient
+  ``P_g = D S_e^-1 G`` maps nodal potentials onto the null space of the
+  scaled curl-curl part — exact to round-off in the interior; the outermost
+  cell layer carries an inherent truncation defect. The Galerkin operator
+  ``A_g = P_g^T A P_g ~= k0^2 P_g^T E P_g`` is a definite nodal Poisson-type
+  operator.
+- Vector auxiliary space: component-wise nodal-to-edge averaging
+  ``P_c = D S_e^-1 Pi_c``; the Galerkin blocks are built from the
+  complex-shifted operator ``A + i beta k0^2 E`` (shifted-Laplacian damping
+  of the indefinite wave part) so multigrid smoothing is stable.
+- Auxiliary solves: geometric multigrid V(1,1) cycles on the structured nodal
+  grid (Galerkin RAP coarsening, damped-Jacobi smoothing, dense solve on the
+  coarsest level).
+- Composition: symmetrized multiplicative Schwarz
+  ``smoother -> gradient -> vector -> gradient -> smoother``. Every stage is
+  complex-symmetric, so the preconditioner is valid for SQMR.
+
+This is the in-repo answer to the plan's 10.4 AMS item: hypre-AMS is
+unavailable on native Windows, and the structured Yee grid lets geometric
+multigrid replace AMG for the auxiliary problems.
+
+STATUS: EXPERIMENTAL — measured outcome on the indefinite time-harmonic
+operator (see plan 10.8): the gradient-space construction is exact (the
+null-space identity holds to round-off) and its multigrid contracts, but
+the overall Hiptmair-Xu composition is non-contractive at these
+wavenumbers even with EXACT auxiliary solves — oblique subspace
+corrections amplify each other on the non-Hermitian indefinite operator.
+This matches hypre-AMS's own scope (definite curl-curl / eddy-current
+regimes). For the indefinite regime use sqmr+ssor+precision='double'
+(converges through 64^3) or the direct backend; large-scale indefinite
+solves need sweeping/DDM-class methods, out of scope here.
+"""
+
+from __future__ import annotations
+
+import cupy as cp
+import cupyx.scipy.sparse as cpsp
+import cupyx.scipy.sparse.linalg as cupy_linalg
+
+# Complex shift applied to the vector auxiliary operator (shifted-Laplacian
+# damping); the gradient space is definite and needs no shift.
+AMS_SHIFT_BETA = 0.5
+# Damped-Jacobi weight for the fine smoother and the multigrid smoother.
+AMS_JACOBI_OMEGA = 0.8
+# Stop coarsening once every axis is at or below this many nodes.
+AMS_COARSEST_NODES = 8
+
+
+def _prolongation_1d(n, dtype):
+    """Vertex-centered 1D linear prolongation (fine n -> coarse (n+1)//2)."""
+    nc = (n + 1) // 2
+    rows, cols, vals = [], [], []
+    for i in range(n):
+        if i % 2 == 0:
+            rows.append(i); cols.append(i // 2); vals.append(1.0)
+        else:
+            left = i // 2
+            right = left + 1
+            if right < nc:
+                rows.extend((i, i)); cols.extend((left, right)); vals.extend((0.5, 0.5))
+            else:
+                rows.append(i); cols.append(left); vals.append(1.0)
+    return cpsp.csr_matrix(
+        (cp.asarray(vals, dtype=dtype), (cp.asarray(rows), cp.asarray(cols))),
+        shape=(n, nc),
+    )
+
+
+def _prolongation_3d(shape, dtype):
+    px = _prolongation_1d(shape[0], dtype)
+    py = _prolongation_1d(shape[1], dtype)
+    pz = _prolongation_1d(shape[2], dtype)
+    return cpsp.kron(cpsp.kron(px, py), pz).tocsr()
+
+
+class _NodalMultigrid:
+    """Geometric V(1,1) multigrid on a structured nodal grid with Galerkin
+    coarsening and damped-Jacobi smoothing. Complex-symmetric throughout
+    (transposes, never conjugates).
+
+    The operator is symmetrically diagonal-scaled first (S A S with
+    S = diag(A)^-1/2), which absorbs the orders-of-magnitude PML weight
+    range into the transfers; the coarsest level is solved densely with a
+    small Tikhonov term (the gradient-space operator has an exact constant
+    null space)."""
+
+    def __init__(self, A, shape, omega=AMS_JACOBI_OMEGA):
+        self.omega = omega
+        A = A.tocsr()
+        self.scale = cp.reciprocal(cp.sqrt(A.diagonal()))
+        scale_matrix = cpsp.diags(self.scale).tocsr()
+        current_A = (scale_matrix @ (A @ scale_matrix)).tocsr()
+        current_shape = tuple(shape)
+        self.levels = []
+        while min(current_shape) > AMS_COARSEST_NODES:
+            P = _prolongation_3d(current_shape, current_A.dtype)
+            R = P.T.tocsr()
+            inv_diag = cp.reciprocal(current_A.diagonal())
+            self.levels.append((current_A, inv_diag, P, R))
+            current_A = (R @ (current_A @ P)).tocsr()
+            current_shape = tuple((s + 1) // 2 for s in current_shape)
+        coarse = current_A.toarray()
+        regularization = 1e-10 * float(cp.abs(coarse.diagonal()).max())
+        coarse = coarse + regularization * cp.eye(coarse.shape[0], dtype=coarse.dtype)
+        self.coarse_dense = coarse
+
+    def solve(self, r):
+        return self.scale * self._cycle(0, self.scale * r)
+
+    def _cycle(self, level, r):
+        if level == len(self.levels):
+            return cp.linalg.solve(self.coarse_dense, r)
+        A, inv_diag, P, R = self.levels[level]
+        z = self.omega * inv_diag * r
+        z = z + P @ self._cycle(level + 1, R @ (r - A @ z))
+        z = z + self.omega * inv_diag * (r - A @ z)
+        return z
+
+
+def _edge_shapes(scene):
+    return (
+        (scene.Nx_ex, scene.Ny_ex, scene.Nz_ex),
+        (scene.Nx_ey, scene.Ny_ey, scene.Nz_ey),
+        (scene.Nx_ez, scene.Ny_ez, scene.Nz_ez),
+    )
+
+
+def _stretch_half_factors(solver, dtype):
+    s_x, s_y, s_z = solver._create_pml_3d()
+    s_x = cp.asarray(s_x, dtype=dtype)
+    s_y = cp.asarray(s_y, dtype=dtype)
+    s_z = cp.asarray(s_z, dtype=dtype)
+    return (
+        (s_x[:-1] + s_x[1:]) / 2,
+        (s_y[:-1] + s_y[1:]) / 2,
+        (s_z[:-1] + s_z[1:]) / 2,
+    )
+
+
+def _edge_stretch_weights(solver, dtype):
+    """1/s along each edge's own axis, per component block (flat)."""
+    scene = solver.scene
+    shapes = _edge_shapes(scene)
+    sxh, syh, szh = _stretch_half_factors(solver, dtype)
+    wx = cp.broadcast_to((1.0 / sxh)[:, None, None], shapes[0]).ravel()
+    wy = cp.broadcast_to((1.0 / syh)[None, :, None], shapes[1]).ravel()
+    wz = cp.broadcast_to((1.0 / szh)[None, None, :], shapes[2]).ravel()
+    return wx, wy, wz
+
+
+def _node_index(scene):
+    return cp.arange(scene.Nx * scene.Ny * scene.Nz).reshape(scene.Nx, scene.Ny, scene.Nz)
+
+
+def _difference_block(node, axis, dx, dtype):
+    """Plain discrete gradient for one component: edges x nodes, +/- 1/dx."""
+    if axis == 0:
+        lo, hi = node[:-1], node[1:]
+    elif axis == 1:
+        lo, hi = node[:, :-1], node[:, 1:]
+    else:
+        lo, hi = node[:, :, :-1], node[:, :, 1:]
+    n_edges = lo.size
+    rows = cp.arange(n_edges)
+    rows2 = cp.concatenate([rows, rows])
+    cols = cp.concatenate([hi.ravel(), lo.ravel()])
+    vals = cp.concatenate([
+        cp.full(n_edges, 1.0 / dx, dtype=dtype),
+        cp.full(n_edges, -1.0 / dx, dtype=dtype),
+    ])
+    return cpsp.csr_matrix((vals, (rows2, cols)), shape=(n_edges, node.size))
+
+
+def _average_block(node, axis, dtype):
+    """Nodal-to-edge tangential averaging for one component: weights 1/2."""
+    if axis == 0:
+        lo, hi = node[:-1], node[1:]
+    elif axis == 1:
+        lo, hi = node[:, :-1], node[:, 1:]
+    else:
+        lo, hi = node[:, :, :-1], node[:, :, 1:]
+    n_edges = lo.size
+    rows = cp.arange(n_edges)
+    rows2 = cp.concatenate([rows, rows])
+    cols = cp.concatenate([hi.ravel(), lo.ravel()])
+    vals = cp.full(2 * n_edges, 0.5, dtype=dtype)
+    return cpsp.csr_matrix((vals, (rows2, cols)), shape=(n_edges, node.size))
+
+
+def _eps_face_vector(solver, dtype):
+    ex = cp.asarray(solver.material_eps_components["x"], dtype=dtype)
+    ey = cp.asarray(solver.material_eps_components["y"], dtype=dtype)
+    ez = cp.asarray(solver.material_eps_components["z"], dtype=dtype)
+    return cp.concatenate([
+        ((ex[:-1] + ex[1:]) / 2).ravel(),
+        ((ey[:, :-1] + ey[:, 1:]) / 2).ravel(),
+        ((ez[:, :, :-1] + ez[:, :, 1:]) / 2).ravel(),
+    ])
+
+
+class AMSPreconditioner:
+    def __init__(self, solver, shift_beta=AMS_SHIFT_BETA):
+        A = solver._iteration_matrix()
+        dtype = A.dtype
+        scene = solver.scene
+        dx = float(scene.dx)
+        node = _node_index(scene)
+        node_shape = (scene.Nx, scene.Ny, scene.Nz)
+        n_nodes = node.size
+        edge_counts = (scene.N_ex, scene.N_ey, scene.N_ez)
+
+        scale = solver._ensure_symmetrization_scale().astype(dtype)
+        stretch = _edge_stretch_weights(solver, dtype)
+        # per-component edge weight: d * (1/s) restricted to that block
+        offsets = (0, edge_counts[0], edge_counts[0] + edge_counts[1])
+        block_weights = [
+            scale[offsets[c]:offsets[c] + edge_counts[c]] * stretch[c]
+            for c in range(3)
+        ]
+
+        # Shifted-Laplacian architecture: every internal stage (smoother,
+        # residual refreshes, auxiliary Galerkin operators) works on the
+        # DAMPED operator A_beta = A + i beta k0^2 E; the whole object then
+        # approximates A_beta^-1, which is the classical robust
+        # preconditioner for the indefinite wave operator A. A fine-level
+        # smoother on the undamped A itself amplifies wave modes.
+        eps_face = _eps_face_vector(solver, dtype)
+        shift = cpsp.diags(1j * shift_beta * solver.k0**2 * eps_face).tocsr()
+        self.A_shift = (A + shift).tocsr()
+        self.inv_diag = cp.reciprocal(self.A_shift.diagonal())
+        self.omega = AMS_JACOBI_OMEGA
+
+        # --- gradient auxiliary space (exact null-space complement) ---
+        grad_blocks = []
+        for c in range(3):
+            block = _difference_block(node, c, dx, dtype)
+            grad_blocks.append(cpsp.diags(block_weights[c]).tocsr() @ block)
+        self.P_grad = cpsp.vstack(grad_blocks).tocsr()
+        A_grad = (self.P_grad.T.tocsr() @ (self.A_shift @ self.P_grad)).tocsr()
+        self.mg_grad = _NodalMultigrid(A_grad, node_shape)
+
+        # --- vector auxiliary spaces ---
+        # The Galerkin block P_c^T A P_c only contains the TRANSVERSE
+        # Laplacian (the Yee (c,c) curl-curl block has no d^2/dc^2 term),
+        # which leaves whole line families unsmoothable. Following the
+        # Hiptmair-Xu identity (curl curl = -Lap + grad div on the
+        # divergence-free complement), use the FULL stretched nodal
+        # Laplacian as the vector surrogate, with the same complex shift
+        # on the mass term.
+        laplacian = (self.P_grad.T.tocsr() @ self.P_grad).tocsr()  # Gs^T D^2 Gs
+        eps_diag = cpsp.diags(eps_face).tocsr()
+        k02 = dtype.type((1.0 + 1j * shift_beta)) * solver.k0**2
+        self.P_vec = []
+        self.mg_vec = []
+        for c in range(3):
+            averaging = _average_block(node, c, dtype)
+            block = cpsp.diags(block_weights[c]).tocsr() @ averaging
+            zero_above = cpsp.csr_matrix((offsets[c], n_nodes), dtype=dtype)
+            below = A.shape[0] - offsets[c] - edge_counts[c]
+            zero_below = cpsp.csr_matrix((below, n_nodes), dtype=dtype)
+            P_c = cpsp.vstack([zero_above, block, zero_below]).tocsr()
+            mass_c = (P_c.T.tocsr() @ (eps_diag @ P_c)).tocsr()
+            A_c = (k02 * mass_c - laplacian).tocsr()
+            self.P_vec.append(P_c)
+            self.mg_vec.append(_NodalMultigrid(A_c, node_shape))
+
+    def _smooth(self, z, r):
+        return z + self.omega * self.inv_diag * (r - self.A_shift @ z)
+
+    def _grad_correct(self, z, r):
+        residual = r - self.A_shift @ z
+        return z + self.P_grad @ self.mg_grad.solve(self.P_grad.T @ residual)
+
+    def _vec_correct(self, z, r):
+        residual = r - self.A_shift @ z
+        for P_c, mg in zip(self.P_vec, self.mg_vec):
+            z = z + P_c @ mg.solve(P_c.T @ residual)
+        return z
+
+    def apply(self, r):
+        # Symmetrized multiplicative Schwarz: S, G, V, G, S
+        z = self.omega * self.inv_diag * r
+        z = self._grad_correct(z, r)
+        z = self._vec_correct(z, r)
+        z = self._grad_correct(z, r)
+        z = self._smooth(z, r)
+        return z
+
+    def as_linear_operator(self):
+        return cupy_linalg.LinearOperator(
+            self.A_shift.shape, matvec=self.apply, dtype=self.A_shift.dtype
+        )
+
+
+def build_ams_preconditioner(solver):
+    return AMSPreconditioner(solver).as_linear_operator()

witwin/maxwell/fdfd/krylov.py

@@ -0,0 +1,280 @@
+"""GPU-native Krylov engines for the FDFD solver.
+
+cupyx.scipy.sparse.linalg ships only gmres/cg/cgs/minres. This module adds
+the short-recurrence engines from the FDFD performance plan (§10.2–10.3):
+
+- ``bicgstab`` — general non-Hermitian, no minimization property.
+- ``tfqmr``    — transpose-free quasi-minimal residual (right-preconditioned).
+- ``idr_s``    — IDR(s) with biorthogonalization (van Gijzen & Sonneveld).
+- ``sqmr``     — simplified QMR for complex-symmetric systems (Freund);
+                 valid for the symmetrized UPML assembly where A = A^T.
+
+All engines run entirely on GPU (sparse matvec plus vector work), accept a
+CSR matrix or LinearOperator for ``A`` and an optional approximate-inverse
+preconditioner ``M`` (applied as ``M @ v``), and follow the cupyx return
+convention: ``(x, info)`` with ``info == 0`` on convergence (relative
+residual ``<= tol``) and the iteration count otherwise.
+"""
+
+from __future__ import annotations
+
+import cupy as cp
+
+
+def _matvec(operator):
+    if operator is None:
+        return lambda v: v
+    if hasattr(operator, "matvec"):
+        return operator.matvec
+    return lambda v: operator @ v
+
+
+def _bilinear(a, b):
+    """Unconjugated bilinear form a^T b used by complex-symmetric methods."""
+    return cp.sum(a * b)
+
+
+def solve(kind, A, b, M=None, tol=1e-6, maxiter=1000, s=4):
+    """Dispatch to one of the engines by name ('bicgstab'|'tfqmr'|'idr'|'sqmr')."""
+    if kind == "bicgstab":
+        return bicgstab(A, b, M=M, tol=tol, maxiter=maxiter)
+    if kind == "tfqmr":
+        return tfqmr(A, b, M=M, tol=tol, maxiter=maxiter)
+    if kind == "idr":
+        return idr_s(A, b, M=M, tol=tol, maxiter=maxiter, s=s)
+    if kind == "sqmr":
+        return sqmr(A, b, M=M, tol=tol, maxiter=maxiter)
+    raise ValueError(f"Unknown Krylov engine {kind!r}.")
+
+
+def bicgstab(A, b, M=None, tol=1e-6, maxiter=1000):
+    matvec = _matvec(A)
+    psolve = _matvec(M)
+    x = cp.zeros_like(b)
+    r = b.copy()
+    r_shadow = r.copy()
+    b_norm = float(cp.linalg.norm(b))
+    if b_norm == 0.0:
+        return x, 0
+    target = tol * b_norm
+
+    rho = alpha = omega = 1.0 + 0.0j
+    v = cp.zeros_like(b)
+    p = cp.zeros_like(b)
+    for iteration in range(1, maxiter + 1):
+        rho_next = complex(cp.vdot(r_shadow, r))
+        if rho_next == 0.0:
+            return x, iteration
+        if iteration == 1:
+            p = r.copy()
+        else:
+            beta = (rho_next / rho) * (alpha / omega)
+            p = r + beta * (p - omega * v)
+        p_hat = psolve(p)
+        v = matvec(p_hat)
+        denominator = complex(cp.vdot(r_shadow, v))
+        if denominator == 0.0:
+            return x, iteration
+        alpha = rho_next / denominator
+        residual_half = r - alpha * v
+        if float(cp.linalg.norm(residual_half)) <= target:
+            x += alpha * p_hat
+            return x, 0
+        s_hat = psolve(residual_half)
+        t = matvec(s_hat)
+        tt = complex(cp.vdot(t, t))
+        if tt == 0.0:
+            return x, iteration
+        omega = complex(cp.vdot(t, residual_half)) / tt
+        x += alpha * p_hat + omega * s_hat
+        r = residual_half - omega * t
+        if float(cp.linalg.norm(r)) <= target:
+            return x, 0
+        if omega == 0.0:
+            return x, iteration
+        rho = rho_next
+    return x, maxiter
+
+
+def sqmr(A, b, M=None, tol=1e-6, maxiter=1000, callback=None):
+    """Freund's simplified QMR for complex-symmetric A (A = A^T, bilinear
+    inner products, no transpose matvec). M should also be symmetric
+    (Jacobi and SSOR of a symmetric matrix are). ``callback``, if given, is
+    invoked as ``callback(iteration, residual_norm)`` once per iteration."""
+    matvec = _matvec(A)
+    psolve = _matvec(M)
+    x = cp.zeros_like(b)
+    r = b.copy()
+    b_norm = float(cp.linalg.norm(b))
+    if b_norm == 0.0:
+        return x, 0
+    target = tol * b_norm
+
+    z = psolve(r)
+    tau = float(cp.linalg.norm(r))
+    rho = complex(_bilinear(r, z))
+    q = z.copy()
+    theta = 0.0
+    d = cp.zeros_like(b)
+    for iteration in range(1, maxiter + 1):
+        Aq = matvec(q)
+        sigma = complex(_bilinear(q, Aq))
+        if sigma == 0.0 or rho == 0.0:
+            return x, iteration
+        alpha = rho / sigma
+        r = r - alpha * Aq
+        r_norm = float(cp.linalg.norm(r))
+        theta_next = r_norm / tau
+        c = 1.0 / (1.0 + theta_next * theta_next) ** 0.5
+        tau = tau * theta_next * c
+        d = (c * c * theta * theta) * d + (c * c * alpha) * q
+        x = x + d
+        if callback is not None:
+            callback(iteration, r_norm)
+        if r_norm <= target:
+            return x, 0
+        z = psolve(r)
+        rho_next = complex(_bilinear(r, z))
+        beta = rho_next / rho
+        q = z + beta * q
+        rho = rho_next
+        theta = theta_next
+    return x, maxiter
+
+
+def tfqmr(A, b, M=None, tol=1e-6, maxiter=1000):
+    """Transpose-free QMR (Freund 1993), right-preconditioned so the
+    recurrence residual matches the true residual."""
+    matvec = _matvec(A)
+    psolve = _matvec(M)
+
+    def op(v):
+        return matvec(psolve(v))
+
+    b_norm = float(cp.linalg.norm(b))
+    if b_norm == 0.0:
+        return cp.zeros_like(b), 0
+    target = tol * b_norm
+
+    z = cp.zeros_like(b)  # solution of (A M^-1) z = b; x = M^-1 z
+    r = b.copy()
+    w = r.copy()
+    u1 = r.copy()
+    r_shadow = r.copy()
+    Au1 = op(u1)
+    v = Au1.copy()
+    d = cp.zeros_like(b)
+    tau = float(cp.linalg.norm(r))
+    theta = 0.0
+    eta = 0.0 + 0.0j
+    rho = complex(cp.vdot(r_shadow, r))
+    iteration = 0
+    while iteration < maxiter:
+        sigma = complex(cp.vdot(r_shadow, v))
+        if sigma == 0.0 or rho == 0.0:
+            break
+        alpha = rho / sigma
+        u2 = u1 - alpha * v
+        Au2 = op(u2)
+        converged = False
+        for u, Au in ((u1, Au1), (u2, Au2)):
+            w = w - alpha * Au
+            theta_next = float(cp.linalg.norm(w)) / tau
+            c = 1.0 / (1.0 + theta_next * theta_next) ** 0.5
+            tau = tau * theta_next * c
+            eta_next = c * c * alpha
+            d = u + (theta * theta * eta / alpha) * d
+            z = z + eta_next * d
+            theta = theta_next
+            eta = eta_next
+            iteration += 1
+            # tau * sqrt(iteration + 1) bounds the true residual
+            if tau * (iteration + 1) ** 0.5 <= target:
+                converged = True
+                break
+        if converged:
+            break
+        rho_next = complex(cp.vdot(r_shadow, w))
+        beta = rho_next / rho
+        u1 = w + beta * u2
+        Au1 = op(u1)
+        v = Au1 + beta * (Au2 + beta * v)
+        rho = rho_next
+    x = psolve(z)
+    residual = float(cp.linalg.norm(b - matvec(x)))
+    return x, 0 if residual <= target else max(iteration, 1)
+
+
+def idr_s(A, b, M=None, tol=1e-6, maxiter=1000, s=4, seed=0):
+    """IDR(s) with biorthogonalization (van Gijzen & Sonneveld, TOMS 2011).
+
+    ``maxiter`` counts matvecs, comparable with the other engines. The shadow
+    space is a fixed orthonormalized complex random block (deterministic
+    ``seed``)."""
+    matvec = _matvec(A)
+    psolve = _matvec(M)
+    n = b.size
+    x = cp.zeros_like(b)
+    r = b.copy()
+    b_norm = float(cp.linalg.norm(b))
+    if b_norm == 0.0:
+        return x, 0
+    target = tol * b_norm
+
+    rng = cp.random.default_rng(seed)
+    shadow = (rng.standard_normal((n, s), dtype=cp.float32)
+              + 1j * rng.standard_normal((n, s), dtype=cp.float32)).astype(b.dtype)
+    shadow, _ = cp.linalg.qr(shadow)
+
+    G = cp.zeros((n, s), dtype=b.dtype)
+    U = cp.zeros((n, s), dtype=b.dtype)
+    Ms = cp.eye(s, dtype=b.dtype)
+    omega = 1.0 + 0.0j
+    kappa = 0.7
+    iteration = 0
+
+    r_norm = float(cp.linalg.norm(r))
+    while r_norm > target and iteration < maxiter:
+        f = shadow.conj().T @ r
+        for k in range(s):
+            c = cp.linalg.solve(Ms[k:, k:], f[k:])
+            v = r - G[:, k:] @ c
+            v = psolve(v)
+            U[:, k] = U[:, k:] @ c + omega * v
+            G[:, k] = matvec(U[:, k])
+            iteration += 1
+            for i in range(k):
+                alpha = complex(cp.vdot(shadow[:, i], G[:, k])) / complex(Ms[i, i])
+                G[:, k] -= alpha * G[:, i]
+                U[:, k] -= alpha * U[:, i]
+            Ms[k:, k] = shadow[:, k:].conj().T @ G[:, k]
+            if complex(Ms[k, k]) == 0.0:
+                return x, iteration
+            beta = complex(f[k]) / complex(Ms[k, k])
+            r = r - beta * G[:, k]
+            x = x + beta * U[:, k]
+            r_norm = float(cp.linalg.norm(r))
+            if r_norm <= target or iteration >= maxiter:
+                break
+            if k + 1 < s:
+                f[k + 1:] = f[k + 1:] - beta * Ms[k + 1:, k]
+        if r_norm <= target or iteration >= maxiter:
+            break
+        # Dimension-reduction step
+        v = psolve(r)
+        t = matvec(v)
+        iteration += 1
+        t_norm = float(cp.linalg.norm(t))
+        if t_norm == 0.0:
+            return x, iteration
+        ts = complex(cp.vdot(t, r))
+        angle = abs(ts) / (t_norm * r_norm)
+        omega = ts / (t_norm * t_norm)
+        if angle < kappa and angle > 0.0:
+            omega = omega * kappa / angle
+        if omega == 0.0:
+            return x, iteration
+        x = x + omega * v
+        r = r - omega * t
+        r_norm = float(cp.linalg.norm(r))
+    return x, 0 if r_norm <= target else iteration

witwin/maxwell/fdfd/plotting.py

@@ -0,0 +1,110 @@
+import numpy as np
+import torch
+
+from ..visualization import (
+    extract_orthogonal_slice,
+    plot_cross_section_panels,
+    plot_orthogonal_views,
+)
+from .postprocess import interpolate_yee_to_center
+
+
+def plot_cross_section(
+    solver,
+    axis="z",
+    position=0.0,
+    component="abs",
+    field_log_scale=False,
+    figsize=(12, 5),
+    save_path=None,
+    verbose=None,
+):
+    if verbose is None:
+        verbose = solver.verbose
+    if solver.E_field is None:
+        print("Run solve() first.")
+        return
+
+    ex, ey, ez = solver.E_field
+    component_str = component.lower()
+
+    if component_str == "abs":
+        field_data, x_c, y_c, z_c = interpolate_yee_to_center(solver)
+        eps_np = solver.scene.permittivity[:-1, :-1, :-1].cpu().numpy()
+        title_prefix = "|E|"
+    else:
+        if component_str == "ex":
+            field_data = torch.abs(ex).cpu().numpy()
+        elif component_str == "ey":
+            field_data = torch.abs(ey).cpu().numpy()
+        elif component_str == "ez":
+            field_data = torch.abs(ez).cpu().numpy()
+        else:
+            raise ValueError(f"Invalid component '{component}'. Use 'abs', 'Ex', 'Ey', or 'Ez'")
+        x_c = solver.scene.x.cpu().numpy()
+        y_c = solver.scene.y.cpu().numpy()
+        z_c = solver.scene.z.cpu().numpy()
+        eps_np = solver.scene.permittivity.cpu().numpy()
+        title_prefix = f"|{component.upper()}|"
+
+    field_info = extract_orthogonal_slice(field_data, axis, position, x_c, y_c, z_c)
+    eps_info = extract_orthogonal_slice(eps_np, axis, position, x_c, y_c, z_c)
+
+    field_slice = field_info["slice"]
+    if field_log_scale:
+        field_slice = np.where(field_slice > 1e-12, field_slice, 1e-12)
+        fmax = np.max(field_slice)
+        if fmax > 0:
+            field_slice = 20 * np.log10(field_slice / fmax)
+        cbar_label = f"{title_prefix} (dB)"
+    else:
+        cbar_label = f"{title_prefix} (V/m)"
+
+    plot_cross_section_panels(
+        field_slice=field_slice,
+        permittivity_slice=eps_info["slice"],
+        x_coords=field_info["x_coords"],
+        y_coords=field_info["y_coords"],
+        xlabel=field_info["xlabel"],
+        ylabel=field_info["ylabel"],
+        field_title=f"{title_prefix} at {axis}={position:.3f}m",
+        field_cbar_label=cbar_label,
+        permittivity_title=f"Permittivity at {axis}={position:.3f}m",
+        figsize=figsize,
+        save_path=save_path,
+        verbose=verbose,
+    )
+
+
+def plot_isotropic_3views(
+    solver,
+    position=0.0,
+    field_log_scale=True,
+    figsize=(18, 5),
+    save_path=None,
+    vmin_db=-60,
+    verbose=None,
+):
+    if verbose is None:
+        verbose = solver.verbose
+    if solver.E_field is None:
+        print("Run solve() first.")
+        return
+
+    field_data, x_c, y_c, z_c = interpolate_yee_to_center(solver)
+    eps = solver.scene.permittivity[:-1, :-1, :-1].cpu().numpy()
+
+    plot_orthogonal_views(
+        field_data=field_data,
+        permittivity=eps,
+        x_coords=x_c,
+        y_coords=y_c,
+        z_coords=z_c,
+        position=position,
+        field_log_scale=field_log_scale,
+        figsize=figsize,
+        save_path=save_path,
+        vmin_db=vmin_db,
+        verbose=verbose,
+        suptitle="Total Electric Field Intensity |E| (FDFD)",
+    )