meanas/meanas/test/test_eme_numerics.py
2026-04-21 19:40:32 -07:00

234 lines
7.6 KiB
Python

import numpy
import pytest
from scipy import sparse
from ..fdmath import vec
from ..fdfd import eme, waveguide_2d, waveguide_cyl
from ._test_builders import complex_ramp, unit_dxes
from .utils import assert_close
SHAPE = (3, 2, 2)
DXES = unit_dxes((2, 2))
WAVENUMBERS_L = numpy.array([1.0, 0.8])
WAVENUMBERS_R = numpy.array([0.9, 0.7])
OMEGA = 1 / 1500
REAL_DXES = unit_dxes((5, 5))
def _mode(scale: float) -> tuple[numpy.ndarray, numpy.ndarray]:
e_field = complex_ramp(SHAPE, offset=1.0 + scale)
h_field = complex_ramp(SHAPE, scale=0.2, offset=2.0, imag_offset=0.05 * scale)
return vec(e_field), vec(h_field)
def _mode_sets() -> tuple[list[tuple[numpy.ndarray, numpy.ndarray]], list[tuple[numpy.ndarray, numpy.ndarray]]]:
left_modes = [_mode(0.0), _mode(0.7)]
right_modes = [_mode(1.4), _mode(2.1)]
return left_modes, right_modes
def _gain_only_tr(*args, **kwargs) -> tuple[numpy.ndarray, numpy.ndarray]:
return numpy.array([[2.0, 0.0], [0.0, 0.5]]), numpy.zeros((2, 2))
def _gain_and_reflection_tr(*args, **kwargs) -> tuple[numpy.ndarray, numpy.ndarray]:
return numpy.array([[2.0, 0.0], [0.0, 0.5]]), numpy.array([[0.0, 1.0], [2.0, 0.0]])
def _nonsymmetric_tr(left_marker: object):
def fake_get_tr(_eh_left, wavenumbers_left, _eh_right, _wavenumbers_right, **kwargs):
if wavenumbers_left is left_marker:
return (
numpy.array([[1.0, 2.0], [0.5, 1.0]]),
numpy.array([[0.0, 1.0], [2.0, 0.0]]),
)
return (
numpy.array([[1.0, -1.0], [0.0, 1.0]]),
numpy.array([[0.0, 0.5], [1.5, 0.0]]),
)
return fake_get_tr
def test_get_tr_returns_finite_bounded_transfer_matrices() -> None:
left_modes, right_modes = _mode_sets()
transmission, reflection = eme.get_tr(
left_modes,
WAVENUMBERS_L,
right_modes,
WAVENUMBERS_R,
dxes=DXES,
)
singular_values = numpy.linalg.svd(transmission, compute_uv=False)
assert transmission.shape == (2, 2)
assert reflection.shape == (2, 2)
assert numpy.isfinite(transmission).all()
assert numpy.isfinite(reflection).all()
assert (singular_values <= 1.0 + 1e-12).all()
def test_get_abcd_matches_explicit_block_formula() -> None:
left_modes, right_modes = _mode_sets()
t12, r12 = eme.get_tr(left_modes, WAVENUMBERS_L, right_modes, WAVENUMBERS_R, dxes=DXES)
t21, r21 = eme.get_tr(right_modes, WAVENUMBERS_R, left_modes, WAVENUMBERS_L, dxes=DXES)
t21_inv = numpy.linalg.pinv(t21)
expected = numpy.block([
[t12 - r21 @ t21_inv @ r12, r21 @ t21_inv],
[-t21_inv @ r12, t21_inv],
])
abcd = eme.get_abcd(left_modes, WAVENUMBERS_L, right_modes, WAVENUMBERS_R, dxes=DXES)
assert sparse.issparse(abcd)
assert abcd.shape == (4, 4)
assert_close(abcd.toarray(), expected)
def test_get_s_plain_matches_block_assembly_from_get_tr() -> None:
left_modes, right_modes = _mode_sets()
t12, r12 = eme.get_tr(left_modes, WAVENUMBERS_L, right_modes, WAVENUMBERS_R, dxes=DXES)
t21, r21 = eme.get_tr(right_modes, WAVENUMBERS_R, left_modes, WAVENUMBERS_L, dxes=DXES)
expected = numpy.block([[r12, t12], [t21, r21]])
ss = eme.get_s(left_modes, WAVENUMBERS_L, right_modes, WAVENUMBERS_R, dxes=DXES)
assert ss.shape == (4, 4)
assert numpy.isfinite(ss).all()
assert_close(ss, expected)
def test_get_s_force_nogain_caps_singular_values(monkeypatch) -> None:
monkeypatch.setattr(eme, 'get_tr', _gain_only_tr)
plain_s = eme.get_s(None, None, None, None)
clipped_s = eme.get_s(None, None, None, None, force_nogain=True)
plain_singular_values = numpy.linalg.svd(plain_s, compute_uv=False)
clipped_singular_values = numpy.linalg.svd(clipped_s, compute_uv=False)
assert plain_singular_values.max() > 1.0
assert (clipped_singular_values <= 1.0 + 1e-12).all()
assert numpy.isfinite(clipped_s).all()
def test_get_s_force_reciprocal_symmetrizes_output(monkeypatch) -> None:
left = object()
right = object()
monkeypatch.setattr(eme, 'get_tr', _nonsymmetric_tr(left))
ss = eme.get_s(None, left, None, right, force_reciprocal=True)
assert_close(ss, ss.T)
def test_get_s_force_nogain_and_reciprocal_returns_finite_output(monkeypatch) -> None:
monkeypatch.setattr(eme, 'get_tr', _gain_and_reflection_tr)
ss = eme.get_s(None, None, None, None, force_nogain=True, force_reciprocal=True)
assert ss.shape == (4, 4)
assert numpy.isfinite(ss).all()
assert_close(ss, ss.T)
assert (numpy.linalg.svd(ss, compute_uv=False) <= 1.0 + 1e-12).all()
def test_get_tr_rejects_length_mismatches() -> None:
left_modes, right_modes = _mode_sets()
with pytest.raises(ValueError, match='same length'):
eme.get_tr(left_modes[:1], WAVENUMBERS_L, right_modes, WAVENUMBERS_R, dxes=DXES)
def test_get_tr_rejects_malformed_mode_tuples() -> None:
bad_modes = [(numpy.ones(4),)]
with pytest.raises(ValueError, match='2-tuple'):
eme.get_tr(bad_modes, [1.0], bad_modes, [1.0], dxes=DXES)
def test_get_tr_rejects_incompatible_field_shapes() -> None:
left_modes = [(numpy.ones(4), numpy.ones(4))]
right_modes = [(numpy.ones(6), numpy.ones(6))]
with pytest.raises(ValueError, match='same E/H shapes'):
eme.get_tr(left_modes, [1.0], right_modes, [1.0], dxes=DXES)
def _build_real_epsilon() -> numpy.ndarray:
epsilon = numpy.ones((3, 5, 5), dtype=float)
epsilon[:, 2, 1] = 2.0
return vec(epsilon)
def _build_straight_mode() -> tuple[tuple[numpy.ndarray, numpy.ndarray], complex, numpy.ndarray]:
epsilon = _build_real_epsilon()
e_xy, wavenumber = waveguide_2d.solve_mode(
0,
omega=OMEGA,
dxes=REAL_DXES,
epsilon=epsilon,
)
e_field, h_field = waveguide_2d.normalized_fields_e(
e_xy,
wavenumber=wavenumber,
omega=OMEGA,
dxes=REAL_DXES,
epsilon=epsilon,
)
return (e_field, h_field), wavenumber, epsilon
def _build_bend_mode() -> tuple[tuple[numpy.ndarray, numpy.ndarray], complex]:
epsilon = vec(numpy.ones((3, 5, 5), dtype=float))
rmin = 10.0
e_xy, angular_wavenumber = waveguide_cyl.solve_mode(
0,
omega=OMEGA,
dxes=REAL_DXES,
epsilon=epsilon,
rmin=rmin,
)
linear_wavenumber = waveguide_cyl.linear_wavenumbers(
[e_xy],
[angular_wavenumber],
epsilon=epsilon,
dxes=REAL_DXES,
rmin=rmin,
)[0]
e_field, h_field = waveguide_cyl.normalized_fields_e(
e_xy,
angular_wavenumber=angular_wavenumber,
omega=OMEGA,
dxes=REAL_DXES,
epsilon=epsilon,
rmin=rmin,
)
return (e_field, h_field), linear_wavenumber
def test_get_s_is_near_identity_for_identical_solved_straight_modes() -> None:
mode, wavenumber, _epsilon = _build_straight_mode()
ss = eme.get_s([mode], [wavenumber], [mode], [wavenumber], dxes=REAL_DXES)
assert ss.shape == (2, 2)
assert numpy.isfinite(ss).all()
assert abs(ss[0, 0]) < 1e-6
assert abs(ss[1, 1]) < 1e-6
assert abs(abs(ss[0, 1]) - 1.0) < 1e-6
assert abs(abs(ss[1, 0]) - 1.0) < 1e-6
assert numpy.linalg.svd(ss, compute_uv=False).max() <= 1.0 + 1e-10
def test_get_s_returns_finite_passive_output_for_small_straight_to_bend_fixture() -> None:
straight_mode, straight_wavenumber, _epsilon = _build_straight_mode()
bend_mode, bend_wavenumber = _build_bend_mode()
ss = eme.get_s([straight_mode], [straight_wavenumber], [bend_mode], [bend_wavenumber], dxes=REAL_DXES)
assert ss.shape == (2, 2)
assert numpy.isfinite(ss).all()
assert numpy.linalg.svd(ss, compute_uv=False).max() <= 1.0 + 1e-10