scipy 1.16.2__cp312-cp312-win_arm64.whl

scipy/special/_support_alternative_backends.py

@@ -0,0 +1,295 @@
+import functools
+import operator
+from collections.abc import Callable
+from dataclasses import dataclass
+from types import ModuleType
+
+import numpy as np
+from scipy._lib._array_api import (
+    array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray,
+    xp_promote, xp_capabilities, SCIPY_ARRAY_API
+)
+import scipy._lib.array_api_extra as xpx
+from . import _ufuncs
+
+
+@dataclass
+class _FuncInfo:
+    # NumPy-only function. IT MUST BE ELEMENTWISE.
+    func: Callable
+    # Number of arguments, not counting out=
+    # This is for testing purposes only, due to the fact that
+    # inspect.signature() just returns *args for ufuncs.
+    n_args: int
+    # @xp_capabilities decorator, for the purpose of
+    # documentation and unit testing. Omit to indicate
+    # full support for all backends.
+    xp_capabilities: Callable[[Callable], Callable] | None = None
+    # Generic implementation to fall back on if there is no native dispatch
+    # available. This is a function that accepts (main namespace, scipy namespace)
+    # and returns the final callable, or None if not available.
+    generic_impl: Callable[
+        [ModuleType, ModuleType | None], Callable | None
+    ] | None = None
+
+    @property
+    def name(self):
+        return self.func.__name__
+
+    # These are needed by @lru_cache below
+    def __hash__(self):
+        return hash(self.func)
+
+    def __eq__(self, other):
+        return isinstance(other, _FuncInfo) and self.func == other.func
+
+    @property
+    def wrapper(self):
+        if self.name in globals():
+            # Already initialised. We are likely in a unit test.
+            # Return function potentially overridden by xpx.testing.lazy_xp_function.
+            import scipy.special
+            return getattr(scipy.special, self.name)
+
+        if SCIPY_ARRAY_API:
+            @functools.wraps(self.func)
+            def wrapped(*args, **kwargs):
+                xp = array_namespace(*args)
+                return self._wrapper_for(xp)(*args, **kwargs)
+
+            # Allow pickling the function. Normally this is done by @wraps,
+            # but in this case it doesn't work because self.func is a ufunc.
+            wrapped.__module__ = "scipy.special"
+            wrapped.__qualname__ = self.name
+            func = wrapped
+        else:
+            func = self.func
+
+        capabilities = self.xp_capabilities or xp_capabilities()
+        # In order to retain a naked ufunc when SCIPY_ARRAY_API is
+        # disabled, xp_capabilities must apply its changes in place.
+        cap_func = capabilities(func)
+        assert cap_func is func
+        return func
+
+    @functools.lru_cache(1000)
+    def _wrapper_for(self, xp):
+        if is_numpy(xp):
+            return self.func
+
+        # If a native implementation is available, use that
+        spx = scipy_namespace_for(xp)
+        f = _get_native_func(xp, spx, self.name)
+        if f is not None:
+            return f
+
+        # If generic Array API implementation is available, use that
+        if self.generic_impl is not None:
+            f = self.generic_impl(xp, spx)
+            if f is not None:
+                return f
+
+        if is_marray(xp):
+            # Unwrap the array, apply the function on the wrapped namespace,
+            # and then re-wrap it.
+            # IMPORTANT: this only works because all functions in this module
+            # are elementwise. Otherwise, we would not be able to define a
+            # general rule for mask propagation.
+
+            _f = globals()[self.name]  # Allow nested wrapping
+            def f(*args, _f=_f, xp=xp, **kwargs):
+                data_args = [arg.data for arg in args]
+                out = _f(*data_args, **kwargs)
+                mask = functools.reduce(operator.or_, (arg.mask for arg in args))
+                return xp.asarray(out, mask=mask)
+
+            return f
+
+        if is_dask(xp):
+            # Apply the function to each block of the Dask array.
+            # IMPORTANT: map_blocks works only because all functions in this module
+            # are elementwise. It would be a grave mistake to apply this to gufuncs
+            # or any other function with reductions, as they would change their
+            # output depending on chunking!
+
+            _f = globals()[self.name]  # Allow nested wrapping
+            def f(*args, _f=_f, xp=xp, **kwargs):
+                # Hide dtype kwarg from map_blocks
+                return xp.map_blocks(functools.partial(_f, **kwargs), *args)
+
+            return f
+
+        # As a final resort, use the NumPy/SciPy implementation
+        _f = self.func
+        def f(*args, _f=_f, xp=xp, **kwargs):
+            # TODO use xpx.lazy_apply to add jax.jit support
+            # (but dtype propagation can be non-trivial)
+            args = [np.asarray(arg) for arg in args]
+            out = _f(*args, **kwargs)
+            return xp.asarray(out)
+
+        return f
+
+
+def _get_native_func(xp, spx, f_name):
+    f = getattr(spx.special, f_name, None) if spx else None
+    if f is None and hasattr(xp, 'special'):
+        # Currently dead branch, in anticipation of 'special' Array API extension
+        # https://github.com/data-apis/array-api/issues/725
+        f = getattr(xp.special, f_name, None)
+    return f
+
+
+def _rel_entr(xp, spx):
+    def __rel_entr(x, y, *, xp=xp):
+        # https://github.com/data-apis/array-api-extra/issues/160
+        mxp = array_namespace(x._meta, y._meta) if is_dask(xp) else xp
+        x, y = xp_promote(x, y, broadcast=True, force_floating=True, xp=xp)
+        xy_pos = (x > 0) & (y > 0)
+        xy_inf = xp.isinf(x) & xp.isinf(y)
+        res = xpx.apply_where(
+            xy_pos & ~xy_inf,
+            (x, y),
+            # Note: for very large x, this can overflow.
+            lambda x, y: x * (mxp.log(x) - mxp.log(y)),
+            fill_value=xp.inf
+        )
+        res = xpx.at(res)[(x == 0) & (y >= 0)].set(0)
+        res = xpx.at(res)[xp.isnan(x) | xp.isnan(y) | (xy_pos & xy_inf)].set(xp.nan)
+        return res
+
+    return __rel_entr
+
+
+def _xlogy(xp, spx):
+    def __xlogy(x, y, *, xp=xp):
+        x, y = xp_promote(x, y, force_floating=True, xp=xp)
+        with np.errstate(divide='ignore', invalid='ignore'):
+            temp = x * xp.log(y)
+        return xp.where(x == 0., 0., temp)
+    return __xlogy
+
+
+
+def _chdtr(xp, spx):
+    # The difference between this and just using `gammainc`
+    # defined by `get_array_special_func` is that if `gammainc`
+    # isn't found, we don't want to use the SciPy version; we'll
+    # return None here and use the SciPy version of `chdtr`.
+    gammainc = _get_native_func(xp, spx, 'gammainc')
+    if gammainc is None:
+        return None
+
+    def __chdtr(v, x):
+        res = gammainc(v / 2, x / 2)  # this is almost all we need
+        # The rest can be removed when google/jax#20507 is resolved
+        mask = (v == 0) & (x > 0)  # JAX returns NaN
+        res = xp.where(mask, 1., res)
+        mask = xp.isinf(v) & xp.isinf(x)  # JAX returns 1.0
+        return xp.where(mask, xp.nan, res)
+    return __chdtr
+
+
+def _chdtrc(xp, spx):
+    # The difference between this and just using `gammaincc`
+    # defined by `get_array_special_func` is that if `gammaincc`
+    # isn't found, we don't want to use the SciPy version; we'll
+    # return None here and use the SciPy version of `chdtrc`.
+    gammaincc = _get_native_func(xp, spx, 'gammaincc')
+    if gammaincc is None:
+        return None
+
+    def __chdtrc(v, x):
+        res = xp.where(x >= 0, gammaincc(v/2, x/2), 1)
+        i_nan = ((x == 0) & (v == 0)) | xp.isnan(x) | xp.isnan(v) | (v <= 0)
+        res = xp.where(i_nan, xp.nan, res)
+        return res
+    return __chdtrc
+
+
+def _betaincc(xp, spx):
+    betainc = _get_native_func(xp, spx, 'betainc')
+    if betainc is None:
+        return None
+
+    def __betaincc(a, b, x):
+        # not perfect; might want to just rely on SciPy
+        return betainc(b, a, 1-x)
+    return __betaincc
+
+
+def _stdtr(xp, spx):
+    betainc = _get_native_func(xp, spx, 'betainc')
+    if betainc is None:
+        return None
+
+    def __stdtr(df, t):
+        x = df / (t ** 2 + df)
+        tail = betainc(df / 2, 0.5, x) / 2
+        return xp.where(t < 0, tail, 1 - tail)
+
+    return __stdtr
+
+
+def _stdtrit(xp, spx):
+    # Need either native stdtr or native betainc
+    stdtr = _get_native_func(xp, spx, 'stdtr') or _stdtr(xp, spx)
+    # If betainc is not defined, the root-finding would be done with `xp`
+    # despite `stdtr` being evaluated with SciPy/NumPy `stdtr`. Save the
+    # conversions: in this case, just evaluate `stdtrit` with SciPy/NumPy.
+    if stdtr is None:
+        return None
+
+    from scipy.optimize.elementwise import bracket_root, find_root
+
+    def __stdtrit(df, p):
+        def fun(t, df, p):  return stdtr(df, t) - p
+        res_bracket = bracket_root(fun, xp.zeros_like(p), args=(df, p))
+        res_root = find_root(fun, res_bracket.bracket, args=(df, p))
+        return res_root.x
+
+    return __stdtrit
+
+
+# Inventory of automatically dispatched functions
+# IMPORTANT: these must all be **elementwise** functions!
+
+# PyTorch doesn't implement `betainc`.
+# On torch CPU we can fall back to NumPy, but on GPU it won't work.
+_needs_betainc = xp_capabilities(cpu_only=True, exceptions=['jax.numpy', 'cupy'])
+
+_special_funcs = (
+    _FuncInfo(_ufuncs.betainc, 3, _needs_betainc),
+    _FuncInfo(_ufuncs.betaincc, 3, _needs_betainc, generic_impl=_betaincc),
+    _FuncInfo(_ufuncs.chdtr, 2, generic_impl=_chdtr),
+    _FuncInfo(_ufuncs.chdtrc, 2, generic_impl=_chdtrc),
+    _FuncInfo(_ufuncs.erf, 1),
+    _FuncInfo(_ufuncs.erfc, 1),
+    _FuncInfo(_ufuncs.entr, 1),
+    _FuncInfo(_ufuncs.expit, 1),
+    _FuncInfo(_ufuncs.i0, 1),
+    _FuncInfo(_ufuncs.i0e, 1),
+    _FuncInfo(_ufuncs.i1, 1),
+    _FuncInfo(_ufuncs.i1e, 1),
+    _FuncInfo(_ufuncs.log_ndtr, 1),
+    _FuncInfo(_ufuncs.logit, 1),
+    _FuncInfo(_ufuncs.gammaln, 1),
+    _FuncInfo(_ufuncs.gammainc, 2),
+    _FuncInfo(_ufuncs.gammaincc, 2),
+    _FuncInfo(_ufuncs.ndtr, 1),
+    _FuncInfo(_ufuncs.ndtri, 1),
+    _FuncInfo(_ufuncs.rel_entr, 2, generic_impl=_rel_entr),
+    _FuncInfo(_ufuncs.stdtr,  2, _needs_betainc, generic_impl=_stdtr),
+    _FuncInfo(_ufuncs.stdtrit, 2,
+              xp_capabilities(
+                  cpu_only=True, exceptions=['cupy'],  # needs betainc
+                  skip_backends=[("jax.numpy", "no scipy.optimize support")]),
+              generic_impl=_stdtrit),
+    _FuncInfo(_ufuncs.xlogy, 2, generic_impl=_xlogy),
+)
+
+# Override ufuncs.
+# When SCIPY_ARRAY_API is disabled, this exclusively updates the docstrings in place
+# and populates the xp_capabilities table, while retaining the original ufuncs.
+globals().update({nfo.func.__name__: nfo.wrapper for nfo in _special_funcs})
+__all__ = [nfo.func.__name__ for nfo in _special_funcs]

