315 lines
9.9 KiB
Python
315 lines
9.9 KiB
Python
import numpy
|
|
import pytest
|
|
from numpy.testing import assert_allclose
|
|
from types import SimpleNamespace
|
|
|
|
from ..fdfd import bloch
|
|
from ._bloch_case import EPSILON, G_MATRIX, H_SIZE, K0_X, Y0, Y0_TWO_MODE, build_overlap_fixture
|
|
from .utils import assert_close
|
|
|
|
|
|
def test_rtrace_atb_matches_real_frobenius_inner_product() -> None:
|
|
a_mat = numpy.array([[1.0 + 2.0j, 3.0 - 1.0j], [2.0j, 4.0]], dtype=complex)
|
|
b_mat = numpy.array([[5.0 - 1.0j, 1.0 + 1.0j], [2.0, 3.0j]], dtype=complex)
|
|
expected = numpy.real(numpy.sum(a_mat.conj() * b_mat))
|
|
|
|
assert bloch._rtrace_AtB(a_mat, b_mat) == expected
|
|
|
|
|
|
def test_symmetrize_returns_hermitian_average() -> None:
|
|
matrix = numpy.array([[1.0 + 2.0j, 3.0 - 1.0j], [2.0j, 4.0]], dtype=complex)
|
|
result = bloch._symmetrize(matrix)
|
|
|
|
assert_close(result, 0.5 * (matrix + matrix.conj().T))
|
|
assert_close(result, result.conj().T)
|
|
|
|
|
|
def test_inner_product_is_nonmutating_and_obeys_sign_symmetry() -> None:
|
|
e_in, h_in, e_out, h_out = build_overlap_fixture()
|
|
originals = (e_in.copy(), h_in.copy(), e_out.copy(), h_out.copy())
|
|
|
|
pp = bloch.inner_product(e_out, h_out, e_in, h_in)
|
|
pn = bloch.inner_product(e_out, h_out, e_in, -h_in)
|
|
np_term = bloch.inner_product(e_out, -h_out, e_in, h_in)
|
|
nn = bloch.inner_product(e_out, -h_out, e_in, -h_in)
|
|
|
|
assert_close(pp, 0.8164965809277263 + 0.0j)
|
|
assert_close(pp, -nn, atol=1e-12, rtol=1e-12)
|
|
assert_close(pn, -np_term, atol=1e-12, rtol=1e-12)
|
|
assert numpy.array_equal(e_in, originals[0])
|
|
assert numpy.array_equal(h_in, originals[1])
|
|
assert numpy.array_equal(e_out, originals[2])
|
|
assert numpy.array_equal(h_out, originals[3])
|
|
|
|
|
|
def test_trq_returns_expected_transmission_and_reflection() -> None:
|
|
e_in, h_in, e_out, h_out = build_overlap_fixture()
|
|
|
|
transmission, reflection = bloch.trq(e_in, h_in, e_out, h_out)
|
|
|
|
assert_close(transmission, 0.9797958971132713 + 0.0j, atol=1e-12, rtol=1e-12)
|
|
assert_close(reflection, 0.2 + 0.0j, atol=1e-12, rtol=1e-12)
|
|
|
|
|
|
def test_eigsolve_returns_finite_modes_with_small_residual() -> None:
|
|
callback_count = 0
|
|
|
|
def callback() -> None:
|
|
nonlocal callback_count
|
|
callback_count += 1
|
|
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=50,
|
|
y0=Y0,
|
|
callback=callback,
|
|
)
|
|
|
|
operator = bloch.maxwell_operator(K0_X, G_MATRIX, EPSILON)
|
|
eigvec = eigvecs[0] / numpy.linalg.norm(eigvecs[0])
|
|
residual = numpy.linalg.norm(operator(eigvec).reshape(-1) - eigvals[0] * eigvec) / numpy.linalg.norm(eigvec)
|
|
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
assert residual < 1e-5
|
|
assert callback_count > 0
|
|
|
|
|
|
def test_eigsolve_without_initial_guess_returns_finite_modes() -> None:
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=20,
|
|
y0=None,
|
|
)
|
|
|
|
operator = bloch.maxwell_operator(K0_X, G_MATRIX, EPSILON)
|
|
eigvec = eigvecs[0] / numpy.linalg.norm(eigvecs[0])
|
|
residual = numpy.linalg.norm(operator(eigvec).reshape(-1) - eigvals[0] * eigvec) / numpy.linalg.norm(eigvec)
|
|
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
assert residual < 1e-5
|
|
|
|
|
|
def test_eigsolve_recovers_from_singular_initial_subspace(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
class FakeRng:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def random(self, shape: tuple[int, ...]) -> numpy.ndarray:
|
|
self.calls += 1
|
|
return numpy.arange(numpy.prod(shape), dtype=float).reshape(shape) + self.calls
|
|
|
|
fake_rng = FakeRng()
|
|
singular_y0 = numpy.vstack([Y0_TWO_MODE[0], Y0_TWO_MODE[0]])
|
|
monkeypatch.setattr(bloch.numpy.random, 'default_rng', lambda: fake_rng)
|
|
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
2,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=20,
|
|
y0=singular_y0,
|
|
)
|
|
|
|
assert fake_rng.calls == 2
|
|
assert eigvals.shape == (2,)
|
|
assert eigvecs.shape == (2, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
|
|
|
|
def test_eigsolve_reconditions_large_trace_initial_subspace(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
original_inv = bloch.numpy.linalg.inv
|
|
original_sqrtm = bloch.scipy.linalg.sqrtm
|
|
sqrtm_calls = 0
|
|
inv_calls = 0
|
|
|
|
def inv_with_large_first_trace(matrix: numpy.ndarray) -> numpy.ndarray:
|
|
nonlocal inv_calls
|
|
inv_calls += 1
|
|
if inv_calls == 1:
|
|
return numpy.eye(matrix.shape[0], dtype=complex) * 1e9
|
|
return original_inv(matrix)
|
|
|
|
def sqrtm_wrapper(matrix: numpy.ndarray) -> numpy.ndarray:
|
|
nonlocal sqrtm_calls
|
|
sqrtm_calls += 1
|
|
return original_sqrtm(matrix)
|
|
|
|
monkeypatch.setattr(bloch.numpy.linalg, 'inv', inv_with_large_first_trace)
|
|
monkeypatch.setattr(bloch.scipy.linalg, 'sqrtm', sqrtm_wrapper)
|
|
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
2,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=20,
|
|
y0=Y0_TWO_MODE,
|
|
)
|
|
|
|
assert sqrtm_calls >= 2
|
|
assert eigvals.shape == (2,)
|
|
assert eigvecs.shape == (2, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
|
|
|
|
def test_eigsolve_qi_memoization_reuses_cached_theta(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
def fake_minimize_scalar(func, method: str, bounds: tuple[float, float], options: dict[str, float]) -> SimpleNamespace:
|
|
theta = 0.3
|
|
first = func(theta)
|
|
second = func(theta)
|
|
assert_allclose(second, first)
|
|
return SimpleNamespace(fun=second, x=theta)
|
|
|
|
monkeypatch.setattr(bloch.scipy.optimize, 'minimize_scalar', fake_minimize_scalar)
|
|
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=1,
|
|
y0=Y0,
|
|
)
|
|
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
|
|
|
|
@pytest.mark.parametrize('theta', [numpy.pi / 2 - 1e-8, 1e-8])
|
|
def test_eigsolve_qi_taylor_expansions_return_finite_modes(monkeypatch: pytest.MonkeyPatch, theta: float) -> None:
|
|
original_inv = bloch.numpy.linalg.inv
|
|
inv_calls = 0
|
|
|
|
def inv_raise_once_for_q(matrix: numpy.ndarray) -> numpy.ndarray:
|
|
nonlocal inv_calls
|
|
inv_calls += 1
|
|
if inv_calls == 3:
|
|
raise numpy.linalg.LinAlgError('forced singular Q')
|
|
return original_inv(matrix)
|
|
|
|
def fake_minimize_scalar(func, method: str, bounds: tuple[float, float], options: dict[str, float]) -> SimpleNamespace:
|
|
value = func(theta)
|
|
return SimpleNamespace(fun=value, x=theta)
|
|
|
|
monkeypatch.setattr(bloch.numpy.linalg, 'inv', inv_raise_once_for_q)
|
|
monkeypatch.setattr(bloch.scipy.optimize, 'minimize_scalar', fake_minimize_scalar)
|
|
|
|
eigvals, eigvecs = bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=1,
|
|
y0=Y0,
|
|
)
|
|
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
|
|
|
|
def test_eigsolve_qi_inexplicable_singularity_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
original_inv = bloch.numpy.linalg.inv
|
|
inv_calls = 0
|
|
|
|
def inv_raise_once_for_q(matrix: numpy.ndarray) -> numpy.ndarray:
|
|
nonlocal inv_calls
|
|
inv_calls += 1
|
|
if inv_calls == 3:
|
|
raise numpy.linalg.LinAlgError('forced singular Q')
|
|
return original_inv(matrix)
|
|
|
|
def fake_minimize_scalar(func, method: str, bounds: tuple[float, float], options: dict[str, float]) -> SimpleNamespace:
|
|
func(numpy.pi / 4)
|
|
raise AssertionError('unreachable after trace_func exception')
|
|
|
|
monkeypatch.setattr(bloch.numpy.linalg, 'inv', inv_raise_once_for_q)
|
|
monkeypatch.setattr(bloch.scipy.optimize, 'minimize_scalar', fake_minimize_scalar)
|
|
|
|
with pytest.raises(Exception, match='Inexplicable singularity in trace_func'):
|
|
bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=1,
|
|
y0=Y0,
|
|
)
|
|
|
|
|
|
def test_find_k_returns_vector_frequency_and_callbacks() -> None:
|
|
target_eigvals, _target_eigvecs = bloch.eigsolve(
|
|
1,
|
|
K0_X,
|
|
G_MATRIX,
|
|
EPSILON,
|
|
tolerance=1e-6,
|
|
max_iters=50,
|
|
y0=Y0,
|
|
)
|
|
target_frequency = float(numpy.sqrt(abs(numpy.real(target_eigvals[0]))))
|
|
|
|
solve_calls = 0
|
|
iter_calls = 0
|
|
|
|
def solve_callback(k_mag: float, eigvals: numpy.ndarray, eigvecs: numpy.ndarray, frequency: float) -> None:
|
|
nonlocal solve_calls
|
|
solve_calls += 1
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert isinstance(k_mag, float)
|
|
assert isinstance(frequency, float)
|
|
|
|
def iter_callback() -> None:
|
|
nonlocal iter_calls
|
|
iter_calls += 1
|
|
|
|
found_k, found_frequency, eigvals, eigvecs = bloch.find_k(
|
|
target_frequency,
|
|
1e-4,
|
|
[1, 0, 0],
|
|
G_MATRIX,
|
|
EPSILON,
|
|
band=0,
|
|
k_bounds=(0.05, 0.15),
|
|
v0=Y0,
|
|
solve_callback=solve_callback,
|
|
iter_callback=iter_callback,
|
|
)
|
|
|
|
assert found_k.shape == (3,)
|
|
assert numpy.isfinite(found_k).all()
|
|
assert_close(numpy.cross(found_k, [1.0, 0.0, 0.0]), 0.0, atol=1e-12, rtol=1e-12)
|
|
assert_close(found_k, K0_X, atol=1e-4, rtol=1e-4)
|
|
assert abs(found_frequency - target_frequency) <= 1e-4
|
|
assert eigvals.shape == (1,)
|
|
assert eigvecs.shape == (1, H_SIZE)
|
|
assert numpy.isfinite(eigvals).all()
|
|
assert numpy.isfinite(eigvecs).all()
|
|
assert solve_calls > 0
|
|
assert iter_calls > 0
|