diff --git a/examples/solid/beam_vibration_analysis/README.md b/examples/solid/beam_vibration_analysis/README.md new file mode 100644 index 0000000..582e7c8 --- /dev/null +++ b/examples/solid/beam_vibration_analysis/README.md @@ -0,0 +1,71 @@ +# Beam vibration analysis — cantilever cylinder (modes + FRF) + +A 3D **structural-vibration** example: a clamped-free circular steel rod +(length $L = 1\,\mathrm{m}$, radius $R = 25\,\mathrm{mm}$) analysed for its +natural **mode shapes** and its harmonic **displacement FRF** (frequency-response +function). It exercises TensorMesh linear elasticity for a generalized +eigenproblem plus a modal harmonic response. Example-only: no public API is added. + +## Model + +Linear elasticity on a tetrahedral mesh gives the stiffness $K$ and consistent +mass $M$; clamping the base face $z = 0$ removes those DOFs. + +**Modes** — generalized eigenproblem (shift-invert Lanczos near $\sigma = 0$): + +$$K\,\phi = \omega^2 M\,\phi, \qquad f_n = \omega_n / (2\pi).$$ + +**FRF** — by modal superposition with the mass-normalized modes +($\phi_r^{\mathsf T} M \phi_r = 1$) and an isotropic (hysteretic) loss factor +$\eta = 2\%$, i.e. a $1\%$ modal damping ratio ($\zeta = \eta/2$): + +$$u(\omega) = \sum_r \frac{\phi_r\,(\phi_r\cdot F)}{\omega_r^2(1 + i\eta) - \omega^2}.$$ + +The tip response to a transverse ($x$) tip force is normalized by its static +($\omega\!\to\!0$) value, giving a **dimensionless dynamic amplification** +$|u_x|/u_\text{static}$ — 1 below the first resonance, peaking on each bending +mode and dropping into antiresonances between them. + +$K$ comes from `LinearElasticityElementAssembler`; the consistent mass is the +scalar `MassElementAssembler` lifted to the three vector components, +$M_\text{vec} = \rho\,(M_\text{scalar} \otimes I_3)$. The eigen solve uses SciPy +sparse (post-processing), keeping the compute core numpy-free. + +```text +mesh (gmsh cylinder, tets) -> assemble (K, M) -> eig (modes) + -> modal FRF -> plot modes + FRF +``` + +## Run + +```bash +python vibration_cylinder.py +``` + +`run_demo(...)` returns diagnostics; `main()` exposes `--n-modes`, `--mesh-h-mm`, +`--frf-points`, `--no-plot`, `--modes-output`, `--frf-output`. + +## What it shows + +* **Mode shapes** (`vibration_cylinder_modes.png`) — deformed 3D cylinders + (colored by displacement magnitude, tip amplitude exaggerated, undeformed axis + drawn as a dashed line), one mode of each kind. Because the section is circular + each **bending** mode is doubled (two orthogonal bending planes), so modes come + in near-equal pairs; above them sit the first **torsion** and **axial** modes. +* **Displacement FRF** (`vibration_cylinder_frf.png`) — tip transverse + displacement amplitude $|u_x|$ (m) vs frequency: sharp peaks at the bending + frequencies with antiresonances between them. The transverse tip force excites + only bending, so the torsion/axial modes show no peak. + +## Accuracy + +The first bending frequency lands at 37.2 Hz vs the Euler-Bernoulli cantilever +value $f_1 = (\beta_1 L)^2/(2\pi)\sqrt{EI/(\rho A L^4)} = 36.2$ Hz. Linear +tetrahedra are a little stiff in bending, so the FEM frequencies sit ~2–3 % +above the slender-beam analytic and drop toward it as the mesh is refined +(`--mesh-h-mm 5`); the axial mode already matches to <0.1 %. + +> Quadratic (`tetra10`) elements are **not** used here: gmsh's second-order +> tet edge-node ordering does not match TensorMesh's, which corrupts the curved +> Jacobian (the mass matrix under-integrates the volume by ~17 %). The example +> therefore stays on linear tets and relies on mesh refinement for accuracy. diff --git a/examples/solid/beam_vibration_analysis/vibration_cylinder.py b/examples/solid/beam_vibration_analysis/vibration_cylinder.py new file mode 100644 index 0000000..29c3535 --- /dev/null +++ b/examples/solid/beam_vibration_analysis/vibration_cylinder.py @@ -0,0 +1,426 @@ +"""Cantilever cylinder eigenmodes + displacement FRF on TensorMesh — 3D elasticity. + +A clamped-free circular steel rod (length ``L`` = 1 m, radius ``R`` = 25 mm) is +analysed for its structural vibration: the natural mode shapes and the harmonic +drive-point frequency-response function (FRF). Linear elasticity on a tetrahedral +mesh gives the stiffness ``K`` and consistent mass ``M``; clamping the base face +removes those degrees of freedom. + +The **modes** solve the generalized eigenproblem + + K phi = omega^2 M phi , f_n = omega_n / (2 pi) , + +for the lowest ``n_modes`` pairs (shift-invert Lanczos near ``sigma = 0``). The +**FRF** is built by modal superposition with the mass-normalized modes and an +isotropic (hysteretic) loss factor ``eta`` = 2 % (a 1 % modal damping ratio, +``zeta = eta / 2``): + + u(omega) = sum_r phi_r (phi_r . F) / ( omega_r^2 (1 + i eta) - omega^2 ) . + +The tip-average ``|u_x|`` is normalized by its static (``omega -> 0``) value, so +the FRF is a **dimensionless** dynamic amplification: 1 below the first resonance, +peaking on each bending mode and dropping into antiresonances between them. + +``K`` comes from :class:`~tensormesh.assemble.LinearElasticityElementAssembler` +and the consistent mass is the scalar +:class:`~tensormesh.assemble.MassElementAssembler` lifted to the 3 vector +components (``M_vec = rho * (M_scalar (x) I_3)``). The eigen solve uses SciPy +sparse (post-processing), keeping the TensorMesh compute core numpy-free. +No public API is added here. + +Workflow: mesh (gmsh cylinder, tets) -> assemble (K, M) -> eig (modes) + -> modal FRF -> plot mode shapes + FRF. + +Run (env with torch + tensormesh + gmsh + scipy): + python "examples/solid/beam_vibration_analysis/vibration_cylinder.py" +""" +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Optional + +import numpy as np +import scipy.sparse as sp +import scipy.sparse.linalg as spla +import torch + +ROOT = Path(__file__).resolve().parents[3] +sys.path.append(str(ROOT)) + +from tensormesh import Mesh +from tensormesh.assemble import (LinearElasticityElementAssembler, + MassElementAssembler) + +HERE = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- # +# Problem definition +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Cylinder: + r"""All parameters of the clamped-free steel cylinder (SI units). + + Fields + ------ + length, radius : cylinder length ``L`` and section radius ``R`` (m). + E, nu, rho : Young's modulus (Pa), Poisson ratio, density (kg/m^3). + n_modes : number of eigenmodes to extract. + eta : isotropic (hysteretic) loss factor; modal damping ratio is ``eta/2``. + f0 : total transverse (x) tip force for the FRF (N). + frf_f0, frf_f1, frf_n : FRF sweep start / stop (Hz) and number of points. + mesh_h : target tetra edge length (m); default ``radius / 4``. + """ + + length: float = 1.0 + radius: float = 0.025 + E: float = 210.0e9 + nu: float = 0.3 + rho: float = 7850.0 + n_modes: int = 12 + eta: float = 0.02 + f0: float = 1.0 + frf_f0: float = 5.0 + frf_f1: float = 1500.0 + frf_n: int = 400 + mesh_h: Optional[float] = None + + @property + def h(self) -> float: + """Resolved tetra edge length (m); default ``radius / 4``.""" + return self.mesh_h or self.radius / 4.0 + + @property + def area(self) -> float: + """Cross-section area ``pi R^2`` (m^2).""" + return float(np.pi * self.radius ** 2) + + @property + def inertia(self) -> float: + """Second moment of area ``pi R^4 / 4`` (m^4).""" + return float(np.pi * self.radius ** 4 / 4.0) + + def analytic_bending(self, n: int = 4) -> np.ndarray: + r"""Euler-Bernoulli cantilever bending frequencies (Hz), first ``n``.""" + beta_l = np.array([1.875104, 4.694091, 7.854757, 10.995541])[:n] + coeff = np.sqrt(self.E * self.inertia / (self.rho * self.area * self.length ** 4)) + return (beta_l ** 2) / (2.0 * np.pi) * coeff + + +# --------------------------------------------------------------------------- # +# 1. Mesh: tetrahedral cylinder, axis = z, base at z = 0 +# --------------------------------------------------------------------------- # +def build_mesh(problem: Cylinder, msh_path: Optional[str] = None) -> Mesh: + """gmsh tetrahedral mesh of the cylinder; returns a tetra-only Mesh.""" + import gmsh + import meshio + + if msh_path is None: + msh_path = str(HERE / "_cylinder.msh") + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + occ = gmsh.model.occ + occ.addCylinder(0, 0, 0, 0, 0, problem.length, problem.radius) + occ.synchronize() + gmsh.option.setNumber("Mesh.MeshSizeMin", problem.h) + gmsh.option.setNumber("Mesh.MeshSizeMax", problem.h) + gmsh.model.mesh.generate(3) + gmsh.write(msh_path) + finally: + gmsh.finalize() + + # keep only the tetrahedra — lower-dim facets confuse a 3D volume assembler. + m = meshio.read(msh_path) + tets = np.vstack([c.data for c in m.cells if c.type == "tetra"]) + clean = meshio.Mesh(points=m.points, cells=[("tetra", tets)]) + return Mesh(clean, reorder=False) + + +# --------------------------------------------------------------------------- # +# 2. Assemble stiffness K and consistent (vector) mass M as SciPy sparse +# --------------------------------------------------------------------------- # +def assemble(problem: Cylinder, mesh: Mesh): + """Return ``(K, M)`` as CSC (3N x 3N) with node-major, component-minor DOFs.""" + pts = mesh.points.to(torch.float64) + + Kasm = LinearElasticityElementAssembler.from_mesh(mesh, E=problem.E, nu=problem.nu) + K = Kasm().to_scipy_coo().tocsc() + + # scalar consistent mass M_ij = int phi_i phi_j ; lift to the 3 components: + # for isotropic density the vector mass is rho * (M_scalar kron I_3), which + # matches the assembler's node-major / component-minor DOF ordering. + Masm = MassElementAssembler.from_mesh(mesh) + Ms = Masm(pts).to_scipy_coo().tocsc() + M = (problem.rho * sp.kron(Ms, sp.eye(3, format="csc"), format="csc")).tocsc() + return K, M + + +# --------------------------------------------------------------------------- # +# 3. Boundary DOFs: clamp the base face z = 0 +# --------------------------------------------------------------------------- # +def dof_masks(problem: Cylinder, mesh: Mesh): + """Return ``(free_dofs, tip_nodes)``: free DOF indices + free-tip node ids.""" + pts = mesh.points.cpu().numpy() + tol = problem.radius * 1e-4 + fixed_node = pts[:, 2] < tol + tip_node = np.where(pts[:, 2] > problem.length - tol)[0] + free = np.where(np.repeat(~fixed_node, 3))[0] + return free, tip_node + + +# --------------------------------------------------------------------------- # +# 4. Eigenmodes: K phi = omega^2 M phi +# --------------------------------------------------------------------------- # +def solve_modes(problem: Cylinder, K, M, free: np.ndarray) -> Dict: + """Lowest ``n_modes`` natural frequencies (Hz) + full-length mode shapes.""" + Kff = K[free][:, free] + Mff = M[free][:, free] + # shift-invert near 0 pulls out the lowest eigenpairs of the GEVP. + lam, vec = spla.eigsh(Kff, k=problem.n_modes, M=Mff, sigma=0.0, which="LM") + order = np.argsort(lam) + lam, vec = lam[order], vec[:, order] + freqs = np.sqrt(np.maximum(lam, 0.0)) / (2.0 * np.pi) + + n = K.shape[0] + shapes = np.zeros((problem.n_modes, n)) + shapes[:, free] = vec.T + return dict(freqs=freqs, shapes=shapes) # shapes[m] is a length-3N vector + + +# --------------------------------------------------------------------------- # +# 5. FRF: drive-point tip receptance by modal superposition +# --------------------------------------------------------------------------- # +def solve_frf(problem: Cylinder, modes: Dict, tip: np.ndarray) -> Dict: + r"""Tip transverse displacement amplitude ``|u_x|`` (m) vs frequency. + + Modal superposition with the mass-normalized modes ``phi_r`` (from + :func:`solve_modes`, where ``phi^T M phi = I``) and hysteretic damping, driven + by the transverse tip force ``F`` = ``f0`` (N): + + u(omega) = sum_r phi_r (phi_r . F) / ( omega_r^2 (1 + i eta) - omega^2 ) . + + Returns the steady-state tip-face-average ``|u_x|`` in meters — the physical + vibration amplitude of the free end at each drive frequency. + """ + shapes, wr = modes["shapes"], 2.0 * np.pi * modes["freqs"] # phi_r, omega_r + tip_x = 3 * tip + 0 # x-DOF per tip node + + # modal force q_r = phi_r . F (total x-force f0 spread over the tip face) + qr = shapes[:, tip_x].sum(axis=1) * (problem.f0 / len(tip)) + phi_tip = shapes[:, tip_x].mean(axis=1) # tip-average phi_r + + freqs = np.linspace(problem.frf_f0, problem.frf_f1, problem.frf_n) + w2 = (2.0 * np.pi * freqs) ** 2 + denom = (wr[:, None] ** 2) * (1.0 + 1j * problem.eta) - w2[None, :] + u_tip = np.abs(((phi_tip * qr)[:, None] / denom).sum(axis=0)) # |u_x| (m) + return dict(frf_hz=freqs, frf_ux=u_tip) + + +# --------------------------------------------------------------------------- # +# 6. Plot: deformed mode shapes + FRF +# --------------------------------------------------------------------------- # +def _mode_kind(u: np.ndarray, pts: np.ndarray, problem: Cylinder) -> str: + """Label a mode bending / torsion / axial from its cross-section motion. + + Grouping nodes into thin ``z``-slices, each slice's rigid motion splits into a + transverse **translation** (bending), a net **twist** about the axis + (torsion) and an axial **stretch** (axial). A slice's *signed* twist cancels + for a bending mode (odd across the section) but accumulates for torsion, so + comparing the slice-summed magnitudes classifies the mode cleanly. + """ + x, y, z = pts[:, 0], pts[:, 1], pts[:, 2] + r2 = x * x + y * y + zbin = np.round(z / problem.length * 200).astype(np.int64) + + def _slice_sum(v): # sum |per-slice mean| + order = np.argsort(zbin) + _, first = np.unique(zbin[order], return_index=True) + sums = np.add.reduceat(v[order], first) + cnts = np.add.reduceat(np.ones_like(v[order]), first) + return np.abs(sums / cnts).sum() + + bend = _slice_sum(u[:, 0]) + _slice_sum(u[:, 1]) # transverse translation + axial = _slice_sum(u[:, 2]) # z-stretch + twist = _slice_sum((x * u[:, 1] - y * u[:, 0]) / (r2.mean() ** 0.5 + 1e-30)) + scores = {"bending": bend, "torsion": twist, "axial": axial} + return max(scores, key=scores.get) + + +def plot_modes(res: Dict, problem: Cylinder, mesh: Mesh, save_path: str, + show: Optional[list] = None) -> None: + """Selected mode shapes as deformed 3D cylinders (colored by |u|). + + The 40:1-slender rod would vanish in a true-aspect 3D box, so each panel + exaggerates the modal deflection (tip amplitude ~ 0.35 L) and stretches the + transverse axes to match, with the undeformed skin drawn faintly for + reference. ``show`` is a list of 0-based mode indices; the default picks one + of each kind (bending / torsion / axial) so the whole family is visible. + """ + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + pts = mesh.points.cpu().numpy() + freqs, shapes = res["freqs"], res["shapes"] + # cylinder skin nodes only -> a clean, light 3D point cloud. + rr = np.hypot(pts[:, 0], pts[:, 1]) + skin = np.where((rr > problem.radius * 0.8) + | (pts[:, 2] < problem.radius * 0.15) + | (pts[:, 2] > problem.length - problem.radius * 0.15))[0] + sp = pts[skin] + + if show is None: + # one representative per kind, ascending, padded with the lowest bendings. + kinds = [_mode_kind(shapes[i].reshape(-1, 3), pts, problem) + for i in range(len(freqs))] + picks, seen = [], set() + for i, k in enumerate(kinds): + if k not in seen: + picks.append(i); seen.add(k) + for i in range(len(freqs)): + if len(picks) >= 6: + break + if i not in picks: + picks.append(i) + show = sorted(picks[:6]) + + amp = 0.30 * problem.length # target tip deflection + R = problem.radius + ncol = 3 + nrow = int(np.ceil(len(show) / ncol)) + fig = plt.figure(figsize=(4.6 * ncol, 3.4 * nrow)) + for panel, m in enumerate(show): + u = shapes[m].reshape(-1, 3) + umag = np.linalg.norm(u, axis=1) + kind = _mode_kind(u, pts, problem) + + # Put the mode's dominant transverse component on the vertical plot axis + # so every bending mode (whichever physical plane it lives in) reads as a + # clean vertical deflection under one fixed camera; the other transverse + # component becomes the (shallow) depth axis. + vax = 0 if np.abs(u[:, 0]).sum() >= np.abs(u[:, 1]).sum() else 1 + dax = 1 - vax + # torsion moves the rim, not the axis; scaling the rim to `amp` would make + # it explode into a fan, so cap it to a few radii -> a visible twisted rod. + target = 3.0 * R if kind == "torsion" else amp + scale = target / (umag.max() + 1e-30) + horiz, depth, vert = sp[:, 2], sp[:, dax], sp[:, vax] + d_h = horiz + scale * u[skin, 2] + d_d = depth + scale * u[skin, dax] + d_v = vert + scale * u[skin, vax] + + # symmetric limits about each transverse range (base axis at 0 included) + # so the deformed rod is centered in its panel; a generous floor keeps a + # no-swing mode (axial) a legible rod rather than an edge-on sliver. + floor = 0.10 * problem.length + + def _center(vals): + lo, hi = min(vals.min(), 0.0), max(vals.max(), 0.0) + mid, half = 0.5 * (lo + hi), max(0.5 * (hi - lo), 0.5 * floor) + return mid, half * 1.15 + + cd, hd = _center(d_d) + cv, hv = _center(d_v) + + ax = fig.add_subplot(nrow, ncol, panel + 1, projection="3d") + # undeformed reference = the neutral axis (thin line, so it never occludes + # the colored rod the way a solid grey tube does in matplotlib's 3D, which + # has no reliable depth sorting between artists). + ax.plot([0.0, problem.length], [0.0, 0.0], [0.0, 0.0], + color="#999999", lw=1.6, ls="--", zorder=1) + ax.scatter(d_h, d_d, d_v, + c=umag[skin] / umag.max(), cmap="turbo", s=11, + linewidths=0, zorder=2) + ax.set_title(f"mode {m + 1}: {freqs[m]:.1f} Hz ({kind})", fontsize=10) + ax.set_xlabel("z (m)", fontsize=8) + ax.set_xlim(0.0, problem.length) + ax.set_ylim(cd - hd, cd + hd) + ax.set_zlim(cv - hv, cv + hv) + ax.set_box_aspect((problem.length, 2 * hd, 2 * hv)) + ax.set_xticks([0, problem.length]) + ax.set_yticks([]); ax.set_zticks([]) + ax.view_init(elev=18, azim=-70) + fig.suptitle("Cantilever cylinder — natural mode shapes", y=1.0, fontsize=13) + fig.tight_layout() + fig.savefig(save_path, dpi=140, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +def plot_frf(res: Dict, problem: Cylinder, save_path: str) -> None: + """Tip transverse displacement amplitude ``|u_x|`` (m) vs frequency.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(figsize=(8.2, 4.6)) + ax.semilogy(res["frf_hz"], res["frf_ux"], color="#0072B2", lw=1.8, zorder=3) + ax.set_xlim(problem.frf_f0, problem.frf_f1) + ax.set_xlabel("frequency (Hz)") + ax.set_ylabel(r"tip displacement amplitude $|u_x|$ (m)") + ax.set_title("Cantilever cylinder — drive-point displacement FRF") + ax.grid(True, which="both", alpha=0.25) + fig.tight_layout() + fig.savefig(save_path, dpi=140, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 7. run_demo / main +# --------------------------------------------------------------------------- # +def run_demo(*, make_plot: bool = True, modes_path: Optional[str] = None, + frf_path: Optional[str] = None, + problem: Optional[Cylinder] = None) -> Dict: + """Solve modes + FRF, save figures, print a summary, return diagnostics.""" + torch.set_default_dtype(torch.float64) + problem = problem or Cylinder() + + mesh = build_mesh(problem) + K, M = assemble(problem, mesh) + free, tip = dof_masks(problem, mesh) + print(f"mesh: {mesh.points.shape[0]} nodes, " + f"{mesh.cells['tetra'].shape[0]} tets; {len(tip)} tip nodes; " + f"{K.shape[0] - len(free)} clamped DOFs", flush=True) + + modes = solve_modes(problem, K, M, free) + frf = solve_frf(problem, modes, tip) + res = dict(mesh=mesh, **modes, **frf) + + print("\nnatural frequencies (Hz):", flush=True) + print(f"{'mode':>4} {'TensorMesh':>12}", flush=True) + for i in range(problem.n_modes): + print(f"{i + 1:>4} {modes['freqs'][i]:12.2f}", flush=True) + # Euler-Bernoulli bending frequencies (each appears twice — two bend planes). + print("Euler-Bernoulli bending (Hz): " + + ", ".join(f"{f:.1f}" for f in problem.analytic_bending()), flush=True) + + if make_plot: + plot_modes(res, problem, mesh, + modes_path or str(HERE / "vibration_cylinder_modes.png")) + plot_frf(res, problem, frf_path or str(HERE / "vibration_cylinder_frf.png")) + return res + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--n-modes", type=int, default=12) + parser.add_argument("--mesh-h-mm", type=float, default=None, help="tetra edge (mm)") + parser.add_argument("--frf-points", type=int, default=400) + parser.add_argument("--no-plot", action="store_true") + parser.add_argument("--modes-output", type=str, default=None) + parser.add_argument("--frf-output", type=str, default=None) + args = parser.parse_args() + + problem = Cylinder(n_modes=args.n_modes, frf_n=args.frf_points, + mesh_h=(args.mesh_h_mm * 1e-3) if args.mesh_h_mm else None) + run_demo(make_plot=not args.no_plot, modes_path=args.modes_output, + frf_path=args.frf_output, problem=problem) + + +if __name__ == "__main__": + main() diff --git a/examples/wave/helmholtz_resonator/README.md b/examples/wave/helmholtz_resonator/README.md new file mode 100644 index 0000000..1dc80bd --- /dev/null +++ b/examples/wave/helmholtz_resonator/README.md @@ -0,0 +1,53 @@ +# Helmholtz resonator + +A 2D **acoustic Helmholtz resonator** — a duct with a necked side cavity, driven +by a plane-wave port — solved in the frequency domain on TensorMesh. It showcases +the library boundary operators [`tensormesh.robin_operator` / +`tensormesh.port_source`](../../../tensormesh/bc.py) (built on +`FacetBilinearAssembler`): the port line mass and the incident-wave load. + +Example-only: no public API is added. It reuses the scalar-Helmholtz assembly +(`LaplaceElementAssembler`, `MassElementAssembler`), the boundary operators, and +the complex sparse solve. + +## Model + +Scalar pressure acoustics $\nabla^2 p + k^2 p = 0$ on the air domain (duct + neck ++ cavity, meshed as one union). All walls are sound-hard (natural Neumann); the +inlet $x=0$ is a first-order plane-wave port that both injects the incident wave +and absorbs the reflected one: + +$$\frac{\partial p}{\partial n} + i k\, p = 2 i k\, p_0 \quad\text{on the inlet.}$$ + +In the weak form the port adds a boundary mass $B=\int_\Gamma \phi_i \phi_j\,\mathrm dS$ +(scaled by $i k$) to the operator and a load $e=\int_\Gamma \phi_i\,\mathrm dS$ +(scaled by $2 i k p_0$) to the right-hand side — both assembled by +`robin_operator` / `port_source` and folded onto the volume layout so the whole +operator stays one complex `SparseMatrix`: + +```text +mesh -> Laplace/Mass assembler -> port operators -> (K - k^2 M + i k B) p = 2 i k p0 e +``` + +The input impedance $Z(f)=p_\text{avg}/u_\text{in}$ is read at the inlet over a +frequency sweep; its peaks are the resonances (the Helmholtz mode plus duct +harmonics). + +## Run + +```bash +python helmholtz_resonator.py +``` + +`run_demo(...)` returns diagnostics and `main()` exposes `--no-plot` / +`--output`. The demo reports the fundamental resonance near **357 Hz** with a +cavity pressure gain of ~12x over the incident amplitude. + +## What it shows + +- **Input impedance** $|Z|/Z_0$ over the sweep (log scale) with the cavity + pressure response — resonance and anti-resonance peaks. +- **Acoustic field** at the resonance: total pressure `Re(p)` and sound-pressure + level, with the pressure node line across the neck. +- **Boundary conditions**: a schematic of the port (red) and sound-hard walls + (blue) over the duct + cavity geometry. diff --git a/examples/wave/helmholtz_resonator/helmholtz_resonator.py b/examples/wave/helmholtz_resonator/helmholtz_resonator.py new file mode 100644 index 0000000..93fc741 --- /dev/null +++ b/examples/wave/helmholtz_resonator/helmholtz_resonator.py @@ -0,0 +1,493 @@ +"""2D Helmholtz-resonator side-branch on a duct — port-driven pressure acoustics. + +A straight air duct carries a plane wave injected by a **port** at the inlet; +a Helmholtz resonator (a narrow neck opening into a larger cavity) hangs off the +duct's far wall. Near the Helmholtz resonance the cavity pressure is strongly +amplified and the port reactance sweeps through zero. Everything is solved on +TensorMesh (torch); no external reference is needed. + +Geometry (meters) — three axis-aligned rectangles, unioned:: + + duct : [0, D] x [0, Dh] D = Dh = 0.1 + neck : [D, D+Ln] x [0, Nh] Ln = 0.01, Nh = 0.002 + cavity : [D+Ln, D+Ln+Lc] x [0, Ch] Lc = Ch = 0.05 + +with ``D = duct_length``, ``Dh = duct_height``, ``Ln = neck_length``, +``Nh = neck_height``, ``Lc = cavity_length``, ``Ch = cavity_height``. The neck +opening (x = D, 0 <= y <= Nh) is interior after the union; the rest of the far +wall (x = D, Nh <= y <= Dh) is rigid. + +Boundary conditions:: + + inlet x = 0 : plane-wave port, incident amplitude p0 = 1 (absorbing + source) + all others : sound-hard, dp/dn = 0 (natural Neumann) + +Helmholtz, time convention ``e^{+iwt}`` (outgoing ~ e^{-ikr}, dp/dn = -ik p): + + (K - k^2 M + i k B_in) p = 2 i k p0 e_in , k = w / c0 = 2 pi f / c0 + + K_ij = int grad phi_i . grad phi_j (LaplaceElementAssembler) + M_ij = int phi_i phi_j (MassElementAssembler) + B_in = int_{x=0} phi_i phi_j ds (hand-rolled inlet line mass) + e_in = int_{x=0} phi_i ds (incident plane-wave load) + +The port line mass ``B_in`` and load ``e_in`` come from the library boundary +operators :func:`tensormesh.robin_operator` and :func:`tensormesh.port_source` +(themselves built on ``FacetBilinearAssembler``); ``B_in`` is folded onto the +volume sparsity pattern so the whole operator stays a single complex +``SparseMatrix`` (``SparseMatrix`` add needs matching layouts). No public API is +added — this example only *uses* the library. + +Workflow: mesh -> assemble (K, M) -> port line operators -> per-f complex solve + -> input impedance Z(f) + pressure field p(x, y; f0). + +Run (env with torch + tensormesh + gmsh): + python examples/wave/helmholtz_resonator/helmholtz_resonator.py +""" + +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import torch + +ROOT = Path(__file__).resolve().parents[3] +sys.path.append(str(ROOT)) + +from tensormesh import (LaplaceElementAssembler, MassElementAssembler, Mesh, + robin_operator, port_source) +from tensormesh.sparse.matrix import SparseMatrix + +HERE = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- # +# Problem definition +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Resonator: + r"""All parameters of the duct + Helmholtz-resonator problem. + + Fields + ------ + duct_length, duct_height : duct dimensions (m). + neck_length, neck_height : neck length (along the duct axis) and aperture (m). + cavity_length, cavity_height : cavity dimensions (m). + rho0, c0 : air density (kg/m^3) and sound speed (m/s) at 20 C, 1 atm. + p0 : incident plane-wave amplitude at the port (Pa). + f_min, f_max, n_freq : frequency sweep (Hz). + mesh_size : target element edge length (m). + f_field : frequency at which the pressure field snapshot is drawn (Hz). + """ + + duct_length: float = 0.1 + duct_height: float = 0.1 + neck_length: float = 0.01 + neck_height: float = 0.002 + cavity_length: float = 0.05 + cavity_height: float = 0.05 + + rho0: float = 1.2041 + c0: float = 343.2 + p0: float = 1.0 + + f_min: float = 20.0 + f_max: float = 500.0 + n_freq: int = 500 + + mesh_size: float = 0.0015 + f_field: float = 357.0 + + @property + def freqs(self) -> torch.Tensor: + """Sweep frequencies (Hz), shape ``[n_freq]``.""" + return torch.linspace(self.f_min, self.f_max, self.n_freq, dtype=torch.float64) + + @property + def total_length(self) -> float: + """Overall x-extent of the geometry (m).""" + return self.duct_length + self.neck_length + self.cavity_length + + @property + def cavity_probe(self) -> Tuple[float, float]: + """Point at the cavity center (m) — a convenient response probe.""" + return (self.duct_length + self.neck_length + 0.5 * self.cavity_length, + 0.5 * self.cavity_height) + + @property + def Z0(self) -> float: + """Characteristic (specific) acoustic impedance of air, ``rho0 c0``.""" + return self.rho0 * self.c0 + + +# --------------------------------------------------------------------------- # +# 1. Mesh: the three rectangles, unioned +# --------------------------------------------------------------------------- # +def build_mesh(problem: Resonator, msh_path: Optional[str] = None) -> Mesh: + """Mesh the duct + neck + cavity union with gmsh, return a TensorMesh ``Mesh``.""" + import gmsh + + if msh_path is None: + msh_path = str(HERE / "_resonator.msh") + h = problem.mesh_size + + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + gmsh.model.add("helmholtz_resonator") + occ = gmsh.model.occ + duct = occ.addRectangle(0.0, 0.0, 0.0, problem.duct_length, problem.duct_height) + neck = occ.addRectangle(problem.duct_length, 0.0, 0.0, + problem.neck_length, problem.neck_height) + cavity = occ.addRectangle(problem.duct_length + problem.neck_length, 0.0, 0.0, + problem.cavity_length, problem.cavity_height) + occ.fuse([(2, duct)], [(2, neck), (2, cavity)]) + occ.synchronize() + gmsh.option.setNumber("Mesh.MeshSizeMin", h) + gmsh.option.setNumber("Mesh.MeshSizeMax", h) + gmsh.model.mesh.generate(2) + gmsh.write(msh_path) + finally: + gmsh.finalize() + + return Mesh.from_file(msh_path, reorder=False) + + +# --------------------------------------------------------------------------- # +# 2. Volume operators (real) — K = grad.grad, M = mass, sharing one layout +# --------------------------------------------------------------------------- # +def assemble_volume(mesh: Mesh) -> Tuple[SparseMatrix, SparseMatrix]: + """Assemble the stiffness ``K`` and mass ``M`` matrices (real, shared layout).""" + K_asm = LaplaceElementAssembler.from_mesh(mesh, quadrature_order=2) + M_asm = MassElementAssembler.from_assembler(K_asm) + K = K_asm(mesh.points) + M = M_asm(mesh.points) + return K, M + + +# --------------------------------------------------------------------------- # +# 3. Port line operators on the inlet x = 0 [hand-rolled, folded onto K layout] +# --------------------------------------------------------------------------- # +def _fold_onto_layout(K: SparseMatrix, rows: torch.Tensor, cols: torch.Tensor, + vals: torch.Tensor) -> torch.Tensor: + """Scatter boundary entries ``(rows, cols, vals)`` onto ``K``'s sparsity pattern. + + Returns a values vector aligned with ``K.row``/``K.col`` so that + ``K.values + `` is a legal same-layout combination. Every boundary + entry must already exist in ``K`` (boundary nodes sharing an edge are + volume-adjacent), otherwise the assertion trips. + """ + n = K.shape[0] + key_t = K.row.long() * n + K.col.long() + order = torch.argsort(key_t) + sorted_keys = key_t[order] + key_b = rows.long() * n + cols.long() + pos = torch.searchsorted(sorted_keys, key_b) + hit = order[pos] + assert torch.equal(key_t[hit], key_b), "boundary entry missing from volume layout" + out = torch.zeros_like(K.values, dtype=torch.float64) + out.index_add_(0, hit, vals.to(torch.float64)) + return out + + +def port_operators(mesh: Mesh, K: SparseMatrix, tol: float = 1e-7 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Inlet line mass ``B_in`` (as K-aligned values), load ``e_in``, inlet indices. + + The inlet is the straight edge ``x = 0``. The boundary bilinear form + ``B_in = int_inlet phi_i phi_j ds`` and the incident-wave load + ``e_in = int_inlet phi_i ds`` are assembled by the library + (:func:`tensormesh.robin_operator` / :func:`tensormesh.port_source`); ``B_in`` + is then folded onto ``K``'s sparsity pattern so the whole operator stays a + single complex ``SparseMatrix`` (``SparseMatrix`` add needs matching layouts). + """ + pts = mesh.points.to(torch.float64) + inlet_mask = pts[:, 0].abs() < tol + inlet = torch.where(inlet_mask)[0] + + B = robin_operator(mesh, inlet_mask, 1.0, points=pts) # int_inlet phi_i phi_j ds + e_in = port_source(mesh, inlet_mask, 1.0, points=pts).to(torch.float64) # int_inlet phi_i ds + B_vals = _fold_onto_layout(K, B.row, B.col, B.values.to(torch.float64)) + return B_vals, e_in, inlet + + +# --------------------------------------------------------------------------- # +# 4. Frequency sweep -> pressure fields, input impedance, cavity spectrum +# --------------------------------------------------------------------------- # +def solve_spectrum(problem: Resonator, verbose: bool = True) -> Dict: + """Solve the port-driven Helmholtz problem over the sweep; return diagnostics. + + Returns a dict with ``freqs``, ``mesh``, ``points``, per-frequency input + impedance ``Z`` (normalized by ``rho0 c0``), the cavity-probe magnitude + spectrum ``probe``, and the complex pressure field ``p_field`` at + ``problem.f_field``. + """ + mesh = build_mesh(problem) + K, M = assemble_volume(mesh) + B_vals, e_in, inlet = port_operators(mesh, K) + pts = mesh.points.to(torch.float64) + + Kv = K.values.to(torch.complex128) + Mv = M.values.to(torch.complex128) + Bv = B_vals.to(torch.complex128) + e_in_c = e_in.to(torch.complex128) + inlet_len = float(e_in.sum()) # = duct_height (int over the inlet) + + freqs = problem.freqs + probe_pt = torch.tensor(problem.cavity_probe, dtype=torch.float64) + jprobe = int(((pts - probe_pt) ** 2).sum(1).argmin()) + + if verbose: + print(f"mesh: {mesh.n_points} nodes, {mesh.n_elements} triangles; " + f"{len(inlet)} inlet nodes", flush=True) + + z_norm = torch.zeros(len(freqs), dtype=torch.complex128) + probe = torch.zeros(len(freqs), dtype=torch.float64) + p_field: Optional[torch.Tensor] = None + jf_field = int((freqs - problem.f_field).abs().argmin()) + + for i, f in enumerate(freqs.tolist()): + k = 2.0 * torch.pi * f / problem.c0 + A_vals = Kv - (k * k) * Mv + (1j * k) * Bv + A = SparseMatrix(A_vals, K.row, K.col, K.shape) + rhs = (2j * k * problem.p0) * e_in_c + p = A.solve(rhs) + + # input impedance at the port from the reflection coefficient: + # total p = p0 (1 + R) at the inlet -> R =

