[bloch] fixup some vectorization and add tests

meanas/fdfd/bloch.py

@@ -136,6 +136,14 @@ except ImportError:
     logger.info('Using numpy fft')
 
 
+def _assemble_hmn_vector(
+        h_m: NDArray[numpy.complex128],
+        h_n: NDArray[numpy.complex128],
+        ) -> NDArray[numpy.complex128]:
+    stacked = numpy.concatenate((numpy.ravel(h_m), numpy.ravel(h_n)))
+    return stacked[:, None]
+
+
 def generate_kmn(
         k0: ArrayLike,
         G_matrix: ArrayLike,
@@ -253,8 +261,8 @@ def maxwell_operator(
             h_m, h_n = b_m, b_n
         else:
             # transform from mn to xyz
-            b_xyz = (m * b_m[:, :, :, None]
+            b_xyz = (m * b_m
-                   + n * b_n[:, :, :, None])
+                   + n * b_n)    # noqa: E128
 
             # divide by mu
             temp = ifftn(b_xyz, axes=range(3))
@@ -265,10 +273,7 @@ def maxwell_operator(
             h_m = numpy.sum(h_xyz * m, axis=3)
             h_n = numpy.sum(h_xyz * n, axis=3)
 
-        h.shape = (h.size,)
+        return _assemble_hmn_vector(h_m, h_n)
-        h = numpy.concatenate((h_m.ravel(), h_n.ravel()), axis=None, out=h)     # ravel and merge
-        h.shape = (h.size, 1)
-        return h
 
     return operator
 
@@ -403,8 +408,8 @@ def inverse_maxwell_operator_approx(
             b_m, b_n = hin_m, hin_n
         else:
             # transform from mn to xyz
-            h_xyz = (m * hin_m[:, :, :, None]
+            h_xyz = (m * hin_m
-                   + n * hin_n[:, :, :, None])
+                   + n * hin_n)  # noqa: E128
 
             # multiply by mu
             temp = ifftn(h_xyz, axes=range(3))
@@ -412,8 +417,8 @@ def inverse_maxwell_operator_approx(
             b_xyz = fftn(temp, axes=range(3))
 
             # transform back to mn
-            b_m = numpy.sum(b_xyz * m, axis=3)
+            b_m = numpy.sum(b_xyz * m, axis=3, keepdims=True)
-            b_n = numpy.sum(b_xyz * n, axis=3)
+            b_n = numpy.sum(b_xyz * n, axis=3, keepdims=True)
 
         # cross product and transform into xyz basis
         e_xyz = (n * b_m
@@ -428,10 +433,7 @@ def inverse_maxwell_operator_approx(
         h_m = numpy.sum(d_xyz * n, axis=3, keepdims=True) / +k_mag
         h_n = numpy.sum(d_xyz * m, axis=3, keepdims=True) / -k_mag
 
-        h.shape = (h.size,)
+        return _assemble_hmn_vector(h_m, h_n)
-        h = numpy.concatenate((h_m, h_n), axis=None, out=h)
-        h.shape = (h.size, 1)
-        return h
 
     return operator
 

meanas/test/test_bloch_foundations.py

@@ -0,0 +1,85 @@
+import numpy
+from numpy.linalg import norm
+
+from ..fdfd import bloch
+
+
+SHAPE = (2, 2, 2)
+K0 = numpy.array([0.1, 0.2, 0.3])
+G_MATRIX = numpy.eye(3)
+EPSILON = numpy.ones((3, *SHAPE), dtype=float)
+MU = numpy.stack([
+    numpy.linspace(2.0, 2.7, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(2.1, 2.8, numpy.prod(SHAPE)).reshape(SHAPE),
+    numpy.linspace(2.2, 2.9, numpy.prod(SHAPE)).reshape(SHAPE),
+])
+H_MN = (numpy.arange(2 * numpy.prod(SHAPE)) + 0.25j).astype(complex)
+ZERO_H_MN = numpy.zeros_like(H_MN)
+
+
+def test_generate_kmn_general_case_returns_orthonormal_basis() -> None:
+    k_mag, m_vecs, n_vecs = bloch.generate_kmn(K0, G_MATRIX, SHAPE)
+
+    assert k_mag.shape == SHAPE + (1,)
+    assert m_vecs.shape == SHAPE + (3,)
+    assert n_vecs.shape == SHAPE + (3,)
+    assert numpy.isfinite(k_mag).all()
+    assert numpy.isfinite(m_vecs).all()
+    assert numpy.isfinite(n_vecs).all()
+
+    numpy.testing.assert_allclose(norm(m_vecs.reshape(-1, 3), axis=1), 1.0)
+    numpy.testing.assert_allclose(norm(n_vecs.reshape(-1, 3), axis=1), 1.0)
+    numpy.testing.assert_allclose(numpy.sum(m_vecs * n_vecs, axis=3), 0.0, atol=1e-12)
+
+
+def test_generate_kmn_z_aligned_uses_default_transverse_basis() -> None:
+    k_mag, m_vecs, n_vecs = bloch.generate_kmn([0.0, 0.0, 0.25], G_MATRIX, (1, 1, 1))
+
+    assert numpy.isfinite(k_mag).all()
+    numpy.testing.assert_allclose(m_vecs[0, 0, 0], [0.0, 1.0, 0.0])
+    numpy.testing.assert_allclose(numpy.sum(m_vecs * n_vecs, axis=3), 0.0, atol=1e-12)
+    numpy.testing.assert_allclose(norm(n_vecs.reshape(-1, 3), axis=1), 1.0)
+
+
+def test_maxwell_operator_returns_finite_column_vector_without_mu() -> None:
+    operator = bloch.maxwell_operator(K0, G_MATRIX, EPSILON)
+
+    result = operator(H_MN.copy())
+    zero_result = operator(ZERO_H_MN.copy())
+
+    assert result.shape == (2 * numpy.prod(SHAPE), 1)
+    assert numpy.isfinite(result).all()
+    numpy.testing.assert_allclose(zero_result, 0.0)
+
+
+def test_maxwell_operator_returns_finite_column_vector_with_mu() -> None:
+    operator = bloch.maxwell_operator(K0, G_MATRIX, EPSILON, MU)
+
+    result = operator(H_MN.copy())
+    zero_result = operator(ZERO_H_MN.copy())
+
+    assert result.shape == (2 * numpy.prod(SHAPE), 1)
+    assert numpy.isfinite(result).all()
+    numpy.testing.assert_allclose(zero_result, 0.0)
+
+
+def test_inverse_maxwell_operator_returns_finite_column_vector_for_both_mu_branches() -> None:
+    for mu in (None, MU):
+        operator = bloch.inverse_maxwell_operator_approx(K0, G_MATRIX, EPSILON, mu)
+
+        result = operator(H_MN.copy())
+        zero_result = operator(ZERO_H_MN.copy())
+
+        assert result.shape == (2 * numpy.prod(SHAPE), 1)
+        assert numpy.isfinite(result).all()
+        numpy.testing.assert_allclose(zero_result, 0.0)
+
+
+def test_bloch_field_converters_return_finite_fields() -> None:
+    e_field = bloch.hmn_2_exyz(K0, G_MATRIX, EPSILON)(H_MN.copy())
+    h_field = bloch.hmn_2_hxyz(K0, G_MATRIX, EPSILON)(H_MN.copy())
+
+    assert e_field.shape == (3, *SHAPE)
+    assert h_field.shape == (3, *SHAPE)
+    assert numpy.isfinite(e_field).all()
+    assert numpy.isfinite(h_field).all()