scipy/special/_test_internal.pyi

@@ -0,0 +1,9 @@
+import numpy as np
+
+def have_fenv() -> bool: ...
+def random_double(size: int, rng: np.random.RandomState) -> np.float64: ...
+def test_add_round(size: int, mode: str, rng: np.random.RandomState): ...
+
+def _dd_exp(xhi: float, xlo: float) -> tuple[float, float]: ...
+def _dd_log(xhi: float, xlo: float) -> tuple[float, float]: ...
+def _dd_expm1(xhi: float, xlo: float) -> tuple[float, float]: ...

scipy/special/_testutils.py

@@ -0,0 +1,321 @@
+import os
+import functools
+import operator
+from scipy._lib import _pep440
+
+import numpy as np
+from numpy.testing import assert_
+import pytest
+
+import scipy.special as sc
+
+__all__ = ['with_special_errors', 'assert_func_equal', 'FuncData']
+
+
+#------------------------------------------------------------------------------
+# Check if a module is present to be used in tests
+#------------------------------------------------------------------------------
+
+class MissingModule:
+    def __init__(self, name):
+        self.name = name
+
+
+def check_version(module, min_ver):
+    if type(module) is MissingModule:
+        return pytest.mark.skip(reason=f"{module.name} is not installed")
+    return pytest.mark.skipif(
+        _pep440.parse(module.__version__) < _pep440.Version(min_ver),
+        reason=f"{module.__name__} version >= {min_ver} required"
+    )
+
+
+#------------------------------------------------------------------------------
+# Enable convergence and loss of precision warnings -- turn off one by one
+#------------------------------------------------------------------------------
+
+def with_special_errors(func):
+    """
+    Enable special function errors (such as underflow, overflow,
+    loss of precision, etc.)
+    """
+    @functools.wraps(func)
+    def wrapper(*a, **kw):
+        with sc.errstate(all='raise'):
+            res = func(*a, **kw)
+        return res
+    return wrapper
+
+
+#------------------------------------------------------------------------------
+# Comparing function values at many data points at once, with helpful
+# error reports
+#------------------------------------------------------------------------------
+
+def assert_func_equal(func, results, points, rtol=None, atol=None,
+                      param_filter=None, knownfailure=None,
+                      vectorized=True, dtype=None, nan_ok=False,
+                      ignore_inf_sign=False, distinguish_nan_and_inf=True):
+    if hasattr(points, 'next'):
+        # it's a generator
+        points = list(points)
+
+    points = np.asarray(points)
+    if points.ndim == 1:
+        points = points[:,None]
+    nparams = points.shape[1]
+
+    if hasattr(results, '__name__'):
+        # function
+        data = points
+        result_columns = None
+        result_func = results
+    else:
+        # dataset
+        data = np.c_[points, results]
+        result_columns = list(range(nparams, data.shape[1]))
+        result_func = None
+
+    fdata = FuncData(func, data, list(range(nparams)),
+                     result_columns=result_columns, result_func=result_func,
+                     rtol=rtol, atol=atol, param_filter=param_filter,
+                     knownfailure=knownfailure, nan_ok=nan_ok, vectorized=vectorized,
+                     ignore_inf_sign=ignore_inf_sign,
+                     distinguish_nan_and_inf=distinguish_nan_and_inf)
+    fdata.check()
+
+
+class FuncData:
+    """
+    Data set for checking a special function.
+
+    Parameters
+    ----------
+    func : function
+        Function to test
+    data : numpy array
+        columnar data to use for testing
+    param_columns : int or tuple of ints
+        Columns indices in which the parameters to `func` lie.
+        Can be imaginary integers to indicate that the parameter
+        should be cast to complex.
+    result_columns : int or tuple of ints, optional
+        Column indices for expected results from `func`.
+    result_func : callable, optional
+        Function to call to obtain results.
+    rtol : float, optional
+        Required relative tolerance. Default is 5*eps.
+    atol : float, optional
+        Required absolute tolerance. Default is 5*tiny.
+    param_filter : function, or tuple of functions/Nones, optional
+        Filter functions to exclude some parameter ranges.
+        If omitted, no filtering is done.
+    knownfailure : str, optional
+        Known failure error message to raise when the test is run.
+        If omitted, no exception is raised.
+    nan_ok : bool, optional
+        If nan is always an accepted result.
+    vectorized : bool, optional
+        Whether all functions passed in are vectorized.
+    ignore_inf_sign : bool, optional
+        Whether to ignore signs of infinities.
+        (Doesn't matter for complex-valued functions.)
+    distinguish_nan_and_inf : bool, optional
+        If True, treat numbers which contain nans or infs as
+        equal. Sets ignore_inf_sign to be True.
+
+    """
+
+    def __init__(self, func, data, param_columns, result_columns=None,
+                 result_func=None, rtol=None, atol=None, param_filter=None,
+                 knownfailure=None, dataname=None, nan_ok=False, vectorized=True,
+                 ignore_inf_sign=False, distinguish_nan_and_inf=True):
+        self.func = func
+        self.data = data
+        self.dataname = dataname
+        if not hasattr(param_columns, '__len__'):
+            param_columns = (param_columns,)
+        self.param_columns = tuple(param_columns)
+        if result_columns is not None:
+            if not hasattr(result_columns, '__len__'):
+                result_columns = (result_columns,)
+            self.result_columns = tuple(result_columns)
+            if result_func is not None:
+                message = "Only result_func or result_columns should be provided"
+                raise ValueError(message)
+        elif result_func is not None:
+            self.result_columns = None
+        else:
+            raise ValueError("Either result_func or result_columns should be provided")
+        self.result_func = result_func
+        self.rtol = rtol
+        self.atol = atol
+        if not hasattr(param_filter, '__len__'):
+            param_filter = (param_filter,)
+        self.param_filter = param_filter
+        self.knownfailure = knownfailure
+        self.nan_ok = nan_ok
+        self.vectorized = vectorized
+        self.ignore_inf_sign = ignore_inf_sign
+        self.distinguish_nan_and_inf = distinguish_nan_and_inf
+        if not self.distinguish_nan_and_inf:
+            self.ignore_inf_sign = True
+
+    def get_tolerances(self, dtype):
+        if not np.issubdtype(dtype, np.inexact):
+            dtype = np.dtype(float)
+        info = np.finfo(dtype)
+        rtol, atol = self.rtol, self.atol
+        if rtol is None:
+            rtol = 5*info.eps
+        if atol is None:
+            atol = 5*info.tiny
+        return rtol, atol
+
+    def check(self, data=None, dtype=None, dtypes=None):
+        """Check the special function against the data."""
+        __tracebackhide__ = operator.methodcaller(
+            'errisinstance', AssertionError
+        )
+
+        if self.knownfailure:
+            pytest.xfail(reason=self.knownfailure)
+
+        if data is None:
+            data = self.data
+
+        if dtype is None:
+            dtype = data.dtype
+        else:
+            data = data.astype(dtype)
+
+        rtol, atol = self.get_tolerances(dtype)
+
+        # Apply given filter functions
+        if self.param_filter:
+            param_mask = np.ones((data.shape[0],), np.bool_)
+            for j, filter in zip(self.param_columns, self.param_filter):
+                if filter:
+                    param_mask &= list(filter(data[:,j]))
+            data = data[param_mask]
+
+        # Pick parameters from the correct columns
+        params = []
+        for idx, j in enumerate(self.param_columns):
+            if np.iscomplexobj(j):
+                j = int(j.imag)
+                params.append(data[:,j].astype(complex))
+            elif dtypes and idx < len(dtypes):
+                params.append(data[:, j].astype(dtypes[idx]))
+            else:
+                params.append(data[:,j])
+
+        # Helper for evaluating results
+        def eval_func_at_params(func, skip_mask=None):
+            if self.vectorized:
+                got = func(*params)
+            else:
+                got = []
+                for j in range(len(params[0])):
+                    if skip_mask is not None and skip_mask[j]:
+                        got.append(np.nan)
+                        continue
+                    got.append(func(*tuple([params[i][j] for i in range(len(params))])))
+                got = np.asarray(got)
+            if not isinstance(got, tuple):
+                got = (got,)
+            return got
+
+        # Evaluate function to be tested
+        got = eval_func_at_params(self.func)
+
+        # Grab the correct results
+        if self.result_columns is not None:
+            # Correct results passed in with the data
+            wanted = tuple([data[:,icol] for icol in self.result_columns])
+        else:
+            # Function producing correct results passed in
+            skip_mask = None
+            if self.nan_ok and len(got) == 1:
+                # Don't spend time evaluating what doesn't need to be evaluated
+                skip_mask = np.isnan(got[0])
+            wanted = eval_func_at_params(self.result_func, skip_mask=skip_mask)
+
+        # Check the validity of each output returned
+        assert_(len(got) == len(wanted))
+
+        for output_num, (x, y) in enumerate(zip(got, wanted)):
+            if np.issubdtype(x.dtype, np.complexfloating) or self.ignore_inf_sign:
+                pinf_x = np.isinf(x)
+                pinf_y = np.isinf(y)
+                minf_x = np.isinf(x)
+                minf_y = np.isinf(y)
+            else:
+                pinf_x = np.isposinf(x)
+                pinf_y = np.isposinf(y)
+                minf_x = np.isneginf(x)
+                minf_y = np.isneginf(y)
+            nan_x = np.isnan(x)
+            nan_y = np.isnan(y)
+
+            with np.errstate(all='ignore'):
+                abs_y = np.absolute(y)
+                abs_y[~np.isfinite(abs_y)] = 0
+                diff = np.absolute(x - y)
+                diff[~np.isfinite(diff)] = 0
+
+                rdiff = diff / np.absolute(y)
+                rdiff[~np.isfinite(rdiff)] = 0
+
+            tol_mask = (diff <= atol + rtol*abs_y)
+            pinf_mask = (pinf_x == pinf_y)
+            minf_mask = (minf_x == minf_y)
+
+            nan_mask = (nan_x == nan_y)
+
+            bad_j = ~(tol_mask & pinf_mask & minf_mask & nan_mask)
+
+            point_count = bad_j.size
+            if self.nan_ok:
+                bad_j &= ~nan_x
+                bad_j &= ~nan_y
+                point_count -= (nan_x | nan_y).sum()
+
+            if not self.distinguish_nan_and_inf and not self.nan_ok:
+                # If nan's are okay we've already covered all these cases
+                inf_x = np.isinf(x)
+                inf_y = np.isinf(y)
+                both_nonfinite = (inf_x & nan_y) | (nan_x & inf_y)
+                bad_j &= ~both_nonfinite
+                point_count -= both_nonfinite.sum()
+
+            if np.any(bad_j):
+                # Some bad results: inform what, where, and how bad
+                msg = [""]
+                msg.append(f"Max |adiff|: {diff[bad_j].max():g}")
+                msg.append(f"Max |rdiff|: {rdiff[bad_j].max():g}")
+                msg.append(f"Bad results ({np.sum(bad_j)} out of "
+                           f"{point_count}) for the following points "
+                           f"(in output {output_num}):")
+                for j in np.nonzero(bad_j)[0]:
+                    j = int(j)
+                    def fmt(x):
+                        return f'{np.array2string(x[j], precision=18):30s}'
+                    a = "  ".join(map(fmt, params))
+                    b = "  ".join(map(fmt, got))
+                    c = "  ".join(map(fmt, wanted))
+                    d = fmt(rdiff)
+                    msg.append(f"{a} => {b} != {c}  (rdiff {d})")
+                assert_(False, "\n".join(msg))
+
+    def __repr__(self):
+        """Pretty-printing"""
+        if np.any(list(map(np.iscomplexobj, self.param_columns))):
+            is_complex = " (complex)"
+        else:
+            is_complex = ""
+        if self.dataname:
+            return (f"<Data for {self.func.__name__}{is_complex}: "
+                    f"{os.path.basename(self.dataname)}>")
+        else:
+            return f"<Data for {self.func.__name__}{is_complex}>"