image-match/python/py/Lib/site-packages/scipy/integrate/tests/test__quad_vec.py

import pytest

import numpy as np
from numpy.testing import assert_allclose

from scipy.integrate import quad_vec
from scipy._lib._array_api import make_xp_test_case

from multiprocessing.dummy import Pool


quadrature_params = pytest.mark.parametrize(
    'quadrature', [None, "gk15", "gk21", "trapezoid"])

def _lorenzian(x):
    return 1 / (1 + x**2)

def _func_with_args(x, a):
    return x * (x + a) * np.arange(3)


@make_xp_test_case(quad_vec)
class TestQuadVec:
    @quadrature_params
    def test_quad_vec_simple(self, quadrature):
        n = np.arange(10)
        def f(x):
            return x ** n
        for epsabs in [0.1, 1e-3, 1e-6]:
            if quadrature == 'trapezoid' and epsabs < 1e-4:
                # slow: skip
                continue

            kwargs = dict(epsabs=epsabs, quadrature=quadrature)

            exact = 2**(n+1)/(n + 1)

            res, err = quad_vec(f, 0, 2, norm='max', **kwargs)
            assert_allclose(res, exact, rtol=0, atol=epsabs)

            res, err = quad_vec(f, 0, 2, norm='2', **kwargs)
            assert np.linalg.norm(res - exact) < epsabs

            res, err = quad_vec(f, 0, 2, norm='max', points=(0.5, 1.0), **kwargs)
            assert_allclose(res, exact, rtol=0, atol=epsabs)

            res, err, *rest = quad_vec(f, 0, 2, norm='max',
                                    epsrel=1e-8,
                                    full_output=True,
                                    limit=10000,
                                    **kwargs)
            assert_allclose(res, exact, rtol=0, atol=epsabs)


    @quadrature_params
    def test_quad_vec_simple_inf(self, quadrature):
        def f(x):
            return 1 / (1 + np.float64(x) ** 2)

        for epsabs in [0.1, 1e-3, 1e-6]:
            if quadrature == 'trapezoid' and epsabs < 1e-4:
                # slow: skip
                continue

            kwargs = dict(norm='max', epsabs=epsabs, quadrature=quadrature)

            res, err = quad_vec(f, 0, np.inf, **kwargs)
            assert_allclose(res, np.pi/2, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, 0, -np.inf, **kwargs)
            assert_allclose(res, -np.pi/2, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, -np.inf, 0, **kwargs)
            assert_allclose(res, np.pi/2, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, np.inf, 0, **kwargs)
            assert_allclose(res, -np.pi/2, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, -np.inf, np.inf, **kwargs)
            assert_allclose(res, np.pi, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, np.inf, -np.inf, **kwargs)
            assert_allclose(res, -np.pi, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, np.inf, np.inf, **kwargs)
            assert_allclose(res, 0, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, -np.inf, -np.inf, **kwargs)
            assert_allclose(res, 0, rtol=0, atol=max(epsabs, err))

            res, err = quad_vec(f, 0, np.inf, points=(1.0, 2.0), **kwargs)
            assert_allclose(res, np.pi/2, rtol=0, atol=max(epsabs, err))

        def f(x):
            return np.sin(x + 2) / (1 + x ** 2)
        exact = np.pi / np.e * np.sin(2)
        epsabs = 1e-5

        res, err, info = quad_vec(f, -np.inf, np.inf, limit=1000, norm='max',
                                  epsabs=epsabs, quadrature=quadrature,
                                  full_output=True)
        assert info.status == 1
        assert_allclose(res, exact, rtol=0, atol=max(epsabs, 1.5 * err))


    def test_quad_vec_args(self):
        def f(x, a):
            return x * (x + a) * np.arange(3)
        a = 2
        exact = np.array([0, 4/3, 8/3])

        res, err = quad_vec(f, 0, 1, args=(a,))
        assert_allclose(res, exact, rtol=0, atol=1e-4)

    @pytest.mark.fail_slow(10)
    def test_quad_vec_pool(self):
        f = _lorenzian
        res, err = quad_vec(f, -np.inf, np.inf, norm='max', epsabs=1e-4, workers=4)
        assert_allclose(res, np.pi, rtol=0, atol=1e-4)

        with Pool(10) as pool:
            def f(x):
                return 1 / (1 + x ** 2)
            res, _ = quad_vec(f, -np.inf, np.inf, norm='max', epsabs=1e-4,
                              workers=pool.map)
            assert_allclose(res, np.pi, rtol=0, atol=1e-4)

    @pytest.mark.fail_slow(10)
    @pytest.mark.parametrize('extra_args', [2, (2,)])
    @pytest.mark.parametrize(
        'workers',
        [1, pytest.param(10, marks=pytest.mark.parallel_threads_limit(4))]
    )
    def test_quad_vec_pool_args(self, extra_args, workers):
        f = _func_with_args
        exact = np.array([0, 4/3, 8/3])

        res, err = quad_vec(f, 0, 1, args=extra_args, workers=workers)
        assert_allclose(res, exact, rtol=0, atol=1e-4)

        with Pool(workers) as pool:
            res, err = quad_vec(f, 0, 1, args=extra_args, workers=pool.map)
            assert_allclose(res, exact, rtol=0, atol=1e-4)

    @quadrature_params
    def test_num_eval(self, quadrature):
        def f(x):
            count[0] += 1
            return x**5

        count = [0]
        res = quad_vec(f, 0, 1, norm='max', full_output=True, quadrature=quadrature)
        assert res[2].neval == count[0]

    def test_info(self):
        def f(x):
            return np.ones((3, 2, 1))

        res, err, info = quad_vec(f, 0, 1, norm='max', full_output=True)

        assert info.success is True
        assert info.status == 0
        assert info.message == 'Target precision reached.'
        assert info.neval > 0
        assert info.intervals.shape[1] == 2
        assert info.integrals.shape == (info.intervals.shape[0], 3, 2, 1)
        assert info.errors.shape == (info.intervals.shape[0],)

    def test_nan_inf(self):
        def f_nan(x):
            return np.nan

        def f_inf(x):
            return np.inf if x < 0.1 else 1/x

        res, err, info = quad_vec(f_nan, 0, 1, full_output=True)
        assert info.status == 3

        res, err, info = quad_vec(f_inf, 0, 1, full_output=True)
        assert info.status == 3


    @pytest.mark.parametrize('a,b', [(0, 1), (0, np.inf), (np.inf, 0),
                                    (-np.inf, np.inf), (np.inf, -np.inf)])
    def test_points(self, a, b):
        # Check that initial interval splitting is done according to
        # `points`, by checking that consecutive sets of 15 point (for
        # gk15) function evaluations lie between `points`

        points = (0, 0.25, 0.5, 0.75, 1.0)
        points += tuple(-x for x in points)

        quadrature_points = 15
        interval_sets = []
        count = 0

        def f(x):
            nonlocal count

            if count % quadrature_points == 0:
                interval_sets.append(set())

            count += 1
            interval_sets[-1].add(float(x))
            return 0.0

        quad_vec(f, a, b, points=points, quadrature='gk15', limit=0)

        # Check that all point sets lie in a single `points` interval
        for p in interval_sets:
            j = np.searchsorted(sorted(points), tuple(p))
            assert np.all(j == j[0])