90 lines
3.5 KiB
Python
90 lines
3.5 KiB
Python
import numpy
|
|
import pytest
|
|
|
|
from ..fdmath import operators, unvec, vec
|
|
from ._test_builders import real_ramp
|
|
from .utils import assert_close
|
|
|
|
|
|
SHAPE = (2, 3, 2)
|
|
SCALAR_FIELD = real_ramp(SHAPE)
|
|
VECTOR_LEFT = real_ramp((3, *SHAPE), offset=0.5)
|
|
VECTOR_RIGHT = real_ramp((3, *SHAPE), scale=1 / 3, offset=2.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_close(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_close(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_close(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_close(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_close(matrix_result, expected)
|
|
|
|
|
|
def test_avg_forward_rejects_invalid_shape() -> None:
|
|
with pytest.raises(Exception, match='Invalid shape'):
|
|
operators.avg_forward(0, (2,))
|