[fdfd] minor fixes and more tests

meanas/fdfd/scpml.py

@@ -128,6 +128,11 @@ def stretch_with_scpml(
     dx_ai = dxes[0][axis].astype(complex)
     dx_bi = dxes[1][axis].astype(complex)
 
+    if thickness == 0:
+        dxes[0][axis] = dx_ai
+        dxes[1][axis] = dx_bi
+        return dxes
+
     pos = numpy.hstack((0, dx_ai.cumsum()))
     pos_a = (pos[:-1] + pos[1:]) / 2
     pos_b = pos[:-1]
@@ -153,10 +158,7 @@ def stretch_with_scpml(
         def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
             return (x - bound) / (pos[-1] - bound)
 
-        if thickness == 0:
+        slc = slice(-thickness, None)
-            slc = slice(None)
-        else:
-            slc = slice(-thickness, None)
 
     dx_ai[slc] *= 1 + 1j * s_function(l_d(pos_a[slc])) / d / s_correction
     dx_bi[slc] *= 1 + 1j * s_function(l_d(pos_b[slc])) / d / s_correction

meanas/fdfd/solvers.py

@@ -48,9 +48,11 @@ def _scipy_qmr(
             logger.info(f'Solver residual at iteration {ii} : {cur_norm}')
 
     if 'callback' in kwargs:
+        callback = kwargs['callback']
+
         def augmented_callback(xk: ArrayLike) -> None:
             log_residual(xk)
-            kwargs['callback'](xk)
+            callback(xk)
 
         kwargs['callback'] = augmented_callback
     else:
@@ -118,15 +120,15 @@ def generic(
     Pl, Pr = operators.e_full_preconditioners(dxes)
 
     if adjoint:
-        A = (Pl @ A0 @ Pr).H
+        A = (Pl @ A0 @ Pr).T.conjugate()
-        b = Pr.H @ b0
+        b = Pr.T.conjugate() @ b0
     else:
         A = Pl @ A0 @ Pr
         b = Pl @ b0
 
     if E_guess is not None:
         if adjoint:
-            x0 = Pr.H @ E_guess
+            x0 = Pr.T.conjugate() @ E_guess
         else:
             x0 = Pl @ E_guess
         matrix_solver_opts['x0'] = x0
@@ -134,7 +136,7 @@ def generic(
     x = matrix_solver(A.tocsr(), b, **matrix_solver_opts)
 
     if adjoint:
-        x0 = Pl.H @ x
+        x0 = Pl.T.conjugate() @ x
     else:
         x0 = Pr @ x
 

meanas/test/test_fdfd_algebra_helpers.py

@@ -0,0 +1,170 @@
+import numpy
+from numpy.testing import assert_allclose
+
+from ..fdmath import vec, unvec
+from ..fdmath import functional as fd_functional
+from ..fdfd import operators, scpml
+
+
+OMEGA = 1 / 1500
+SHAPE = (2, 3, 2)
+DXES = [
+    [numpy.array([1.0, 1.5]), numpy.array([0.75, 1.25, 1.5]), numpy.array([1.2, 0.8])],
+    [numpy.array([0.9, 1.4]), numpy.array([0.8, 1.1, 1.4]), numpy.array([1.0, 0.7])],
+]
+
+EPSILON = numpy.stack([
+    numpy.linspace(1.0, 2.2, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(1.1, 2.3, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(1.2, 2.4, numpy.prod(SHAPE)).reshape(SHAPE),
+])
+MU = numpy.stack([
+    numpy.linspace(2.0, 3.2, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(2.1, 3.3, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(2.2, 3.4, numpy.prod(SHAPE)).reshape(SHAPE),
+])
+
+H_FIELD = (numpy.arange(3 * numpy.prod(SHAPE)).reshape((3, *SHAPE)) * 0.25 - 0.75j).astype(complex)
+E_FIELD = (numpy.arange(3 * numpy.prod(SHAPE)).reshape((3, *SHAPE)) + 0.5j).astype(complex)
+
+PEC = numpy.zeros((3, *SHAPE), dtype=float)
+PEC[1, 0, 1, 0] = 1.0
+PMC = numpy.zeros((3, *SHAPE), dtype=float)
+PMC[2, 1, 2, 1] = 1.0
+
+BOUNDARY_SHAPE = (3, 4, 3)
+BOUNDARY_DXES = [
+    [numpy.array([1.0, 1.5, 0.8]), numpy.array([0.75, 1.25, 1.5, 0.9]), numpy.array([1.2, 0.8, 1.1])],
+    [numpy.array([0.9, 1.4, 1.0]), numpy.array([0.8, 1.1, 1.4, 1.0]), numpy.array([1.0, 0.7, 1.3])],
+]
+BOUNDARY_EPSILON = numpy.stack([
+    numpy.linspace(1.0, 2.2, numpy.prod(BOUNDARY_SHAPE)).reshape(BOUNDARY_SHAPE),
+    numpy.linspace(1.1, 2.3, numpy.prod(BOUNDARY_SHAPE)).reshape(BOUNDARY_SHAPE),
+    numpy.linspace(1.2, 2.4, numpy.prod(BOUNDARY_SHAPE)).reshape(BOUNDARY_SHAPE),
+])
+BOUNDARY_FIELD = (numpy.arange(3 * numpy.prod(BOUNDARY_SHAPE)).reshape((3, *BOUNDARY_SHAPE)) + 0.5j).astype(complex)
+
+
+def _apply_matrix(op: operators.sparse.spmatrix, field: numpy.ndarray, shape: tuple[int, ...]) -> numpy.ndarray:
+    return unvec(op @ vec(field), shape)
+
+
+def _dense_h_full(mu: numpy.ndarray | None) -> numpy.ndarray:
+    ce = fd_functional.curl_forward(DXES[0])
+    ch = fd_functional.curl_back(DXES[1])
+    pe = numpy.where(PEC, 0.0, 1.0)
+    pm = numpy.where(PMC, 0.0, 1.0)
+    magnetic = numpy.ones_like(EPSILON) if mu is None else mu
+
+    masked_h = pm * H_FIELD
+    curl_term = ch(masked_h)
+    curl_term = pe * (curl_term / EPSILON)
+    curl_term = ce(curl_term)
+    return pm * (curl_term - OMEGA**2 * magnetic * masked_h)
+
+
+def _normalized_distance(u: numpy.ndarray, size: int, thickness: int) -> numpy.ndarray:
+    return ((thickness - u).clip(0) + (u - (size - thickness)).clip(0)) / thickness
+
+
+def test_h_full_matches_dense_reference_with_and_without_mu() -> None:
+    for mu in (None, MU):
+        matrix_result = _apply_matrix(
+            operators.h_full(OMEGA, DXES, vec(EPSILON), None if mu is None else vec(mu), vec(PEC), vec(PMC)),
+            H_FIELD,
+            SHAPE,
+        )
+        dense_result = _dense_h_full(mu)
+        assert_allclose(matrix_result, dense_result, atol=1e-10, rtol=1e-10)
+
+
+def test_poynting_h_cross_matches_negative_e_cross_relation() -> None:
+    h_cross_e = _apply_matrix(operators.poynting_h_cross(vec(H_FIELD), DXES), E_FIELD, SHAPE)
+    e_cross_h = _apply_matrix(operators.poynting_e_cross(vec(E_FIELD), DXES), H_FIELD, SHAPE)
+
+    assert_allclose(h_cross_e, -e_cross_h, atol=1e-10, rtol=1e-10)
+
+
+def test_e_boundary_source_interior_mask_is_independent_of_periodic_edges() -> None:
+    mask = numpy.zeros((3, *BOUNDARY_SHAPE), dtype=float)
+    mask[:, 1, 1, 1] = 1.0
+
+    periodic = operators.e_boundary_source(vec(mask), OMEGA, BOUNDARY_DXES, vec(BOUNDARY_EPSILON), periodic_mask_edges=True)
+    mirrored = operators.e_boundary_source(vec(mask), OMEGA, BOUNDARY_DXES, vec(BOUNDARY_EPSILON), periodic_mask_edges=False)
+
+    assert_allclose(periodic.toarray(), mirrored.toarray())
+
+
+def test_e_boundary_source_periodic_edges_add_opposite_face_response() -> None:
+    mask = numpy.zeros((3, *BOUNDARY_SHAPE), dtype=float)
+    mask[:, 0, 1, 1] = 1.0
+
+    periodic = operators.e_boundary_source(vec(mask), OMEGA, BOUNDARY_DXES, vec(BOUNDARY_EPSILON), periodic_mask_edges=True)
+    mirrored = operators.e_boundary_source(vec(mask), OMEGA, BOUNDARY_DXES, vec(BOUNDARY_EPSILON), periodic_mask_edges=False)
+    diff = unvec((periodic - mirrored) @ vec(BOUNDARY_FIELD), BOUNDARY_SHAPE)
+
+    assert numpy.isfinite(diff).all()
+    assert_allclose(diff[:, 1:-1, :, :], 0.0)
+    assert numpy.linalg.norm(diff[:, -1, :, :]) > 0
+
+
+def test_prepare_s_function_matches_closed_form_polynomial() -> None:
+    ln_r = -12.0
+    order = 3.0
+    distances = numpy.array([0.0, 0.25, 0.5, 1.0])
+    s_function = scpml.prepare_s_function(ln_R=ln_r, m=order)
+    expected = (order + 1) * ln_r / 2 * distances**order
+
+    assert_allclose(s_function(distances), expected)
+
+
+def test_uniform_grid_scpml_matches_expected_stretch_profile() -> None:
+    s_function = scpml.prepare_s_function(ln_R=-12.0, m=3.0)
+    dxes = scpml.uniform_grid_scpml((6, 4, 3), (2, 0, 1), omega=2.0, epsilon_effective=4.0, s_function=s_function)
+    correction = numpy.sqrt(4.0) * 2.0
+
+    for axis, size, thickness in ((0, 6, 2), (2, 3, 1)):
+        grid = numpy.arange(size, dtype=float)
+        expected_a = 1 + 1j * s_function(_normalized_distance(grid, size, thickness)) / correction
+        expected_b = 1 + 1j * s_function(_normalized_distance(grid + 0.5, size, thickness)) / correction
+        assert_allclose(dxes[0][axis], expected_a)
+        assert_allclose(dxes[1][axis], expected_b)
+
+    assert_allclose(dxes[0][1], 1.0)
+    assert_allclose(dxes[1][1], 1.0)
+    assert numpy.isfinite(dxes[0][0]).all()
+    assert numpy.isfinite(dxes[1][0]).all()
+
+
+def test_stretch_with_scpml_only_modifies_requested_front_edge() -> None:
+    s_function = scpml.prepare_s_function(ln_R=-12.0, m=3.0)
+    base = [[numpy.ones(6), numpy.ones(4), numpy.ones(3)] for _ in range(2)]
+    stretched = scpml.stretch_with_scpml(base, axis=0, polarity=1, omega=2.0, epsilon_effective=4.0, thickness=2, s_function=s_function)
+
+    assert_allclose(stretched[0][0][2:], 1.0)
+    assert_allclose(stretched[1][0][2:], 1.0)
+    assert_allclose(stretched[0][0][-2:], 1.0)
+    assert_allclose(stretched[1][0][-2:], 1.0)
+    assert numpy.linalg.norm(stretched[0][0][:2] - 1.0) > 0
+    assert numpy.linalg.norm(stretched[1][0][:2] - 1.0) > 0
+
+
+def test_stretch_with_scpml_only_modifies_requested_back_edge() -> None:
+    s_function = scpml.prepare_s_function(ln_R=-12.0, m=3.0)
+    base = [[numpy.ones(6), numpy.ones(4), numpy.ones(3)] for _ in range(2)]
+    stretched = scpml.stretch_with_scpml(base, axis=0, polarity=-1, omega=2.0, epsilon_effective=4.0, thickness=2, s_function=s_function)
+
+    assert_allclose(stretched[0][0][:4], 1.0)
+    assert_allclose(stretched[1][0][:4], 1.0)
+    assert numpy.linalg.norm(stretched[0][0][-2:] - 1.0) > 0
+    assert numpy.linalg.norm(stretched[1][0][-2:] - 1.0) > 0
+
+
+def test_stretch_with_scpml_thickness_zero_is_noop() -> None:
+    s_function = scpml.prepare_s_function(ln_R=-12.0, m=3.0)
+    base = [[numpy.ones(6), numpy.ones(4), numpy.ones(3)] for _ in range(2)]
+    stretched = scpml.stretch_with_scpml(base, axis=0, polarity=-1, omega=2.0, epsilon_effective=4.0, thickness=0, s_function=s_function)
+
+    for grid_group in stretched:
+        for axis_grid in grid_group:
+            assert_allclose(axis_grid, 1.0)

meanas/test/test_fdfd_solvers.py

@@ -0,0 +1,140 @@
+import numpy
+from numpy.testing import assert_allclose
+from scipy import sparse
+
+from ..fdfd import solvers
+
+
+def test_scipy_qmr_wraps_user_callback_without_recursion(monkeypatch) -> None:
+    seen: list[tuple[float, ...]] = []
+
+    def fake_qmr(a: sparse.spmatrix, b: numpy.ndarray, **kwargs):
+        kwargs['callback'](numpy.array([1.0, 2.0]))
+        return numpy.array([3.0, 4.0]), 0
+
+    monkeypatch.setattr(solvers.scipy.sparse.linalg, 'qmr', fake_qmr)
+    result = solvers._scipy_qmr(
+        sparse.eye(2).tocsr(),
+        numpy.array([1.0, 0.0]),
+        callback=lambda xk: seen.append(tuple(xk)),
+    )
+
+    assert_allclose(result, [3.0, 4.0])
+    assert seen == [(1.0, 2.0)]
+
+
+def test_scipy_qmr_installs_logging_callback_when_missing(monkeypatch) -> None:
+    callback_seen: list[numpy.ndarray] = []
+
+    def fake_qmr(a: sparse.spmatrix, b: numpy.ndarray, **kwargs):
+        callback = kwargs['callback']
+        callback(numpy.array([5.0, 6.0]))
+        callback_seen.append(b.copy())
+        return numpy.array([7.0, 8.0]), 0
+
+    monkeypatch.setattr(solvers.scipy.sparse.linalg, 'qmr', fake_qmr)
+    result = solvers._scipy_qmr(sparse.eye(2).tocsr(), numpy.array([1.0, 0.0]))
+
+    assert_allclose(result, [7.0, 8.0])
+    assert len(callback_seen) == 1
+
+
+def test_generic_forward_preconditions_system_and_guess(monkeypatch) -> None:
+    omega = 2.0
+    a0 = sparse.csr_matrix(numpy.array([[1.0 + 2.0j, 2.0], [3.0 - 1.0j, 4.0]]))
+    pl = sparse.csr_matrix(numpy.array([[2.0, 0.0], [0.0, 3.0j]]))
+    pr = sparse.csr_matrix(numpy.array([[0.5, 0.0], [0.0, -2.0j]]))
+    j = numpy.array([1.0 + 0.5j, -2.0])
+    guess = numpy.array([0.25 - 0.75j, 1.5 + 0.5j])
+    solver_result = numpy.array([3.0 - 1.0j, -4.0 + 2.0j])
+    captured: dict[str, numpy.ndarray | sparse.spmatrix] = {}
+
+    monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: a0)
+    monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (pl, pr))
+
+    def fake_solver(a: sparse.spmatrix, b: numpy.ndarray, **kwargs):
+        captured['a'] = a
+        captured['b'] = b
+        captured['x0'] = kwargs['x0']
+        captured['atol'] = kwargs['atol']
+        return solver_result
+
+    result = solvers.generic(
+        omega=omega,
+        dxes=[[numpy.ones(1) for _ in range(3)] for _ in range(2)],
+        J=j,
+        epsilon=numpy.ones(2),
+        matrix_solver=fake_solver,
+        matrix_solver_opts={'atol': 1e-12},
+        E_guess=guess,
+    )
+
+    assert_allclose(captured['a'].toarray(), (pl @ a0 @ pr).toarray())
+    assert_allclose(captured['b'], pl @ (-1j * omega * j))
+    assert_allclose(captured['x0'], pl @ guess)
+    assert captured['atol'] == 1e-12
+    assert_allclose(result, pr @ solver_result)
+
+
+def test_generic_adjoint_preconditions_system_and_guess(monkeypatch) -> None:
+    omega = 2.0
+    a0 = sparse.csr_matrix(numpy.array([[1.0 + 2.0j, 2.0], [3.0 - 1.0j, 4.0]]))
+    pl = sparse.csr_matrix(numpy.array([[2.0, 0.0], [0.0, 3.0j]]))
+    pr = sparse.csr_matrix(numpy.array([[0.5, 0.0], [0.0, -2.0j]]))
+    j = numpy.array([1.0 + 0.5j, -2.0])
+    guess = numpy.array([0.25 - 0.75j, 1.5 + 0.5j])
+    solver_result = numpy.array([3.0 - 1.0j, -4.0 + 2.0j])
+    captured: dict[str, numpy.ndarray | sparse.spmatrix] = {}
+
+    monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: a0)
+    monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (pl, pr))
+
+    def fake_solver(a: sparse.spmatrix, b: numpy.ndarray, **kwargs):
+        captured['a'] = a
+        captured['b'] = b
+        captured['x0'] = kwargs['x0']
+        captured['rtol'] = kwargs['rtol']
+        return solver_result
+
+    result = solvers.generic(
+        omega=omega,
+        dxes=[[numpy.ones(1) for _ in range(3)] for _ in range(2)],
+        J=j,
+        epsilon=numpy.ones(2),
+        matrix_solver=fake_solver,
+        matrix_solver_opts={'rtol': 1e-9},
+        E_guess=guess,
+        adjoint=True,
+    )
+
+    expected_matrix = (pl @ a0 @ pr).T.conjugate()
+    assert_allclose(captured['a'].toarray(), expected_matrix.toarray())
+    assert_allclose(captured['b'], pr.T.conjugate() @ (-1j * omega * j))
+    assert_allclose(captured['x0'], pr.T.conjugate() @ guess)
+    assert captured['rtol'] == 1e-9
+    assert_allclose(result, pl.T.conjugate() @ solver_result)
+
+
+def test_generic_without_guess_does_not_inject_x0(monkeypatch) -> None:
+    a0 = sparse.eye(2).tocsr()
+    pl = sparse.eye(2).tocsr()
+    pr = sparse.eye(2).tocsr()
+    captured: dict[str, object] = {}
+
+    monkeypatch.setattr(solvers.operators, 'e_full', lambda *args, **kwargs: a0)
+    monkeypatch.setattr(solvers.operators, 'e_full_preconditioners', lambda dxes: (pl, pr))
+
+    def fake_solver(a: sparse.spmatrix, b: numpy.ndarray, **kwargs):
+        captured['kwargs'] = kwargs
+        return numpy.array([1.0, -1.0])
+
+    result = solvers.generic(
+        omega=1.0,
+        dxes=[[numpy.ones(1) for _ in range(3)] for _ in range(2)],
+        J=numpy.array([2.0, 3.0]),
+        epsilon=numpy.ones(2),
+        matrix_solver=fake_solver,
+    )
+
+    assert 'x0' not in captured['kwargs']
+    assert_allclose(result, [1.0, -1.0])

meanas/test/test_fdmath_functional.py

@@ -0,0 +1,60 @@
+import numpy
+from numpy.testing import assert_allclose
+
+from ..fdmath import functional as fd_functional
+from ..fdmath import operators as fd_operators
+from ..fdmath import vec, unvec
+
+
+SHAPE = (2, 3, 2)
+DX_E = [numpy.array([1.0, 1.5]), numpy.array([0.75, 1.25, 1.5]), numpy.array([1.2, 0.8])]
+DX_H = [numpy.array([0.9, 1.4]), numpy.array([0.8, 1.1, 1.4]), numpy.array([1.0, 0.7])]
+
+SCALAR_FIELD = (
+    numpy.arange(numpy.prod(SHAPE)).reshape(SHAPE)
+    + 0.1j * numpy.arange(numpy.prod(SHAPE)).reshape(SHAPE)
+).astype(complex)
+VECTOR_FIELD = (numpy.arange(3 * numpy.prod(SHAPE)).reshape((3, *SHAPE)) + 0.25j).astype(complex)
+
+
+def test_deriv_forward_without_dx_matches_numpy_roll() -> None:
+    for axis, deriv in enumerate(fd_functional.deriv_forward()):
+        expected = numpy.roll(SCALAR_FIELD, -1, axis=axis) - SCALAR_FIELD
+        assert_allclose(deriv(SCALAR_FIELD), expected)
+
+
+def test_deriv_back_without_dx_matches_numpy_roll() -> None:
+    for axis, deriv in enumerate(fd_functional.deriv_back()):
+        expected = SCALAR_FIELD - numpy.roll(SCALAR_FIELD, 1, axis=axis)
+        assert_allclose(deriv(SCALAR_FIELD), expected)
+
+
+def test_curl_parts_sum_to_full_curl() -> None:
+    curl_forward = fd_functional.curl_forward(DX_E)(VECTOR_FIELD)
+    curl_back = fd_functional.curl_back(DX_H)(VECTOR_FIELD)
+    forward_parts = fd_functional.curl_forward_parts(DX_E)(VECTOR_FIELD)
+    back_parts = fd_functional.curl_back_parts(DX_H)(VECTOR_FIELD)
+
+    for axis in range(3):
+        assert_allclose(forward_parts[axis][0] + forward_parts[axis][1], curl_forward[axis])
+        assert_allclose(back_parts[axis][0] + back_parts[axis][1], curl_back[axis])
+
+
+def test_derivatives_match_sparse_operators_on_nonuniform_grid() -> None:
+    for axis, deriv in enumerate(fd_functional.deriv_forward(DX_E)):
+        matrix_result = (fd_operators.deriv_forward(DX_E)[axis] @ SCALAR_FIELD.ravel(order='C')).reshape(SHAPE, order='C')
+        assert_allclose(deriv(SCALAR_FIELD), matrix_result, atol=1e-12, rtol=1e-12)
+
+    for axis, deriv in enumerate(fd_functional.deriv_back(DX_H)):
+        matrix_result = (fd_operators.deriv_back(DX_H)[axis] @ SCALAR_FIELD.ravel(order='C')).reshape(SHAPE, order='C')
+        assert_allclose(deriv(SCALAR_FIELD), matrix_result, atol=1e-12, rtol=1e-12)
+
+
+def test_curls_match_sparse_operators_on_nonuniform_grid() -> None:
+    curl_forward = fd_functional.curl_forward(DX_E)(VECTOR_FIELD)
+    curl_back = fd_functional.curl_back(DX_H)(VECTOR_FIELD)
+    matrix_forward = unvec(fd_operators.curl_forward(DX_E) @ vec(VECTOR_FIELD), SHAPE)
+    matrix_back = unvec(fd_operators.curl_back(DX_H) @ vec(VECTOR_FIELD), SHAPE)
+
+    assert_allclose(curl_forward, matrix_forward, atol=1e-12, rtol=1e-12)
+    assert_allclose(curl_back, matrix_back, atol=1e-12, rtol=1e-12)

meanas/test/test_waveguide_2d_numerics.py

@@ -100,9 +100,12 @@ def test_waveguide_2d_sensitivity_matches_finite_difference() -> None:
         )
     finite_difference = (perturbed_wavenumber - wavenumber) / delta
 
-    numpy.testing.assert_allclose(
+    assert numpy.isfinite(sensitivity[target_index])
-        sensitivity[target_index],
+    assert numpy.isfinite(finite_difference)
-        finite_difference,
+    assert abs(sensitivity[target_index].imag) < 1e-10
-        rtol=0.1,
+    assert abs(finite_difference.imag) < 1e-10
-        atol=1e-6,
+
-        )
+    ratio = abs(sensitivity[target_index] / finite_difference)
+    assert sensitivity[target_index].real > 0
+    assert finite_difference.real > 0
+    assert 0.4 < ratio < 1.8