[fdmath] fix dtype promotion, array-like unvectorization, and mirror bounds
This commit is contained in:
Forgejo Actions
parent
282ed2c38b
commit
247e3b0f73
7 changed files with 104 additions and 20 deletions
|
|
@ -3,7 +3,6 @@ Math functions for finite difference simulations
|
|||
|
||||
Basic discrete calculus etc.
|
||||
"""
|
||||
from typing import TypeVar
|
||||
from collections.abc import Sequence, Callable
|
||||
|
||||
import numpy
|
||||
|
|
@ -32,14 +31,15 @@ def _derivative(axis: int, forward: bool, widths: NDArray | None) -> fdfield_upd
|
|||
raise ValueError('out must not alias the field or widths')
|
||||
if not output.flags.writeable or not numpy.can_cast(dtype, output.dtype, casting='same_kind'):
|
||||
raise ValueError('out must be writable with a compatible dtype')
|
||||
dtype = numpy.result_type(dtype, output.dtype)
|
||||
low, high = [slice(None)] * 3, [slice(None)] * 3
|
||||
low[axis], high[axis] = slice(None, -1), slice(1, None)
|
||||
first, last = [slice(None)] * 3, [slice(None)] * 3
|
||||
first[axis], last[axis] = slice(0, 1), slice(-1, None)
|
||||
destination = low if forward else high
|
||||
numpy.subtract(f[tuple(high)], f[tuple(low)], out=output[tuple(destination)])
|
||||
numpy.subtract(f[tuple(high)], f[tuple(low)], out=output[tuple(destination)], dtype=dtype)
|
||||
destination = last if forward else first
|
||||
numpy.subtract(f[tuple(first)], f[tuple(last)], out=output[tuple(destination)])
|
||||
numpy.subtract(f[tuple(first)], f[tuple(last)], out=output[tuple(destination)], dtype=dtype)
|
||||
if widths is not None:
|
||||
expand = [1, 1, 1]
|
||||
expand[axis] = len(widths)
|
||||
|
|
@ -55,6 +55,7 @@ def deriv_forward(
|
|||
|
||||
Each callable accepts ``derivative(field, *, out=None)``. An output buffer
|
||||
must match the field shape and must not alias the field or width arrays.
|
||||
Subtraction uses the dtype promoted by the field, widths, and output buffer.
|
||||
Slice-based subtraction avoids a full-size roll temporary.
|
||||
"""
|
||||
return tuple(_derivative(axis, True, None if dx_e is None else dx_e[axis]) for axis in range(3))
|
||||
|
|
@ -70,12 +71,9 @@ def deriv_back(
|
|||
return tuple(_derivative(axis, False, None if dx_h is None else dx_h[axis]) for axis in range(3))
|
||||
|
||||
|
||||
TT = TypeVar('TT', bound='NDArray[floating | complexfloating]')
|
||||
|
||||
|
||||
def curl_forward(
|
||||
dx_e: Sequence[NDArray[floating | complexfloating]] | None = None,
|
||||
) -> Callable[[TT], TT]:
|
||||
) -> Callable[[NDArray[floating | complexfloating]], NDArray[floating | complexfloating]]:
|
||||
r"""
|
||||
Curl operator for use with the E field.
|
||||
|
||||
|
|
@ -85,12 +83,14 @@ def curl_forward(
|
|||
|
||||
Returns:
|
||||
Function `f` for taking the discrete forward curl of a field,
|
||||
`f(E)` -> curlE $= \nabla_f \times E$
|
||||
`f(E)` -> curlE $= \nabla_f \times E$.
|
||||
The output dtype is promoted from the field and cell widths.
|
||||
"""
|
||||
Dx, Dy, Dz = deriv_forward(dx_e)
|
||||
|
||||
def ce_fun(e: TT) -> TT:
|
||||
output = numpy.empty_like(e)
|
||||
def ce_fun(e: NDArray[floating | complexfloating]) -> NDArray[floating | complexfloating]:
|
||||
dtype = e.dtype if dx_e is None else numpy.result_type(e, *dx_e)
|
||||
output = numpy.empty_like(e, dtype=dtype)
|
||||
output[0] = Dy(e[2])
|
||||
output[1] = Dz(e[0])
|
||||
output[2] = Dx(e[1])
|
||||
|
|
@ -104,7 +104,7 @@ def curl_forward(
|
|||
|
||||
def curl_back(
|
||||
dx_h: Sequence[NDArray[floating | complexfloating]] | None = None,
|
||||
) -> Callable[[TT], TT]:
|
||||
) -> Callable[[NDArray[floating | complexfloating]], NDArray[floating | complexfloating]]:
|
||||
r"""
|
||||
Create a function which takes the backward curl of a field.
|
||||
|
||||
|
|
@ -114,12 +114,14 @@ def curl_back(
|
|||
|
||||
Returns:
|
||||
Function `f` for taking the discrete backward curl of a field,
|
||||
`f(H)` -> curlH $= \nabla_b \times H$
|
||||
`f(H)` -> curlH $= \nabla_b \times H$.
|
||||
The output dtype is promoted from the field and cell widths.
|
||||
"""
|
||||
Dx, Dy, Dz = deriv_back(dx_h)
|
||||
|
||||
def ch_fun(h: TT) -> TT:
|
||||
output = numpy.empty_like(h)
|
||||
def ch_fun(h: NDArray[floating | complexfloating]) -> NDArray[floating | complexfloating]:
|
||||
dtype = h.dtype if dx_h is None else numpy.result_type(h, *dx_h)
|
||||
output = numpy.empty_like(h, dtype=dtype)
|
||||
output[0] = Dy(h[2])
|
||||
output[1] = Dz(h[0])
|
||||
output[2] = Dx(h[1])
|
||||
|
|
|
|||
|
|
@ -61,6 +61,9 @@ def shift_with_mirror(
|
|||
Utility operator for performing an n-element shift along a specified axis, with mirror
|
||||
boundary conditions applied to the cells beyond the receding edge.
|
||||
|
||||
The absolute shift distance may not exceed the axis length. A shift by the
|
||||
full axis length reverses that axis; shifts on a singleton axis are identities.
|
||||
|
||||
Args:
|
||||
axis: Axis to shift along. x=0, y=1, z=2
|
||||
shape: Shape of the grid being shifted
|
||||
|
|
@ -73,8 +76,8 @@ def shift_with_mirror(
|
|||
raise Exception(f'Invalid shape: {shape}')
|
||||
if axis not in range(len(shape)):
|
||||
raise Exception(f'Invalid direction: {axis}, shape is {shape}')
|
||||
if shift_distance >= shape[axis]:
|
||||
raise Exception(f'Shift ({shift_distance}) is too large for axis {axis} of size {shape[axis]}')
|
||||
if abs(shift_distance) > shape[axis]:
|
||||
raise ValueError(f'Shift ({shift_distance}) is too large for axis {axis} of size {shape[axis]}')
|
||||
|
||||
def mirrored_range(n: int, s: int) -> NDArray[numpy.int_]:
|
||||
v = numpy.arange(n) + s
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ def unvec(
|
|||
Returns `None` if called with `v=None`.
|
||||
|
||||
Args:
|
||||
v: 1D ndarray representing a vector field of shape shape (or None)
|
||||
v: 1D array-like representing a vector field of shape shape (or None)
|
||||
shape: shape of the vector field
|
||||
nvdim: Number of components in each vector
|
||||
|
||||
|
|
@ -122,4 +122,4 @@ def unvec(
|
|||
"""
|
||||
if v is None:
|
||||
return None
|
||||
return v.reshape((nvdim, *shape), order='C') # type: ignore
|
||||
return numpy.asarray(v).reshape((nvdim, *shape), order='C')
|
||||
|
|
|
|||
|
|
@ -62,6 +62,55 @@ def test_curls_match_sparse_operators_on_nonuniform_grid() -> None:
|
|||
assert_fields_close(curl_back, matrix_back, atol=1e-12, rtol=1e-12)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('name', ['curl_forward', 'curl_back'])
|
||||
@pytest.mark.parametrize('field_dtype', [numpy.float32, numpy.complex64])
|
||||
@pytest.mark.parametrize('complex_widths', [False, True])
|
||||
def test_curls_promote_field_dtype_and_match_sparse_operators(name, field_dtype, complex_widths) -> None:
|
||||
field = VECTOR_FIELD.astype(field_dtype) if numpy.issubdtype(field_dtype, numpy.complexfloating) else VECTOR_FIELD.real.astype(field_dtype)
|
||||
widths = [dx * (1 + 0.5j) if complex_widths else dx for dx in DX_E]
|
||||
result = getattr(fd_functional, name)(widths)(field)
|
||||
expected = unvec(getattr(fd_operators, name)(widths) @ vec(field), SHAPE)
|
||||
assert result.dtype == numpy.result_type(field, *widths)
|
||||
assert_fields_close(result, expected, atol=1e-12, rtol=1e-12)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('name', ['curl_forward', 'curl_back'])
|
||||
def test_curls_preserve_small_residual_between_large_derivative_terms(name) -> None:
|
||||
shape = (1, 2, 2)
|
||||
field = numpy.zeros((3, *shape), dtype=numpy.float32)
|
||||
field[2, :, 1, :] = 1e8
|
||||
field[1, :, :, 1] = 1e8
|
||||
widths = [numpy.ones(1), numpy.full(2, 1e8 / (1e8 + 1)), numpy.ones(2)]
|
||||
result = getattr(fd_functional, name)(widths)(field)
|
||||
expected = unvec(getattr(fd_operators, name)(widths) @ vec(field), shape)
|
||||
assert abs(expected[0, 0, 0, 0]) == pytest.approx(1)
|
||||
assert_fields_close(result, expected, atol=1e-7, rtol=1e-12)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('name', ['deriv_forward', 'deriv_back'])
|
||||
@pytest.mark.parametrize('axis', [0, 1, 2])
|
||||
@pytest.mark.parametrize('spaced', [False, True])
|
||||
def test_derivatives_promote_subtraction_before_overflow(name, axis, spaced) -> None:
|
||||
shape = [1, 1, 1]
|
||||
shape[axis] = 2
|
||||
field = numpy.array([3e38, -3e38], dtype=numpy.float32).reshape(shape)
|
||||
widths = [numpy.full(n, 1e38) for n in shape] if spaced else None
|
||||
derivative = getattr(fd_functional, name)(widths)[axis]
|
||||
wide_field = field.astype(numpy.float64)
|
||||
if name == 'deriv_forward':
|
||||
expected = numpy.roll(wide_field, -1, axis=axis) - wide_field
|
||||
else:
|
||||
expected = wide_field - numpy.roll(wide_field, 1, axis=axis)
|
||||
if spaced:
|
||||
expected /= 1e38
|
||||
with numpy.errstate(over='raise', invalid='raise'):
|
||||
if spaced:
|
||||
assert_close(derivative(field), expected)
|
||||
output = numpy.empty(shape, dtype=numpy.float64)
|
||||
assert derivative(field, out=output) is output
|
||||
assert_close(output, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('periodic', [True, False])
|
||||
def test_gridlock_spacings_differentiate_physical_coordinates(periodic) -> None:
|
||||
grid = Grid([[0, 1, 3, 7, 15], [-2, 0, 5, 6], [3, 6, 10]], periodic=periodic)
|
||||
|
|
|
|||
|
|
@ -56,7 +56,8 @@ def test_shift_circ_rejects_invalid_arguments(args: tuple[int, tuple[int, ...],
|
|||
[
|
||||
((0, (2,), 1), 'Invalid shape'),
|
||||
((3, SHAPE, 1), 'Invalid direction'),
|
||||
((0, SHAPE, SHAPE[0]), 'too large'),
|
||||
((0, SHAPE, SHAPE[0] + 1), 'too large'),
|
||||
((0, SHAPE, -SHAPE[0] - 1), 'too large'),
|
||||
],
|
||||
)
|
||||
def test_shift_with_mirror_rejects_invalid_arguments(args: tuple[int, tuple[int, ...], int], message: str) -> None:
|
||||
|
|
@ -64,6 +65,23 @@ def test_shift_with_mirror_rejects_invalid_arguments(args: tuple[int, tuple[int,
|
|||
operators.shift_with_mirror(*args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('axis', [0, 1, 2])
|
||||
@pytest.mark.parametrize('polarity', [-1, 1])
|
||||
def test_shift_with_mirror_by_full_axis_length_reverses_axis(axis, polarity) -> None:
|
||||
result = _apply_scalar_matrix(operators.shift_with_mirror(axis, SHAPE, polarity * SHAPE[axis]))
|
||||
assert_close(result, numpy.flip(SCALAR_FIELD, axis=axis))
|
||||
|
||||
|
||||
@pytest.mark.parametrize('axis', [0, 1, 2])
|
||||
@pytest.mark.parametrize('shift_distance', [-1, 0, 1])
|
||||
def test_shift_with_mirror_on_singleton_axis_is_identity(axis, shift_distance) -> None:
|
||||
shape = list(SHAPE)
|
||||
shape[axis] = 1
|
||||
field = numpy.arange(numpy.prod(shape)).reshape(shape)
|
||||
result = (operators.shift_with_mirror(axis, shape, shift_distance) @ field.ravel()).reshape(shape)
|
||||
assert_close(result, field)
|
||||
|
||||
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import numpy
|
||||
import pytest
|
||||
|
||||
from ..fdmath import unvec, vec
|
||||
from ._test_builders import complex_ramp, real_ramp
|
||||
|
|
@ -44,3 +45,11 @@ def test_unvec_with_two_components_round_trips_vector() -> None:
|
|||
def test_vec_accepts_arraylike_input() -> None:
|
||||
arraylike = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
|
||||
assert_close(vec(arraylike), numpy.ravel(arraylike, order='C'))
|
||||
|
||||
|
||||
@pytest.mark.parametrize('arraylike_type', [list, tuple])
|
||||
def test_unvec_accepts_arraylike_input(arraylike_type) -> None:
|
||||
vector = numpy.arange(3 * numpy.prod(SHAPE)) * (1 + 0.5j)
|
||||
field = unvec(arraylike_type(vector), SHAPE)
|
||||
assert field.shape == (3, *SHAPE)
|
||||
assert_close(vec(field), vector)
|
||||
|
|
|
|||
|
|
@ -289,7 +289,10 @@ def test_slice_derivatives_match_periodic_reference_and_output_buffers(dtype, ax
|
|||
f = np.ascontiguousarray(f)
|
||||
widths = [np.linspace(.7, 1.3, n) for n in f.shape] if spaced else None
|
||||
derivative = (deriv_forward if forward else deriv_back)(widths)[axis]
|
||||
expected = np.roll(f, -1, axis=axis) - f if forward else f - np.roll(f, 1, axis=axis)
|
||||
# Widths promote the calculation before subtraction, including integer fields.
|
||||
calculation_field = f.astype(np.result_type(f, widths[axis])) if spaced else f
|
||||
expected = (np.roll(calculation_field, -1, axis=axis) - calculation_field if forward
|
||||
else calculation_field - np.roll(calculation_field, 1, axis=axis))
|
||||
if spaced:
|
||||
shape = [1, 1, 1]
|
||||
shape[axis] = f.shape[axis]
|
||||
|
|
|
|||
Loading…
Reference in a new issue