[bloch] add some more tests and clean up solves

meanas/fdfd/bloch.py

@@ -451,7 +451,7 @@ def find_k(
         solve_callback: Callable[..., None] | None = None,
         iter_callback: Callable[..., None] | None = None,
         v0: NDArray[numpy.complex128] | None = None,
-        ) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
+        ) -> tuple[NDArray[numpy.float64], float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
     """
     Search for a bloch vector that has a given frequency.
 
@@ -496,15 +496,15 @@ def find_k(
 
     res = scipy.optimize.minimize_scalar(
         lambda x: abs(get_f(x, band) - frequency),
-        k_guess,
+        method='bounded',
-        method='Bounded',
         bounds=k_bounds,
         options={'xatol': abs(tolerance)},
         )
 
     assert n is not None
     assert v is not None
-    return float(res.x * direction), float(res.fun + frequency), n, v
+    actual_frequency = get_f(float(res.x), band)
+    return direction * float(res.x), float(actual_frequency), n, v
 
 
 def eigsolve(
@@ -725,7 +725,12 @@ def eigsolve(
                 amax=pi,
                 )
 
-        result = scipy.optimize.minimize_scalar(trace_func, bounds=(0, pi), tol=tolerance)
+        result = scipy.optimize.minimize_scalar(
+            trace_func,
+            method='bounded',
+            bounds=(0, pi),
+            options={'xatol': tolerance},
+            )
         new_E = result.fun
         theta = result.x
 
@@ -754,7 +759,7 @@ def eigsolve(
         v = eigvecs[:, i]
         n = eigvals[i]
         v /= norm(v)
-        Av = (scipy_op @ v.copy())[:, 0]
+        Av = numpy.asarray(scipy_op @ v.copy()).reshape(-1)
         eigness = norm(Av - (v.conj() @ Av) * v)
         f = numpy.sqrt(-numpy.real(n))
         df = numpy.sqrt(-numpy.real(n) + eigness)
@@ -823,18 +828,18 @@ def inner_product(
     # eRxhR_x = numpy.cross(eR.reshape(3, -1), hR.reshape(3, -1), axis=0).reshape(eR.shape)[0] / normR
     # logger.info(f'power {eRxhR_x.sum() / 2})
 
-    eR /= numpy.sqrt(norm2R)
+    eR_norm = eR / numpy.sqrt(abs(norm2R))
-    hR /= numpy.sqrt(norm2R)
+    hR_norm = hR / numpy.sqrt(abs(norm2R))
-    eL /= numpy.sqrt(norm2L)
+    eL_norm = eL / numpy.sqrt(abs(norm2L))
-    hL /= numpy.sqrt(norm2L)
+    hL_norm = hL / numpy.sqrt(abs(norm2L))
 
     # (eR x hL)[0] and (eL x hR)[0]
-    eRxhL_x = eR[1] * hL[2] - eR[2] - hL[1]
+    eRxhL_x = eR_norm[1] * hL_norm[2] - eR_norm[2] * hL_norm[1]
-    eLxhR_x = eL[1] * hR[2] - eL[2] - hR[1]
+    eLxhR_x = eL_norm[1] * hR_norm[2] - eL_norm[2] * hR_norm[1]
 
     #return 1j *  (eRxhL_x - eLxhR_x).sum() / numpy.sqrt(norm2R * norm2L)
     #return (eRxhL_x.sum() - eLxhR_x.sum()) / numpy.sqrt(norm2R * norm2L)
-    return eRxhL_x.sum() - eLxhR_x.sum()
+    return eLxhR_x.sum() - eRxhL_x.sum()
 
 
 def trq(
@@ -848,8 +853,8 @@ def trq(
     np = inner_product(eO, -hO, eI,  hI)
     nn = inner_product(eO, -hO, eI, -hI)
 
-    assert pp == -nn
+    assert numpy.allclose(pp, -nn, atol=1e-12, rtol=1e-12)
-    assert pn == -np
+    assert numpy.allclose(pn, -np, atol=1e-12, rtol=1e-12)
 
     logger.info(f'''
         {pp=:4g} {pn=:4g}

meanas/test/test_bloch_interactions.py

@@ -0,0 +1,151 @@
+import numpy
+from numpy.testing import assert_allclose
+
+from ..fdfd import bloch
+
+
+SHAPE = (2, 2, 2)
+G_MATRIX = numpy.eye(3)
+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, :]
+
+
+def build_overlap_fixture() -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+    e_in = numpy.zeros((3, *SHAPE), dtype=complex)
+    h_in = numpy.zeros_like(e_in)
+    e_out = numpy.zeros_like(e_in)
+    h_out = numpy.zeros_like(e_in)
+
+    e_in[1] = 1.0
+    h_in[2] = 2.0
+    e_out[1] = 3.0
+    h_out[2] = 4.0
+    return e_in, h_in, e_out, h_out
+
+
+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_allclose(result, 0.5 * (matrix + matrix.conj().T))
+    assert_allclose(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_allclose(pp, 0.8164965809277263 + 0.0j)
+    assert_allclose(pp, -nn, atol=1e-12, rtol=1e-12)
+    assert_allclose(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_allclose(transmission, 0.9797958971132713 + 0.0j, atol=1e-12, rtol=1e-12)
+    assert_allclose(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,
+        G_MATRIX,
+        EPSILON,
+        tolerance=1e-6,
+        max_iters=50,
+        y0=Y0,
+        callback=callback,
+    )
+
+    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
+    assert callback_count > 0
+
+
+def test_find_k_returns_vector_frequency_and_callbacks() -> None:
+    target_eigvals, _target_eigvecs = bloch.eigsolve(
+        1,
+        K0,
+        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_allclose(numpy.cross(found_k, [1.0, 0.0, 0.0]), 0.0, atol=1e-12, rtol=1e-12)
+    assert_allclose(found_k, K0, 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