Fix complex interpolation and normalized period edge cases

Author: DeepCogNeural

src/array_api_extra/_agnostic/_interpolation.py

@@ -59,6 +59,39 @@ def _safe_divide(
     return xp.where(overflow & ~positive_overflow, -infinity, out)
 
 
+def _safe_multiply(
+    x1: Array, x2: Array, /, *, zero_times_infinity: float, xp: ArrayNamespace
+) -> Array:
+    """Multiply arrays without warnings from overflow or zero times infinity."""
+    zero = xp.zeros_like(x1)
+    one = xp.ones_like(x2)
+    maximum = xp.asarray(
+        xp.finfo(x1.dtype).max, dtype=x1.dtype, device=_compat.device(x1)
+    )
+    absolute_x1 = xp.abs(x1)
+    absolute_x2 = xp.abs(x2)
+    large_x2 = absolute_x2 > 1
+    overflow_threshold = maximum / xp.where(large_x2, absolute_x2, one)
+    overflow = (
+        xp.isfinite(x1)
+        & xp.isfinite(x2)
+        & large_x2
+        & (absolute_x1 > overflow_threshold)
+    )
+    zero_inf = ((x1 == 0) & xp.isinf(x2)) | (xp.isinf(x1) & (x2 == 0))
+    suppressed = overflow | zero_inf
+
+    out = xp.where(suppressed, zero, x1) * xp.where(suppressed, zero, x2)
+    infinity = xp.asarray(math.inf, dtype=x1.dtype, device=_compat.device(x1))
+    positive_overflow = overflow & ((x1 > 0) == (x2 > 0))
+    out = xp.where(positive_overflow, infinity, out)
+    out = xp.where(overflow & ~positive_overflow, -infinity, out)
+    zero_inf_value = xp.asarray(
+        zero_times_infinity, dtype=x1.dtype, device=_compat.device(x1)
+    )
+    return xp.where(zero_inf, zero_inf_value, out)
+
+
 def _interp_component(
     x: Array,
     x_lo: Array,
@@ -68,6 +101,7 @@ def _interp_component(
     /,
     *,
     inactive: Array,
+    reciprocal_first: bool = False,
     xp: ArrayNamespace,
 ) -> Array:
     """Interpolate one real component without invalid arithmetic."""
@@ -87,17 +121,28 @@ def _interp_component(
     coordinate_difference = _safe_difference(safe_x_hi, safe_x_lo, xp=xp)
     value_difference = _safe_difference(safe_y_hi, safe_y_lo, xp=xp)
     indeterminate_slope = xp.isinf(coordinate_difference) & xp.isinf(value_difference)
-    slope = _safe_divide(
-        xp.where(
-            indeterminate_slope, xp.zeros_like(value_difference), value_difference
-        ),
-        xp.where(
-            indeterminate_slope,
-            xp.ones_like(coordinate_difference),
-            coordinate_difference,
-        ),
-        xp=xp,
+    safe_value_difference = xp.where(
+        indeterminate_slope, xp.zeros_like(value_difference), value_difference
+    )
+    safe_coordinate_difference = xp.where(
+        indeterminate_slope,
+        xp.ones_like(coordinate_difference),
+        coordinate_difference,
     )
+    if reciprocal_first:
+        inverse_coordinate_difference = _safe_divide(
+            xp.ones_like(safe_coordinate_difference),
+            safe_coordinate_difference,
+            xp=xp,
+        )
+        slope = _safe_multiply(
+            safe_value_difference,
+            inverse_coordinate_difference,
+            zero_times_infinity=math.nan,
+            xp=xp,
+        )
+    else:
+        slope = _safe_divide(safe_value_difference, safe_coordinate_difference, xp=xp)
     slope = xp.where(indeterminate_slope, nan, slope)
 
     left_delta = _safe_difference(safe_x, safe_x_lo, xp=xp)
@@ -131,13 +176,22 @@ def _interp_component(
 
     only_lo_coordinate_infinite = xp.isinf(x_lo) & ~xp.isinf(x_hi)
     only_hi_coordinate_infinite = ~xp.isinf(x_lo) & xp.isinf(x_hi)
+    finite_y_lo = xp.where(values_finite, y_lo, xp.zeros_like(y_lo))
+    finite_y_hi = xp.where(values_finite, y_hi, xp.zeros_like(y_hi))
+    infinite_coordinate_value_overflow = values_finite & xp.isinf(
+        _safe_difference(finite_y_hi, finite_y_lo, xp=xp)
+    )
     infinite_coordinate_out = xp.where(
-        values_finite & only_lo_coordinate_infinite,
-        y_hi,
+        infinite_coordinate_value_overflow,
+        nan,
         xp.where(
-            values_finite & only_hi_coordinate_infinite,
-            y_lo,
-            xp.where(equal_values, y_lo, nan),
+            values_finite & only_lo_coordinate_infinite,
+            y_hi,
+            xp.where(
+                values_finite & only_hi_coordinate_infinite,
+                y_lo,
+                xp.where(equal_values, y_lo, nan),
+            ),
         ),
     )
     return xp.where(coordinate_infinite, infinite_coordinate_out, out)
@@ -234,6 +288,7 @@ def interp(
             y_lo_real,
             y_hi_real,
             inactive=inactive,
+            reciprocal_first=True,
             xp=xp,
         )
         out_imag = _interp_component(
@@ -243,6 +298,7 @@ def interp(
             y_lo_imag,
             y_hi_imag,
             inactive=inactive,
+            reciprocal_first=True,
             xp=xp,
         )
         left_array = values[0] if left is None else left

src/array_api_extra/_interpolation.py

@@ -148,10 +148,10 @@ def interp(
         the last element of `values` is used. Ignored, without validation, when
         `period` is provided.
     period : real scalar, optional
-        Finite period for the sample and query coordinates. A negative value is
-        treated as its absolute value and zero is invalid. When provided,
-        coordinates are normalized to the period and samples are sorted by
-        normalized coordinate.
+        Period for the sample and query coordinates. It must remain finite and
+        nonzero when converted to ``float64``. A negative value is treated as its
+        absolute value. When provided, coordinates are normalized to the period
+        and samples are sorted by normalized coordinate.
     xp : array_namespace, optional
         The standard-compatible namespace for the array arguments. Default: infer.
 
@@ -197,13 +197,17 @@ def interp(
         if not _is_real_scalar(period):
             msg = "`period` must be a finite real scalar or None."
             raise TypeError(msg)
+        try:
+            period = float(abs(period))
+        except (OverflowError, ValueError) as error:
+            msg = "`period` must be representable as a finite float."
+            raise ValueError(msg) from error
         if not math.isfinite(period):
-            msg = "`period` must be finite."
+            msg = "`period` must be finite after conversion to float."
             raise ValueError(msg)
         if period == 0:
-            msg = "`period` must be nonzero."
+            msg = "`period` must be nonzero after conversion to float."
             raise ValueError(msg)
-        period = float(abs(period))
 
     namespace_args: list[Array] = [x_points, values]
     if _compat.is_array_api_obj(x):

tests/main/test_interpolation.py

@@ -1,3 +1,4 @@
+from fractions import Fraction
 from typing import Any, Literal
 
 import numpy as np
@@ -151,13 +152,16 @@ def test_reference_extreme_finite_coordinates(
         )
         assert_equal(actual, xp.asarray(expected, dtype=xp.float64))
 
+    @pytest.mark.parametrize("values", [[0.0, 1.0], [0.0, 1e-320]])
     @implementations
     def test_reference_subnormal_interval(
-        self, xp: ArrayNamespace, implementation: Implementation
+        self,
+        xp: ArrayNamespace,
+        implementation: Implementation,
+        values: list[float],
     ):
         x_points = np.asarray([0.0, 1e-320])
         x = np.asarray([5e-321])
-        values = np.asarray([0.0, 1.0])
         expected = np.interp(x, x_points, values)
 
         actual = _interp(
@@ -169,6 +173,67 @@ def test_reference_subnormal_interval(
         )
         assert_equal(actual, xp.asarray(expected, dtype=xp.float64))
 
+    @pytest.mark.parametrize(
+        "end_value", [complex(1e-320, 1e-320), complex(1e-320, 0.0)]
+    )
+    @implementations
+    def test_reference_subnormal_complex_slope(
+        self,
+        xp: ArrayNamespace,
+        implementation: Implementation,
+        end_value: complex,
+    ):
+        x_points = np.asarray([0.0, 1e-320])
+        x = np.asarray([5e-321])
+        values = np.asarray([0j, end_value])
+        expected = np.interp(x, x_points, values)
+
+        actual = _interp(
+            implementation,
+            xp.asarray(x, dtype=xp.float64),
+            xp.asarray(x_points, dtype=xp.float64),
+            xp.asarray(values, dtype=xp.complex128),
+            left=0,
+            xp=xp,
+        )
+        assert_equal(xp.real(actual), xp.asarray(expected.real, dtype=xp.float64))
+        assert_equal(xp.imag(actual), xp.asarray(expected.imag, dtype=xp.float64))
+
+    @pytest.mark.parametrize(
+        ("x_points", "x"),
+        [([-np.inf, 0.0], [-1.0]), ([0.0, np.inf], [1.0])],
+        ids=["left", "right"],
+    )
+    @pytest.mark.parametrize("complex_values", [False, True], ids=["real", "complex"])
+    @implementations
+    def test_reference_infinite_coordinate_value_overflow(
+        self,
+        xp: ArrayNamespace,
+        implementation: Implementation,
+        x_points: list[float],
+        x: list[float],
+        complex_values: bool,
+    ):
+        values = np.asarray([-1e308, 1e308])
+        dtype = xp.float64
+        if complex_values:
+            values = values + values * 1j
+            dtype = xp.complex128
+        expected = np.interp(x, x_points, values)
+
+        actual = _interp(
+            implementation,
+            xp.asarray(x, dtype=xp.float64),
+            xp.asarray(x_points, dtype=xp.float64),
+            xp.asarray(values, dtype=dtype),
+            xp=xp,
+        )
+        if complex_values:
+            assert_equal(xp.real(actual), xp.asarray(expected.real, dtype=xp.float64))
+            assert_equal(xp.imag(actual), xp.asarray(expected.imag, dtype=xp.float64))
+        else:
+            assert_equal(actual, xp.asarray(expected, dtype=xp.float64))
+
     # NumPy recommends strictly increasing sample coordinates. array-api-extra
     # deliberately defines these repeated-knot cases, including which value wins at
     # the exact knot, rather than exposing searchsorted clipping or a 0/0 division.
@@ -239,7 +304,10 @@ def test_complex_nan_query(
             xp.asarray(values, dtype=xp.complex128),
             xp=xp,
         )
-        assert_equal(actual, xp.asarray(expected, dtype=xp.complex128))
+        assert_equal(xp.real(actual), xp.asarray(expected.real, dtype=xp.float64))
+        assert_equal(xp.imag(actual), xp.asarray(expected.imag, dtype=xp.float64))
+        with pytest.raises(AssertionError):
+            assert_equal(xp.imag(actual), xp.asarray([123.0], dtype=xp.float64))
 
     @implementations
     def test_left_right_and_complex_values(
@@ -426,9 +494,11 @@ def test_bound_and_period_validation(self, xp: ArrayNamespace):
         with pytest.raises(TypeError):
             _ = xpx_interp(x, x_points, values, left=1j)
 
-        for period in (0.0, np.inf, -np.inf, np.nan):
+        for period in (0.0, np.inf, -np.inf, np.nan, Fraction(1, 10**400)):
             with pytest.raises(ValueError, match=r"must be (finite|nonzero)"):
                 _ = xpx_interp(x, x_points, values, period=period)
+        with pytest.raises(ValueError, match="representable as a finite float"):
+            _ = xpx_interp(x, x_points, values, period=Fraction(10**400, 1))
         with pytest.raises(TypeError):
             _ = xpx_interp(x, x_points, values, period=True)
         with pytest.raises(TypeError):