[fdfd.operators] fix eh_full for non-None mu

meanas/test/test_fdmath_operators.py

@@ -0,0 +1,89 @@
+import numpy
+import pytest
+from numpy.testing import assert_allclose, assert_array_equal
+
+from ..fdmath import operators, unvec, vec
+
+
+SHAPE = (2, 3, 2)
+SCALAR_FIELD = numpy.arange(numpy.prod(SHAPE), dtype=float).reshape(SHAPE, order='C')
+VECTOR_LEFT = (numpy.arange(3 * numpy.prod(SHAPE), dtype=float).reshape((3, *SHAPE), order='C') + 0.5)
+VECTOR_RIGHT = (2.0 + numpy.arange(3 * numpy.prod(SHAPE), dtype=float).reshape((3, *SHAPE), order='C') / 3.0)
+
+
+def _apply_scalar_matrix(op: operators.sparse.spmatrix) -> numpy.ndarray:
+    return (op @ SCALAR_FIELD.ravel(order='C')).reshape(SHAPE, order='C')
+
+
+def _mirrored_indices(size: int, shift_distance: int) -> numpy.ndarray:
+    indices = numpy.arange(size) + shift_distance
+    indices = numpy.where(indices >= size, 2 * size - indices - 1, indices)
+    indices = numpy.where(indices < 0, -1 - indices, indices)
+    return indices
+
+
+@pytest.mark.parametrize(('axis', 'shift_distance'), [(0, 1), (1, -1), (2, 1)])
+def test_shift_circ_matches_numpy_roll(axis: int, shift_distance: int) -> None:
+    matrix_result = _apply_scalar_matrix(operators.shift_circ(axis, SHAPE, shift_distance))
+    expected = numpy.roll(SCALAR_FIELD, -shift_distance, axis=axis)
+    assert_array_equal(matrix_result, expected)
+
+
+@pytest.mark.parametrize(('axis', 'shift_distance'), [(0, 1), (1, -1), (2, 1)])
+def test_shift_with_mirror_matches_explicit_mirrored_indices(axis: int, shift_distance: int) -> None:
+    matrix_result = _apply_scalar_matrix(operators.shift_with_mirror(axis, SHAPE, shift_distance))
+    indices = [numpy.arange(length) for length in SHAPE]
+    indices[axis] = _mirrored_indices(SHAPE[axis], shift_distance)
+    expected = SCALAR_FIELD[numpy.ix_(*indices)]
+    assert_array_equal(matrix_result, expected)
+
+
+@pytest.mark.parametrize(
+    ('args', 'message'),
+    [
+        ((0, (2,), 1), 'Invalid shape'),
+        ((3, SHAPE, 1), 'Invalid direction'),
+    ],
+    )
+def test_shift_circ_rejects_invalid_arguments(args: tuple[int, tuple[int, ...], int], message: str) -> None:
+    with pytest.raises(Exception, match=message):
+        operators.shift_circ(*args)
+
+
+@pytest.mark.parametrize(
+    ('args', 'message'),
+    [
+        ((0, (2,), 1), 'Invalid shape'),
+        ((3, SHAPE, 1), 'Invalid direction'),
+        ((0, SHAPE, SHAPE[0]), 'too large'),
+    ],
+    )
+def test_shift_with_mirror_rejects_invalid_arguments(args: tuple[int, tuple[int, ...], int], message: str) -> None:
+    with pytest.raises(Exception, match=message):
+        operators.shift_with_mirror(*args)
+
+
+def test_vec_cross_matches_pointwise_cross_product() -> None:
+    matrix_result = unvec(operators.vec_cross(vec(VECTOR_LEFT)) @ vec(VECTOR_RIGHT), SHAPE)
+    expected = numpy.empty_like(VECTOR_LEFT)
+    expected[0] = VECTOR_LEFT[1] * VECTOR_RIGHT[2] - VECTOR_LEFT[2] * VECTOR_RIGHT[1]
+    expected[1] = VECTOR_LEFT[2] * VECTOR_RIGHT[0] - VECTOR_LEFT[0] * VECTOR_RIGHT[2]
+    expected[2] = VECTOR_LEFT[0] * VECTOR_RIGHT[1] - VECTOR_LEFT[1] * VECTOR_RIGHT[0]
+    assert_allclose(matrix_result, expected)
+
+
+def test_avg_forward_matches_half_sum_with_forward_neighbor() -> None:
+    matrix_result = _apply_scalar_matrix(operators.avg_forward(1, SHAPE))
+    expected = 0.5 * (SCALAR_FIELD + numpy.roll(SCALAR_FIELD, -1, axis=1))
+    assert_allclose(matrix_result, expected)
+
+
+def test_avg_back_matches_half_sum_with_backward_neighbor() -> None:
+    matrix_result = _apply_scalar_matrix(operators.avg_back(1, SHAPE))
+    expected = 0.5 * (SCALAR_FIELD + numpy.roll(SCALAR_FIELD, 1, axis=1))
+    assert_allclose(matrix_result, expected)
+
+
+def test_avg_forward_rejects_invalid_shape() -> None:
+    with pytest.raises(Exception, match='Invalid shape'):
+        operators.avg_forward(0, (2,))