diff --git a/python/micromode/beamz.py b/python/micromode/beamz.py index 9d13112..116517f 100644 --- a/python/micromode/beamz.py +++ b/python/micromode/beamz.py @@ -8,11 +8,12 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Literal, TypedDict, cast +from typing import Any, Literal, TypedDict, cast import numpy as np from .raster import solve_grid +from .yee import refine_x_mode_at_fixed_beta, validate_x_mode_refinement AxisName = Literal["x", "y", "z"] DirectionName = Literal["+x", "-x", "+y", "-y", "+z", "-z"] @@ -72,6 +73,8 @@ class ModePlaneSpec: component_offsets: dict[str, dict[str, float]] = field(default_factory=lambda: dict(_YEE_OFFSETS_3D)) component_permittivity: dict[str, np.ndarray] = field(default_factory=dict) component_permeability: dict[str, np.ndarray] = field(default_factory=dict) + # Guarded by joint field, power, energy, and discrete-Maxwell validation. + yee_refinement: bool = True boundary: str = "beamz-finite-aperture" def __post_init__(self) -> None: @@ -134,6 +137,7 @@ def __post_init__(self) -> None: object.__setattr__(self, "num_modes", None if self.num_modes is None else int(self.num_modes)) object.__setattr__(self, "aperture_pad_cells", int(self.aperture_pad_cells)) object.__setattr__(self, "aperture_window_alpha", float(self.aperture_window_alpha)) + object.__setattr__(self, "yee_refinement", bool(self.yee_refinement)) if self.polarization is not None: pol = str(self.polarization).lower() @@ -213,6 +217,60 @@ def solve_beamz_mode(spec: ModePlaneSpec) -> DiscreteMode: ) omega = 2.0 * np.pi * spec.frequency k_num = _solve_numeric_k_axis(omega, spec.dt, spec.resolution, selected["neff"]) + boundary_neff = _boundary_refractive_index(spec.scalar_permittivity) + yee_refinement_eligible = ( + spec.axis == "x" and bool(spec.component_permittivity) and float(np.real(selected["neff"])) > boundary_neff + ) + yee_refinement_requested = bool(spec.yee_refinement) + yee_refinement_attempted = yee_refinement_requested and yee_refinement_eligible + yee_refinement_accepted = False + yee_refinement_rejection_reason = "" + yee_validation: dict[str, Any] = {} + yee_residual = np.nan + yee_frequency_ratio = np.nan + yee_initial_frequency_ratio = np.nan + if yee_refinement_requested and not yee_refinement_eligible: + yee_refinement_rejection_reason = "mode is not eligible for x-normal guided-mode refinement" + if yee_refinement_attempted: + seed_profiles = {name: np.asarray(value, dtype=np.complex128).copy() for name, value in profiles.items()} + seed_k_num = float(k_num) + try: + candidate, yee_residual, yee_frequency_ratio, candidate_k_num, yee_initial_frequency_ratio = ( + refine_x_mode_at_fixed_beta( + seed_profiles, + indices, + component_permittivity=spec.component_permittivity, + component_permeability=spec.component_permeability, + omega=omega, + dt=spec.dt, + resolution=spec.resolution, + k_num=seed_k_num, + direction_sign=_direction_sign(spec.direction), + ) + ) + yee_refinement_accepted, yee_validation = validate_x_mode_refinement( + seed_profiles, + candidate, + indices, + component_permittivity=spec.component_permittivity, + component_permeability=spec.component_permeability, + omega=omega, + dt=spec.dt, + resolution=spec.resolution, + k_num=candidate_k_num, + direction_sign=_direction_sign(spec.direction), + ) + yee_refinement_rejection_reason = str(yee_validation.get("rejection_reason", "")) + if yee_refinement_accepted: + profiles = candidate + k_num = float(candidate_k_num) + else: + profiles = seed_profiles + k_num = seed_k_num + except Exception as exc: + profiles = seed_profiles + k_num = seed_k_num + yee_refinement_rejection_reason = f"{type(exc).__name__}: {exc}" profiles, power_scale = _normalize_profiles_by_phase_referenced_flux( profiles, indices, @@ -238,6 +296,17 @@ def solve_beamz_mode(spec: ModePlaneSpec) -> DiscreteMode: "phase_reference": spec.phase_reference, "time_convention": spec.time_convention, "aperture_window_alpha": spec.aperture_window_alpha, + "yee_refinement": yee_refinement_accepted, + "yee_refinement_requested": yee_refinement_requested, + "yee_refinement_eligible": yee_refinement_eligible, + "yee_refinement_attempted": yee_refinement_attempted, + "yee_refinement_accepted": yee_refinement_accepted, + "yee_refinement_rejection_reason": yee_refinement_rejection_reason, + "yee_refinement_validation": yee_validation, + "boundary_neff": float(boundary_neff), + "yee_residual": float(yee_residual), + "yee_frequency_ratio": float(yee_frequency_ratio), + "yee_initial_frequency_ratio": float(yee_initial_frequency_ratio), "power_before_phase_reference": float(extra.get("initial_power", np.nan)), "power_after_phase_reference": float( _modal_power_from_profiles( @@ -274,6 +343,15 @@ def solve_beamz_mode(spec: ModePlaneSpec) -> DiscreteMode: ) +def _boundary_refractive_index(permittivity: np.ndarray) -> float: + """Return the largest refractive index touching the mode-plane boundary.""" + eps = np.asarray(permittivity, dtype=np.complex128) + if eps.ndim != 2 or min(eps.shape) == 0: + return 0.0 + boundary = np.concatenate((eps[0], eps[-1], eps[1:-1, 0], eps[1:-1, -1])) + return float(np.sqrt(max(float(np.max(np.real(boundary))), 0.0))) + + def _candidate_modes(result, spec: ModePlaneSpec) -> list[_ModeCandidate]: candidates = [] count = int(result.n_complex.shape[1]) @@ -579,8 +657,8 @@ def _normalize_profiles_by_flux( direction_sign: float, ) -> float: flux = _modal_power_from_profiles(profiles, axis=axis, d_area=d_area, direction_sign=direction_sign) - if np.isfinite(flux) and abs(flux) > 1e-18: - scale = float(np.clip(np.sqrt(1.0 / abs(flux)), 1e-6, 1e6)) + if np.isfinite(flux) and abs(flux) > np.finfo(float).tiny: + scale = float(np.sqrt(1.0 / abs(flux))) for key, value in profiles.items(): profiles[key] = np.asarray(value, dtype=np.complex128) * scale return float(flux) @@ -613,9 +691,9 @@ def _normalize_profiles_by_phase_referenced_flux( d_area=d_area, direction_sign=direction_sign, ) - if (not np.isfinite(flux)) or abs(flux) <= 1e-18: + if (not np.isfinite(flux)) or abs(flux) <= np.finfo(float).tiny: return profiles, 1.0 - scale = float(np.clip(np.sqrt(1.0 / abs(flux)), 1e-6, 1e6)) + scale = float(np.sqrt(1.0 / abs(flux))) return ( {key: np.asarray(value, dtype=np.complex128) * scale for key, value in profiles.items()}, scale, @@ -660,23 +738,12 @@ def _modal_power_from_profiles( terms = ("Ez", "Ex", "Hx", "Hz") else: terms = ("Ex", "Ey", "Hy", "Hx") - arrays = [ - np.asarray( - profiles.get(name, np.zeros((0,), dtype=np.complex128)), - dtype=np.complex128, - ) - for name in terms - ] + arrays = [np.asarray(profiles.get(name, ()), dtype=np.complex128) for name in terms] if any(arr.size == 0 for arr in arrays): return 0.0 - arrays = [arr[:, None] if arr.ndim == 1 else arr for arr in arrays] - rows = min(arr.shape[0] for arr in arrays) - cols = min(arr.shape[1] for arr in arrays) - if rows <= 0 or cols <= 0: - return 0.0 - a0, a1, b0, b1 = [arr[:rows, :cols] for arr in arrays] - s_axis = a0 * np.conjugate(b0) - a1 * np.conjugate(b1) - return float(0.5 * direction_sign * np.real(np.sum(s_axis) * float(d_area))) + a0, a1, b0, b1 = arrays + flux = np.vdot(b0.reshape(-1), a0.reshape(-1)) - np.vdot(b1.reshape(-1), a1.reshape(-1)) + return float(0.5 * direction_sign * np.real(flux * float(d_area))) def _axis_index_from_component_indices(indices: ComponentIndex | None, axis: AxisName) -> int | None: diff --git a/python/micromode/yee.py b/python/micromode/yee.py new file mode 100644 index 0000000..58331b7 --- /dev/null +++ b/python/micromode/yee.py @@ -0,0 +1,349 @@ +"""Fixed-beta refinement on BeamZ's x-normal Yee cross section.""" + +from __future__ import annotations + +from typing import Any, cast + +import numpy as np + +_C0 = 299_792_458.0 +_EPS0 = 8.854_187_812_8e-12 +_MU0 = 1.256_637_062_12e-6 + +_E_COMPONENTS = ("Ex", "Ey", "Ez") +_H_COMPONENTS = ("Hx", "Hy", "Hz") + + +def refine_x_mode_at_fixed_beta( + profiles, + indices, + *, + component_permittivity, + component_permeability, + omega, + dt, + resolution, + k_num, + direction_sign, + _correct_beta=True, +): + """Refine E at the selected beta, then reconstruct H from Faraday's law.""" + from scipy import sparse + from scipy.sparse import linalg as spla + + out = {name: np.asarray(value, dtype=np.complex128) for name, value in profiles.items()} + nz, ny = out["Ex"].shape + d = float(resolution) + sign = 1.0 if float(direction_sign) >= 0.0 else -1.0 + delta = -2j * sign * np.sin(0.5 * float(k_num) * d) / d + curl_e, curl_h, curl_e_beta, curl_h_beta = _yee_curl_operators( + sparse, + out, + resolution=d, + delta=delta, + ) + eps, mu = _component_material_vectors( + out, + indices, + component_permittivity=component_permittivity, + component_permeability=component_permeability, + ) + ex_n, ey_n, _ez_n = (out[name].size for name in _E_COMPONENTS) + hx_n, hy_n, _hz_n = (out[name].size for name in _H_COMPONENTS) + operator = curl_h @ sparse.diags(1.0 / mu) @ curl_e + mass = sparse.diags(eps) + omega_d = float(omega) if dt is None else 2.0 * np.sin(0.5 * float(omega) * float(dt)) / float(dt) + target = (omega_d / _C0) ** 2 + seed = np.concatenate([out[name].reshape(-1) for name in ("Ex", "Ey", "Ez")]) + count = min(4, operator.shape[0] - 2) + eigenpairs = spla.eigs( + operator, + M=mass, + k=count, + sigma=target, + v0=seed, + tol=1e-9, # pyright: ignore[reportArgumentType] -- SciPy infers int from its untyped default. + ) + values, vectors = cast(tuple[np.ndarray, np.ndarray], eigenpairs) + overlaps = np.abs(np.conjugate(seed) @ vectors) + selected = int(np.argmax(overlaps)) + electric = np.asarray(vectors[:, selected], dtype=np.complex128) + phase = np.vdot(seed, electric) + if abs(phase) > 1e-30: + electric *= np.exp(-1j * np.angle(phase)) + ex, ey, ez = np.split(electric, (ex_n, ex_n + ey_n)) + out["Ex"], out["Ey"], out["Ez"] = ex.reshape(nz, ny), ey.reshape(nz, ny - 1), ez.reshape(nz - 1, ny) + + magnetic = 1j * (curl_e @ electric) / (_MU0 * omega_d * mu) + hx, hy, hz = np.split(magnetic, (hx_n, hx_n + hy_n)) + out["Hx"], out["Hy"], out["Hz"] = hx.reshape(nz - 1, ny - 1), hy.reshape(nz - 1, ny), hz.reshape(nz, ny - 1) + residual = np.linalg.norm(operator @ electric - values[selected] * (mass @ electric)) + residual /= max(np.linalg.norm(values[selected] * (mass @ electric)), np.finfo(float).eps) + eigenvalue = float(np.real(values[selected])) + frequency_ratio = float(np.sqrt(max(eigenvalue / target, 0.0))) + frequency_ratio_initial = frequency_ratio + if _correct_beta: + q_axis = 2.0 * np.sin(0.5 * float(k_num) * d) / d + q_probe = 0.98 * q_axis + delta_probe = -1j * sign * q_probe + delta_change = delta_probe - delta + operator_probe = ( + (curl_h + delta_change * curl_h_beta) @ sparse.diags(1.0 / mu) @ (curl_e + delta_change * curl_e_beta) + ) + denominator = np.vdot(electric, mass @ electric) + eigenvalue_probe = float(np.real(np.vdot(electric, operator_probe @ electric) / denominator)) + slope = (eigenvalue_probe - eigenvalue) / (q_probe**2 - q_axis**2) + if not np.isfinite(slope) or abs(slope) <= np.finfo(float).eps: + slope = 0.5 + q_corrected = np.sqrt(max(q_axis**2 - (eigenvalue - target) / slope, 0.0)) + k_corrected = 2.0 * np.arcsin(np.clip(0.5 * q_corrected * d, -1.0, 1.0)) / d + if np.isfinite(k_corrected) and k_corrected > 0.0: + refined, residual, frequency_ratio, _, _ = refine_x_mode_at_fixed_beta( + out, + indices, + component_permittivity=component_permittivity, + component_permeability=component_permeability, + omega=omega, + dt=dt, + resolution=resolution, + k_num=k_corrected, + direction_sign=direction_sign, + _correct_beta=False, + ) + return refined, residual, frequency_ratio, float(k_corrected), frequency_ratio_initial + return out, float(residual), frequency_ratio, float(k_num), frequency_ratio + + +def validate_x_mode_refinement( + seed_profiles, + candidate_profiles, + indices, + *, + component_permittivity, + component_permeability, + omega, + dt, + resolution, + k_num, + direction_sign, + minimum_electric_overlap=None, + minimum_magnetic_overlap=0.5, + maximum_impedance_change=4.0, + maximum_energy_change=4.0, + maximum_energy_imbalance=4.0, + maximum_maxwell_residual=5e-3, +): + """Validate that a refined x-normal profile remains the same physical mode. + + The refinement eigen-residual alone only establishes that the selected + electric vector solves the assembled curl-curl problem. These checks also + require mode identity, a consistent E/H balance, forward signed power, and + both first-order discrete Maxwell equations. Electric overlap remains a + diagnostic by default because Yee refinement can legitimately redistribute + longitudinal electric components; callers may opt into a hard threshold. + """ + from scipy import sparse + + seed = {name: np.asarray(value, dtype=np.complex128) for name, value in seed_profiles.items()} + candidate = {name: np.asarray(value, dtype=np.complex128) for name, value in candidate_profiles.items()} + diagnostics: dict[str, Any] = { + "electric_overlap": _normalized_component_overlap(seed, candidate, _E_COMPONENTS), + "magnetic_overlap": _normalized_component_overlap(seed, candidate, _H_COMPONENTS), + } + seed_impedance = _rms_impedance(seed) + candidate_impedance = _rms_impedance(candidate) + diagnostics["seed_rms_impedance"] = seed_impedance + diagnostics["candidate_rms_impedance"] = candidate_impedance + diagnostics["impedance_change"] = _ratio_change(candidate_impedance, seed_impedance) + + eps, mu = _component_material_vectors( + candidate, + indices, + component_permittivity=component_permittivity, + component_permeability=component_permeability, + ) + seed_energy_ratio = _energy_ratio(seed, eps, mu) + candidate_energy_ratio = _energy_ratio(candidate, eps, mu) + diagnostics["seed_energy_ratio"] = seed_energy_ratio + diagnostics["candidate_energy_ratio"] = candidate_energy_ratio + diagnostics["energy_ratio_change"] = _ratio_change(candidate_energy_ratio, seed_energy_ratio) + diagnostics["candidate_energy_imbalance"] = _ratio_change(candidate_energy_ratio, 1.0) + + seed_power = _signed_x_power(seed, resolution=resolution, direction_sign=direction_sign) + candidate_power = _signed_x_power(candidate, resolution=resolution, direction_sign=direction_sign) + diagnostics["seed_signed_power"] = seed_power + diagnostics["candidate_signed_power"] = candidate_power + + sign = 1.0 if float(direction_sign) >= 0.0 else -1.0 + delta = -2j * sign * np.sin(0.5 * float(k_num) * float(resolution)) / float(resolution) + curl_e, curl_h, _curl_e_beta, _curl_h_beta = _yee_curl_operators( + sparse, + candidate, + resolution=float(resolution), + delta=delta, + ) + omega_d = float(omega) if dt is None else 2.0 * np.sin(0.5 * float(omega) * float(dt)) / float(dt) + electric = _component_vector(candidate, _E_COMPONENTS) + magnetic = _component_vector(candidate, _H_COMPONENTS) + faraday_left = curl_e @ electric + faraday_right = -1j * _MU0 * omega_d * mu * magnetic + ampere_left = curl_h @ magnetic + ampere_right = 1j * _EPS0 * omega_d * eps * electric + diagnostics["faraday_residual"] = _relative_pair_residual(faraday_left, faraday_right) + diagnostics["ampere_residual"] = _relative_pair_residual(ampere_left, ampere_right) + + rejection_reasons = [] + if minimum_electric_overlap is not None and diagnostics["electric_overlap"] < float(minimum_electric_overlap): + rejection_reasons.append("electric overlap") + if diagnostics["magnetic_overlap"] < float(minimum_magnetic_overlap): + rejection_reasons.append("magnetic overlap") + if diagnostics["impedance_change"] > float(maximum_impedance_change): + rejection_reasons.append("E/H impedance change") + if diagnostics["energy_ratio_change"] > float(maximum_energy_change): + rejection_reasons.append("electric/magnetic energy change") + if diagnostics["candidate_energy_imbalance"] > float(maximum_energy_imbalance): + rejection_reasons.append("electric/magnetic energy imbalance") + if not np.isfinite(candidate_power) or candidate_power <= np.finfo(float).tiny: + rejection_reasons.append("non-forward signed power") + if diagnostics["faraday_residual"] > float(maximum_maxwell_residual): + rejection_reasons.append("Faraday residual") + if diagnostics["ampere_residual"] > float(maximum_maxwell_residual): + rejection_reasons.append("Ampere residual") + diagnostics["accepted"] = not rejection_reasons + diagnostics["rejection_reason"] = ", ".join(rejection_reasons) + return not rejection_reasons, diagnostics + + +def _yee_curl_operators(sparse, profiles, *, resolution, delta): + nz, ny = profiles["Ex"].shape + dz = _forward_difference(sparse, nz, resolution) + dy = _forward_difference(sparse, ny, resolution) + bz, by = -dz.T, -dy.T + iz = sparse.eye(nz, format="csc") + iy = sparse.eye(ny, format="csc") + izm = sparse.eye(nz - 1, format="csc") + iym = sparse.eye(ny - 1, format="csc") + dy_ez = sparse.kron(izm, dy, format="csc") + dz_ey = sparse.kron(dz, iym, format="csc") + dz_ex = sparse.kron(dz, iy, format="csc") + dy_ex = sparse.kron(iz, dy, format="csc") + by_hz = sparse.kron(iz, by, format="csc") + bz_hy = sparse.kron(bz, iy, format="csc") + bz_hx = sparse.kron(bz, iym, format="csc") + by_hx = sparse.kron(izm, by, format="csc") + ex_n, ey_n, ez_n = (profiles[name].size for name in _E_COMPONENTS) + hx_n, hy_n, hz_n = (profiles[name].size for name in _H_COMPONENTS) + z = sparse.csc_matrix + curl_e = sparse.bmat( + [ + [z((hx_n, ex_n)), -dz_ey, dy_ez], + [dz_ex, z((hy_n, ey_n)), -delta * sparse.eye(ez_n)], + [-dy_ex, delta * sparse.eye(ey_n), z((hz_n, ez_n))], + ], + format="csc", + ) + curl_h = sparse.bmat( + [ + [z((ex_n, hx_n)), -bz_hy, by_hz], + [bz_hx, z((ey_n, hy_n)), -delta * sparse.eye(hz_n)], + [-by_hx, delta * sparse.eye(hy_n), z((ez_n, hz_n))], + ], + format="csc", + ) + curl_e_beta = sparse.bmat( + [ + [z((hx_n, ex_n)), z((hx_n, ey_n)), z((hx_n, ez_n))], + [z((hy_n, ex_n)), z((hy_n, ey_n)), -sparse.eye(ez_n)], + [z((hz_n, ex_n)), sparse.eye(ey_n), z((hz_n, ez_n))], + ], + format="csc", + ) + curl_h_beta = sparse.bmat( + [ + [z((ex_n, hx_n)), z((ex_n, hy_n)), z((ex_n, hz_n))], + [z((ey_n, hx_n)), z((ey_n, hy_n)), -sparse.eye(hz_n)], + [z((ez_n, hx_n)), sparse.eye(hy_n), z((ez_n, hz_n))], + ], + format="csc", + ) + return curl_e, curl_h, curl_e_beta, curl_h_beta + + +def _component_material_vectors( + profiles, + indices, + *, + component_permittivity, + component_permeability, +): + eps = np.concatenate( + [_material(component_permittivity[name], indices[name], profiles[name].shape) for name in _E_COMPONENTS] + ) + mu = np.concatenate( + [_material(component_permeability[name], indices[name], profiles[name].shape) for name in _H_COMPONENTS] + ) + return eps, mu + + +def _component_vector(profiles, names): + return np.concatenate([np.asarray(profiles[name], dtype=np.complex128).reshape(-1) for name in names]) + + +def _normalized_component_overlap(seed, candidate, names): + left = _component_vector(seed, names) + right = _component_vector(candidate, names) + denominator = np.linalg.norm(left) * np.linalg.norm(right) + if not np.isfinite(denominator) or denominator <= np.finfo(float).tiny: + return 0.0 + return float(abs(np.vdot(left, right)) / denominator) + + +def _rms_impedance(profiles): + electric_norm = np.linalg.norm(_component_vector(profiles, _E_COMPONENTS)) + magnetic_norm = np.linalg.norm(_component_vector(profiles, _H_COMPONENTS)) + if not np.isfinite(electric_norm) or not np.isfinite(magnetic_norm) or magnetic_norm <= np.finfo(float).tiny: + return np.inf + return float(electric_norm / magnetic_norm) + + +def _energy_ratio(profiles, eps, mu): + electric = _component_vector(profiles, _E_COMPONENTS) + magnetic = _component_vector(profiles, _H_COMPONENTS) + electric_energy = _EPS0 * float(np.real(np.vdot(electric, eps * electric))) + magnetic_energy = _MU0 * float(np.real(np.vdot(magnetic, mu * magnetic))) + if not np.isfinite(electric_energy) or not np.isfinite(magnetic_energy) or magnetic_energy <= np.finfo(float).tiny: + return np.inf + return float(electric_energy / magnetic_energy) + + +def _ratio_change(candidate, seed): + if not np.isfinite(candidate) or not np.isfinite(seed) or min(candidate, seed) <= np.finfo(float).tiny: + return np.inf + return float(max(candidate / seed, seed / candidate)) + + +def _signed_x_power(profiles, *, resolution, direction_sign): + flux = np.vdot(profiles["Hz"].reshape(-1), profiles["Ey"].reshape(-1)) + flux -= np.vdot(profiles["Hy"].reshape(-1), profiles["Ez"].reshape(-1)) + return float(0.5 * float(direction_sign) * np.real(flux) * float(resolution) ** 2) + + +def _relative_pair_residual(left, right): + denominator = max(np.linalg.norm(left), np.linalg.norm(right), np.finfo(float).eps) + return float(np.linalg.norm(left - right) / denominator) + + +def _material(values, index, shape): + arr = np.asarray(values, dtype=np.complex128) + profile = arr[index] if arr.ndim == 3 else arr + profile = np.asarray(profile, dtype=np.complex128).squeeze() + if profile.shape != shape: + raise ValueError(f"Material shape {profile.shape} does not match {shape}") + return profile.reshape(-1) + + +def _forward_difference(sparse, count, spacing): + rows = np.repeat(np.arange(count - 1), 2) + cols = np.column_stack((np.arange(count - 1), np.arange(1, count))).reshape(-1) + data = np.tile(np.asarray([-1.0, 1.0]) / spacing, count - 1) + return sparse.csc_matrix((data, (rows, cols)), shape=(count - 1, count)) diff --git a/tests/test_micromode_api.py b/tests/test_micromode_api.py index 19a7b73..be9c790 100644 --- a/tests/test_micromode_api.py +++ b/tests/test_micromode_api.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Sequence +from dataclasses import replace from typing import Any import numpy as np @@ -10,6 +11,7 @@ import xarray as xr import micromode as mm +from micromode.beamz import _boundary_refractive_index def _linspace_edges(start: float, stop: float, count: int) -> tuple[float, ...]: @@ -117,6 +119,22 @@ def _beamz_mode_plane_spec(axis: str = "x") -> mm.ModePlaneSpec: ) +def _with_component_materials(spec: mm.ModePlaneSpec, *, yee_refinement: bool | None = None) -> mm.ModePlaneSpec: + component_permittivity = { + name: np.full(spec.component_shapes[name], 2.25, dtype=np.float64) for name in ("Ex", "Ey", "Ez") + } + component_permeability = { + name: np.ones(spec.component_shapes[name], dtype=np.float64) for name in ("Hx", "Hy", "Hz") + } + updates: dict[str, object] = { + "component_permittivity": component_permittivity, + "component_permeability": component_permeability, + } + if yee_refinement is not None: + updates["yee_refinement"] = yee_refinement + return replace(spec, **updates) + + def test_grid_api_solves_with_scipy_solver(): """Verify grid api solves with scipy solver.""" eps, x_edges, y_edges = _strip_grid(6, 5) @@ -178,6 +196,63 @@ def test_beamz_mode_plane_contract_returns_component_local_profiles(axis: str): assert abs(float(discrete.diagnostics["power_after_phase_reference"])) == pytest.approx(1.0, rel=2e-6, abs=2e-6) +def test_beamz_yee_refinement_can_be_disabled(monkeypatch): + """Verify callers can explicitly retain the unrefined mode profile.""" + import micromode.beamz as beamz_module + + spec = _with_component_materials(_beamz_mode_plane_spec("x"), yee_refinement=False) + + def unexpected_refinement(*_args, **_kwargs): + raise AssertionError("disabled mode solve must not call Yee refinement") + + monkeypatch.setattr(beamz_module, "refine_x_mode_at_fixed_beta", unexpected_refinement) + discrete = mm.solve_beamz_mode(spec) + + assert discrete.diagnostics["yee_refinement_requested"] is False + assert discrete.diagnostics["yee_refinement_attempted"] is False + assert discrete.diagnostics["yee_refinement_accepted"] is False + assert abs(float(discrete.diagnostics["power_after_phase_reference"])) == pytest.approx(1.0, rel=2e-6) + + +def test_beamz_rejects_unphysical_refinement_and_falls_back_atomically(monkeypatch): + """Verify a broken E/H candidate cannot replace the original profiles or k.""" + import micromode.beamz as beamz_module + + base_spec = _with_component_materials(_beamz_mode_plane_spec("x"), yee_refinement=False) + baseline = mm.solve_beamz_mode(base_spec) + + def broken_refinement(profiles, _indices, **kwargs): + candidate = {name: np.asarray(value, dtype=np.complex128).copy() for name, value in profiles.items()} + for component in ("Ex", "Ey", "Ez"): + candidate[component] *= 1e11 + return candidate, 1e-15, 1.0, 1.2 * float(kwargs["k_num"]), 1.0 + + monkeypatch.setattr(beamz_module, "refine_x_mode_at_fixed_beta", broken_refinement) + refined_spec = _with_component_materials(_beamz_mode_plane_spec("x")) + assert refined_spec.yee_refinement is True + guarded = mm.solve_beamz_mode(refined_spec) + + assert guarded.diagnostics["yee_refinement_requested"] is True + assert guarded.diagnostics["yee_refinement_attempted"] is True + assert guarded.diagnostics["yee_refinement_accepted"] is False + assert "E/H impedance change" in guarded.diagnostics["yee_refinement_rejection_reason"] + assert guarded.k_num_axis == pytest.approx(baseline.k_num_axis) + for component in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"): + np.testing.assert_allclose(guarded.profiles[component], baseline.profiles[component], rtol=1e-9, atol=1e-9) + assert abs(float(guarded.diagnostics["power_after_phase_reference"])) == pytest.approx(1.0, rel=2e-6) + + +def test_beamz_boundary_index_uses_all_mode_plane_edges(): + eps = np.ones((5, 6), dtype=float) + eps[0, 2] = 1.4**2 + eps[-1, 3] = 1.5**2 + eps[2, 0] = 1.6**2 + eps[3, -1] = 1.7**2 + eps[2, 3] = 3.4**2 + + assert _boundary_refractive_index(eps) == pytest.approx(1.7) + + def test_beamz_mode_plane_contract_validates_direction_axis(): """Verify BEAMZ contract rejects mismatched direction metadata.""" spec = _beamz_mode_plane_spec("x") diff --git a/tests/test_yee.py b/tests/test_yee.py new file mode 100644 index 0000000..34ea12b --- /dev/null +++ b/tests/test_yee.py @@ -0,0 +1,122 @@ +import numpy as np +import pytest + +from micromode.yee import refine_x_mode_at_fixed_beta, validate_x_mode_refinement + + +def _refinement_fixture(): + nz, ny = 8, 9 + shapes = { + "Ex": (nz, ny), + "Ey": (nz, ny - 1), + "Ez": (nz - 1, ny), + "Hx": (nz - 1, ny - 1), + "Hy": (nz - 1, ny), + "Hz": (nz, ny - 1), + } + z = np.sin(np.pi * (np.arange(nz) + 0.5) / nz)[:, None] + y = np.sin(np.pi * (np.arange(ny - 1) + 1.0) / ny)[None, :] + profiles = {name: np.zeros(shape, dtype=np.complex128) for name, shape in shapes.items()} + profiles["Ey"] = z * y + profiles["Hz"] = profiles["Ey"] / 250.0 + indices = {name: (slice(None), slice(None)) for name in shapes} + eps = {name: np.full(shape, 2.25) for name, shape in shapes.items() if name.startswith("E")} + mu = {name: np.ones(shape) for name, shape in shapes.items() if name.startswith("H")} + kwargs = { + "component_permittivity": eps, + "component_permeability": mu, + "omega": 2.0 * np.pi * 193.0e12, + "dt": 8.0e-17, + "resolution": 80.0e-9, + "k_num": 1.1e7, + "direction_sign": 1.0, + } + return shapes, profiles, indices, kwargs + + +def test_fixed_beta_refinement_returns_normalized_yee_shapes(): + shapes, profiles, indices, kwargs = _refinement_fixture() + + refined, residual, ratio, k_num, initial_ratio = refine_x_mode_at_fixed_beta( + profiles, + indices, + **kwargs, + ) + + assert {name: value.shape for name, value in refined.items()} == shapes + assert residual < 1e-8 + assert abs(ratio - 1.0) < abs(initial_ratio - 1.0) + assert k_num > 0.0 + assert initial_ratio > 0.0 + + +def test_refinement_validator_accepts_balanced_maxwell_mode(): + _shapes, profiles, indices, kwargs = _refinement_fixture() + refined, _residual, _ratio, k_num, _initial_ratio = refine_x_mode_at_fixed_beta(profiles, indices, **kwargs) + + accepted, diagnostics = validate_x_mode_refinement( + profiles, + refined, + indices, + **{**kwargs, "k_num": k_num}, + ) + + assert accepted + assert diagnostics["magnetic_overlap"] > 0.9 + assert diagnostics["impedance_change"] < 1.01 + assert diagnostics["energy_ratio_change"] < 1.01 + assert diagnostics["faraday_residual"] < 1e-12 + assert diagnostics["ampere_residual"] < 5e-3 + + +def test_refinement_validator_rejects_broken_e_h_balance(): + _shapes, profiles, indices, kwargs = _refinement_fixture() + refined, _residual, _ratio, k_num, _initial_ratio = refine_x_mode_at_fixed_beta(profiles, indices, **kwargs) + broken = dict(refined) + for component in ("Ex", "Ey", "Ez"): + broken[component] = refined[component] * 1e11 + + accepted, diagnostics = validate_x_mode_refinement( + profiles, + broken, + indices, + **{**kwargs, "k_num": k_num}, + ) + + assert not accepted + assert diagnostics["impedance_change"] > 1e10 + assert diagnostics["energy_ratio_change"] > 1e20 + assert "E/H impedance change" in diagnostics["rejection_reason"] + assert "electric/magnetic energy change" in diagnostics["rejection_reason"] + + +def test_refinement_validator_treats_electric_overlap_as_diagnostic_by_default(): + _shapes, profiles, indices, kwargs = _refinement_fixture() + refined, _residual, _ratio, k_num, _initial_ratio = refine_x_mode_at_fixed_beta(profiles, indices, **kwargs) + electric = np.concatenate([refined[name].reshape(-1) for name in ("Ex", "Ey", "Ez")]) + orthogonal = np.roll(electric, 1) + orthogonal -= electric * (np.vdot(electric, orthogonal) / np.vdot(electric, electric)) + orthogonal *= np.linalg.norm(electric) / np.linalg.norm(orthogonal) + low_overlap = 0.1 * electric + np.sqrt(1.0 - 0.1**2) * orthogonal + seed = {name: np.asarray(value).copy() for name, value in refined.items()} + split = np.cumsum([refined[name].size for name in ("Ex", "Ey")]) + for name, values in zip(("Ex", "Ey", "Ez"), np.split(low_overlap, split), strict=True): + seed[name] = values.reshape(refined[name].shape) + + accepted, diagnostics = validate_x_mode_refinement( + seed, + refined, + indices, + **{**kwargs, "k_num": k_num}, + ) + explicitly_rejected, _ = validate_x_mode_refinement( + seed, + refined, + indices, + **{**kwargs, "k_num": k_num}, + minimum_electric_overlap=0.2, + ) + + assert diagnostics["electric_overlap"] == pytest.approx(0.1) + assert accepted + assert not explicitly_rejected