# mypy: disable-error-code="attr-defined" import pytest import numpy as np from numpy.testing import assert_allclose, assert_equal import scipy._lib._elementwise_iterative_method as eim from scipy import special, stats from scipy.integrate import quad_vec from scipy.integrate._tanhsinh import _tanhsinh, _pair_cache, _nsum from scipy.stats._discrete_distns import _gen_harmonic_gt1 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): err = abs(res - ref) if not log: return err with np.errstate(divide='ignore'): return np.log10(err) def test_input_validation(self): f = self.f1 message = '`f` must be callable.' with pytest.raises(ValueError, match=message): _tanhsinh(42, 0, f.b) message = '...must be True or False.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, log=2) message = '...must be real numbers.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 1+1j, f.b) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, atol='ekki') with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, rtol=pytest) message = '...must be non-negative and finite.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, rtol=-1) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, atol=np.inf) message = '...may not be positive infinity.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, rtol=np.inf, log=True) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, atol=np.inf, log=True) message = '...must be integers.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, maxlevel=object()) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, maxfun=1+1j) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, minlevel="migratory coconut") message = '...must be non-negative.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, maxlevel=-1) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, maxfun=-1) with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, minlevel=-1) message = '...must be True or False.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, preserve_shape=2) message = '...must be callable.' with pytest.raises(ValueError, match=message): _tanhsinh(f, 0, f.b, callback='elderberry') @pytest.mark.parametrize("limits, ref", [ [(0, np.inf), 0.5], # b infinite [(-np.inf, 0), 0.5], # a infinite [(-np.inf, np.inf), 1], # a and b infinite [(np.inf, -np.inf), -1], # flipped limits [(1, -1), stats.norm.cdf(-1) - stats.norm.cdf(1)], # flipped limits ]) def test_integral_transforms(self, limits, ref): # Check that the integral transforms are behaving for both normal and # log integration dist = stats.norm() res = _tanhsinh(dist.pdf, *limits) assert_allclose(res.integral, ref) logres = _tanhsinh(dist.logpdf, *limits, log=True) assert_allclose(np.exp(logres.integral), ref) # Transformation should not make the result complex unnecessarily assert (np.issubdtype(logres.integral.dtype, np.floating) if ref > 0 else np.issubdtype(logres.integral.dtype, np.complexfloating)) assert_allclose(np.exp(logres.error), res.error, atol=1e-16) # 15 skipped intentionally; it's very difficult numerically @pytest.mark.parametrize('f_number', range(1, 15)) def test_basic(self, f_number): 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.parametrize('ref', (0.5, [0.4, 0.6])) @pytest.mark.parametrize('case', stats._distr_params.distcont) def test_accuracy(self, ref, case): 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') 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): # Test for correct functionality, output shapes, and dtypes for various # input shapes. rng = np.random.default_rng(82456839535679456794) a = rng.random(shape) b = rng.random(shape) p = rng.random(shape) n = np.prod(shape) def f(x, p): f.ncall += 1 f.feval += 1 if (x.size == n or x.ndim <=1) else x.shape[-1] return x**p f.ncall = 0 f.feval = 0 @np.vectorize 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).ravel() attrs = ['integral', 'error', 'success', 'status', 'nfev', 'maxlevel'] for attr in attrs: ref_attr = [getattr(ref, attr) for ref in refs] res_attr = getattr(res, attr) assert_allclose(res_attr.ravel(), ref_attr, rtol=1e-15) assert_equal(res_attr.shape, shape) assert np.issubdtype(res.success.dtype, np.bool_) assert np.issubdtype(res.status.dtype, np.integer) assert np.issubdtype(res.nfev.dtype, np.integer) assert np.issubdtype(res.maxlevel.dtype, np.integer) assert_equal(np.max(res.nfev), f.feval) # maxlevel = 2 -> 3 function calls (2 initialization, 1 work) assert np.max(res.maxlevel) >= 2 assert_equal(np.max(res.maxlevel), f.ncall) def test_flags(self): # 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: np.exp(-x**2), # converges lambda x: np.exp(x), # reaches maxiter due to order=2 lambda x: np.full_like(x, np.nan)[()]] # stops due to NaN res = [funcs[j](x) for x, j in zip(xs, js.ravel())] return res f.nit = 0 args = (np.arange(3, dtype=np.int64),) res = _tanhsinh(f, [np.inf]*3, [-np.inf]*3, maxlevel=5, args=args) ref_flags = np.array([0, -2, -3]) assert_equal(res.status, ref_flags) def test_flags_preserve_shape(self): # Same test as above but using `preserve_shape` option to simplify. def f(x): return [np.exp(-x[0]**2), # converges np.exp(x[1]), # reaches maxiter due to order=2 np.full_like(x[2], np.nan)[()]] # stops due to NaN res = _tanhsinh(f, [np.inf]*3, [-np.inf]*3, maxlevel=5, preserve_shape=True) ref_flags = np.array([0, -2, -3]) assert_equal(res.status, ref_flags) def test_preserve_shape(self): # Test `preserve_shape` option def f(x): return np.asarray([[x, np.sin(10 * x)], [np.cos(30 * x), x * np.sin(100 * x)]]) ref = quad_vec(f, 0, 1) res = _tanhsinh(f, 0, 1, preserve_shape=True) assert_allclose(res.integral, ref[0]) def test_convergence(self): # demonstrate that number of accurate digits doubles each iteration f = self.f1 last_logerr = 0 for i in range(4): res = _tanhsinh(f, 0, f.b, minlevel=0, maxlevel=i) logerr = self.error(res.integral, f.ref, log=True) assert (logerr < last_logerr * 2 or logerr < -15.5) last_logerr = logerr def test_options_and_result_attributes(self): # demonstrate that options are behaving as advertised and status # messages are as intended def f(x): f.calls += 1 f.feval += np.size(x) return self.f2(x) f.ref = self.f2.ref f.b = self.f2.b default_rtol = 1e-12 default_atol = f.ref * default_rtol # effective default absolute tol # Test default options f.feval, f.calls = 0, 0 ref = _tanhsinh(f, 0, f.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 = ref.maxlevel res = _tanhsinh(f, 0, f.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, 0, f.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(ref.error, np.inf) res = _tanhsinh(f, 0, f.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(ref.error, -np.inf) res = _tanhsinh(f, 0, f.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(ref.error/ref.integral, np.inf) res = _tanhsinh(f, 0, f.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(ref.error/ref.integral, -np.inf) res = _tanhsinh(f, 0, f.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.parametrize('rtol', [1e-4, 1e-14]) def test_log(self, rtol): # Test equivalence of log-integration and regular integration dist = stats.norm() test_tols = dict(atol=1e-18, rtol=1e-15) # Positive integrand (real log-integrand) res = _tanhsinh(dist.logpdf, -1, 2, log=True, rtol=np.log(rtol)) ref = _tanhsinh(dist.pdf, -1, 2, rtol=rtol) assert_allclose(np.exp(res.integral), ref.integral, **test_tols) assert_allclose(np.exp(res.error), ref.error, **test_tols) assert res.nfev == ref.nfev # Real integrand (complex log-integrand) def f(x): return -dist.logpdf(x)*dist.pdf(x) def logf(x): return np.log(dist.logpdf(x) + 0j) + dist.logpdf(x) + np.pi * 1j res = _tanhsinh(logf, -np.inf, np.inf, log=True) ref = _tanhsinh(f, -np.inf, np.inf) # 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'): assert_allclose(np.exp(res.integral), ref.integral, **test_tols) assert_allclose(np.exp(res.error), ref.error, **test_tols) assert res.nfev == ref.nfev def test_complex(self): # Test integration of complex integrand # Finite limits def f(x): return np.exp(1j * x) res = _tanhsinh(f, 0, np.pi/4) ref = np.sqrt(2)/2 + (1-np.sqrt(2)/2)*1j assert_allclose(res.integral, ref) # Infinite limits dist1 = stats.norm(scale=1) dist2 = stats.norm(scale=2) def f(x): return dist1.pdf(x) + 1j*dist2.pdf(x) res = _tanhsinh(f, np.inf, -np.inf) assert_allclose(res.integral, -(1+1j)) @pytest.mark.parametrize("maxlevel", range(4)) def test_minlevel(self, maxlevel): # 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 def f(x): f.calls += 1 f.feval += np.size(x) f.x = np.concatenate((f.x, x.ravel())) return self.f2(x) f.feval, f.calls, f.x = 0, 0, np.array([]) ref = _tanhsinh(f, 0, self.f2.b, minlevel=0, maxlevel=maxlevel) ref_x = np.sort(f.x) for minlevel in range(0, maxlevel + 1): f.feval, f.calls, f.x = 0, 0, np.array([]) options = dict(minlevel=minlevel, maxlevel=maxlevel) res = _tanhsinh(f, 0, self.f2.b, **options) # Should be very close; all that has changed is the order of values assert_allclose(res.integral, ref.integral, rtol=4e-16) # Difference in absolute errors << magnitude of integral assert_allclose(res.error, ref.error, atol=4e-16 * ref.integral) assert res.nfev == f.feval == len(f.x) assert f.calls == maxlevel - minlevel + 1 + 1 # 1 validation call assert res.status == ref.status assert_equal(ref_x, np.sort(f.x)) def test_improper_integrals(self): # Test handling of infinite limits of integration (mixed with finite limits) def f(x): x[np.isinf(x)] = np.nan return np.exp(-x**2) a = [-np.inf, 0, -np.inf, np.inf, -20, -np.inf, -20] b = [np.inf, np.inf, 0, -np.inf, 20, 20, np.inf] ref = np.sqrt(np.pi) res = _tanhsinh(f, a, b) assert_allclose(res.integral, [ref, ref/2, ref/2, -ref, ref, ref, ref]) @pytest.mark.parametrize("limits", ((0, 3), ([-np.inf, 0], [3, 3]))) @pytest.mark.parametrize("dtype", (np.float32, np.float64)) def test_dtype(self, limits, dtype): # Test that dtypes are preserved a, b = np.asarray(limits, dtype=dtype)[()] def f(x): assert x.dtype == dtype return np.exp(x) rtol = 1e-12 if dtype == np.float64 else 1e-5 res = _tanhsinh(f, a, b, rtol=rtol) assert res.integral.dtype == dtype assert res.error.dtype == dtype assert np.all(res.success) assert_allclose(res.integral, np.exp(b)-np.exp(a), rtol=rtol) def test_maxiter_callback(self): # Test behavior of `maxiter` parameter and `callback` interface a, b = -np.inf, np.inf def f(x): return np.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 callback.res[key] == 1 assert res[key] == -2 assert res2[key] == -4 else: assert res2[key] == callback.res[key] == res[key] def test_jumpstart(self): # 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, b = -np.inf, np.inf def f(x): return np.exp(-x*x) def callback(res): callback.integrals.append(res.integral) callback.errors.append(res.error) callback.integrals = [] callback.errors = [] maxlevel = 4 _tanhsinh(f, a, b, minlevel=0, maxlevel=maxlevel, callback=callback) integrals = [] errors = [] for i in range(maxlevel + 1): res = _tanhsinh(f, a, b, minlevel=i, maxlevel=i) integrals.append(res.integral) errors.append(res.error) assert_allclose(callback.integrals[1:], integrals, rtol=1e-15) assert_allclose(callback.errors[1:], errors, rtol=1e-15, atol=1e-16) def test_special_cases(self): # Test edge cases and other special cases # Test that integers are not passed to `f` # (otherwise this would overflow) def f(x): assert np.issubdtype(x.dtype, np.floating) return x ** 99 res = _tanhsinh(f, 0, 1) assert res.success assert_allclose(res.integral, 1/100) # Test levels 0 and 1; error is NaN res = _tanhsinh(f, 0, 1, maxlevel=0) assert res.integral > 0 assert_equal(res.error, np.nan) res = _tanhsinh(f, 0, 1, maxlevel=1) assert res.integral > 0 assert_equal(res.error, np.nan) # Tes equal left and right integration limits res = _tanhsinh(f, 1, 1) assert res.success assert res.maxlevel == -1 assert_allclose(res.integral, 0) # Test scalar `args` (not in tuple) def f(x, c): return x**c res = _tanhsinh(f, 0, 1, args=99) assert_allclose(res.integral, 1/100) # Test NaNs a = [np.nan, 0, 0, 0] b = [1, np.nan, 1, 1] c = [1, 1, np.nan, 1] res = _tanhsinh(f, a, b, args=(c,)) assert_allclose(res.integral, [np.nan, np.nan, np.nan, 0.5]) assert_allclose(res.error[:3], np.nan) assert_equal(res.status, [-3, -3, -3, 0]) assert_equal(res.success, [False, False, False, True]) assert_equal(res.nfev[:3], 1) # 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 = np.empty(0) _pair_cache.wj = np.empty(0) _pair_cache.indices = [0] _pair_cache.h0 = None res = _tanhsinh(lambda x: x*1j, 0, 1) assert_allclose(res.integral, 0.5*1j) res = _tanhsinh(lambda x: x, 0, 1) assert_allclose(res.integral, 0.5) # Test zero-size shape = (0, 3) res = _tanhsinh(lambda x: x, 0, np.zeros(shape)) attrs = ['integral', 'error', 'success', 'status', 'nfev', 'maxlevel'] for attr in attrs: assert_equal(res[attr].shape, shape) class TestNSum: rng = np.random.default_rng(5895448232066142650) p = rng.uniform(1, 10, size=10) 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): f = self.f1 message = '`f` must be callable.' with pytest.raises(ValueError, match=message): _nsum(42, f.a, f.b) message = '...must be True or False.' with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, log=2) message = '...must be real numbers.' with pytest.raises(ValueError, match=message): _nsum(f, 1+1j, f.b) with pytest.raises(ValueError, match=message): _nsum(f, f.a, None) with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, step=object()) with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, atol='ekki') with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, rtol=pytest) with np.errstate(all='ignore'): res = _nsum(f, [np.nan, -np.inf, np.inf], 1) assert np.all((res.status == -1) & np.isnan(res.sum) & np.isnan(res.error) & ~res.success & res.nfev == 1) res = _nsum(f, 10, [np.nan, 1]) assert np.all((res.status == -1) & np.isnan(res.sum) & np.isnan(res.error) & ~res.success & res.nfev == 1) res = _nsum(f, 1, 10, step=[np.nan, -np.inf, np.inf, -1, 0]) assert np.all((res.status == -1) & np.isnan(res.sum) & np.isnan(res.error) & ~res.success & res.nfev == 1) message = '...must be non-negative and finite.' with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, rtol=-1) with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, atol=np.inf) message = '...may not be positive infinity.' with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, rtol=np.inf, log=True) with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, atol=np.inf, log=True) message = '...must be a non-negative integer.' with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, maxterms=3.5) with pytest.raises(ValueError, match=message): _nsum(f, f.a, f.b, maxterms=-2) @pytest.mark.parametrize('f_number', range(1, 4)) def test_basic(self, f_number): f = getattr(self, f"f{f_number}") res = _nsum(f, f.a, f.b, args=f.args) assert_allclose(res.sum, f.ref) assert_equal(res.status, 0) assert_equal(res.success, True) with np.errstate(divide='ignore'): logres = _nsum(lambda *args: np.log(f(*args)), f.a, f.b, log=True, args=f.args) assert_allclose(np.exp(logres.sum), res.sum) assert_allclose(np.exp(logres.error), res.error) assert_equal(logres.status, 0) assert_equal(logres.success, True) @pytest.mark.parametrize('maxterms', [0, 1, 10, 20, 100]) def test_integral(self, maxterms): # test precise behavior of integral approximation f = self.f1 def logf(x): return -2*np.log(x) def F(x): return -1 / x a = np.asarray([1, 5])[:, np.newaxis] b = np.asarray([20, 100, np.inf])[:, np.newaxis, np.newaxis] step = np.asarray([0.5, 1, 2]).reshape((-1, 1, 1, 1)) nsteps = np.floor((b - a)/step) b_original = b b = a + nsteps*step k = a + maxterms*step # partial sum direct = f(a + np.arange(maxterms)*step).sum(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 a, b, step = np.broadcast_arrays(a, b, step) for i in np.ndindex(a.shape): ai, bi, stepi = a[i], b[i], step[i] if (bi - ai)/stepi + 1 <= maxterms: direct = f(np.arange(ai, bi+stepi, stepi)).sum() ref_sum[i] = direct ref_err[i] = direct * np.finfo(direct).eps rtol = 1e-12 res = _nsum(f, a, b_original, step=step, maxterms=maxterms, rtol=rtol) assert_allclose(res.sum, ref_sum, rtol=10*rtol) assert_allclose(res.error, ref_err, rtol=100*rtol) assert_equal(res.status, 0) assert_equal(res.success, True) i = ((b_original - a)/step + 1 <= maxterms) assert_allclose(res.sum[i], ref_sum[i], rtol=1e-15) assert_allclose(res.error[i], ref_err[i], rtol=1e-15) logres = _nsum(logf, a, b_original, step=step, log=True, rtol=np.log(rtol), maxterms=maxterms) assert_allclose(np.exp(logres.sum), res.sum) assert_allclose(np.exp(logres.error), res.error) assert_equal(logres.status, 0) assert_equal(logres.success, True) @pytest.mark.parametrize('shape', [tuple(), (12,), (3, 4), (3, 2, 2)]) def test_vectorization(self, shape): # 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 = np.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, a, b, maxterms=1000, args=(p,)) refs = _nsum_single(a, b, p, maxterms=1000).ravel() attrs = ['sum', 'error', 'success', 'status', 'nfev'] for attr in attrs: ref_attr = [getattr(ref, attr) for ref in refs] res_attr = getattr(res, attr) assert_allclose(res_attr.ravel(), ref_attr, rtol=1e-15) assert_equal(res_attr.shape, shape) assert np.issubdtype(res.success.dtype, np.bool_) assert np.issubdtype(res.status.dtype, np.integer) assert np.issubdtype(res.nfev.dtype, np.integer) assert_equal(np.max(res.nfev), f.feval) def test_status(self): f = self.f2 p = [2, 2, 0.9, 1.1] a = [0, 0, 1, 1] b = [10, np.inf, np.inf, np.inf] ref = special.zeta(p, 1) with np.errstate(divide='ignore'): # intentionally dividing by zero res = _nsum(f, a, b, args=(p,)) assert_equal(res.success, [False, False, False, True]) assert_equal(res.status, [-3, -3, -2, 0]) assert_allclose(res.sum[res.success], ref[res.success]) def test_nfev(self): def f(x): f.nfev += np.size(x) return 1 / x**2 f.nfev = 0 res = _nsum(f, 1, 10) assert_equal(res.nfev, f.nfev) f.nfev = 0 res = _nsum(f, 1, np.inf, atol=1e-6) assert_equal(res.nfev, f.nfev) def test_inclusive(self): # There was an edge case off-by one bug when `_direct` was called with # `inclusive=True`. Check that this is resolved. res = _nsum(lambda k: 1 / k ** 2, [1, 4], np.inf, maxterms=500, atol=0.1) ref = _nsum(lambda k: 1 / k ** 2, [1, 4], np.inf) assert np.all(res.sum > (ref.sum - res.error)) assert np.all(res.sum < (ref.sum + res.error)) def test_special_case(self): # test equal lower/upper limit f = self.f1 a = b = 2 res = _nsum(f, a, b) assert_equal(res.sum, f(a)) # Test scalar `args` (not in tuple) res = _nsum(self.f2, 1, np.inf, args=2) assert_allclose(res.sum, self.f1.ref) # f1.ref is correct w/ args=2 # Test 0 size input a = np.empty((3, 1, 1)) # arbitrary broadcastable shapes b = np.empty((0, 1)) # could use Hypothesis p = np.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, 0, 10, maxterms=0) assert np.isnan(res.sum) assert np.isnan(res.error) assert res.status == -2 res = _nsum(f, 0, 10, maxterms=1) assert np.isnan(res.sum) assert np.isnan(res.error) assert res.status == -3 # Test NaNs # should skip both direct and integral methods if there are NaNs a = [np.nan, 1, 1, 1] b = [np.inf, np.nan, np.inf, np.inf] p = [2, 2, np.nan, 2] res = _nsum(self.f2, a, b, args=(p,)) assert_allclose(res.sum, [np.nan, np.nan, np.nan, self.f1.ref]) assert_allclose(res.error[:3], np.nan) assert_equal(res.status, [-1, -1, -3, 0]) assert_equal(res.success, [False, False, False, True]) # Ideally res.nfev[2] would be 1, but `tanhsinh` has some function evals assert_equal(res.nfev[:2], 1) @pytest.mark.parametrize('dtype', [np.float32, np.float64]) def test_dtype(self, dtype): def f(k): assert k.dtype == dtype return 1 / k ** np.asarray(2, dtype=dtype)[()] a = np.asarray(1, dtype=dtype) b = np.asarray([10, np.inf], dtype=dtype) res = _nsum(f, a, b) assert res.sum.dtype == dtype assert res.error.dtype == dtype rtol = 1e-12 if dtype == np.float64 else 1e-6 ref = _gen_harmonic_gt1(b, 2) assert_allclose(res.sum, ref, rtol=rtol)