ENH: add one_minus_square

docs/sphinx/api-elementwise.md

@@ -11,6 +11,7 @@
     deg2rad
     isclose
     nan_to_num
+    one_minus_square
     rad2deg
     sinc
 ```

src/array_api_extra/__init__.py

@@ -5,7 +5,15 @@
 from ._agnostic._inspection import default_dtype
 from ._at import at
 from ._creation import create_diagonal, one_hot
-from ._elementwise import angle, deg2rad, isclose, nan_to_num, rad2deg, sinc
+from ._elementwise import (
+    angle,
+    deg2rad,
+    isclose,
+    nan_to_num,
+    one_minus_square,
+    rad2deg,
+    sinc,
+)
 from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index
 from ._lazy import lazy_apply
 from ._linalg import kron
@@ -42,6 +50,7 @@
     "nansum",
     "nunique",
     "one_hot",
+    "one_minus_square",
     "pad",
     "partition",
     "rad2deg",

src/array_api_extra/_agnostic/_elementwise.py

@@ -14,6 +14,7 @@
     "deg2rad",
     "isclose",
     "nan_to_num",
+    "one_minus_square",
     "rad2deg",
     "sinc",
 ]
@@ -315,3 +316,9 @@ def rad2deg(x: Array, /, *, xp: ArrayNamespace) -> Array:
     # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._elementwise`."""
     return x * 180 / xp.pi
+
+
+def one_minus_square(x: Array, /) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._elementwise`."""
+    return (1 + x) * (1 - x)

src/array_api_extra/_elementwise.py

@@ -4,7 +4,15 @@
 from ._lib import _compat, _helpers
 from ._lib._typing import Array, ArrayNamespace
 
-__all__ = ["angle", "deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"]
+__all__ = [
+    "angle",
+    "deg2rad",
+    "isclose",
+    "nan_to_num",
+    "one_minus_square",
+    "rad2deg",
+    "sinc",
+]
 
 
 def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array:
@@ -340,6 +348,63 @@
     return _agnostic._elementwise.nan_to_num(y, fill_value=fill_value, xp=xp)
 
 
+def one_minus_square(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
+    r"""
+    Compute :math:`1 - x^2` element-wise.
+
+    The result is evaluated as :math:`(1 + x)(1 - x)`, which avoids the
+    catastrophic cancellation of the naive formula when :math:`|x|` is close to
+    1.
+
+    Parameters
+    ----------
+    x : array
+        Input array. Must have a numeric dtype.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        ``1 - x**2`` calculated element-wise, with the same shape and dtype as `x`.
+
+    See Also
+    --------
+    jax.lax.one_minus_square : JAX equivalent, which additionally defines its
+        derivative as :math:`-2x` for accurate gradients near zero.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> xpx.one_minus_square(xp.asarray([0.0, 0.5, 1.0, 2.0]), xp=xp)
+    Array([ 1.  ,  0.75,  0.  , -3.  ], dtype=array_api_strict.float64)
+
+    Near 1, the result keeps full relative precision:
+
+    >>> x = xp.asarray(1 - 2.0**-30)
+    >>> float(xpx.one_minus_square(x, xp=xp)) == 2.0**-29 - 2.0**-60
+    True
+    >>> float(1 - x**2) == 2.0**-29 - 2.0**-60
+    False
+    """
+    if xp is None:
+        xp = _compat.array_namespace(x)
+
+    if not xp.isdtype(x.dtype, "numeric"):
+        msg = "`x` must have a numeric dtype."
+        raise TypeError(msg)
+
+    if _compat.is_jax_namespace(xp):
+        import jax
+
+        # Added in JAX 0.11.2.
+        if hasattr(jax.lax, "one_minus_square"):
+            return jax.lax.one_minus_square(x)
+
+    return _agnostic._elementwise.one_minus_square(x)
+
+
 def sinc(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
     r"""
     Return the normalized sinc function.

tests/main/test_elementwise.py

@@ -15,6 +15,7 @@
     deg2rad,
     isclose,
     nan_to_num,
+    one_minus_square,
     rad2deg,
     sinc,
 )
@@ -28,6 +29,7 @@
 lazy_xp_function(deg2rad)
 lazy_xp_function(isclose)
 lazy_xp_function(nan_to_num)
+lazy_xp_function(one_minus_square)
 lazy_xp_function(rad2deg)
 lazy_xp_function(sinc)
 
@@ -697,6 +699,45 @@ def test_device(self, xp: ArrayNamespace, device: Device):
         assert get_device(angle(a)) == device
 
 
+class TestOneMinusSquare:
+    def test_basic(self, xp: ArrayNamespace):
+        x = xp.asarray([-2.0, -1.0, 0.0, 0.5, 1.0, 3.0])
+        expected = xp.asarray([-3.0, 0.0, 1.0, 0.75, 0.0, -8.0])
+        assert_equal(one_minus_square(x), expected)
+
+    def test_near_one(self, xp: ArrayNamespace):
+        x = xp.asarray([1 - 2.0**-30, -(1 - 2.0**-30)], dtype=xp.float64)
+        expected = xp.asarray([2.0**-29 - 2.0**-60] * 2, dtype=xp.float64)
+        assert_equal(one_minus_square(x), expected)
+
+    @pytest.mark.parametrize("dtype_name", ["float32", "float64"])
+    def test_dtype(self, xp: ArrayNamespace, dtype_name: str):
+        dtype = getattr(xp, dtype_name)
+        x = xp.asarray([0.5], dtype=dtype)
+        actual = one_minus_square(x, xp=xp)
+        assert actual.dtype == dtype
+        assert_equal(actual, xp.asarray([0.75], dtype=dtype))
+
+    def test_integral(self, xp: ArrayNamespace):
+        x = xp.asarray([0, 1, 2, -3], dtype=xp.int32)
+        actual = one_minus_square(x, xp=xp)
+        assert_equal(actual, xp.asarray([1, 0, -3, -8], dtype=xp.int32))
+
+    def test_complex(self, xp: ArrayNamespace):
+        x = xp.asarray([0.5 + 1j], dtype=xp.complex128)
+        actual = one_minus_square(x, xp=xp)
+        assert_close(actual, xp.asarray([1.75 - 1j], dtype=xp.complex128))
+
+    def test_bool(self, xp: ArrayNamespace):
+        x = xp.asarray([True])
+        with pytest.raises(TypeError, match="numeric dtype"):
+            _ = one_minus_square(x, xp=xp)
+
+    def test_device(self, xp: ArrayNamespace, device: Device):
+        x = xp.asarray([0.5], device=device)
+        assert get_device(one_minus_square(x)) == device
+
+
 class TestDeg2Rad:
     def test_basic(self, xp: ArrayNamespace):
         x = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0])