venv/lib/python3.11/site-packages/scipy/integrate/tests/test_tanhsinh.py

# mypy: disable-error-code="attr-defined"
import os
import pytest
import math

import numpy as np
from numpy.testing import assert_allclose

from scipy.conftest import array_api_compatible
import scipy._lib._elementwise_iterative_method as eim
from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal
from scipy._lib._array_api import array_namespace, xp_size, xp_ravel, xp_copy, is_numpy
from scipy import special, stats
from scipy.integrate import quad_vec, nsum, tanhsinh as _tanhsinh
from scipy.integrate._tanhsinh import _pair_cache
from scipy.stats._discrete_distns import _gen_harmonic_gt1


def norm_pdf(x, xp=None):
    xp = array_namespace(x) if xp is None else xp
    return 1/(2*xp.pi)**0.5 * xp.exp(-x**2/2)

def norm_logpdf(x, xp=None):
    xp = array_namespace(x) if xp is None else xp
    return -0.5*math.log(2*xp.pi) - x**2/2


def _vectorize(xp):
    # xp-compatible version of np.vectorize
    # assumes arguments are all arrays of the same shape
    def decorator(f):
        def wrapped(*arg_arrays):
            shape = arg_arrays[0].shape
            arg_arrays = [xp_ravel(arg_array) for arg_array in arg_arrays]
            res = []
            for i in range(math.prod(shape)):
                arg_scalars = [arg_array[i] for arg_array in arg_arrays]
                res.append(f(*arg_scalars))
            return res

        return wrapped

    return decorator


