From b9aa06355b4ca59de63fe0bb2a787020b352ecb3 Mon Sep 17 00:00:00 2001 From: Forgejo Actions Date: Wed, 28 Oct 2026 18:00:34 -0700 Subject: [PATCH] [fdmath] widen curl calculations and sparse reciprocals before arithmetic --- meanas/fdmath/functional.py | 16 ++++++++++++---- meanas/fdmath/operators.py | 4 ++-- meanas/test/test_fdmath_functional.py | 18 ++++++++++++++++++ meanas/test/test_fdmath_operators.py | 16 ++++++++++++++++ 4 files changed, 48 insertions(+), 6 deletions(-) diff --git a/meanas/fdmath/functional.py b/meanas/fdmath/functional.py index c48848a..329198e 100644 --- a/meanas/fdmath/functional.py +++ b/meanas/fdmath/functional.py @@ -71,6 +71,12 @@ def deriv_back( return tuple(_derivative(axis, False, None if dx_h is None else dx_h[axis]) for axis in range(3)) +def _promote_field(field: NDArray, widths: Sequence[NDArray] | None) -> NDArray: + """Use one calculation dtype for all components of a curl.""" + dtype = field.dtype if widths is None else numpy.result_type(field, *widths) + return numpy.asarray(field, dtype=dtype) + + def curl_forward( dx_e: Sequence[NDArray[floating | complexfloating]] | None = None, ) -> Callable[[NDArray[floating | complexfloating]], NDArray[floating | complexfloating]]: @@ -89,8 +95,8 @@ def curl_forward( Dx, Dy, Dz = deriv_forward(dx_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) + e = _promote_field(e, dx_e) + output = numpy.empty_like(e) output[0] = Dy(e[2]) output[1] = Dz(e[0]) output[2] = Dx(e[1]) @@ -120,8 +126,8 @@ def curl_back( Dx, Dy, Dz = deriv_back(dx_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) + h = _promote_field(h, dx_h) + output = numpy.empty_like(h) output[0] = Dy(h[2]) output[1] = Dz(h[0]) output[2] = Dx(h[1]) @@ -139,6 +145,7 @@ def curl_forward_parts( Dx, Dy, Dz = deriv_forward(dx_e) def mkparts_fwd(e: fdfield) -> tuple[tuple[fdfield, fdfield], ...]: + e = _promote_field(e, dx_e) return ((-Dz(e[1]), Dy(e[2])), ( Dz(e[0]), -Dx(e[2])), (-Dy(e[0]), Dx(e[1]))) @@ -152,6 +159,7 @@ def curl_back_parts( Dx, Dy, Dz = deriv_back(dx_h) def mkparts_back(h: fdfield) -> tuple[tuple[fdfield, fdfield], ...]: + h = _promote_field(h, dx_h) return ((-Dz(h[1]), Dy(h[2])), ( Dz(h[0]), -Dx(h[2])), (-Dy(h[0]), Dx(h[1]))) diff --git a/meanas/fdmath/operators.py b/meanas/fdmath/operators.py index d9348ba..a9fd732 100644 --- a/meanas/fdmath/operators.py +++ b/meanas/fdmath/operators.py @@ -120,7 +120,7 @@ def deriv_forward( def deriv(axis: int) -> sparse.sparray: return shift_circ(axis, shape, 1) - sparse.eye_array(n) - Ds = [sparse.diags_array(+1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags_array(numpy.reciprocal(dx.ravel(order='C'), dtype=numpy.result_type(dx, numpy.float64))) @ deriv(a) for a, dx in enumerate(dx_e_expanded)] return Ds @@ -147,7 +147,7 @@ def deriv_back( def deriv(axis: int) -> sparse.sparray: return shift_circ(axis, shape, -1) - sparse.eye_array(n) - Ds = [sparse.diags_array(-1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags_array(-numpy.reciprocal(dx.ravel(order='C'), dtype=numpy.result_type(dx, numpy.float64))) @ deriv(a) for a, dx in enumerate(dx_h_expanded)] return Ds diff --git a/meanas/test/test_fdmath_functional.py b/meanas/test/test_fdmath_functional.py index 3a1f3cf..03a0d7d 100644 --- a/meanas/test/test_fdmath_functional.py +++ b/meanas/test/test_fdmath_functional.py @@ -87,6 +87,24 @@ def test_curls_preserve_small_residual_between_large_derivative_terms(name) -> N assert_fields_close(result, expected, atol=1e-7, rtol=1e-12) +@pytest.mark.parametrize('name', ['curl_forward', 'curl_back']) +@pytest.mark.parametrize('complex_field', [False, True]) +def test_curls_use_common_precision_for_mixed_width_dtypes(name, complex_field) -> None: + shape = (1, 2, 2) + field = numpy.zeros((3, *shape), dtype=numpy.complex64 if complex_field else numpy.float32) + field[2, :, 1, :] = 1e8 + field[1, :, :, 1] = 1e8 + widths = [numpy.ones(1), numpy.full(2, 0.7, dtype=numpy.float32), numpy.full(2, float(numpy.float32(0.7)))] + wide_widths = [dx.astype(numpy.float64) for dx in widths] + expected = unvec(getattr(fd_operators, name)(wide_widths) @ vec(field), shape) + assert expected[0, 0, 0, 0] == 0 + result = getattr(fd_functional, name)(widths)(field) + assert_fields_close(result, expected, atol=1e-8, rtol=1e-12) + parts = getattr(fd_functional, name + '_parts')(widths)(field) + for component, (left, right) in enumerate(parts): + assert_close(left + right, expected[component], atol=1e-8, rtol=1e-12) + + @pytest.mark.parametrize('name', ['deriv_forward', 'deriv_back']) @pytest.mark.parametrize('axis', [0, 1, 2]) @pytest.mark.parametrize('spaced', [False, True]) diff --git a/meanas/test/test_fdmath_operators.py b/meanas/test/test_fdmath_operators.py index 2ba7755..0ef88a2 100644 --- a/meanas/test/test_fdmath_operators.py +++ b/meanas/test/test_fdmath_operators.py @@ -91,6 +91,22 @@ def test_vec_cross_matches_pointwise_cross_product() -> None: assert_close(matrix_result, expected) +@pytest.mark.parametrize('name', ['deriv_forward', 'deriv_back']) +@pytest.mark.parametrize('complex_widths', [False, True]) +def test_sparse_derivatives_widen_reciprocals_before_overflow(name, complex_widths) -> None: + dtype = numpy.complex64 if complex_widths else numpy.float32 + width = numpy.array(1e-39 * (1 + 1j) if complex_widths else 1e-39, dtype=dtype) + widths = [numpy.full(2, width, dtype=dtype), numpy.ones(1), numpy.ones(1)] + wide_widths = [dx.astype(numpy.complex128 if complex_widths else numpy.float64) for dx in widths] + field = numpy.array([0, 1e-39]) + with numpy.errstate(over='raise', invalid='raise'): + matrix = getattr(operators, name)(widths)[0] + result = matrix @ field + assert numpy.isfinite(matrix.data).all() + expected = getattr(operators, name)(wide_widths)[0] @ field + assert_close(result, expected, atol=1e-12, rtol=1e-12) + + 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))