_inlet / p0 - 1, + # normalized specific impedance z = (1 + R) / (1 - R). + p_avg = (e_in_c @ p) / inlet_len + R = p_avg / problem.p0 - 1.0 + z_norm[i] = (1.0 + R) / (1.0 - R) + probe[i] = float(p[jprobe].abs()) + if i == jf_field: + p_field = p.clone() + + return dict(freqs=freqs, mesh=mesh, points=pts, Z=z_norm, probe=probe, + p_field=p_field, f_field=float(freqs[jf_field]), + Z0=problem.Z0, jprobe=jprobe) + + +# --------------------------------------------------------------------------- # +# 5. Plots — (a) input impedance Z(f), (b) pressure field at f0 +# --------------------------------------------------------------------------- # +def plot_impedance(res: Dict, save_path: str) -> None: + """Port impedance magnitude ``|Z|/Z0`` (log scale) + cavity response. + + The system is lossless (``|R| = 1``), so ``Z`` is purely reactive: ``|Z|`` + peaks at the **Helmholtz resonance** and dips to zero at the + **anti-resonance**. Both extrema show cleanly on a logarithmic axis. + """ + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import numpy as np + + f = res["freqs"].cpu().numpy() + Z = res["Z"].cpu().numpy() + Zmag = 1.0 / np.maximum(np.abs(Z), 1e-30) # peaks at resonance + probe = res["probe"].cpu().numpy() + + BLUE, GREEN = "#0072B2", "#009E73" + fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(8.5, 7.2), sharex=True) + + ax0.semilogy(f, Zmag, "-", color=BLUE, lw=1.5) + ax0.set_xlim(f.min(), f.max()) + ax0.set_ylabel(r"$|Z| / \rho_0 c_0$") + ax0.set_title("Helmholtz-resonator port impedance (log scale)") + ax0.grid(True, which="both", alpha=0.25) + + ax1.plot(f, probe, "-", color=GREEN, lw=1.8) + ax1.set_xlabel("frequency (Hz)") + ax1.set_ylabel(r"$|p|$ at cavity center (Pa)") + ax1.set_title("Cavity pressure response") + ax1.grid(True, alpha=0.25) + + fig.tight_layout() + fig.savefig(save_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +def _refine(triang, values, subdiv: int): + """Cubic-interpolated uniform refinement -> smooth fields (e.g. node lines).""" + import matplotlib.tri as mtri + + refiner = mtri.UniformTriRefiner(triang) + interp = mtri.CubicTriInterpolator(triang, values) + return refiner.refine_field(values, interp, subdiv=subdiv) + + +def plot_field(res: Dict, problem: Resonator, save_path: str, + p_ref: float = 20e-6, subdiv: int = 3) -> None: + """Acoustic field at ``f_field``: total acoustic pressure ``Re(p)`` (Pa) and + total sound pressure level (dB). + + Both panels are drawn on a cubically-refined triangulation so the pressure + node (|p| -> 0, where the SPL plunges) renders as a smooth line instead of a + jagged one. SPL uses the RMS pressure against ``p_ref`` (20 uPa in air): + ``Lp = 20 log10(|p| / (sqrt(2) p_ref))``. + """ + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import matplotlib.tri as mtri + import numpy as np + + pts = res["points"].cpu().numpy() + tris = res["mesh"].cells["triangle"].cpu().numpy() + triang = mtri.Triangulation(pts[:, 0], pts[:, 1], triangles=tris) + p = res["p_field"].cpu().numpy() + f0 = res["f_field"] + + # refine Re and Im (smooth) on a shared finer mesh, then recombine + fine, p_re = _refine(triang, p.real, subdiv) + _, p_im = _refine(triang, p.imag, subdiv) + amp = np.hypot(p_re, p_im) + spl = 20.0 * np.log10(np.maximum(amp / np.sqrt(2.0), 1e-30) / p_ref) + lim = float(np.abs(p_re).max()) + + xmax = problem.total_length + panels = [ + ("Total acoustic pressure (Pa)", p_re, "RdBu_r", "Pa", + dict(vmin=-lim, vmax=lim)), + ("Total sound pressure level (dB)", spl, "jet", "dB", + dict(vmin=float(spl.min()), vmax=float(spl.max()))), + ] + fig, axes = plt.subplots(2, 1, figsize=(8.5, 6.8)) + for ax, (title, data, cmap, unit, norm) in zip(axes, panels): + tpc = ax.tripcolor(fine, data, shading="gouraud", cmap=cmap, + rasterized=True, **norm) + ax.set_title(title) + ax.set_aspect("equal") + ax.set_xlim(-0.005, xmax + 0.025) + ax.set_ylim(-0.005, problem.duct_height + 0.005) + ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)") + cb = plt.colorbar(tpc, ax=ax, shrink=0.9, pad=0.02) + cb.ax.set_title(unit, fontsize=9) + + fig.tight_layout(rect=(0, 0, 1, 0.96)) + # center the suptitle over the plot axes (not the figure, which the right + # color-bars would otherwise skew left) + pos = axes[0].get_position() + fig.suptitle(f"Helmholtz resonator @ {f0:.0f} Hz", + x=pos.x0 + 0.5 * pos.width, y=0.99) + fig.savefig(save_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 6. Boundary-condition schematic +# --------------------------------------------------------------------------- # +def plot_boundary_conditions(problem: Resonator, save_path: str) -> None: + """Schematic of the geometry with the boundary conditions labeled. + + Inlet (x = 0): plane-wave port (absorbing + incident source). Every other + outer wall: sound-hard (rigid, ``dp/dn = 0``). The neck opening is an + interior interface (continuity), not a boundary. + """ + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.patches import Rectangle + from matplotlib.lines import Line2D + + D, Dh = problem.duct_length, problem.duct_height + Ln, Nh = problem.neck_length, problem.neck_height + Lc, Ch = problem.cavity_length, problem.cavity_height + PORT, HARD = "#D62728", "#1F5FBF" # red = port, blue = sound-hard + xmax = D + Ln + Lc + + fig, ax = plt.subplots(figsize=(10, 6.4)) + for x, y, w, h in ((0, 0, D, Dh), (D, 0, Ln, Nh), (D + Ln, 0, Lc, Ch)): + ax.add_patch(Rectangle((x, y), w, h, fc="#e5eefb", ec="none", zorder=1)) + + segs = [ + ((0, 0), (0, Dh), "port"), # inlet + ((0, Dh), (D, Dh), "hard"), # duct top + ((D, Dh), (D, Nh), "hard"), # duct right, above neck + ((D, Nh), (D + Ln, Nh), "hard"), # neck top + ((D + Ln, Nh), (D + Ln, Ch), "hard"), # cavity left, above neck + ((D + Ln, Ch), (D + Ln + Lc, Ch), "hard"), # cavity top + ((D + Ln + Lc, Ch), (D + Ln + Lc, 0), "hard"), # cavity right + ((D + Ln + Lc, 0), (0, 0), "hard"), # bottom + ] + for (x0, y0), (x1, y1), kind in segs: + ax.plot([x0, x1], [y0, y1], "-", + color=PORT if kind == "port" else HARD, lw=4, zorder=5, + solid_capstyle="round") + + ax.text(0.5 * D, 0.5 * Dh, "duct\n(air)", ha="center", va="center", + fontsize=12, color="#1e293b") + ax.text(D + Ln + 0.5 * Lc, 0.5 * Ch, "cavity", ha="center", va="center", + fontsize=12, color="#1e293b") + + ax.legend(handles=[ + Line2D([0], [0], color=PORT, lw=4, label="plane-wave port (inlet $x=0$)"), + Line2D([0], [0], color=HARD, lw=4, label=r"sound-hard walls ($\partial p/\partial n = 0$)"), + ], loc="upper right", fontsize=9, framealpha=0.96) + + ax.set_aspect("equal") + ax.set_xlim(-0.04 * xmax, xmax * 1.06) + ax.set_ylim(-0.06 * Dh, Dh * 1.12) + ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)") + ax.set_title("Helmholtz resonator — port-driven duct with side cavity", fontsize=13) + ax.grid(True, color="#f1f5f9", lw=0.6); ax.set_axisbelow(True) + fig.tight_layout() + fig.savefig(save_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 7. run_demo / main +# --------------------------------------------------------------------------- # +def run_demo(*, make_plot: bool = True, + impedance_path: Optional[str] = None, + field_path: Optional[str] = None, + bc_path: Optional[str] = None, + problem: Optional[Resonator] = None) -> Dict: + """Run the case, save the two figures, print a summary, return diagnostics.""" + problem = problem or Resonator() + res = solve_spectrum(problem) + + import numpy as np + f = res["freqs"].cpu().numpy() + probe = res["probe"].cpu().numpy() + ipk = int(np.argmax(probe)) + print(f"resonance: f = {f[ipk]:.0f} Hz, " + f"|p|_cavity = {probe[ipk]:.2f} Pa (incident p0 = {problem.p0:g})", + flush=True) + + if make_plot: + plot_boundary_conditions(problem, bc_path or str(HERE / "boundary_conditions.png")) + plot_impedance(res, impedance_path or str(HERE / "impedance.png")) + plot_field(res, problem, field_path or str(HERE / "acoustic_field.png")) + return res + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--no-plot", action="store_true") + parser.add_argument("--mesh-size", type=float, default=Resonator.mesh_size) + parser.add_argument("--f-field", type=float, default=Resonator.f_field, + help="frequency (Hz) for the field snapshot") + parser.add_argument("--impedance-output", type=str, default=None) + parser.add_argument("--field-output", type=str, default=None) + args = parser.parse_args() + + torch.set_default_dtype(torch.float64) + problem = Resonator(mesh_size=args.mesh_size, f_field=args.f_field) + run_demo(make_plot=not args.no_plot, + impedance_path=args.impedance_output, + field_path=args.field_output, + problem=problem) + + +if __name__ == "__main__": + main() diff --git a/examples/wave/optical_ring_resonator/README.md b/examples/wave/optical_ring_resonator/README.md new file mode 100644 index 0000000..7226123 --- /dev/null +++ b/examples/wave/optical_ring_resonator/README.md @@ -0,0 +1,65 @@ +# Optical ring resonator (waveguide-coupled microdisk) + +A 2D **photonic** example: a silicon bus waveguide evanescently couples light +into a silicon **microdisk** (both $n=3.48$) embedded in silica ($n=1.44$), with +a stretched-coordinate **PML** frame making the domain open. On resonance the +disk lights up in a **whispering-gallery mode** (WGM) — a bright ring of +azimuthal lobes on the rim. + +It exercises the open-domain wave operators +[`tensormesh.cartesian_pml`](../../../tensormesh/pml.py) + +`AnisotropicLaplaceElementAssembler` + `ScaledMassElementAssembler`. Example-only: +no public API is added. + +## Model + +In TM polarization the out-of-plane field $E_z$ obeys a scalar Helmholtz +equation with the PML's complex coordinate stretch $s_x, s_y$: + +$$\nabla\cdot(\mathbf\Lambda\,\nabla E_z) + k_0^2\,\varepsilon_r\, s_x s_y\, E_z = \text{source}, +\qquad \mathbf\Lambda = \mathrm{diag}(s_y/s_x,\ s_x/s_y).$$ + +Outside the PML $s_x=s_y=1$ and this is the ordinary Helmholtz equation; +$\varepsilon_r$ is $n_\text{core}^2$ in the waveguide and disk, $n_\text{clad}^2$ +elsewhere. A **modal soft source** launches the guided mode: the analytic +fundamental transverse profile of the bus slab (effective index $n_\text{eff}$ +from the even-mode dispersion $k_x\tan(k_x w/2)=\gamma$) is imprinted on a thin +launch plane and phased as a two-plane directional launch, so power travels +toward the disk instead of radiating both ways from a current blob. + +```text +mesh (gmsh, conforming) -> assemble (K_PML - k0^2 M) -> complex solve -> E_z +``` + +The example uses an inline `ElementAssembler` that evaluates $\varepsilon_r$ and +the stretch **per quadrature point** (sharper at the Si/SiO2 interfaces than +nodal coefficients); the reusable library path is +`cartesian_pml` + `AnisotropicLaplaceElementAssembler` + +`ScaledMassElementAssembler`, which reproduces the same operator and matches the +2D Hankel Green function to correlation 0.994. + +The PML validity was checked against a purely-absorbing reference: without it a +PEC/hard box traps the field into a standing wave (~8x the amplitude). + +## Run + +```bash +python optical_ring_resonator.py +``` + +`run_demo(...)` returns diagnostics; `main()` exposes `--lam0-nm`, `--order` +(1 or 2), `--mesh-h-nm`, `--no-plot`, `--output`. + +**On resonance vs off resonance.** The disk WGM is a sharp, high-Q resonance, so +the drive wavelength must sit on it. With linear (P1) elements numerical +dispersion blue-shifts the physical ~1.55 µm resonance to ~1.512 µm (the default +`lam0`); driving there gives a clean rim mode, while off resonance the disk fills +with a messy radial mix. Refining the mesh or `--order 2` moves the numerical +resonance back toward 1.55 µm. + +## What it shows + +A two-panel figure of the on-resonance field: `Re(E_z)` (the guided mode feeding +a ring of alternating rim lobes) and `|E_z|` (the bright whispering-gallery ring, +the waveguide depleted as power couples into the disk), both decaying into the +PML frame. diff --git a/examples/wave/optical_ring_resonator/optical_ring_resonator.py b/examples/wave/optical_ring_resonator/optical_ring_resonator.py new file mode 100644 index 0000000..256f568 --- /dev/null +++ b/examples/wave/optical_ring_resonator/optical_ring_resonator.py @@ -0,0 +1,463 @@ +"""2D photonic waveguide + microdisk coupler on TensorMesh — TM / E_z. + +A silicon bus waveguide evanescently couples light into a silicon microdisk +(both n = 3.48) embedded in silica (n = 1.44); a stretched-coordinate PML frame +makes the domain open. In TM polarization the out-of-plane field E_z obeys the +scalar Helmholtz equation + + div( Lambda grad E_z ) + k0^2 eps_r sx sy E_z = source , + +where the PML gives the complex diagonal tensor ``Lambda = diag(sy/sx, sx/sy)`` +and mass scaling ``sx sy`` (both 1 outside the PML). This is a plain scalar +Helmholtz with a spatially-varying complex coefficient, assembled natively by a +custom :class:`~tensormesh.ElementAssembler` whose ``forward`` reads the +quadrature coordinate and evaluates ``eps_r`` (Si in the waveguide + disk, SiO2 +elsewhere) and the PML stretch on the fly. The waveguide is fed by a **modal +soft source**: the analytic fundamental transverse profile of the bus slab is +imprinted on a thin launch plane and phased as a two-plane directional launch, +so it injects the guided mode travelling toward the disk (not an isotropic +current blob). On resonance it feeds a whispering-gallery mode (WGM) of the +disk — a bright ring of azimuthal lobes on the rim. + +TensorMesh now ships a native PML (``tensormesh.cartesian_pml`` + +``AnisotropicLaplaceElementAssembler`` + ``ScaledMassElementAssembler``); this +example keeps an equivalent inline assembler that evaluates ``eps_r`` and the +stretch *per quadrature point*, which resolves the sharp Si/SiO2 interfaces a +touch better than nodal coefficients. No public API is added here. + +The disk WGM is a sharp resonance, so the drive wavelength must sit on it: with +linear (P1) elements the physical ~1.55 um resonance blue-shifts to ~1.512 um +(numerical dispersion), which is the default ``lam0``. Off resonance the disk +fills with a messy radial mix instead of a clean rim mode; refining the mesh or +using ``mesh_order=2`` moves the numerical resonance back toward 1.55 um. + +Workflow: mesh (gmsh, conforming) -> assemble (K_PML - k0^2 M) -> complex solve + -> E_z field (guided mode + disk WGM). + +Run (env with torch + tensormesh + gmsh): + python examples/wave/optical_ring_resonator/optical_ring_resonator.py +""" +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Optional, Tuple + +import numpy as np +import torch + +ROOT = Path(__file__).resolve().parents[3] +sys.path.append(str(ROOT)) + +from tensormesh import ElementAssembler, MassElementAssembler, Mesh +from tensormesh.sparse.matrix import SparseMatrix + +HERE = Path(__file__).resolve().parent +C0 = 299792458.0 + + +# --------------------------------------------------------------------------- # +# Problem definition +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Coupler: + r"""All parameters of the waveguide + microdisk coupler (SI, meters). + + Fields + ------ + domain, pml : full domain size and PML-frame thickness (m). + wg_width, wg_x : bus-waveguide width and center x (m). + disk_r, disk_x, disk_y : microdisk radius and center (m). + src_y : y of the modal launch plane on the waveguide (m). + launch_sigma : longitudinal 1/e half-width of the soft-source window (m); + defaults to ``1.5 * h``. + directional : if True, use a two-plane launch (a second plane a quarter + guided-wavelength ahead, phased by +i) so the mode radiates toward the + disk (+y) and the backward wave cancels; if False, a single symmetric + plane launches both ways. + n_core, n_clad : Si and SiO2 refractive indices. + lam0 : drive wavelength (m). Default 1.512 um — the disk's whispering-gallery + resonance for the default mesh. The ~1.55 um physical resonance + blue-shifts numerically with linear (P1) elements; drive at the numerical + resonance (or refine / use ``mesh_order=2``) to see a clean rim mode. + pml_strength : imaginary coordinate-stretch amplitude of the PML. + mesh_h : target element edge length (m); default lam0/n_core/12. + mesh_order : 1 (linear) or 2 (quadratic ``triangle6``) elements. + """ + + domain: float = 6.0e-6 + pml: float = 0.6e-6 + wg_width: float = 0.25e-6 + wg_x: float = 1.6e-6 + disk_r: float = 1.0e-6 + disk_x: float = 2.75e-6 + disk_y: float = 3.0e-6 + src_y: float = 0.6e-6 # launch at the PML/cladding interface (= pml) + launch_sigma: Optional[float] = None + directional: bool = True + n_core: float = 3.48 + n_clad: float = 1.44 + lam0: float = 1.512e-6 # disk WGM resonance for the default P1 mesh + pml_strength: float = 10.0 + mesh_h: Optional[float] = None + mesh_order: int = 1 + + @property + def k0(self) -> float: + """Vacuum wavenumber ``2 pi / lam0`` (rad/m).""" + return 2.0 * np.pi / self.lam0 + + @property + def freq(self) -> float: + """Frequency ``c0 / lam0`` (Hz).""" + return C0 / self.lam0 + + @property + def h(self) -> float: + """Resolved mesh edge length (m).""" + return self.mesh_h or self.lam0 / self.n_core / 12.0 + + @property + def sigma(self) -> float: + """Resolved longitudinal 1/e half-width of the launch window (m).""" + return self.launch_sigma or 1.5 * self.h + + +# --------------------------------------------------------------------------- # +# Custom TM-Maxwell (scalar Helmholtz) assembler with an on-the-fly PML +# --------------------------------------------------------------------------- # +class MaxwellTMAssembler(ElementAssembler): + r"""Integrand ``(Lambda grad u).grad v - k0^2 eps_r sx sy u v`` per point. + + ``eps_r`` and the PML stretch are evaluated from the quadrature coordinate + ``x``, so material interfaces and the PML profile are resolved at quadrature + level. Geometry / PML parameters are stashed on the instance by + :meth:`__post_init__`. + """ + + def __post_init__(self, geo: dict): + for k, v in geo.items(): + setattr(self, k, v) + + def _stretch(self, t): + """Complex SC-PML stretch ``s(t)``: 1 inside, ``1 - i*strength*xi^2`` in the PML.""" + d, L = self.pml, self.domain + xi = torch.clamp((d - t) / d, min=0.0) + torch.clamp((t - (L - d)) / d, min=0.0) + return 1.0 - 1j * self.pml_strength * xi * xi + + def forward(self, gradu, gradv, u, v, x): + xx, yy = x[0].real, x[1].real + in_core = (torch.abs(xx - self.wg_x) <= self.wg_width / 2) | \ + ((xx - self.disk_x) ** 2 + (yy - self.disk_y) ** 2 <= self.disk_r ** 2) + eps = torch.where(in_core, + torch.as_tensor(self.eps_core, dtype=torch.complex128), + torch.as_tensor(self.eps_clad, dtype=torch.complex128)) + sx, sy = self._stretch(xx), self._stretch(yy) + z = torch.zeros_like(sx) + lam = torch.stack([torch.stack([sy / sx, z]), torch.stack([z, sx / sy])]) + eps_s = (self.k0 ** 2) * eps * sx * sy + return gradu.to(torch.complex128) @ lam @ gradv.to(torch.complex128) - eps_s * u * v + + +# --------------------------------------------------------------------------- # +# 1. Mesh: square domain conforming to the waveguide / disk interfaces +# --------------------------------------------------------------------------- # +def build_mesh(problem: Coupler, msh_path: Optional[str] = None) -> Mesh: + """gmsh mesh conforming to the waveguide/disk/source interfaces.""" + import gmsh + + if msh_path is None: + msh_path = str(HERE / "_coupler.msh") + L, h = problem.domain, problem.h + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + occ = gmsh.model.occ + sq = occ.addRectangle(0, 0, 0, L, L) + wg = occ.addRectangle(problem.wg_x - problem.wg_width / 2, 0, 0, problem.wg_width, L) + dk = occ.addDisk(problem.disk_x, problem.disk_y, 0, problem.disk_r, problem.disk_r) + occ.fragment([(2, sq)], [(2, wg), (2, dk)]) + occ.synchronize() + gmsh.option.setNumber("Mesh.MeshSizeMin", h) + gmsh.option.setNumber("Mesh.MeshSizeMax", h) + gmsh.model.mesh.generate(2) + if problem.mesh_order >= 2: + gmsh.model.mesh.setOrder(problem.mesh_order) + gmsh.write(msh_path) + finally: + gmsh.finalize() + return Mesh.from_file(msh_path, reorder=False) + + +# --------------------------------------------------------------------------- # +# 1b. Modal soft source: fundamental transverse mode of the bus waveguide +# --------------------------------------------------------------------------- # +def slab_mode(problem: Coupler) -> Tuple[float, float, float]: + r"""Fundamental even :math:`E_z` mode of the bus waveguide (symmetric slab). + + Solves the even-mode dispersion :math:`k_x \tan(k_x w/2) = \gamma` for the + guided effective index :math:`n_\mathrm{eff}\in(n_\mathrm{clad},n_\mathrm{core})`, + with transverse core wavenumber :math:`k_x = k_0\sqrt{n_\mathrm{core}^2-n_\mathrm{eff}^2}` + and cladding decay :math:`\gamma = k_0\sqrt{n_\mathrm{eff}^2-n_\mathrm{clad}^2}`. + Returns ``(n_eff, kx, gamma)`` (``kx, gamma`` in rad/m). + """ + import math + + k0, w = problem.k0, problem.wg_width + n1, n2 = problem.n_core, problem.n_clad + + def f(neff: float) -> float: + kx = k0 * math.sqrt(max(n1 * n1 - neff * neff, 0.0)) + ga = k0 * math.sqrt(max(neff * neff - n2 * n2, 0.0)) + # even-mode condition times cos(kx w/2): avoids the tan asymptote. + return kx * math.sin(kx * w / 2) - ga * math.cos(kx * w / 2) + + # Scan from just below n_core downward for the first (fundamental) root. + lo, hi, steps = n2 + 1e-9, n1 - 1e-9, 4000 + prev_n, prev_f, root = hi, f(hi), None + for i in range(1, steps + 1): + n = hi - (hi - lo) * i / steps + fn = f(n) + if prev_f * fn <= 0.0: + a, fa, b = n, fn, prev_n + for _ in range(100): # bisection + m = 0.5 * (a + b) + if fa * f(m) <= 0.0: + b = m + else: + a, fa = m, f(m) + root = 0.5 * (a + b) + break + prev_n, prev_f = n, fn + if root is None: + raise RuntimeError("no guided even mode found for the bus waveguide") + + kx = k0 * math.sqrt(n1 * n1 - root * root) + ga = k0 * math.sqrt(root * root - n2 * n2) + return root, kx, ga + + +def launch_field(problem: Coupler, pts: torch.Tensor, + neff: float, kx: float, ga: float) -> torch.Tensor: + r"""Complex nodal soft source imprinting the fundamental bus mode. + + The transverse profile :math:`\psi(x)` (cosine in the core, evanescent tails + in the cladding) is imprinted on a thin Gaussian window in ``y`` at the + launch plane ``src_y``. When ``directional`` is set, a second window a + quarter guided-wavelength ahead and phased by ``+i`` makes the ``+y`` + (toward-disk) wave add and the backward wave cancel. The consistent FEM + load is then ``rhs = M @ launch_field``. + """ + import math + + x = pts[:, 0] - problem.wg_x + y = pts[:, 1] + w, sig = problem.wg_width, problem.sigma + ax = torch.abs(x) + psi = torch.where(ax <= w / 2, + torch.cos(kx * x), + math.cos(kx * w / 2) * torch.exp(-ga * (ax - w / 2))) + + def band(yc: float) -> torch.Tensor: + return torch.exp(-((y - yc) / sig) ** 2) + + lam_g = 2.0 * math.pi / (problem.k0 * neff) # guided wavelength + win = band(problem.src_y).to(torch.complex128) + if problem.directional: + # second plane a quarter guided-wavelength ahead, phased by -i so the + # outgoing wave (e^{-i beta y} for this solver's convention) adds in +y + # (toward the disk) and cancels in -y. + win = win - 1j * band(problem.src_y + lam_g / 4.0) + return psi.to(torch.complex128) * win + + +# --------------------------------------------------------------------------- # +# 2. Assemble the PML Helmholtz operator, add the source, solve +# --------------------------------------------------------------------------- # +def solve(problem: Coupler, verbose: bool = True) -> Dict: + """Solve the TM Helmholtz coupler; return the complex field ``E_z``.""" + torch.set_default_dtype(torch.float64) + mesh = build_mesh(problem) + pts = mesh.points.to(torch.float64) + geo = dict(k0=problem.k0, eps_core=problem.n_core ** 2, eps_clad=problem.n_clad ** 2, + wg_x=problem.wg_x, wg_width=problem.wg_width, disk_x=problem.disk_x, + disk_y=problem.disk_y, disk_r=problem.disk_r, domain=problem.domain, + pml=problem.pml, pml_strength=problem.pml_strength) + + asm = MaxwellTMAssembler.from_mesh(mesh, quadrature_order=4, geo=geo) + asm.type(torch.float64) + H = asm(points=pts) + + # Modal soft source: imprint the fundamental bus-waveguide mode on a thin + # launch plane and apply it as the consistent FEM load rhs = M @ psi. + neff, kx, ga = slab_mode(problem) + Masm = MassElementAssembler.from_mesh(mesh, quadrature_order=4) + Msp = Masm(pts) + src = launch_field(problem, pts, neff, kx, ga) + rhs = Msp @ src + + if verbose: + print(f"mesh: {mesh.n_points} nodes, {mesh.n_elements} triangles; " + f"bus mode n_eff = {neff:.4f} ({'directional' if problem.directional else 'symmetric'} " + f"launch); lam0 = {problem.lam0*1e9:.0f} nm", flush=True) + + ez = SparseMatrix(H.values, H.row, H.col, H.shape).solve(rhs) + return dict(mesh=mesh, points=pts.cpu().numpy(), Ez=ez.cpu().numpy()) + + +# --------------------------------------------------------------------------- # +# 3. Plot the field +# --------------------------------------------------------------------------- # +def plot_field(res: Dict, problem: Coupler, save_path: str) -> None: + """Draw ``Re(E_z)`` and ``|E_z|`` of the coupler (guided mode + disk WGM).""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.tri import Triangulation + + pts = res["points"] + cells = res["mesh"].cells + key = "triangle6" if "triangle6" in list(cells.keys()) else "triangle" + tris = cells[key].cpu().numpy()[:, :3] + tri = Triangulation(pts[:, 0] * 1e6, pts[:, 1] * 1e6, triangles=tris) + ez = res["Ez"] + ez = ez / np.abs(ez).max() # normalize (source strength arbitrary) + + fig, axes = plt.subplots(1, 2, figsize=(12, 5.2)) + panels = [("Re($E_z$)", ez.real, "RdBu_r", dict(vmin=-1, vmax=1)), + ("|$E_z$|", np.abs(ez), "inferno", dict(vmin=0, vmax=0.5))] + for ax, (title, data, cmap, norm) in zip(axes, panels): + tpc = ax.tripcolor(tri, data, shading="gouraud", cmap=cmap, rasterized=True, **norm) + ax.set_aspect("equal") + ax.set_title(f"{title} @ {problem.lam0*1e9:.0f} nm") + ax.set_xlabel("x (µm)"); ax.set_ylabel("y (µm)") + cb = plt.colorbar(tpc, ax=ax, shrink=0.88) + cb.ax.set_title("$E_z$\n(norm.)", fontsize=9) + fig.suptitle("Waveguide-coupled silicon microdisk", y=1.0) + fig.tight_layout() + fig.savefig(save_path, dpi=140, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 3b. Plot the setup: materials, PML frame, source, boundary conditions +# --------------------------------------------------------------------------- # +def plot_setup(problem: Coupler, save_path: str) -> None: + """Draw the physical setup: materials (Si / SiO2), the PML frame, the + line-current launch, and the applied boundary conditions.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.patches import Circle, Rectangle + from matplotlib.lines import Line2D + + # Wong colorblind-safe qualitative palette. + c_clad = "#56B4E9" # SiO2 cladding (sky blue) + c_core = "#E69F00" # Si core (orange) + c_pml = "#999999" # PML frame (grey) + c_src = "#D55E00" # source (vermillion) + + L, d = problem.domain * 1e6, problem.pml * 1e6 + fig, ax = plt.subplots(figsize=(6.4, 6.4)) + + # SiO2 cladding fills the whole domain. + ax.add_patch(Rectangle((0, 0), L, L, facecolor=c_clad, edgecolor="none", + alpha=0.35, zorder=0)) + # PML frame: solid grey border of thickness `pml` on all four sides. + for xy, w, h in [((0, 0), L, d), ((0, L - d), L, d), + ((0, d), d, L - 2 * d), ((L - d, d), d, L - 2 * d)]: + ax.add_patch(Rectangle(xy, w, h, facecolor=c_pml, edgecolor="none", + alpha=0.30, zorder=1)) + + # Si core: bus waveguide (full-height strip) + microdisk, with a thin edge. + core_edge = "#8a6100" + wg_x0 = (problem.wg_x - problem.wg_width / 2) * 1e6 + ax.add_patch(Rectangle((wg_x0, 0), problem.wg_width * 1e6, L, + facecolor=c_core, edgecolor=core_edge, linewidth=0.8, + alpha=0.95, zorder=2)) + ax.add_patch(Circle((problem.disk_x * 1e6, problem.disk_y * 1e6), + problem.disk_r * 1e6, facecolor=c_core, edgecolor=core_edge, + linewidth=0.8, alpha=0.95, zorder=2)) + + # Modal launch plane (thin line across the waveguide) + toward-disk arrow. + xc, ys = problem.wg_x * 1e6, problem.src_y * 1e6 + hw = problem.wg_width * 1e6 / 2 # span exactly the waveguide width + ax.plot([xc - hw, xc + hw], [ys, ys], color=c_src, linewidth=2.4, zorder=4) + if problem.directional: + ax.annotate("", xy=(xc, ys + 0.55), xytext=(xc, ys), + arrowprops=dict(arrowstyle="-|>", color=c_src, linewidth=2.0), + zorder=5) + + # Labels. + ax.annotate("mode\nlaunch", (xc + hw + 0.15, ys), ha="left", va="center", + fontsize=8, color=c_src, zorder=5) + + ax.set_xlim(0, L); ax.set_ylim(0, L) + ax.set_aspect("equal") + ax.set_xlabel("x (µm)"); ax.set_ylabel("y (µm)") + ax.set_title("Simulation setup: materials & PML") + + legend = [Line2D([0], [0], marker="s", color="none", markerfacecolor=c_core, + markersize=11, label="Si core ($n=%.2f$)" % problem.n_core), + Line2D([0], [0], marker="s", color="none", markerfacecolor=c_clad, + markersize=11, alpha=0.5, label="SiO$_2$ ($n=%.2f$)" % problem.n_clad), + Line2D([0], [0], marker="s", color="none", markerfacecolor=c_pml, + markersize=11, alpha=0.5, label="PML (radiating BC)"), + Line2D([0], [0], color=c_src, linewidth=2.4, label="modal soft source")] + ax.legend(handles=legend, loc="upper right", fontsize=8, framealpha=0.9, + labelspacing=1.0, handletextpad=0.9, borderpad=0.9) + + fig.tight_layout() + fig.savefig(save_path, dpi=140, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 4. run_demo / main +# --------------------------------------------------------------------------- # +def run_demo(*, make_plot: bool = True, field_path: Optional[str] = None, + setup_path: Optional[str] = None, + problem: Optional[Coupler] = None) -> Dict: + """Solve the coupler, save the setup + field figures, return diagnostics.""" + problem = problem or Coupler() + if make_plot: + plot_setup(problem, setup_path or str(HERE / "optical_ring_resonator_setup.png")) + res = solve(problem) + ez = res["Ez"] + disk = (((res["points"][:, 0] - problem.disk_x) ** 2 + + (res["points"][:, 1] - problem.disk_y) ** 2) < problem.disk_r ** 2) + print(f"|E_z| max {np.abs(ez).max():.2e}; mean |E_z| in disk / whole = " + f"{np.abs(ez[disk]).mean() / np.abs(ez).mean():.2f} (coupling into the WGM)", + flush=True) + if make_plot: + plot_field(res, problem, field_path or str(HERE / "optical_ring_resonator.png")) + return res + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--lam0-nm", type=float, default=1512.0, + help="drive wavelength (nm); default = disk WGM resonance") + parser.add_argument("--mesh-h-nm", type=float, default=None, help="mesh edge (nm)") + parser.add_argument("--order", type=int, default=1, choices=(1, 2)) + parser.add_argument("--pml-strength", type=float, default=10.0) + parser.add_argument("--no-plot", action="store_true") + parser.add_argument("--output", type=str, default=None, help="field figure path") + parser.add_argument("--setup-output", type=str, default=None, + help="setup (materials/BC) figure path") + args = parser.parse_args() + + torch.set_default_dtype(torch.float64) + problem = Coupler(lam0=args.lam0_nm * 1e-9, + mesh_h=(args.mesh_h_nm * 1e-9) if args.mesh_h_nm else None, + mesh_order=args.order, pml_strength=args.pml_strength) + run_demo(make_plot=not args.no_plot, field_path=args.output, + setup_path=args.setup_output, problem=problem) + + +if __name__ == "__main__": + main() diff --git a/examples/wave/phononic_crystal_2d/README.md b/examples/wave/phononic_crystal_2d/README.md index 197c4a8..a3623ec 100644 --- a/examples/wave/phononic_crystal_2d/README.md +++ b/examples/wave/phononic_crystal_2d/README.md @@ -32,27 +32,27 @@ bands: mesh -> Laplace/Mass assembler -> BlochReducer -> generalized eig ``` -- Square (rigid): $K_{ij}=\int\nabla\phi_i\!\cdot\!\nabla\phi_j$, +- Square (rigid): $K_{ij}=\int\nabla\phi_i\cdot\nabla\phi_j$, $M_{ij}=\int\phi_i\phi_j$, $K p=(\omega/c)^2 M p$, $f=c\sqrt{\mu}/2\pi$; - path $M\!-\!\Gamma\!-\!X\!-\!M$. + path $M\to\Gamma\to X\to M$. - Triangular (penetrable): material varies in space, so the operators are - **weighted** — $K_{ij}=\int\frac1\rho\nabla\phi_i\!\cdot\!\nabla\phi_j$, + **weighted** — $K_{ij}=\int\frac1\rho\nabla\phi_i\cdot\nabla\phi_j$, $M_{ij}=\int\frac1{\rho c^2}\phi_i\phi_j$, $K p=\omega^2 M p$ — assembled with a per-element `ElementAssembler` carrying $1/\rho,\ 1/(\rho c^2)$ over a - conformal steel/water mesh; path $M\!-\!\Gamma\!-\!K\!-\!M$. + conformal steel/water mesh; path $M\to\Gamma\to K\to M$. Unit-cell meshes use gmsh `setPeriodic` (and `fragment` for the two-domain case) so opposite edges carry matching nodes — the precondition `BlochReducer` needs. **Transmission** — a normally-incident plane wave ($p_0=1$ Pa) crosses a finite -slab of rigid cylinders; the power transmission $T(f)=\langle|p|^2\rangle_{\rm -out}/|p_0|^2$ is swept over frequency: +slab of rigid cylinders; the power transmission $T(f)=\langle|p|^2\rangle_{\mathrm{out}}/|p_0|^2$ is swept +over frequency: ```text mesh -> Laplace/Mass assembler -> (K - k^2 M - i k B) p = -2 i k p0 e_in ``` -The first-order radiation term $B$ and the incident load $e_{\rm in}$ are +The first-order radiation term $B$ and the incident load $e_{\mathrm{in}}$ are hand-rolled boundary line integrals (no PML needed — this matches COMSOL's first-order "Plane Wave Radiation"). At normal incidence the lateral periodic walls are mirror-symmetry planes, equivalent to natural Neumann. @@ -65,8 +65,8 @@ next to the script, so the comparison reproduces **offline** — no COMSOL neede | Script | Reference | Agreement | | --- | --- | --- | -| `band_structure_square.py` | square $M\!-\!\Gamma\!-\!X\!-\!M$, 31 k-points | mean **0.08 %**, p95 0.16 % | -| `band_structure_triangular.py` | triangular $M\!-\!\Gamma\!-\!K\!-\!M$, 31 k-points | mean **0.51 %**, p95 1.45 % | +| `band_structure_square.py` | square $M\to\Gamma\to X\to M$, 31 k-points | mean **0.08 %**, p95 0.16 % | +| `band_structure_triangular.py` | triangular $M\to\Gamma\to K\to M$, 31 k-points | mean **0.51 %**, p95 1.45 % | | `transmission_slab.py` | 30–120 kHz, 46 frequencies | mean $|\Delta T|$ **0.007** | Bands are compared k-point by k-point at COMSOL's exact wavevectors (lowest 10 diff --git a/examples/wave/wave_guide_mode/README.md b/examples/wave/wave_guide_mode/README.md new file mode 100644 index 0000000..d27dbad --- /dev/null +++ b/examples/wave/wave_guide_mode/README.md @@ -0,0 +1,78 @@ +# Waveguide mode analysis — rectangular dielectric waveguide + +A 2D **photonic** example: the guided modes of a rectangular dielectric +waveguide, found from its cross-section — a high-index rectangular core +(width $w$ × height $h$, index $n_\text{core}$) in a lower-index cladding +($n_\text{clad}$). Example-only: no public API is added. + +## Model + +In the weakly-guiding (scalar) approximation the transverse field $E(x,y)$ and +propagation constant $\beta$ satisfy the transverse Helmholtz equation + +$$\nabla_t^2 E + \big(k_0^2\,n(x,y)^2 - \beta^2\big)E = 0,\qquad k_0 = 2\pi/\lambda_0,$$ + +a generalized eigenproblem for $\beta^2$. With the builtin stiffness $K$ +(`LaplaceElementAssembler`), mass $M$ (`MassElementAssembler`) and the +index-weighted mass $M_\varepsilon = \int n^2\,\phi_i\phi_j$ +(`ScaledMassElementAssembler`): + +$$\big(K - k_0^2 M_\varepsilon\big)\,E = -\beta^2 M\,E.$$ + +The field is clamped to zero on the outer box (guided fields decay in the +cladding), so the modes come from a shift-invert Lanczos solve near the core +light line $\beta^2 = (k_0 n_\text{core})^2$. A mode is **guided** when its +effective index $n_\text{eff} = \beta/k_0$ satisfies +$n_\text{clad} < n_\text{eff} < n_\text{core}$. + +```text +mesh (gmsh, core in box) -> assemble (K, M, M_eps) -> eig (beta^2) + -> guided modes (n_eff, fields) +``` + +## Run + +```bash +python waveguide_modes.py +``` + +`run_demo(...)` returns diagnostics; `main()` exposes `--core-w-um`, +`--core-h-um`, `--n-core`, `--n-clad`, `--lam0-nm`, `--n-modes`, `--mesh-h-um`, +`--no-plot`, `--output`. + +The mode count is set by the two normalized frequencies +$V_x = k_0\,(w/2)\,\mathrm{NA}$ and $V_y = k_0\,(h/2)\,\mathrm{NA}$ with +$\mathrm{NA}=\sqrt{n_\text{core}^2-n_\text{clad}^2}$. The default +($w=18\,\mu\mathrm{m}$, $h=14\,\mu\mathrm{m}$, $n_\text{core}=1.450$, +$n_\text{clad}=1.444$, $\lambda_0=1.55\,\mu\mathrm{m}$) gives +$V_x\approx4.81$, $V_y\approx3.74$ and seven guided modes (the gallery shows the +first six). + +## What it shows + +`waveguide_modes.png` — the guided transverse mode fields $E(x,y)$ (red/blue = +field sign, core outlined). Modes are labelled $E_{pq}$ by their lobe counts +along $x$ and $y$: $E_{11}$ (single lobe), $E_{21}$ / $E_{12}$ (dipoles along +$x$ / $y$), $E_{22}$ (quadrupole), $E_{31}$ / $E_{13}$ (three lobes across the +width / height). + +## Accuracy + +Each computed effective index is compared to the separable **Marcatili** +approximation: solve the symmetric-slab dispersion $U\tan U = W$ (even) / +$-U\cot U = W$ (odd), $U^2+W^2=(a\,k_0\,\mathrm{NA})^2$, independently across the +width and height, then combine $\beta^2 = k_0^2 n_\text{core}^2 - k_x^2 - k_y^2$. +FEM and Marcatili $n_\text{eff}$ agree to ~$10^{-4}$ (Marcatili is itself +approximate — it neglects the field in the core corners), and both give the same +guided-mode set: + +| mode | $n_\text{eff}$ (FEM) | $n_\text{eff}$ (Marcatili) | +|-----:|---------------------:|---------------------------:| +| $E_{11}$ | 1.44902 | 1.44891 | +| $E_{21}$ | 1.44784 | 1.44763 | +| $E_{12}$ | 1.44732 | 1.44703 | +| $E_{22}$ | 1.44614 | 1.44574 | +| $E_{31}$ | 1.44595 | 1.44560 | +| $E_{13}$ | 1.44478 | 1.44430 | + +Refining the mesh (`--mesh-h-um`) or box (`box_factor`) tightens the FEM values. diff --git a/examples/wave/wave_guide_mode/waveguide_modes.py b/examples/wave/wave_guide_mode/waveguide_modes.py new file mode 100644 index 0000000..89a47d8 --- /dev/null +++ b/examples/wave/wave_guide_mode/waveguide_modes.py @@ -0,0 +1,343 @@ +"""Rectangular dielectric waveguide mode analysis on TensorMesh — scalar Helmholtz. + +The guided modes of a rectangular dielectric waveguide are found from its 2D +cross-section: a high-index rectangular core (width ``w`` x height ``h``, index +``n_core``) embedded in a lower-index cladding (``n_clad``). In the +weakly-guiding (scalar) approximation the modal field ``E(x, y)`` and propagation +constant ``beta`` obey the transverse Helmholtz equation + + lap_t E + ( k0^2 n(x,y)^2 - beta^2 ) E = 0 , k0 = 2 pi / lam0 , + +a generalized eigenproblem for ``beta^2``. Assembled with the builtin +:class:`~tensormesh.assemble.LaplaceElementAssembler` (stiffness ``K``), +:class:`~tensormesh.assemble.MassElementAssembler` (mass ``M``) and +:class:`~tensormesh.assemble.ScaledMassElementAssembler` (index-weighted mass +``M_eps = int n^2 phi_i phi_j``) it reads + + ( K - k0^2 M_eps ) E = -beta^2 M E . + +The field is clamped to zero on the outer box (guided modes decay in the +cladding), so the eigenpairs come from a shift-invert Lanczos solve near the core +light line ``beta^2 = (k0 n_core)^2``. A mode is **guided** when its effective +index ``n_eff = beta / k0`` lies between the cladding and core indices, +``n_clad < n_eff < n_core``. Each is compared to the separable (Marcatili) slab +approximation. No public API is added here. + +Workflow: mesh (gmsh, rectangular core in box) -> assemble (K, M, M_eps) + -> generalized eig (beta^2) -> guided modes (n_eff + fields). + +Run (env with torch + tensormesh + gmsh + scipy): + python examples/wave/wave_guide_mode/waveguide_modes.py +""" +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +import numpy as np +import scipy.sparse.linalg as spla +from scipy.optimize import brentq +import torch + +ROOT = Path(__file__).resolve().parents[3] +sys.path.append(str(ROOT)) + +from tensormesh import Mesh +from tensormesh.assemble import (LaplaceElementAssembler, MassElementAssembler, + ScaledMassElementAssembler) + +HERE = Path(__file__).resolve().parent + + +# --------------------------------------------------------------------------- # +# Problem definition +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class Waveguide: + r"""Rectangular step-index waveguide parameters (SI, meters). + + Fields + ------ + core_w, core_h : core width and height (m). + n_core, n_clad : core / cladding refractive indices. + lam0 : vacuum wavelength (m). + box_factor : outer (square) box half-width in units of ``max(w, h)``. + n_modes : number of eigenpairs to request from the solver. + mesh_h : target element edge length (m); default ``min(w, h) / 12``. + """ + + core_w: float = 18.0e-6 + core_h: float = 14.0e-6 + n_core: float = 1.450 + n_clad: float = 1.444 + lam0: float = 1.55e-6 + box_factor: float = 2.0 + n_modes: int = 12 + mesh_h: Optional[float] = None + + @property + def k0(self) -> float: + """Vacuum wavenumber ``2 pi / lam0`` (rad/m).""" + return 2.0 * np.pi / self.lam0 + + @property + def NA(self) -> float: + """Numerical aperture ``sqrt(n_core^2 - n_clad^2)``.""" + return float(np.sqrt(self.n_core ** 2 - self.n_clad ** 2)) + + @property + def box(self) -> float: + """Outer box half-width (m).""" + return self.box_factor * max(self.core_w, self.core_h) + + @property + def Vx(self) -> float: + """Normalized frequency across the width ``k0 (w/2) NA``.""" + return self.k0 * (self.core_w / 2.0) * self.NA + + @property + def Vy(self) -> float: + """Normalized frequency across the height ``k0 (h/2) NA``.""" + return self.k0 * (self.core_h / 2.0) * self.NA + + @property + def h(self) -> float: + """Resolved mesh edge length (m).""" + return self.mesh_h or min(self.core_w, self.core_h) / 10.0 + + +# --------------------------------------------------------------------------- # +# 1. Mesh: rectangular core centered in a square cladding box +# --------------------------------------------------------------------------- # +def build_mesh(problem: Waveguide, msh_path: Optional[str] = None) -> Mesh: + """gmsh mesh of the square box with the core rectangle as an interior interface.""" + import gmsh + + if msh_path is None: + msh_path = str(HERE / "_waveguide.msh") + b, w, ht, hh = problem.box, problem.core_w, problem.core_h, problem.h + gmsh.initialize() + try: + gmsh.option.setNumber("General.Terminal", 0) + occ = gmsh.model.occ + box = occ.addRectangle(-b, -b, 0, 2 * b, 2 * b) + core = occ.addRectangle(-w / 2, -ht / 2, 0, w, ht) + occ.fragment([(2, box)], [(2, core)]) + occ.synchronize() + gmsh.option.setNumber("Mesh.MeshSizeMin", hh) + gmsh.option.setNumber("Mesh.MeshSizeMax", hh) + gmsh.model.mesh.generate(2) + gmsh.write(msh_path) + finally: + gmsh.finalize() + return Mesh.from_file(msh_path, reorder=False) + + +# --------------------------------------------------------------------------- # +# 2. Assemble K, M and the index-weighted mass M_eps = int n^2 phi phi +# --------------------------------------------------------------------------- # +def assemble(problem: Waveguide, mesh: Mesh): + """Return CSC ``(A, B)`` with ``A = K - k0^2 M_eps``, ``B = M``, + free mask.""" + pts = mesh.points.to(torch.float64) + xy = pts.cpu().numpy() + + K = LaplaceElementAssembler.from_mesh(mesh)(pts).to_scipy_coo().tocsc() + M = MassElementAssembler.from_mesh(mesh, quadrature_order=4)(pts).to_scipy_coo().tocsc() + + # nodal n^2: n_core^2 inside the core rectangle, n_clad^2 in the cladding. + in_core = (np.abs(xy[:, 0]) <= problem.core_w / 2) & \ + (np.abs(xy[:, 1]) <= problem.core_h / 2) + n2 = np.where(in_core, problem.n_core ** 2, problem.n_clad ** 2) + c = torch.as_tensor(n2, dtype=torch.float64) + Meps = ScaledMassElementAssembler.from_mesh(mesh, quadrature_order=4)( + pts, point_data={"c": c}).to_scipy_coo().tocsc() + + A = (K - problem.k0 ** 2 * Meps).tocsc() + + # clamp the outer box boundary (guided fields decay to zero there). + tol = problem.h * 0.5 + on_box = (np.abs(np.abs(xy[:, 0]) - problem.box) < tol) | \ + (np.abs(np.abs(xy[:, 1]) - problem.box) < tol) + free = np.where(~on_box)[0] + return A, M, free + + +# --------------------------------------------------------------------------- # +# 3. Solve the generalized eigenproblem for the guided modes +# --------------------------------------------------------------------------- # +def solve_modes(problem: Waveguide, A, B, free: np.ndarray) -> Dict: + """Guided modes: effective indices ``n_eff`` (desc) + full-length fields.""" + Aff = A[free][:, free] + Bff = B[free][:, free] + # (K - k0^2 M_eps) E = -beta^2 M E ; the most-guided modes sit just below the + # core light line, so shift-invert near lambda = -(k0 n_core)^2. + sigma = -(problem.k0 * problem.n_core) ** 2 + lam, vec = spla.eigsh(Aff, k=problem.n_modes, M=Bff, sigma=sigma, which="LM") + + beta2 = -lam + n_eff = np.sqrt(np.clip(beta2, 0.0, None)) / problem.k0 + guided = (n_eff > problem.n_clad) & (n_eff < problem.n_core) + + order = np.argsort(n_eff[guided])[::-1] # most-confined first + idx = np.where(guided)[0][order] + n = A.shape[0] + fields = np.zeros((len(idx), n)) + for j, i in enumerate(idx): + f = np.zeros(n) + f[free] = vec[:, i] + f /= np.abs(f).max() # normalize sign/amplitude + fields[j] = f + return dict(n_eff=n_eff[idx], fields=fields) + + +# --------------------------------------------------------------------------- # +# 4. Analytic reference: separable (Marcatili) symmetric-slab approximation +# --------------------------------------------------------------------------- # +def _slab_kx(half_width: float, k0: float, NA: float) -> List[float]: + r"""Transverse wavenumbers ``kx`` of a symmetric slab of full width ``2 a``. + + Even / odd guided modes satisfy ``U tan U = W`` / ``-U cot U = W`` with + ``U = kx a``, ``W = gamma a`` and ``U^2 + W^2 = (a k0 NA)^2``. + """ + Vd = half_width * k0 * NA + if Vd <= 0: + return [] + grid = np.linspace(1e-6, Vd - 1e-9, 6000) + + def g(u): # even (+) and odd (-) merged + w = np.sqrt(max(Vd * Vd - u * u, 0.0)) + return u * np.tan(u) - w, u / np.tan(u) + w # (even, odd) residuals + + kxs: List[float] = [] + for sel in (0, 1): # even then odd branch + vals = np.array([g(u)[sel] for u in grid]) + for i in range(len(grid) - 1): + a0, a1 = vals[i], vals[i + 1] + if np.isfinite(a0) and np.isfinite(a1) and a0 * a1 < 0 \ + and abs(a0) < 10 and abs(a1) < 10: # skip tan/cot poles + try: + u = brentq(lambda u: g(u)[sel], grid[i], grid[i + 1], maxiter=200) + kxs.append(u / half_width) + except ValueError: + pass + return sorted(kxs) + + +def analytic_slab(problem: Waveguide) -> List[dict]: + """Guided ``E_{pq}`` modes from the separable slab product ``kx * ky``.""" + kx = _slab_kx(problem.core_w / 2, problem.k0, problem.NA) + ky = _slab_kx(problem.core_h / 2, problem.k0, problem.NA) + kco = problem.k0 * problem.n_core + kcl = problem.k0 * problem.n_clad + out: List[dict] = [] + for p, kxi in enumerate(kx, 1): + for q, kyj in enumerate(ky, 1): + beta2 = kco ** 2 - kxi ** 2 - kyj ** 2 + if beta2 > kcl ** 2: # still guided + out.append(dict(p=p, q=q, n_eff=np.sqrt(beta2) / problem.k0)) + return sorted(out, key=lambda d: -d["n_eff"]) + + +# --------------------------------------------------------------------------- # +# 5. Plot the guided mode-field gallery +# --------------------------------------------------------------------------- # +def plot_modes(res: Dict, problem: Waveguide, mesh: Mesh, save_path: str, + labels: Optional[List[str]] = None, n_show: int = 6) -> None: + """Gallery of the guided transverse mode fields ``E(x, y)``.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.patches import Rectangle + from matplotlib.tri import Triangulation + + pts = mesh.points.cpu().numpy() + tris = mesh.cells["triangle"].cpu().numpy()[:, :3] + tri = Triangulation(pts[:, 0] * 1e6, pts[:, 1] * 1e6, triangles=tris) + + n_show = min(n_show, len(res["fields"])) + ncol = 3 + nrow = int(np.ceil(n_show / ncol)) + fig, axes = plt.subplots(nrow, ncol, figsize=(3.8 * ncol, 3.5 * nrow), + squeeze=False) + w_um, h_um = problem.core_w * 1e6, problem.core_h * 1e6 + lim = 1.6 * max(w_um, h_um) / 2 + for k in range(nrow * ncol): + ax = axes[k // ncol][k % ncol] + if k >= n_show: + ax.axis("off"); continue + e = res["fields"][k] + ax.tripcolor(tri, e, shading="gouraud", cmap="RdBu_r", + vmin=-1, vmax=1, rasterized=True) + ax.add_patch(Rectangle((-w_um / 2, -h_um / 2), w_um, h_um, + fill=False, ec="k", lw=0.8)) + ax.set_aspect("equal") + ax.set_xlim(-lim, lim); ax.set_ylim(-lim, lim) + tag = f"${labels[k]}$" if labels and k < len(labels) and labels[k] else "" + ax.set_title(f"mode {k + 1}: {tag}\n$n_{{eff}}$ = {res['n_eff'][k]:.5f}", + fontsize=10) + ax.set_xlabel("x (µm)", fontsize=8); ax.set_ylabel("y (µm)", fontsize=8) + fig.suptitle("Rectangular waveguide modes", y=1.0, fontsize=13) + fig.tight_layout() + fig.savefig(save_path, dpi=140, bbox_inches="tight") + plt.close(fig) + print(f"saved {save_path}", flush=True) + + +# --------------------------------------------------------------------------- # +# 6. run_demo / main +# --------------------------------------------------------------------------- # +def run_demo(*, make_plot: bool = True, field_path: Optional[str] = None, + problem: Optional[Waveguide] = None) -> Dict: + """Solve the guided modes, save the field gallery, print a summary.""" + torch.set_default_dtype(torch.float64) + problem = problem or Waveguide() + + mesh = build_mesh(problem) + A, B, free = assemble(problem, mesh) + print(f"mesh: {mesh.n_points} nodes, {mesh.n_elements} triangles; " + f"Vx = {problem.Vx:.3f}, Vy = {problem.Vy:.3f}", flush=True) + + modes = solve_modes(problem, A, B, free) + res = dict(mesh=mesh, **modes) + ana = analytic_slab(problem) + + print(f"\nguided modes: {len(modes['n_eff'])} FEM / " + f"{len(ana)} analytic (Marcatili)", flush=True) + print(f"{'#':>3} {'n_eff (FEM)':>12} {'Marcatili':>11} mode", flush=True) + for i, ne in enumerate(modes["n_eff"]): + tag = f"{ana[i]['n_eff']:11.5f} E{ana[i]['p']}{ana[i]['q']}" \ + if i < len(ana) else "" + print(f"{i + 1:>3} {ne:12.5f} {tag}", flush=True) + + labels = [f"E_{{{d['p']}{d['q']}}}" for d in ana] + if make_plot: + plot_modes(res, problem, mesh, + field_path or str(HERE / "waveguide_modes.png"), labels=labels) + return res + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--core-w-um", type=float, default=18.0, help="core width (µm)") + parser.add_argument("--core-h-um", type=float, default=14.0, help="core height (µm)") + parser.add_argument("--n-core", type=float, default=1.450) + parser.add_argument("--n-clad", type=float, default=1.444) + parser.add_argument("--lam0-nm", type=float, default=1550.0) + parser.add_argument("--n-modes", type=int, default=12) + parser.add_argument("--mesh-h-um", type=float, default=None) + parser.add_argument("--no-plot", action="store_true") + parser.add_argument("--output", type=str, default=None) + args = parser.parse_args() + + problem = Waveguide(core_w=args.core_w_um * 1e-6, core_h=args.core_h_um * 1e-6, + n_core=args.n_core, n_clad=args.n_clad, + lam0=args.lam0_nm * 1e-9, n_modes=args.n_modes, + mesh_h=(args.mesh_h_um * 1e-6) if args.mesh_h_um else None) + run_demo(make_plot=not args.no_plot, field_path=args.output, problem=problem) + + +if __name__ == "__main__": + main() diff --git a/tensormesh/__init__.py b/tensormesh/__init__.py index b14b725..45fc6d9 100644 --- a/tensormesh/__init__.py +++ b/tensormesh/__init__.py @@ -1,5 +1,5 @@ from .mesh import Mesh -from .operator import Condenser, BlochReducer +from .operator import Condenser, BlochReducer, robin_operator, port_source # from .element import get_shape_val, get_shape_grad, get_basis from .element import Transformation,\ Element,\ @@ -14,8 +14,8 @@ from .element.element_type2dimension import element_type2dimension from .element import element_type2element,\ element_types -from .assemble import ElementAssembler, MixedElementAssembler, Field, BlockLayout, NodeAssembler, FacetAssembler -from .assemble import LaplaceElementAssembler, MassElementAssembler, LinearElasticityElementAssembler, const_node_assembler, func_node_assembler +from .assemble import ElementAssembler, MixedElementAssembler, Field, BlockLayout, NodeAssembler, FacetAssembler, FacetBilinearAssembler +from .assemble import LaplaceElementAssembler, AnisotropicLaplaceElementAssembler, MassElementAssembler, ScaledMassElementAssembler, LinearElasticityElementAssembler, const_node_assembler, func_node_assembler from .functional import * from .dataset import MeshGen from .distributed import ( diff --git a/tensormesh/assemble/__init__.py b/tensormesh/assemble/__init__.py index 5edccc1..8269f2e 100644 --- a/tensormesh/assemble/__init__.py +++ b/tensormesh/assemble/__init__.py @@ -12,9 +12,12 @@ from .mixed_assembler import MixedElementAssembler, Field, BlockLayout from .node_assembler import NodeAssembler from .facet_assembler import FacetAssembler +from .facet_bilinear import FacetBilinearAssembler from .builtin import ( LaplaceElementAssembler, + AnisotropicLaplaceElementAssembler, MassElementAssembler, + ScaledMassElementAssembler, LinearElasticityElementAssembler, NeoHookeanModel, J2Plasticity, @@ -31,8 +34,11 @@ "BlockLayout", "NodeAssembler", "FacetAssembler", + "FacetBilinearAssembler", "LaplaceElementAssembler", + "AnisotropicLaplaceElementAssembler", "MassElementAssembler", + "ScaledMassElementAssembler", "LinearElasticityElementAssembler", "NeoHookeanModel", "J2Plasticity", diff --git a/tensormesh/assemble/builtin.py b/tensormesh/assemble/builtin.py index 0da050d..10e9729 100644 --- a/tensormesh/assemble/builtin.py +++ b/tensormesh/assemble/builtin.py @@ -61,7 +61,41 @@ class LaplaceElementAssembler(ElementAssembler): """ def forward(self, gradu, gradv): return gradu @ gradv - + + +class AnisotropicLaplaceElementAssembler(ElementAssembler): + r"""Tensor-coefficient (anisotropic) diffusion / stiffness assembler. + + Generalizes :class:`LaplaceElementAssembler` to a full material tensor + :math:`\mathbf{A}`: + + .. math:: + + K_{ij}^e = \int_{\Omega^e} (\mathbf{A}\, \nabla N_i) \cdot \nabla N_j \, \mathrm{d}\Omega + = \int_{\Omega^e} \nabla N_i^\top \mathbf{A}\, \nabla N_j \, \mathrm{d}\Omega . + + The coefficient ``A`` is a nodal :math:`[|\mathcal V|, D, D]` field passed + through ``point_data`` (interpolated to each quadrature point), so it may be + spatially varying, anisotropic, and **complex** — e.g. the diagonal + coordinate-stretch tensor :math:`\mathrm{diag}(s_y/s_x,\, s_x/s_y)` of a + stretched-coordinate PML, or an anisotropic conductivity/permittivity. With + ``A`` the identity this reduces exactly to :class:`LaplaceElementAssembler`. + + Examples + -------- + .. code-block:: python + + mesh = Mesh.gen_rectangle(chara_length=0.1) + A = torch.eye(2).expand(mesh.n_points, 2, 2).clone() # [N, D, D] + K = AnisotropicLaplaceElementAssembler.from_mesh(mesh)( + mesh.points, point_data={"A": A}) + """ + def forward(self, gradu, gradv, A): + # promote the (real) shape gradients to A's dtype so a complex tensor + # coefficient (e.g. a PML stretch) works without a manual cast + return gradu.to(A.dtype) @ A @ gradv.to(A.dtype) + + class MassElementAssembler(ElementAssembler): r"""Mass Element Assembler. @@ -99,7 +133,28 @@ class MassElementAssembler(ElementAssembler): """ def forward(self, u, v): return u * v - + + +class ScaledMassElementAssembler(ElementAssembler): + r"""Coefficient-weighted mass assembler :math:`\int_\Omega c\, N_i N_j\, d\Omega`. + + Like :class:`MassElementAssembler` but with a scalar nodal coefficient ``c`` + (from ``point_data``, possibly spatially varying and **complex**). Handy for + the reaction / potential term of a Helmholtz problem — e.g. the PML mass + scaling :math:`k_0^2\,\varepsilon\, s_{\mathrm{prod}}` (see + :func:`tensormesh.cartesian_pml`), or a variable wave speed. + + Examples + -------- + .. code-block:: python + + c = k0**2 * eps_r * s_prod # complex nodal field [N] + M = ScaledMassElementAssembler.from_mesh(mesh)(mesh.points, point_data={"c": c}) + """ + def forward(self, u, v, c): + return c * u * v + + class LinearElasticityElementAssembler(ElementAssembler): r"""Linear Elasticity Element Assembler. diff --git a/tensormesh/assemble/facet_bilinear.py b/tensormesh/assemble/facet_bilinear.py new file mode 100644 index 0000000..15ad245 --- /dev/null +++ b/tensormesh/assemble/facet_bilinear.py @@ -0,0 +1,152 @@ +r"""Facet **bilinear** assembler — boundary matrices for Robin/impedance/port BCs. + +:class:`~tensormesh.FacetAssembler` integrates a *linear* form over the boundary +and returns a node **vector** (tractions, loads, energy). Robin / impedance / +absorbing / plane-wave-port conditions instead need a boundary *bilinear* form + +.. math:: + + B_{ij} = \int_{\Gamma} c(\mathbf x)\, N_i N_j \, \mathrm dS + +i.e. a boundary **mass matrix** (optionally with a coefficient), which adds to +the volume operator. :class:`FacetBilinearAssembler` fills exactly that gap: its +``forward(u, v, ...)`` is evaluated per basis pair (like +:class:`~tensormesh.ElementAssembler`, but over facets) and the result is +scattered into a :class:`~tensormesh.sparse.matrix.SparseMatrix`. + +It shares all of :class:`FacetAssembler`'s topology construction +(:meth:`from_mesh` / :meth:`from_elements`); only the assembly returns a matrix. +""" +from __future__ import annotations + +import inspect +from typing import Callable, Mapping, Optional + +import torch + +from .facet_assembler import FacetAssembler +from ..sparse.matrix import SparseMatrix +from ..vmap import vmap + + +class FacetBilinearAssembler(FacetAssembler): + r"""Assemble a bilinear form over boundary facets into a sparse matrix. + + Override :meth:`forward` with the integrand of a boundary bilinear form at a + single facet-quadrature point and basis pair — e.g. ``u * v`` for a boundary + mass matrix, or ``c * u * v`` with a nodal coefficient ``c`` (impedance, + radiation). ``gradu`` / ``gradv`` and any ``point_data`` key are available + just as in :class:`~tensormesh.ElementAssembler`. + + Calling the assembler returns a :class:`~tensormesh.sparse.matrix.SparseMatrix` + of shape ``[n_points, n_points]``. + + Examples + -------- + Boundary mass matrix on the selected boundary: + + .. code-block:: python + + class BoundaryMass(FacetBilinearAssembler): + def forward(self, u, v): + return u * v + + B = BoundaryMass.from_mesh(mesh, boundary_mask=inlet)(mesh.points) + """ + + __autodoc__ = ["__call__", "forward"] + + def __call__(self, points: Optional[torch.Tensor] = None, + func: Optional[Callable] = None, + point_data: Optional[Mapping[str, torch.Tensor]] = None, + ) -> SparseMatrix: + r"""Integrate the facet bilinear form; return a ``[N, N]`` SparseMatrix.""" + if point_data is None: + point_data = {} + point_data = dict(point_data) + + if points is not None: + self = self.type(points.dtype).to(points.device) + for et in self.element_types: + self.transformation[et].update_points(points) + else: + points = next(iter(self.transformation.values())).points # type: ignore + point_data["x"] = points # type: ignore + for key, value in point_data.items(): + assert value.shape[0] == self.n_points, \ + f"point_data['{key}'] must be [n_points, ...], got {tuple(value.shape)}" + + fn = self.forward if func is None else func + signature = inspect.signature(fn) + + # Every facet argument carries a leading (facet, quadrature) pair, so both + # of those vmap levels map dim 0 for *all* args; only the two basis levels + # differ. in_dims per arg at (facet, quad, u_basis, v_basis): + def _dims(key): + if key in ("u", "gradu"): + return (0, 0, 0, None) + if key in ("v", "gradv"): + return (0, 0, None, 0) + if key in point_data or (key.startswith("grad") and key[4:] in point_data): + return (0, 0, None, None) # coefficient: facet+quad only + raise ValueError(f"key {key} is not supported (use u/v/gradu/gradv " + f"or a point_data key)") + dims = [_dims(k) for k in signature.parameters] + facet_d = tuple(d[0] for d in dims) + quad_d = tuple(d[1] for d in dims) + u_d = tuple(d[2] for d in dims) + v_d = tuple(d[3] for d in dims) + parallel_fn = vmap(vmap(vmap(vmap(fn, in_dims=v_d), in_dims=u_d), + in_dims=quad_d), in_dims=facet_d) + + N = self.n_points + rows_all, cols_all, vals_all = [], [], [] + for et in self.element_types: + trans = self.transformation[et] # type: ignore + proj = self.projector[et] # type: ignore + m: torch.Tensor = self.facet_mask[et].item() # type: ignore [n_elem, n_facet] + n_basis = trans.n_basis + + # selected-facet shape values / gradients / area-weights + sv = trans.facet_shape_val.repeat(trans.n_elements, 1, 1, 1)[m] # [F, Q, B] + sg = trans.facet_shape_grad[m] # [F, Q, B, D] + fxw = trans.FxW[m] # [F, Q] + n_facet = sv.shape[0] + ele_pd = {k: v[trans.elements] for k, v in point_data.items()} + + args = [] + for key in signature.parameters: + if key in ("u", "v"): + args.append(sv) + elif key in ("gradu", "gradv"): + args.append(sg) + elif key in ele_pd: + fsv = trans.facet_shape_val + if ele_pd[key].is_complex() and not fsv.is_complex(): + fsv = fsv.to(ele_pd[key].dtype) # promote real shapes + pd = torch.einsum("eb...,fqb->efq...", ele_pd[key], fsv)[m] + args.append(pd) # [F, Q, ...] + elif key.startswith("grad") and key[4:] in ele_pd: + coeff = ele_pd[key[4:]] + fsg = trans.facet_shape_grad + if coeff.is_complex() and not fsg.is_complex(): + fsg = fsg.to(coeff.dtype) + gd = torch.einsum("eb...,efqbd->efq...d", coeff, fsg)[m] + args.append(gd) + else: + raise NotImplementedError(f"key {key} not implemented") + + integ = parallel_fn(*args) # [F, Q, B, B] + local = torch.einsum("fqab,fq->fab", integ, fxw.to(integ.dtype)) # [F, B, B] + + # scatter: entry (a, b) of facet f -> global (dof[f,a], dof[f,b]) + cell_dofs = proj.indices.reshape(n_facet, n_basis).to(local.device) # [F, B] + rows = cell_dofs[:, :, None].expand(n_facet, n_basis, n_basis).reshape(-1) + cols = cell_dofs[:, None, :].expand(n_facet, n_basis, n_basis).reshape(-1) + rows_all.append(rows); cols_all.append(cols); vals_all.append(local.reshape(-1)) + + row = torch.cat(rows_all); col = torch.cat(cols_all); val = torch.cat(vals_all) + # coalesce duplicate (row, col) entries (nodes shared by several facets) + coo = torch.sparse_coo_tensor(torch.stack([row, col]), val, (N, N)).coalesce() + idx = coo.indices() + return SparseMatrix(coo.values(), idx[0], idx[1], (N, N)) diff --git a/tensormesh/functional/__init__.py b/tensormesh/functional/__init__.py index 294a832..885af8a 100644 --- a/tensormesh/functional/__init__.py +++ b/tensormesh/functional/__init__.py @@ -1,3 +1,4 @@ from .ops import * from .elasticity import * -from .plastic import * \ No newline at end of file +from .plastic import * +from .pml import cartesian_pml diff --git a/tensormesh/functional/pml.py b/tensormesh/functional/pml.py new file mode 100644 index 0000000..72bee06 --- /dev/null +++ b/tensormesh/functional/pml.py @@ -0,0 +1,98 @@ +r"""Stretched-coordinate Perfectly Matched Layer (PML) coefficients. + +A PML makes a truncated FEM domain behave as if it were open: outgoing waves +are absorbed in a thin frame instead of reflecting off the outer boundary. In +the frequency domain this is realized by an analytic **complex coordinate +stretch** :math:`x \to \tilde x` with :math:`\partial\tilde x/\partial x = s_x`, +which turns the scalar Helmholtz operator into an anisotropic one: + +.. math:: + + \nabla\cdot(\mathbf\Lambda\,\nabla u) + k_0^2\, \varepsilon\, s_{\mathrm{prod}}\, u = f , + \qquad + \mathbf\Lambda = \operatorname{diag}\!\Big(\tfrac{s_y s_z}{s_x},\, + \tfrac{s_x s_z}{s_y},\, + \tfrac{s_x s_y}{s_z}\Big), + \quad s_{\mathrm{prod}} = s_x s_y s_z , + +so it plugs straight into :class:`~tensormesh.AnisotropicLaplaceElementAssembler` +(the ``\mathbf\Lambda`` tensor) plus a mass term scaled by +:math:`k_0^2 \varepsilon\, s_{\mathrm{prod}}`. + +:func:`cartesian_pml` returns those two nodal fields for a rectangular PML frame. + + Lam, sprod = cartesian_pml(mesh.points, bounds=[(0, L), (0, L)], thickness=d) + K = AnisotropicLaplaceElementAssembler.from_mesh(mesh)(pts, point_data={"A": Lam}) + # M scaled by k0**2 * eps_r * sprod (e.g. ScaledMassElementAssembler) +""" +from __future__ import annotations + +from typing import List, Sequence, Tuple, Union + +import torch + + +def cartesian_pml(points: torch.Tensor, + *, + bounds: Sequence[Tuple[float, float]], + thickness: Union[float, Sequence[float]], + strength: float = 5.0, + order: int = 2, + ) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Complex coordinate-stretch tensor field for a Cartesian PML frame. + + The stretch on each axis is unity in the interior and ramps into the PML as + :math:`s(\xi) = 1 - i\,\sigma\,\xi^{p}`, where :math:`\xi \in [0, 1]` is the + normalized depth into the PML, :math:`\sigma` is ``strength`` and :math:`p` + is ``order``. Larger ``strength`` / thicker layers absorb more strongly; + too steep a ramp for the mesh re-introduces discretization reflection. + + Parameters + ---------- + points : torch.Tensor + Nodal coordinates, shape :math:`[|\mathcal V|, D]` with ``D`` in ``{1,2,3}``. + bounds : sequence of (lo, hi) + Outer domain extent per axis; the PML occupies ``[lo, lo+thickness]`` and + ``[hi-thickness, hi]`` on every axis. + thickness : float or sequence of float + PML thickness, scalar (same on all axes) or one value per axis. + strength : float, optional + Imaginary stretch amplitude :math:`\sigma`; default ``5.0``. + order : int, optional + Polynomial ramp order :math:`p`; default ``2``. + + Returns + ------- + Lambda : torch.Tensor + Complex diagonal stretch tensor, shape :math:`[|\mathcal V|, D, D]`, for + :class:`~tensormesh.AnisotropicLaplaceElementAssembler`. + s_prod : torch.Tensor + Complex product :math:`\prod_a s_a`, shape :math:`[|\mathcal V|]`, the + mass-term scaling. + """ + if points.dim() != 2: + raise ValueError(f"points must be [N, D], got {tuple(points.shape)}") + n, dim = points.shape + if len(bounds) != dim: + raise ValueError(f"bounds must have {dim} (lo, hi) pairs, got {len(bounds)}") + if isinstance(thickness, (int, float)): + thickness = [float(thickness)] * dim + dev = points.device + + s_list: List[torch.Tensor] = [] + for a in range(dim): + lo, hi = bounds[a] + d = float(thickness[a]) + t = points[:, a].to(torch.float64) + below = torch.clamp((lo + d - t) / d, min=0.0, max=1.0) # ramp near lo + above = torch.clamp((t - (hi - d)) / d, min=0.0, max=1.0) # ramp near hi + xi = below + above # 0 in the interior + s_list.append(1.0 - 1j * strength * xi ** order) + s = torch.stack(s_list, dim=1).to(torch.complex128).to(dev) # [N, D] + s_prod = torch.prod(s, dim=1) # [N] + + Lambda = torch.zeros(n, dim, dim, dtype=torch.complex128, device=dev) + for a in range(dim): + # diag_a = (prod_{b != a} s_b) / s_a = s_prod / s_a^2 + Lambda[:, a, a] = s_prod / (s[:, a] * s[:, a]) + return Lambda, s_prod diff --git a/tensormesh/operator/__init__.py b/tensormesh/operator/__init__.py index f0e2170..9bd8c30 100644 --- a/tensormesh/operator/__init__.py +++ b/tensormesh/operator/__init__.py @@ -3,11 +3,14 @@ Exposes :class:`Condenser`, which applies Dirichlet boundary conditions to an assembled FEM system via static condensation, and :class:`BlochReducer`, the periodic counterpart that ties opposite unit-cell faces with a wavevector- -dependent Floquet phase (band structure / Bloch-periodic solves). See the class -docstrings and the User Guide chapter on boundary conditions. +dependent Floquet phase (band structure / Bloch-periodic solves); plus +:func:`robin_operator` / :func:`port_source`, the assembled boundary matrix and +load of first-order absorbing / plane-wave-port conditions for wave problems. +See the class docstrings and the User Guide chapter on boundary conditions. """ from .condense import Condenser from .bloch import BlochReducer +from .boundary import robin_operator, port_source -__all__ = ["Condenser", "BlochReducer"] +__all__ = ["Condenser", "BlochReducer", "robin_operator", "port_source"] diff --git a/tensormesh/operator/boundary.py b/tensormesh/operator/boundary.py new file mode 100644 index 0000000..a83a7e5 --- /dev/null +++ b/tensormesh/operator/boundary.py @@ -0,0 +1,121 @@ +r"""Boundary operators for wave problems — Robin / impedance / absorbing ports. + +Built on :class:`~tensormesh.FacetBilinearAssembler`, these helpers produce the +boundary **matrix** and **load** of a first-order absorbing / plane-wave-port +condition for the scalar Helmholtz equation. + +For an outgoing/absorbing (Sommerfeld) boundary and an incident plane wave +:math:`u_{\text{inc}}` the condition + +.. math:: + + \frac{\partial u}{\partial n} + i k\, u = 2 i k\, u_{\text{inc}} + \quad\text{on } \Gamma_{\text{port}} + +contributes, in the weak form, + +* to the **operator**: :math:`+\,i k \int_\Gamma N_i N_j\,\mathrm dS = i k\,\mathbf B` + — see :func:`robin_operator`; +* to the **right-hand side**: :math:`+\,2 i k \int_\Gamma u_{\text{inc}} N_i\,\mathrm dS` + — see :func:`port_source`. + +Recipe for a driven, non-reflecting Helmholtz solve:: + + A = K - k**2 * M + 1j*k * robin_operator(mesh, port_and_open_boundaries) + b = 2j*k * port_source(mesh, inlet, incident=1.0) + u = A.solve(b) + +Dropping the incident term (``port_source``) leaves a purely absorbing (open) +boundary; a wall is simply a boundary left out of ``robin_operator`` (natural +Neumann / sound-hard). +""" +from __future__ import annotations + +from typing import Optional, Union + +import torch + +from ..assemble.facet_bilinear import FacetBilinearAssembler +from ..mesh import Mesh +from ..sparse.matrix import SparseMatrix + + +class _RobinMass(FacetBilinearAssembler): + """Boundary mass ``\\int_\\Gamma c\\, N_i N_j`` (coefficient ``c`` from point_data).""" + def forward(self, u, v, c): + return c * u * v + + +def robin_operator(mesh: Mesh, + boundary_mask: Optional[Union[str, torch.Tensor]] = None, + coeff: Union[float, complex, torch.Tensor] = 1.0, + *, + points: Optional[torch.Tensor] = None, + quadrature_order: int = 2, + ) -> SparseMatrix: + r"""Boundary (Robin / impedance) matrix :math:`\int_\Gamma c\, N_i N_j\,\mathrm dS`. + + Parameters + ---------- + mesh : Mesh + boundary_mask : str, torch.Tensor, or None + Which boundary to integrate over (a per-node boolean tensor, a named + mask, or ``None`` for the full boundary) — passed to + :meth:`FacetBilinearAssembler.from_mesh`. + coeff : float, complex, or torch.Tensor + Impedance/absorption coefficient :math:`c`; a scalar (e.g. ``1j*k`` for a + first-order absorbing boundary) or a nodal ``[n_points]`` field. + points : torch.Tensor, optional + Node coordinates (defaults to ``mesh.points``). + quadrature_order : int, optional + Facet quadrature order (default ``2``). + + Returns + ------- + SparseMatrix + The ``[n_points, n_points]`` boundary matrix. + """ + pts = mesh.points if points is None else points + n = pts.shape[0] + if not torch.is_tensor(coeff): + # scalar coefficient: follow the mesh precision (float64 mesh -> + # complex128 / float64 coefficient) instead of torch's float32 default + if isinstance(coeff, complex): + cdtype = torch.complex128 if pts.dtype == torch.float64 else torch.complex64 + else: + cdtype = pts.dtype + c = torch.full((n,), coeff, dtype=cdtype, device=pts.device) + else: + c = coeff + asm = _RobinMass.from_mesh(mesh, boundary_mask=boundary_mask, + quadrature_order=quadrature_order) + return asm(pts, point_data={"c": c}) + + +def port_source(mesh: Mesh, + boundary_mask: Optional[Union[str, torch.Tensor]] = None, + incident: Union[float, complex, torch.Tensor] = 1.0, + *, + points: Optional[torch.Tensor] = None, + quadrature_order: int = 2, + ) -> torch.Tensor: + r"""Port load vector :math:`\int_\Gamma u_{\text{inc}}\, N_i\,\mathrm dS`. + + This is the consistent boundary load of an incident field ``incident`` (a + scalar amplitude for a uniform plane-wave port, or a nodal ``[n_points]`` + field for a shaped one). Multiply by ``2j*k`` for the plane-wave-port RHS. + + Returns + ------- + torch.Tensor + A ``[n_points]`` load vector (nonzero only on the selected boundary). + """ + pts = mesh.points if points is None else points + n = pts.shape[0] + B = robin_operator(mesh, boundary_mask, 1.0, points=pts, + quadrature_order=quadrature_order) + if not torch.is_tensor(incident): + inc = torch.full((n,), incident, dtype=B.values.dtype) + else: + inc = incident.to(B.values.dtype) + return B @ inc # int_Gamma incident * N_i (Sum_j N_j = 1 recovers int N_i) diff --git a/tests/assemble/test_wave.py b/tests/assemble/test_wave.py new file mode 100644 index 0000000..1074769 --- /dev/null +++ b/tests/assemble/test_wave.py @@ -0,0 +1,186 @@ +"""Tests for the open-domain wave operators: anisotropic / scaled-mass / +facet-bilinear assemblers, the Cartesian PML, and the Robin / port boundary +operators.""" + +import sys +sys.path.append("../..") + +import numpy as np +import pytest +import torch + +from tensormesh import ( + Mesh, + LaplaceElementAssembler, + AnisotropicLaplaceElementAssembler, + MassElementAssembler, + ScaledMassElementAssembler, + FacetBilinearAssembler, + cartesian_pml, + robin_operator, + port_source, +) + + +@pytest.fixture(autouse=True) +def _default_float64(): + """Use double precision for the exact-equality asserts, restoring the global + default afterward so this module never leaks its dtype into other tests.""" + prev = torch.get_default_dtype() + torch.set_default_dtype(torch.float64) + yield + torch.set_default_dtype(prev) + + +def _unit_square(h=0.2): + return Mesh.gen_rectangle(chara_length=h, element_type="tri", + left=0, right=1, bottom=0, top=1) + + +def _dense(sp): + return sp.to_scipy_coo().tocsr().toarray() + + +# --------------------------------------------------------------------------- # +# AnisotropicLaplaceElementAssembler +# --------------------------------------------------------------------------- # +def test_anisotropic_laplace_identity_equals_laplace(): + mesh = _unit_square() + K = LaplaceElementAssembler.from_mesh(mesh, quadrature_order=2)(mesh.points) + A = torch.eye(2).expand(mesh.n_points, 2, 2).clone() + Ka = AnisotropicLaplaceElementAssembler.from_mesh(mesh, quadrature_order=2)( + mesh.points, point_data={"A": A}) + assert np.abs(_dense(K) - _dense(Ka)).max() < 1e-12 + + +def test_anisotropic_laplace_diagonal_scales_direction(): + # A = diag(2, 1) doubles the xx-contribution vs plain Laplace + mesh = _unit_square() + A = torch.diag(torch.tensor([2.0, 1.0])).expand(mesh.n_points, 2, 2).clone() + Ka = AnisotropicLaplaceElementAssembler.from_mesh(mesh, quadrature_order=2)( + mesh.points, point_data={"A": A}) + # symmetric, real, and not equal to the isotropic operator + d = _dense(Ka) + assert np.allclose(d, d.T, atol=1e-10) + + +def test_anisotropic_laplace_complex_coefficient(): + mesh = _unit_square() + A = (torch.eye(2) * (1.0 - 0.5j)).expand(mesh.n_points, 2, 2).clone() + Ka = AnisotropicLaplaceElementAssembler.from_mesh(mesh, quadrature_order=2)( + mesh.points, point_data={"A": A}) + assert Ka.values.is_complex() + # (1 - 0.5j) * Laplace + K = LaplaceElementAssembler.from_mesh(mesh, quadrature_order=2)(mesh.points) + assert np.abs(_dense(Ka) - (1.0 - 0.5j) * _dense(K)).max() < 1e-12 + + +# --------------------------------------------------------------------------- # +# ScaledMassElementAssembler +# --------------------------------------------------------------------------- # +def test_scaled_mass_unit_coefficient_equals_mass(): + mesh = _unit_square() + M = MassElementAssembler.from_mesh(mesh, quadrature_order=2)(mesh.points) + c = torch.ones(mesh.n_points) + Ms = ScaledMassElementAssembler.from_mesh(mesh, quadrature_order=2)( + mesh.points, point_data={"c": c}) + assert np.abs(_dense(M) - _dense(Ms)).max() < 1e-12 + + +def test_scaled_mass_constant_coefficient_scales(): + mesh = _unit_square() + M = MassElementAssembler.from_mesh(mesh, quadrature_order=2)(mesh.points) + c = 3.0 * torch.ones(mesh.n_points) + Ms = ScaledMassElementAssembler.from_mesh(mesh, quadrature_order=2)( + mesh.points, point_data={"c": c}) + assert np.abs(_dense(Ms) - 3.0 * _dense(M)).max() < 1e-12 + + +# --------------------------------------------------------------------------- # +# FacetBilinearAssembler (boundary mass) +# --------------------------------------------------------------------------- # +class _BoundaryMass(FacetBilinearAssembler): + def forward(self, u, v): + return u * v + + +class _CoeffBoundaryMass(FacetBilinearAssembler): + def forward(self, u, v, c): + return c * u * v + + +def test_facet_bilinear_full_boundary_is_perimeter(): + # sum over B = int_Gamma (sum_i N_i)(sum_j N_j) dS = int_Gamma 1 dS = |Gamma| + mesh = _unit_square(0.1) + B = _BoundaryMass.from_mesh(mesh, boundary_mask=None)(mesh.points) + assert abs(B.values.sum().item() - 4.0) < 1e-9 # unit-square perimeter + + +def test_facet_bilinear_single_edge_length(): + mesh = _unit_square(0.1) + bottom = mesh.points[:, 1] < 1e-9 + B = _BoundaryMass.from_mesh(mesh, boundary_mask=bottom)(mesh.points) + assert abs(B.values.sum().item() - 1.0) < 1e-9 # bottom edge length + + +def test_facet_bilinear_coefficient(): + mesh = _unit_square(0.1) + bottom = mesh.points[:, 1] < 1e-9 + c = 2.0 * torch.ones(mesh.n_points) + B = _CoeffBoundaryMass.from_mesh(mesh, boundary_mask=bottom)( + mesh.points, point_data={"c": c}) + assert abs(B.values.sum().item() - 2.0) < 1e-9 + + +# --------------------------------------------------------------------------- # +# cartesian_pml +# --------------------------------------------------------------------------- # +def test_cartesian_pml_interior_is_identity(): + # points well inside the domain -> stretch 1 -> Lambda = I, s_prod = 1 + pts = torch.tensor([[3.0, 3.0], [3.0, 2.0], [2.5, 3.5]]) # interior of [0,6]^2, tpml=0.6 + Lam, sprod = cartesian_pml(pts, bounds=[(0.0, 6.0), (0.0, 6.0)], + thickness=0.6, strength=5.0, order=2) + assert torch.allclose(Lam, torch.eye(2, dtype=torch.complex128).expand(3, 2, 2)) + assert torch.allclose(sprod, torch.ones(3, dtype=torch.complex128)) + + +def test_cartesian_pml_tensor_and_product(): + # inside the PML the diagonal is (sy/sx, sx/sy) and s_prod = sx*sy + pts = torch.tensor([[0.1, 3.0]]) # deep in the -x PML + L, d, sig, p = 6.0, 0.6, 5.0, 2 + Lam, sprod = cartesian_pml(pts, bounds=[(0.0, L), (0.0, L)], + thickness=d, strength=sig, order=p) + xi = (d - 0.1) / d + sx = torch.tensor(1.0 - 1j * sig * xi ** p, dtype=torch.complex128) + sy = torch.tensor(1.0 + 0j, dtype=torch.complex128) # y interior + assert torch.allclose(Lam[0, 0, 0], sy / sx) + assert torch.allclose(Lam[0, 1, 1], sx / sy) + assert torch.allclose(Lam[0, 0, 1], torch.zeros((), dtype=torch.complex128)) + assert torch.allclose(sprod[0], sx * sy) + + +# --------------------------------------------------------------------------- # +# robin_operator / port_source +# --------------------------------------------------------------------------- # +def test_robin_operator_full_boundary_is_perimeter(): + mesh = _unit_square(0.1) + B = robin_operator(mesh, boundary_mask=None, coeff=1.0) + assert abs(B.values.sum().item() - 4.0) < 1e-9 + + +def test_robin_operator_complex_coefficient(): + mesh = _unit_square(0.1) + B = robin_operator(mesh, boundary_mask=None, coeff=1j) + assert B.values.is_complex() + assert abs(B.values.sum().imag.item() - 4.0) < 1e-9 # i * perimeter + + +def test_port_source_integrates_shape_functions(): + # sum_i int_Gamma N_i dS = int_Gamma 1 dS = perimeter + mesh = _unit_square(0.1) + e = port_source(mesh, boundary_mask=None, incident=1.0) + assert abs(e.real.sum().item() - 4.0) < 1e-9 + # a single edge integrates to its length + bottom = mesh.points[:, 1] < 1e-9 + e_edge = port_source(mesh, boundary_mask=bottom, incident=1.0) + assert abs(e_edge.real.sum().item() - 1.0) < 1e-9