@array_api_compatible
@pytest.mark.usefixtures("skip_xp_backends")
@pytest.mark.skip_xp_backends(
    'array_api_strict', reason='Currently uses fancy indexing assignment.'
)
@pytest.mark.skip_xp_backends(
    'jax.numpy', reason='JAX arrays do not support item assignment.'
)
class TestTanhSinh:

    # Test problems from [1] Section 6
    def f1(self, t):
        return t * np.log(1 + t)

    f1.ref = 0.25
    f1.b = 1

    def f2(self, t):
        return t ** 2 * np.arctan(t)

    f2.ref = (np.pi - 2 + 2 * np.log(2)) / 12
    f2.b = 1

    def f3(self, t):
        return np.exp(t) * np.cos(t)

    f3.ref = (np.exp(np.pi / 2) - 1) / 2
    f3.b = np.pi / 2

    def f4(self, t):
        a = np.sqrt(2 + t ** 2)
        return np.arctan(a) / ((1 + t ** 2) * a)

    f4.ref = 5 * np.pi ** 2 / 96
    f4.b = 1

    def f5(self, t):
        return np.sqrt(t) * np.log(t)

    f5.ref = -4 / 9
    f5.b = 1

    def f6(self, t):
        return np.sqrt(1 - t ** 2)

    f6.ref = np.pi / 4
    f6.b = 1

    def f7(self, t):
        return np.sqrt(t) / np.sqrt(1 - t ** 2)

    f7.ref = 2 * np.sqrt(np.pi) * special.gamma(3 / 4) / special.gamma(1 / 4)
    f7.b = 1

    def f8(self, t):
        return np.log(t) ** 2

    f8.ref = 2
    f8.b = 1

    def f9(self, t):
        return np.log(np.cos(t))

    f9.ref = -np.pi * np.log(2) / 2
    f9.b = np.pi / 2

    def f10(self, t):
        return np.sqrt(np.tan(t))

    f10.ref = np.pi * np.sqrt(2) / 2
    f10.b = np.pi / 2

    def f11(self, t):
        return 1 / (1 + t ** 2)

    f11.ref = np.pi / 2
    f11.b = np.inf

    def f12(self, t):
        return np.exp(-t) / np.sqrt(t)

    f12.ref = np.sqrt(np.pi)
    f12.b = np.inf

    def f13(self, t):
        return np.exp(-t ** 2 / 2)

    f13.ref = np.sqrt(np.pi / 2)
    f13.b = np.inf

    def f14(self, t):
        return np.exp(-t) * np.cos(t)

    f14.ref = 0.5
    f14.b = np.inf

    def f15(self, t):
        return np.sin(t) / t

    f15.ref = np.pi / 2
    f15.b = np.inf

    def error(self, res, ref, log=False, xp=None):
        xp = array_namespace(res, ref) if xp is None else xp
        err = abs(res - ref)

        if not log:
            return err

        with np.errstate(divide='ignore'):
            return xp.log10(err)

    def test_input_validation(self, xp):
        f = self.f1

        zero = xp.asarray(0)
        f_b = xp.asarray(f.b)

        message = '`f` must be callable.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(42, zero, f_b)

        message = '...must be True or False.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, log=2)

        message = '...must be real numbers.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, xp.asarray(1+1j), f_b)
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, atol='ekki')
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, rtol=pytest)

        message = '...must be non-negative and finite.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, rtol=-1)
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, atol=xp.inf)

        message = '...may not be positive infinity.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, rtol=xp.inf, log=True)
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, atol=xp.inf, log=True)

        message = '...must be integers.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, maxlevel=object())
        # with pytest.raises(ValueError, match=message):  # unused for now
        #     _tanhsinh(f, zero, f_b, maxfun=1+1j)
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, minlevel="migratory coconut")

        message = '...must be non-negative.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, maxlevel=-1)
        # with pytest.raises(ValueError, match=message):  # unused for now
        #     _tanhsinh(f, zero, f_b, maxfun=-1)
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, minlevel=-1)

        message = '...must be True or False.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, preserve_shape=2)

        message = '...must be callable.'
        with pytest.raises(ValueError, match=message):
            _tanhsinh(f, zero, f_b, callback='elderberry')

    @pytest.mark.parametrize("limits, ref", [
        [(0, math.inf), 0.5],  # b infinite
        [(-math.inf, 0), 0.5],  # a infinite
        [(-math.inf, math.inf), 1.],  # a and b infinite
        [(math.inf, -math.inf), -1.],  # flipped limits
        [(1, -1), stats.norm.cdf(-1.) -  stats.norm.cdf(1.)],  # flipped limits
    ])
    def test_integral_transforms(self, limits, ref, xp):
        # Check that the integral transforms are behaving for both normal and
        # log integration
        limits = [xp.asarray(limit) for limit in limits]
        dtype = xp.asarray(float(limits[0])).dtype
        ref = xp.asarray(ref, dtype=dtype)

        res = _tanhsinh(norm_pdf, *limits)
        xp_assert_close(res.integral, ref)

        logres = _tanhsinh(norm_logpdf, *limits, log=True)
        xp_assert_close(xp.exp(logres.integral), ref, check_dtype=False)
        # Transformation should not make the result complex unnecessarily
        xp_test = array_namespace(*limits)  # we need xp.isdtype
        assert (xp_test.isdtype(logres.integral.dtype, "real floating") if ref > 0
                else xp_test.isdtype(logres.integral.dtype, "complex floating"))

        xp_assert_close(xp.exp(logres.error), res.error, atol=1e-16, check_dtype=False)

    # 15 skipped intentionally; it's very difficult numerically
    @pytest.mark.skip_xp_backends(np_only=True,
                                  reason='Cumbersome to convert everything.')
    @pytest.mark.parametrize('f_number', range(1, 15))
    def test_basic(self, f_number, xp):
        f = getattr(self, f"f{f_number}")
        rtol = 2e-8
        res = _tanhsinh(f, 0, f.b, rtol=rtol)
        assert_allclose(res.integral, f.ref, rtol=rtol)
        if f_number not in {14}:  # mildly underestimates error here
            true_error = abs(self.error(res.integral, f.ref)/res.integral)
            assert true_error < res.error

        if f_number in {7, 10, 12}:  # succeeds, but doesn't know it
            return

        assert res.success
        assert res.status == 0

    @pytest.mark.skip_xp_backends(np_only=True,
                                  reason="Distributions aren't xp-compatible.")
    @pytest.mark.parametrize('ref', (0.5, [0.4, 0.6]))
    @pytest.mark.parametrize('case', stats._distr_params.distcont)
    def test_accuracy(self, ref, case, xp):
        distname, params = case
        if distname in {'dgamma', 'dweibull', 'laplace', 'kstwo'}:
            # should split up interval at first-derivative discontinuity
            pytest.skip('tanh-sinh is not great for non-smooth integrands')
        if (distname in {'studentized_range', 'levy_stable'}
                and not int(os.getenv('SCIPY_XSLOW', 0))):
            pytest.skip('This case passes, but it is too slow.')
        dist = getattr(stats, distname)(*params)
        x = dist.interval(ref)
        res = _tanhsinh(dist.pdf, *x)
        assert_allclose(res.integral, ref)

    @pytest.mark.parametrize('shape', [tuple(), (12,), (3, 4), (3, 2, 2)])
    def test_vectorization(self, shape, xp):
        # Test for correct functionality, output shapes, and dtypes for various
        # input shapes.
        rng = np.random.default_rng(82456839535679456794)
        a = xp.asarray(rng.random(shape))
        b = xp.asarray(rng.random(shape))
        p = xp.asarray(rng.random(shape))
        n = math.prod(shape)

        def f(x, p):
            f.ncall += 1
            f.feval += 1 if (xp_size(x) == n or x.ndim <= 1) else x.shape[-1]
            return x**p
        f.ncall = 0
        f.feval = 0

        @_vectorize(xp)
        def _tanhsinh_single(a, b, p):
            return _tanhsinh(lambda x: x**p, a, b)

        res = _tanhsinh(f, a, b, args=(p,))
        refs = _tanhsinh_single(a, b, p)

        xp_test = array_namespace(a)  # need xp.stack, isdtype
        attrs = ['integral', 'error', 'success', 'status', 'nfev', 'maxlevel']
        for attr in attrs:
            ref_attr = xp_test.stack([getattr(ref, attr) for ref in refs])
            res_attr = xp_ravel(getattr(res, attr))
            xp_assert_close(res_attr, ref_attr, rtol=1e-15)
            assert getattr(res, attr).shape == shape

        assert xp_test.isdtype(res.success.dtype, 'bool')
        assert xp_test.isdtype(res.status.dtype, 'integral')
        assert xp_test.isdtype(res.nfev.dtype, 'integral')
        assert xp_test.isdtype(res.maxlevel.dtype, 'integral')
        assert xp.max(res.nfev) == f.feval
        # maxlevel = 2 -> 3 function calls (2 initialization, 1 work)
        assert xp.max(res.maxlevel) >= 2
        assert xp.max(res.maxlevel) == f.ncall

    def test_flags(self, xp):
        # Test cases that should produce different status flags; show that all
        # can be produced simultaneously.
        def f(xs, js):
            f.nit += 1
            funcs = [lambda x: xp.exp(-x**2),  # converges
                     lambda x: xp.exp(x),  # reaches maxiter due to order=2
                     lambda x: xp.full_like(x, xp.nan)]  # stops due to NaN
            res = []
            for i in range(xp_size(js)):
                x = xs[i, ...]
                j = int(xp_ravel(js)[i])
                res.append(funcs[j](x))
            return xp.stack(res)
        f.nit = 0

        args = (xp.arange(3, dtype=xp.int64),)
        a = xp.asarray([xp.inf]*3)
        b = xp.asarray([-xp.inf] * 3)
        res = _tanhsinh(f, a, b, maxlevel=5, args=args)
        ref_flags = xp.asarray([0, -2, -3], dtype=xp.int32)
        xp_assert_equal(res.status, ref_flags)

    def test_flags_preserve_shape(self, xp):
        # Same test as above but using `preserve_shape` option to simplify.
        def f(x):
            res = [xp.exp(-x[0]**2),  # converges
                   xp.exp(x[1]),  # reaches maxiter due to order=2
                   xp.full_like(x[2], xp.nan)]  # stops due to NaN
            return xp.stack(res)

        a = xp.asarray([xp.inf] * 3)
        b = xp.asarray([-xp.inf] * 3)
        res = _tanhsinh(f, a, b, maxlevel=5, preserve_shape=True)
        ref_flags = xp.asarray([0, -2, -3], dtype=xp.int32)
        xp_assert_equal(res.status, ref_flags)

    def test_preserve_shape(self, xp):
        # Test `preserve_shape` option
        def f(x, xp):
            return xp.stack([xp.stack([x, xp.sin(10 * x)]),
                             xp.stack([xp.cos(30 * x), x * xp.sin(100 * x)])])

        ref = quad_vec(lambda x: f(x, np), 0, 1)
        res = _tanhsinh(lambda x: f(x, xp), xp.asarray(0), xp.asarray(1),
                        preserve_shape=True)
        dtype = xp.asarray(0.).dtype
        xp_assert_close(res.integral, xp.asarray(ref[0], dtype=dtype))

    def test_convergence(self, xp):
        # demonstrate that number of accurate digits doubles each iteration
        dtype = xp.float64  # this only works with good precision
        def f(t):
            return t * xp.log(1 + t)
        ref = xp.asarray(0.25, dtype=dtype)
        a, b = xp.asarray(0., dtype=dtype), xp.asarray(1., dtype=dtype)

        last_logerr = 0
        for i in range(4):
            res = _tanhsinh(f, a, b, minlevel=0, maxlevel=i)
            logerr = self.error(res.integral, ref, log=True, xp=xp)
            assert (logerr < last_logerr * 2 or logerr < -15.5)
            last_logerr = logerr

    def test_options_and_result_attributes(self, xp):
        # demonstrate that options are behaving as advertised and status
        # messages are as intended
        xp_test = array_namespace(xp.asarray(1.))  # need xp.atan

        def f(x):
            f.calls += 1
            f.feval += xp_size(xp.asarray(x))
            return x**2 * xp_test.atan(x)

        f.ref = xp.asarray((math.pi - 2 + 2 * math.log(2)) / 12, dtype=xp.float64)

        default_rtol = 1e-12
        default_atol = f.ref * default_rtol  # effective default absolute tol

        # Keep things simpler by leaving tolerances fixed rather than
        # having to make them dtype-dependent
        a = xp.asarray(0., dtype=xp.float64)
        b = xp.asarray(1., dtype=xp.float64)

        # Test default options
        f.feval, f.calls = 0, 0
        ref = _tanhsinh(f, a, b)
        assert self.error(ref.integral, f.ref) < ref.error < default_atol
        assert ref.nfev == f.feval
        ref.calls = f.calls  # reference number of function calls
        assert ref.success
        assert ref.status == 0

        # Test `maxlevel` equal to required max level
        # We should get all the same results
        f.feval, f.calls = 0, 0
        maxlevel = int(ref.maxlevel)
        res = _tanhsinh(f, a, b, maxlevel=maxlevel)
        res.calls = f.calls
        assert res == ref

        # Now reduce the maximum level. We won't meet tolerances.
        f.feval, f.calls = 0, 0
        maxlevel -= 1
        assert maxlevel >= 2  # can't compare errors otherwise
        res = _tanhsinh(f, a, b, maxlevel=maxlevel)
        assert self.error(res.integral, f.ref) < res.error > default_atol
        assert res.nfev == f.feval < ref.nfev
        assert f.calls == ref.calls - 1
        assert not res.success
        assert res.status == eim._ECONVERR

        # `maxfun` is currently not enforced

        # # Test `maxfun` equal to required number of function evaluations
        # # We should get all the same results
        # f.feval, f.calls = 0, 0
        # maxfun = ref.nfev
        # res = _tanhsinh(f, 0, f.b, maxfun = maxfun)
        # assert res == ref
        #
        # # Now reduce `maxfun`. We won't meet tolerances.
        # f.feval, f.calls = 0, 0
        # maxfun -= 1
        # res = _tanhsinh(f, 0, f.b, maxfun=maxfun)
        # assert self.error(res.integral, f.ref) < res.error > default_atol
        # assert res.nfev == f.feval < ref.nfev
        # assert f.calls == ref.calls - 1
        # assert not res.success
        # assert res.status == 2

        # Take this result to be the new reference
        ref = res
        ref.calls = f.calls

        # Test `atol`
        f.feval, f.calls = 0, 0
        # With this tolerance, we should get the exact same result as ref
        atol = np.nextafter(float(ref.error), np.inf)
        res = _tanhsinh(f, a, b, rtol=0, atol=atol)
        assert res.integral == ref.integral
        assert res.error == ref.error
        assert res.nfev == f.feval == ref.nfev
        assert f.calls == ref.calls
        # Except the result is considered to be successful
        assert res.success
        assert res.status == 0

        f.feval, f.calls = 0, 0
        # With a tighter tolerance, we should get a more accurate result
        atol = np.nextafter(float(ref.error), -np.inf)
        res = _tanhsinh(f, a, b, rtol=0, atol=atol)
        assert self.error(res.integral, f.ref) < res.error < atol
        assert res.nfev == f.feval > ref.nfev
        assert f.calls > ref.calls
        assert res.success
        assert res.status == 0

        # Test `rtol`
        f.feval, f.calls = 0, 0
        # With this tolerance, we should get the exact same result as ref
        rtol = np.nextafter(float(ref.error/ref.integral), np.inf)
        res = _tanhsinh(f, a, b, rtol=rtol)
        assert res.integral == ref.integral
        assert res.error == ref.error
        assert res.nfev == f.feval == ref.nfev
        assert f.calls == ref.calls
        # Except the result is considered to be successful
        assert res.success
        assert res.status == 0

        f.feval, f.calls = 0, 0
        # With a tighter tolerance, we should get a more accurate result
        rtol = np.nextafter(float(ref.error/ref.integral), -np.inf)
        res = _tanhsinh(f, a, b, rtol=rtol)
        assert self.error(res.integral, f.ref)/f.ref < res.error/res.integral < rtol
        assert res.nfev == f.feval > ref.nfev
        assert f.calls > ref.calls
        assert res.success
        assert res.status == 0

    @pytest.mark.skip_xp_backends('torch', reason=
            'https://github.com/scipy/scipy/pull/21149#issuecomment-2330477359',
    )
    @pytest.mark.parametrize('rtol', [1e-4, 1e-14])
    def test_log(self, rtol, xp):
        # Test equivalence of log-integration and regular integration
        test_tols = dict(atol=1e-18, rtol=1e-15)

        # Positive integrand (real log-integrand)
        a = xp.asarray(-1., dtype=xp.float64)
        b = xp.asarray(2., dtype=xp.float64)
        res = _tanhsinh(norm_logpdf, a, b, log=True, rtol=math.log(rtol))
        ref = _tanhsinh(norm_pdf, a, b, rtol=rtol)
        xp_assert_close(xp.exp(res.integral), ref.integral, **test_tols)
        xp_assert_close(xp.exp(res.error), ref.error, **test_tols)
        assert res.nfev == ref.nfev

        # Real integrand (complex log-integrand)
        def f(x):
            return -norm_logpdf(x)*norm_pdf(x)

        def logf(x):
            return xp.log(norm_logpdf(x) + 0j) + norm_logpdf(x) + xp.pi * 1j

        a = xp.asarray(-xp.inf, dtype=xp.float64)
        b = xp.asarray(xp.inf, dtype=xp.float64)
        res = _tanhsinh(logf, a, b, log=True)
        ref = _tanhsinh(f, a, b)
        # In gh-19173, we saw `invalid` warnings on one CI platform.
        # Silencing `all` because I can't reproduce locally and don't want
        # to risk the need to run CI again.
        with np.errstate(all='ignore'):
            xp_assert_close(xp.exp(res.integral), ref.integral, **test_tols,
                            check_dtype=False)
            xp_assert_close(xp.exp(res.error), ref.error, **test_tols,
                            check_dtype=False)
        assert res.nfev == ref.nfev

    def test_complex(self, xp):
        # Test integration of complex integrand
        # Finite limits
        def f(x):
            return xp.exp(1j * x)

        a, b = xp.asarray(0.), xp.asarray(xp.pi/4)
        res = _tanhsinh(f, a, b)
        ref = math.sqrt(2)/2 + (1-math.sqrt(2)/2)*1j
        xp_assert_close(res.integral, xp.asarray(ref))

        # Infinite limits
        def f(x):
            return norm_pdf(x) + 1j/2*norm_pdf(x/2)

        a, b = xp.asarray(xp.inf), xp.asarray(-xp.inf)
        res = _tanhsinh(f, a, b)
        xp_assert_close(res.integral, xp.asarray(-(1+1j)))

    @pytest.mark.parametrize("maxlevel", range(4))
    def test_minlevel(self, maxlevel, xp):
        # Verify that minlevel does not change the values at which the
        # integrand is evaluated or the integral/error estimates, only the
        # number of function calls

        # need `xp.concat`, `xp.atan`, and `xp.sort`
        xp_test = array_namespace(xp.asarray(1.))

        def f(x):
            f.calls += 1
            f.feval += xp_size(xp.asarray(x))
            f.x = xp_test.concat((f.x, xp_ravel(x)))
            return x**2 * xp_test.atan(x)

        f.feval, f.calls, f.x = 0, 0, xp.asarray([])

        a = xp.asarray(0, dtype=xp.float64)
        b = xp.asarray(1, dtype=xp.float64)
        ref = _tanhsinh(f, a, b, minlevel=0, maxlevel=maxlevel)
        ref_x = xp_test.sort(f.x)

        for minlevel in range(0, maxlevel + 1):
            f.feval, f.calls, f.x = 0, 0, xp.asarray([])
            options = dict(minlevel=minlevel, maxlevel=maxlevel)
            res = _tanhsinh(f, a, b, **options)
            # Should be very close; all that has changed is the order of values
            xp_assert_close(res.integral, ref.integral, rtol=4e-16)
            # Difference in absolute errors << magnitude of integral
            xp_assert_close(res.error, ref.error, atol=4e-16 * ref.integral)
            assert res.nfev == f.feval == f.x.shape[0]
            assert f.calls == maxlevel - minlevel + 1 + 1  # 1 validation call
            assert res.status == ref.status
            xp_assert_equal(ref_x, xp_test.sort(f.x))

    def test_improper_integrals(self, xp):
        # Test handling of infinite limits of integration (mixed with finite limits)
        def f(x):
            x[xp.isinf(x)] = xp.nan
            return xp.exp(-x**2)
        a = xp.asarray([-xp.inf, 0, -xp.inf, xp.inf, -20, -xp.inf, -20])
        b = xp.asarray([xp.inf, xp.inf, 0, -xp.inf, 20, 20, xp.inf])
        ref = math.sqrt(math.pi)
        ref = xp.asarray([ref, ref/2, ref/2, -ref, ref, ref, ref])
        res = _tanhsinh(f, a, b)
        xp_assert_close(res.integral, ref)

    @pytest.mark.parametrize("limits", ((0, 3), ([-math.inf, 0], [3, 3])))
    @pytest.mark.parametrize("dtype", ('float32', 'float64'))
    def test_dtype(self, limits, dtype, xp):
        # Test that dtypes are preserved
        dtype = getattr(xp, dtype)
        a, b = xp.asarray(limits, dtype=dtype)

        def f(x):
            assert x.dtype == dtype
            return xp.exp(x)

        rtol = 1e-12 if dtype == xp.float64 else 1e-5
        res = _tanhsinh(f, a, b, rtol=rtol)
        assert res.integral.dtype == dtype
        assert res.error.dtype == dtype
        assert xp.all(res.success)
        xp_assert_close(res.integral, xp.exp(b)-xp.exp(a))

    def test_maxiter_callback(self, xp):
        # Test behavior of `maxiter` parameter and `callback` interface
        a, b = xp.asarray(-xp.inf), xp.asarray(xp.inf)
        def f(x):
            return xp.exp(-x*x)

        minlevel, maxlevel = 0, 2
        maxiter = maxlevel - minlevel + 1
        kwargs = dict(minlevel=minlevel, maxlevel=maxlevel, rtol=1e-15)
        res = _tanhsinh(f, a, b, **kwargs)
        assert not res.success
        assert res.maxlevel == maxlevel

        def callback(res):
            callback.iter += 1
            callback.res = res
            assert hasattr(res, 'integral')
            assert res.status == 1
            if callback.iter == maxiter:
                raise StopIteration
        callback.iter = -1  # callback called once before first iteration
        callback.res = None

        del kwargs['maxlevel']
        res2 = _tanhsinh(f, a, b, **kwargs, callback=callback)
        # terminating with callback is identical to terminating due to maxiter
        # (except for `status`)
        for key in res.keys():
            if key == 'status':
                assert res[key] == -2
                assert res2[key] == -4
            else:
                assert res2[key] == callback.res[key] == res[key]

    def test_jumpstart(self, xp):
        # The intermediate results at each level i should be the same as the
        # final results when jumpstarting at level i; i.e. minlevel=maxlevel=i
        a = xp.asarray(-xp.inf, dtype=xp.float64)
        b = xp.asarray(xp.inf, dtype=xp.float64)

        def f(x):
            return xp.exp(-x*x)

        def callback(res):
            callback.integrals.append(xp_copy(res.integral)[()])
            callback.errors.append(xp_copy(res.error)[()])
        callback.integrals = []
        callback.errors = []

        maxlevel = 4
        _tanhsinh(f, a, b, minlevel=0, maxlevel=maxlevel, callback=callback)

        for i in range(maxlevel + 1):
            res = _tanhsinh(f, a, b, minlevel=i, maxlevel=i)
            xp_assert_close(callback.integrals[1+i], res.integral, rtol=1e-15)
            xp_assert_close(callback.errors[1+i], res.error, rtol=1e-15, atol=1e-16)

    def test_special_cases(self, xp):
        # Test edge cases and other special cases
        a, b = xp.asarray(0), xp.asarray(1)
        xp_test = array_namespace(a, b)  # need `xp.isdtype`

        def f(x):
            assert xp_test.isdtype(x.dtype, "real floating")
            return x

        res = _tanhsinh(f, a, b)
        assert res.success
        xp_assert_close(res.integral, xp.asarray(0.5))

        # Test levels 0 and 1; error is NaN
        res = _tanhsinh(f, a, b, maxlevel=0)
        assert res.integral > 0
        xp_assert_equal(res.error, xp.asarray(xp.nan))
        res = _tanhsinh(f, a, b, maxlevel=1)
        assert res.integral > 0
        xp_assert_equal(res.error, xp.asarray(xp.nan))

        # Test equal left and right integration limits
        res = _tanhsinh(f, b, b)
        assert res.success
        assert res.maxlevel == -1
        xp_assert_close(res.integral, xp.asarray(0.))

        # Test scalar `args` (not in tuple)
        def f(x, c):
            return x**c

        res = _tanhsinh(f, a, b, args=29)
        xp_assert_close(res.integral, xp.asarray(1/30))

        # Test NaNs
        a = xp.asarray([xp.nan, 0, 0, 0])
        b = xp.asarray([1, xp.nan, 1, 1])
        c = xp.asarray([1, 1, xp.nan, 1])
        res = _tanhsinh(f, a, b, args=(c,))
        xp_assert_close(res.integral, xp.asarray([xp.nan, xp.nan, xp.nan, 0.5]))
        xp_assert_equal(res.error[:3], xp.full((3,), xp.nan))
        xp_assert_equal(res.status, xp.asarray([-3, -3, -3, 0], dtype=xp.int32))
        xp_assert_equal(res.success, xp.asarray([False, False, False, True]))
        xp_assert_equal(res.nfev[:3], xp.full((3,), 1, dtype=xp.int32))

        # Test complex integral followed by real integral
        # Previously, h0 was of the result dtype. If the `dtype` were complex,
        # this could lead to complex cached abscissae/weights. If these get
        # cast to real dtype for a subsequent real integral, we would get a
        # ComplexWarning. Check that this is avoided.
        _pair_cache.xjc = xp.empty(0)
        _pair_cache.wj = xp.empty(0)
        _pair_cache.indices = [0]
        _pair_cache.h0 = None
        a, b = xp.asarray(0), xp.asarray(1)
        res = _tanhsinh(lambda x: xp.asarray(x*1j), a, b)
        xp_assert_close(res.integral, xp.asarray(0.5*1j))
        res = _tanhsinh(lambda x: x, a, b)
        xp_assert_close(res.integral, xp.asarray(0.5))

        # Test zero-size
        shape = (0, 3)
        res = _tanhsinh(lambda x: x, xp.asarray(0), xp.zeros(shape))
        attrs = ['integral', 'error', 'success', 'status', 'nfev', 'maxlevel']
        for attr in attrs:
            assert res[attr].shape == shape

    @pytest.mark.skip_xp_backends(np_only=True)
    def test_compress_nodes_weights_gh21496(self, xp):
        # See discussion in:
        # https://github.com/scipy/scipy/pull/21496#discussion_r1878681049
        # This would cause "ValueError: attempt to get argmax of an empty sequence"
        # Check that this has been resolved.
        x = np.full(65, 3)
        x[-1] = 1000
        _tanhsinh(np.sin, 1, x)


