image-match/python/py/Lib/site-packages/scipy/stats/tests/test_continued_fraction.py

import math

import pytest
import numpy as np

from scipy._lib._array_api import array_namespace
from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_less, xp_assert_equal
from scipy.stats._continued_fraction import _continued_fraction


@pytest.mark.skip_xp_backends('array_api_strict', reason='No fancy indexing assignment')
@pytest.mark.skip_xp_backends('jax.numpy', reason="Don't support mutation")
# dask doesn't like lines like this
# n = int(xp.real(xp_ravel(n))[0])
# (at some point in here the shape becomes nan)
@pytest.mark.skip_xp_backends('dask.array', reason="dask has issues with the shapes")
class TestContinuedFraction:
    rng = np.random.default_rng(5895448232066142650)
    p = rng.uniform(1, 10, size=10)

    def a1(self, n, x=1.5):
        if n == 0:
            y = 0*x
        elif n == 1:
            y = x
        else:
            y = -x**2
        if np.isscalar(y) and np.__version__ < "2.0":
            y = np.full_like(x, y)  # preserve dtype pre NEP 50
        return y

    def b1(self, n, x=1.5):
        if n == 0:
            y = 0*x
        else:
            one = x/x  # gets array of correct type, dtype, and shape
            y = one * (2*n - 1)
        if np.isscalar(y) and np.__version__ < "2.0":
            y = np.full_like(x, y)  # preserve dtype pre NEP 50
        return y

    def log_a1(self, n, x):
        xp = array_namespace(x)
        if n == 0:
            y = xp.full_like(x, -xp.asarray(math.inf, dtype=x.dtype))
        elif n == 1:
            y = xp.log(x)
        else:
            y = 2 * xp.log(x) + math.pi * 1j
        return y

    def log_b1(self, n, x):
        xp = array_namespace(x)
        if n == 0:
            y = xp.full_like(x, -xp.asarray(math.inf, dtype=x.dtype))
        else:
            one = x - x  # gets array of correct type, dtype, and shape
            y = one + math.log(2 * n - 1)
        return y

    def test_input_validation(self, xp):
        a1 = self.a1
        b1 = self.b1

        message = '`a` and `b` must be callable.'
        with pytest.raises(ValueError, match=message):
            _continued_fraction(1, b1)
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, 1)

        message = r'`eps` and `tiny` must be \(or represent the logarithm of\)...'
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'eps': -10})
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'eps': np.nan})
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'eps': 1+1j}, log=True)
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'tiny': 0})
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'tiny': np.inf})
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, tolerances={'tiny': np.inf}, log=True)
        # this should not raise
        kwargs = dict(args=xp.asarray(1.5+0j), log=True, maxiter=0)
        _continued_fraction(a1, b1, tolerances={'eps': -10}, **kwargs)
        _continued_fraction(a1, b1, tolerances={'tiny': -10}, **kwargs)

        message = '`maxiter` must be a non-negative integer.'
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, maxiter=-1)

        message = '`log` must be boolean.'
        with pytest.raises(ValueError, match=message):
            _continued_fraction(a1, b1, log=2)

    @pytest.mark.parametrize('dtype', ['float32', 'float64', 'complex64', 'complex128'])
    @pytest.mark.parametrize('shape', [(), (1,), (3,), (3, 2)])
    def test_basic(self, shape, dtype, xp):
        np_dtype = getattr(np, dtype)
        xp_dtype = getattr(xp, dtype)
        rng = np.random.default_rng(2435908729190400)

        x = rng.random(shape).astype(np_dtype)
        x = x + rng.random(shape).astype(np_dtype)*1j if dtype.startswith('c') else x
        x = xp.asarray(x, dtype=xp_dtype)

        res = _continued_fraction(self.a1, self.b1, args=(x,))
        ref = xp.tan(x)
        xp_assert_close(res.f, ref)

    @pytest.mark.skip_xp_backends('torch', reason='pytorch/pytorch#136063')
    @pytest.mark.parametrize('dtype', ['float32', 'float64'])
    @pytest.mark.parametrize('shape', [(), (1,), (3,), (3, 2)])
    def test_log(self, shape, dtype, xp):
        if (np.__version__ < "2") and (dtype == 'float32'):
            pytest.skip("Scalar dtypes only respected after NEP 50.")
        np_dtype = getattr(np, dtype)
        rng = np.random.default_rng(2435908729190400)
        x = rng.random(shape).astype(np_dtype)
        x = xp.asarray(x)

        res = _continued_fraction(self.log_a1, self.log_b1, args=(x + 0j,), log=True)
        ref = xp.tan(x)
        xp_assert_close(xp.exp(xp.real(res.f)), ref)

    def test_maxiter(self, xp):
        rng = np.random.default_rng(2435908729190400)
        x = xp.asarray(rng.random(), dtype=xp.float64)
        ref = xp.tan(x)

        res1 = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=3)
        assert res1.nit == 3

        res2 = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=6)
        assert res2.nit == 6

        xp_assert_less(xp.abs(res2.f - ref), xp.abs(res1.f - ref))

    def test_eps(self, xp):
        x = xp.asarray(1.5, dtype=xp.float64)  # x = 1.5 is the default defined above
        ref = xp.tan(x)
        res1 = _continued_fraction(self.a1, self.b1, args=(x,),
                                   tolerances={'eps': 1e-6})
        res2 = _continued_fraction(self.a1, self.b1, args=(x,))
        xp_assert_less(res1.nit, res2.nit)
        xp_assert_less(xp.abs(res2.f - ref), xp.abs(res1.f - ref))

    def test_feval(self, xp):
        def a(n, x):
            a.nfev += 1
            return n * x

        def b(n, x):
            b.nfev += 1
            return n * x

        a.nfev, b.nfev = 0, 0

        res = _continued_fraction(a, b, args=(xp.asarray(1.),))
        assert res.nfev == a.nfev == b.nfev == res.nit + 1

    def test_status(self, xp):
        x = xp.asarray([1, 10, np.nan], dtype=xp.float64)
        res = _continued_fraction(self.a1, self.b1, args=(x,), maxiter=15)
        xp_assert_equal(res.success, xp.asarray([True, False, False]))
        xp_assert_equal(res.status, xp.asarray([0, -2, -3], dtype=xp.int32))

    def test_special_cases(self, xp):
        one = xp.asarray(1)
        res = _continued_fraction(lambda x: one, lambda x: one, maxiter=0)
        xp_assert_close(res.f, xp.asarray(1.))
        assert res.nit == res.nfev - 1 == 0