[fdmath] widen curl calculations and sparse reciprocals before arithmetic

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])))

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

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])

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))