@array_api_compatible
@pytest.mark.usefixtures("skip_xp_backends")
@pytest.mark.skip_xp_backends('array_api_strict', reason='No fancy indexing.')
@pytest.mark.skip_xp_backends('jax.numpy', reason='No mutation.')
class TestNSum:
    rng = np.random.default_rng(5895448232066142650)
    p = rng.uniform(1, 10, size=10).tolist()

    def f1(self, k):
        # Integers are never passed to `f1`; if they were, we'd get
        # integer to negative integer power error
        return k**(-2)

    f1.ref = np.pi**2/6
    f1.a = 1
    f1.b = np.inf
    f1.args = tuple()

    def f2(self, k, p):
        return 1 / k**p

    f2.ref = special.zeta(p, 1)
    f2.a = 1.
    f2.b = np.inf
    f2.args = (p,)

    def f3(self, k, p):
        return 1 / k**p

    f3.a = 1
    f3.b = rng.integers(5, 15, size=(3, 1))
    f3.ref = _gen_harmonic_gt1(f3.b, p)
    f3.args = (p,)

    def test_input_validation(self, xp):
        f = self.f1
        a, b = xp.asarray(f.a), xp.asarray(f.b)

        message = '`f` must be callable.'
        with pytest.raises(ValueError, match=message):
            nsum(42, a, b)

        message = '...must be True or False.'
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, log=2)

        message = '...must be real numbers.'
        with pytest.raises(ValueError, match=message):
            nsum(f, xp.asarray(1+1j), b)
        with pytest.raises(ValueError, match=message):
            nsum(f, a, xp.asarray(1+1j))
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, step=xp.asarray(1+1j))
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(atol='ekki'))
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(rtol=pytest))

        with np.errstate(all='ignore'):
            res = nsum(f, xp.asarray([np.nan, np.inf]), xp.asarray(1.))
            assert xp.all((res.status == -1) & xp.isnan(res.sum)
                          & xp.isnan(res.error) & ~res.success & res.nfev == 1)
            res = nsum(f, xp.asarray(10.), xp.asarray([np.nan, 1]))
            assert xp.all((res.status == -1) & xp.isnan(res.sum)
                          & xp.isnan(res.error) & ~res.success & res.nfev == 1)
            res = nsum(f, xp.asarray(1.), xp.asarray(10.),
                       step=xp.asarray([xp.nan, -xp.inf, xp.inf, -1, 0]))
            assert xp.all((res.status == -1) & xp.isnan(res.sum)
                          & xp.isnan(res.error) & ~res.success & res.nfev == 1)

        message = '...must be non-negative and finite.'
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(rtol=-1))
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(atol=np.inf))

        message = '...may not be positive infinity.'
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(rtol=np.inf), log=True)
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, tolerances=dict(atol=np.inf), log=True)

        message = '...must be a non-negative integer.'
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, maxterms=3.5)
        with pytest.raises(ValueError, match=message):
            nsum(f, a, b, maxterms=-2)

    @pytest.mark.parametrize('f_number', range(1, 4))
    def test_basic(self, f_number, xp):
        dtype = xp.asarray(1.).dtype
        f = getattr(self, f"f{f_number}")
        a, b = xp.asarray(f.a), xp.asarray(f.b),
        args = tuple(xp.asarray(arg) for arg in f.args)
        ref = xp.asarray(f.ref, dtype=dtype)
        res = nsum(f, a, b, args=args)
        xp_assert_close(res.sum, ref)
        xp_assert_equal(res.status, xp.zeros(ref.shape, dtype=xp.int32))
        xp_test = array_namespace(a)  # CuPy doesn't have `bool`
        xp_assert_equal(res.success, xp.ones(ref.shape, dtype=xp_test.bool))

        with np.errstate(divide='ignore'):
            logres = nsum(lambda *args: xp.log(f(*args)),
                           a, b, log=True, args=args)
        xp_assert_close(xp.exp(logres.sum), res.sum)
        xp_assert_close(xp.exp(logres.error), res.error, atol=1e-15)
        xp_assert_equal(logres.status, res.status)
        xp_assert_equal(logres.success, res.success)

    @pytest.mark.parametrize('maxterms', [0, 1, 10, 20, 100])
    def test_integral(self, maxterms, xp):
        # test precise behavior of integral approximation
        f = self.f1

        def logf(x):
            return -2*xp.log(x)

        def F(x):
            return -1 / x

        a = xp.asarray([1, 5], dtype=xp.float64)[:, xp.newaxis]
        b = xp.asarray([20, 100, xp.inf], dtype=xp.float64)[:, xp.newaxis, xp.newaxis]
        step = xp.asarray([0.5, 1, 2], dtype=xp.float64).reshape((-1, 1, 1, 1))
        nsteps = xp.floor((b - a)/step)
        b_original = b
        b = a + nsteps*step

        k = a + maxterms*step
        # partial sum
        direct = xp.sum(f(a + xp.arange(maxterms)*step), axis=-1, keepdims=True)
        integral = (F(b) - F(k))/step  # integral approximation of remainder
        low = direct + integral + f(b)  # theoretical lower bound
        high = direct + integral + f(k)  # theoretical upper bound
        ref_sum = (low + high)/2  # nsum uses average of the two
        ref_err = (high - low)/2  # error (assuming perfect quadrature)

        # correct reference values where number of terms < maxterms
        xp_test = array_namespace(a)  # torch needs broadcast_arrays
        a, b, step = xp_test.broadcast_arrays(a, b, step)
        for i in np.ndindex(a.shape):
            ai, bi, stepi = float(a[i]), float(b[i]), float(step[i])
            if (bi - ai)/stepi + 1 <= maxterms:
                direct = xp.sum(f(xp.arange(ai, bi+stepi, stepi, dtype=xp.float64)))
                ref_sum[i] = direct
                ref_err[i] = direct * xp.finfo(direct.dtype).eps

        rtol = 1e-12
        res = nsum(f, a, b_original, step=step, maxterms=maxterms,
                   tolerances=dict(rtol=rtol))
        xp_assert_close(res.sum, ref_sum, rtol=10*rtol)
        xp_assert_close(res.error, ref_err, rtol=100*rtol)

        i = ((b_original - a)/step + 1 <= maxterms)
        xp_assert_close(res.sum[i], ref_sum[i], rtol=1e-15)
        xp_assert_close(res.error[i], ref_err[i], rtol=1e-15)

        logres = nsum(logf, a, b_original, step=step, log=True,
                      tolerances=dict(rtol=math.log(rtol)), maxterms=maxterms)
        xp_assert_close(xp.exp(logres.sum), res.sum)
        xp_assert_close(xp.exp(logres.error), res.error)

    @pytest.mark.parametrize('shape', [tuple(), (12,), (3, 4), (3, 2, 2)])
    def test_vectorization(self, shape, xp):
        # Test for correct functionality, output shapes, and dtypes for various
        # input shapes.
        rng = np.random.default_rng(82456839535679456794)
        a = rng.integers(1, 10, size=shape)
        # when the sum can be computed directly or `maxterms` is large enough
        # to meet `atol`, there are slight differences (for good reason)
        # between vectorized call and looping.
        b = np.inf
        p = rng.random(shape) + 1
        n = math.prod(shape)

        def f(x, p):
            f.feval += 1 if (x.size == n or x.ndim <= 1) else x.shape[-1]
            return 1 / x ** p

        f.feval = 0

        @np.vectorize
        def nsum_single(a, b, p, maxterms):
            return nsum(lambda x: 1 / x**p, a, b, maxterms=maxterms)

        res = nsum(f, xp.asarray(a), xp.asarray(b), maxterms=1000,
                   args=(xp.asarray(p),))
        refs = nsum_single(a, b, p, maxterms=1000).ravel()

        attrs = ['sum', 'error', 'success', 'status', 'nfev']
        for attr in attrs:
            ref_attr = [xp.asarray(getattr(ref, attr)) for ref in refs]
            res_attr = getattr(res, attr)
            xp_assert_close(xp_ravel(res_attr), xp.asarray(ref_attr), rtol=1e-15)
            assert res_attr.shape == shape

        xp_test = array_namespace(xp.asarray(1.))
        assert xp_test.isdtype(res.success.dtype, 'bool')
        assert xp_test.isdtype(res.status.dtype, 'integral')
        assert xp_test.isdtype(res.nfev.dtype, 'integral')
        if is_numpy(xp):  # other libraries might have different number
            assert int(xp.max(res.nfev)) == f.feval

    def test_status(self, xp):
        f = self.f2

        p = [2, 2, 0.9, 1.1, 2, 2]
        a = xp.asarray([0, 0, 1, 1, 1, np.nan], dtype=xp.float64)
        b = xp.asarray([10, np.inf, np.inf, np.inf, np.inf, np.inf], dtype=xp.float64)
        ref = special.zeta(p, 1)
        p = xp.asarray(p, dtype=xp.float64)

        with np.errstate(divide='ignore'):  # intentionally dividing by zero
            res = nsum(f, a, b, args=(p,))

        ref_success = xp.asarray([False, False, False, False, True, False])
        ref_status = xp.asarray([-3, -3, -2, -4, 0, -1], dtype=xp.int32)
        xp_assert_equal(res.success, ref_success)
        xp_assert_equal(res.status, ref_status)
        xp_assert_close(res.sum[res.success], xp.asarray(ref)[res.success])

    def test_nfev(self, xp):
        def f(x):
            f.nfev += xp_size(x)
            return 1 / x**2

        f.nfev = 0
        res = nsum(f, xp.asarray(1), xp.asarray(10))
        assert res.nfev == f.nfev

        f.nfev = 0
        res = nsum(f, xp.asarray(1), xp.asarray(xp.inf), tolerances=dict(atol=1e-6))
        assert res.nfev == f.nfev

    def test_inclusive(self, xp):
        # There was an edge case off-by one bug when `_direct` was called with
        # `inclusive=True`. Check that this is resolved.
        a = xp.asarray([1, 4])
        b = xp.asarray(xp.inf)
        res = nsum(lambda k: 1 / k ** 2, a, b,
                   maxterms=500, tolerances=dict(atol=0.1))
        ref = nsum(lambda k: 1 / k ** 2, a, b)
        assert xp.all(res.sum > (ref.sum - res.error))
        assert xp.all(res.sum < (ref.sum + res.error))

    @pytest.mark.parametrize('log', [True, False])
    def test_infinite_bounds(self, log, xp):
        a = xp.asarray([1, -np.inf, -np.inf])
        b = xp.asarray([np.inf, -1, np.inf])
        c = xp.asarray([1, 2, 3])

        def f(x, a):
            return (xp.log(xp.tanh(a / 2)) - a*xp.abs(x) if log
                    else xp.tanh(a/2) * xp.exp(-a*xp.abs(x)))

        res = nsum(f, a, b, args=(c,), log=log)
        ref = xp.asarray([stats.dlaplace.sf(0, 1), stats.dlaplace.sf(0, 2), 1])
        ref = xp.log(ref) if log else ref
        atol = (1e-10 if a.dtype==xp.float64 else 1e-5) if log else 0
        xp_assert_close(res.sum, xp.asarray(ref, dtype=a.dtype), atol=atol)

        # # Make sure the sign of `x` passed into `f` is correct.
        def f(x, c):
            return -3*xp.log(c*x) if log else 1 / (c*x)**3

        a = xp.asarray([1, -np.inf])
        b = xp.asarray([np.inf, -1])
        arg = xp.asarray([1, -1])
        res = nsum(f, a, b, args=(arg,), log=log)
        ref = np.log(special.zeta(3)) if log else special.zeta(3)
        xp_assert_close(res.sum, xp.full(a.shape, ref, dtype=a.dtype))

    def test_decreasing_check(self, xp):
        # Test accuracy when we start sum on an uphill slope.
        # Without the decreasing check, the terms would look small enough to
        # use the integral approximation. Because the function is not decreasing,
        # the error is not bounded by the magnitude of the last term of the
        # partial sum. In this case, the error would be  ~1e-4, causing the test
        # to fail.
        def f(x):
            return xp.exp(-x ** 2)

        a, b = xp.asarray(-25, dtype=xp.float64), xp.asarray(np.inf, dtype=xp.float64)
        res = nsum(f, a, b)

        # Reference computed with mpmath:
        # from mpmath import mp
        # mp.dps = 50
        # def fmp(x): return mp.exp(-x**2)
        # ref = mp.nsum(fmp, (-25, 0)) + mp.nsum(fmp, (1, mp.inf))
        ref = xp.asarray(1.772637204826652, dtype=xp.float64)

        xp_assert_close(res.sum, ref, rtol=1e-15)

    def test_special_case(self, xp):
        # test equal lower/upper limit
        f = self.f1
        a = b = xp.asarray(2)
        res = nsum(f, a, b)
        xp_assert_equal(res.sum, xp.asarray(f(2)))

        # Test scalar `args` (not in tuple)
        res = nsum(self.f2, xp.asarray(1), xp.asarray(np.inf), args=xp.asarray(2))
        xp_assert_close(res.sum, xp.asarray(self.f1.ref))  # f1.ref is correct w/ args=2

        # Test 0 size input
        a = xp.empty((3, 1, 1))  # arbitrary broadcastable shapes
        b = xp.empty((0, 1))  # could use Hypothesis
        p = xp.empty(4)  # but it's overkill
        shape = np.broadcast_shapes(a.shape, b.shape, p.shape)
        res = nsum(self.f2, a, b, args=(p,))
        assert res.sum.shape == shape
        assert res.status.shape == shape
        assert res.nfev.shape == shape

        # Test maxterms=0
        def f(x):
            with np.errstate(divide='ignore'):
                return 1 / x

        res = nsum(f, xp.asarray(0), xp.asarray(10), maxterms=0)
        assert xp.isnan(res.sum)
        assert xp.isnan(res.error)
        assert res.status == -2

        res = nsum(f, xp.asarray(0), xp.asarray(10), maxterms=1)
        assert xp.isnan(res.sum)
        assert xp.isnan(res.error)
        assert res.status == -3

        # Test NaNs
        # should skip both direct and integral methods if there are NaNs
        a = xp.asarray([xp.nan, 1, 1, 1])
        b = xp.asarray([xp.inf, xp.nan, xp.inf, xp.inf])
        p = xp.asarray([2, 2, xp.nan, 2])
        res = nsum(self.f2, a, b, args=(p,))
        xp_assert_close(res.sum, xp.asarray([xp.nan, xp.nan, xp.nan, self.f1.ref]))
        xp_assert_close(res.error[:3], xp.full((3,), xp.nan))
        xp_assert_equal(res.status, xp.asarray([-1, -1, -3, 0], dtype=xp.int32))
        xp_assert_equal(res.success, xp.asarray([False, False, False, True]))
        # Ideally res.nfev[2] would be 1, but `tanhsinh` has some function evals
        xp_assert_equal(res.nfev[:2], xp.full((2,), 1, dtype=xp.int32))

    @pytest.mark.parametrize('dtype', ['float32', 'float64'])
    def test_dtype(self, dtype, xp):
        dtype = getattr(xp, dtype)

        def f(k):
            assert k.dtype == dtype
            return 1 / k ** xp.asarray(2, dtype=dtype)

        a = xp.asarray(1, dtype=dtype)
        b = xp.asarray([10, xp.inf], dtype=dtype)
        res = nsum(f, a, b)
        assert res.sum.dtype == dtype
        assert res.error.dtype == dtype

        rtol = 1e-12 if dtype == xp.float64 else 1e-6
        ref = _gen_harmonic_gt1(np.asarray([10, xp.inf]), 2)
        xp_assert_close(res.sum, xp.asarray(ref, dtype=dtype), rtol=rtol)

    @pytest.mark.parametrize('case', [(10, 100), (100, 10)])
    def test_nondivisible_interval(self, case, xp):
        # When the limits of the sum are such that (b - a)/step
        # is not exactly integral, check that only floor((b - a)/step)
        # terms are included.
        n, maxterms = case

        def f(k):
            return 1 / k ** 2

        a = np.e
        step = 1 / 3
        b0 = a + n * step
        i = np.arange(-2, 3)
        b = b0 + i * np.spacing(b0)
        ns = np.floor((b - a) / step)
        assert len(set(ns)) == 2

        a, b = xp.asarray(a, dtype=xp.float64), xp.asarray(b, dtype=xp.float64)
        step, ns = xp.asarray(step, dtype=xp.float64), xp.asarray(ns, dtype=xp.float64)
        res = nsum(f, a, b, step=step, maxterms=maxterms)
        xp_assert_equal(xp.diff(ns) > 0, xp.diff(res.sum) > 0)
        xp_assert_close(res.sum[-1], res.sum[0] + f(b0))

    @pytest.mark.skip_xp_backends(np_only=True, reason='Needs beta function.')
    def test_logser_kurtosis_gh20648(self, xp):
        # Some functions return NaN at infinity rather than 0 like they should.
        # Check that this is accounted for.
        ref = stats.yulesimon.moment(4, 5)
        def f(x):
            return stats.yulesimon._pmf(x, 5) * x**4

        with np.errstate(invalid='ignore'):
            assert np.isnan(f(np.inf))

        res = nsum(f, 1, np.inf)
        assert_allclose(res.sum, ref)