[tests] more test coverage

meanas/test/test_bloch_interactions.py

@@ -1,5 +1,7 @@
 import numpy
+import pytest
 from numpy.testing import assert_allclose
+from types import SimpleNamespace
 
 from ..fdfd import bloch
 
@@ -10,6 +12,12 @@ EPSILON = numpy.ones((3, *SHAPE), dtype=float)
 K0 = numpy.array([0.1, 0.0, 0.0], dtype=float)
 H_SIZE = 2 * numpy.prod(SHAPE)
 Y0 = (numpy.arange(H_SIZE, dtype=float) + 1j * numpy.linspace(0.1, 0.9, H_SIZE))[None, :]
+Y0_TWO_MODE = numpy.vstack(
+    [
+        numpy.arange(H_SIZE, dtype=float) + 1j * numpy.linspace(0.1, 0.9, H_SIZE),
+        numpy.linspace(2.0, 3.5, H_SIZE) - 0.5j * numpy.arange(H_SIZE, dtype=float),
+    ],
+)
 
 
 def build_overlap_fixture() -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
@@ -98,6 +106,187 @@ def test_eigsolve_returns_finite_modes_with_small_residual() -> None:
     assert callback_count > 0
 
 
+def test_eigsolve_without_initial_guess_returns_finite_modes() -> None:
+    eigvals, eigvecs = bloch.eigsolve(
+        1,
+        K0,
+        G_MATRIX,
+        EPSILON,
+        tolerance=1e-6,
+        max_iters=20,
+        y0=None,
+    )
+
+    operator = bloch.maxwell_operator(K0, 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,
+        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,
+        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,
+        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,
+        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,
+            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,