import numpy as np from scipy._lib._array_api import ( array_namespace, xp_ravel, xp_copy, xp_promote ) import scipy._lib._elementwise_iterative_method as eim from scipy._lib._util import _RichResult from scipy import special # Todo: # Avoid special-casing key 'n' in _lib._elementwise_iterative_method::_check_termination # Rearrange termination condition to allow absolute and relative tolerances? # Interpret/return |f_n - f_{n-1}| as an error estimate? # Return gracefully for size=0 arrays def _logaddexp(x, y, xp=None): # logaddexp that supports complex numbers xp = array_namespace(x, y) if xp is None else xp x, y = xp.broadcast_arrays(x, y) xy = xp.stack((x, y), axis=0) return special.logsumexp(xy, axis=0) def _continued_fraction_iv(a, b, args, tolerances, maxiter, log): # Input validation for `_continued_fraction` if not callable(a) or not callable(b): raise ValueError('`a` and `b` must be callable.') if not np.iterable(args): args = (args,) # Call each callable once to determine namespace and dtypes a0, b0 = a(0, *args), b(0, *args) xp = array_namespace(a0, b0, *args) a0, b0, *args = xp_promote(a0, b0, *args, force_floating=True, broadcast=True, xp=xp) shape, dtype = a0.shape, a0.dtype a0, b0, *args = (xp_ravel(arg) for arg in (a0, b0) + tuple(args)) tolerances = {} if tolerances is None else tolerances eps = tolerances.get('eps', None) tiny = tolerances.get('tiny', None) # tolerances are floats, not arrays, so it's OK to use NumPy message = ('`eps` and `tiny` must be (or represent the logarithm of) ' 'finite, positive, real scalars.') tols = np.asarray([eps if eps is not None else 1, tiny if tiny is not None else 1]) not_real = (not np.issubdtype(tols.dtype, np.number) or np.issubdtype(tols.dtype, np.complexfloating)) not_positive = np.any(tols <= 0) if not log else False not_finite = not np.all(np.isfinite(tols)) not_scalar = tols.shape != (2,) if not_real or not_positive or not_finite or not_scalar: raise ValueError(message) maxiter_int = int(maxiter) if maxiter != maxiter_int or maxiter < 0: raise ValueError('`maxiter` must be a non-negative integer.') if not isinstance(log, bool): raise ValueError('`log` must be boolean.') return a, b, args, eps, tiny, maxiter, log, a0, b0, shape, dtype, xp def _continued_fraction(a, b, *, args=(), tolerances=None, maxiter=100, log=False): r"""Evaluate a generalized continued fraction numerically. `_continued_fraction` iteratively evaluates convergents of a continued fraction given coefficients returned by callables `a` and `b`. Iteration terminates when `maxiter` terms have been evaluated or a termination criterion controlled by `tolerances` is satisfied, and the final convergent is returned as the ``f`` attribute of the result object. This function works elementwise when `args` contains (broadcastable) arrays. Parameters ---------- a, b: callable Functions that provide the *numerator* and *denominator* coefficients of the continued fraction, respectively. The signature of each must be:: a(n: int, *argsj) -> ndarray where ``n`` is the coefficient number and ``argsj`` is a tuple, which may contain an arbitrary number of arrays of any shape. `a` and `b` must be elementwise functions: each scalar element ``a(n, *argsj)[i]`` must equal ``a(n, *[argj[i] for argj in argsj])`` for valid indices ``i``. `a` and `b` must not mutate the arrays in ``argsj``. The result shape is the broadcasted shape of ``a(0, *args)`` and ``b(0, *args)``. The dtype used throughout computation is the result dtype of these terms if it is a float, and the default float of the array library otherwise. The numerical value of ``a(0, *args)`` is ignored, and the value of the leading term ``b(0, *args)`` is the so-called "integer" part of the continued fraction (although it need not be integral). args : tuple of array_like, optional Additional positional *array* arguments to be passed to `a` and `b`. Arrays must be broadcastable with one another. If the coefficient callables require additional arguments that are not broadcastable with one another, wrap them with callables `a` and `b` such that `a` and `b` accept only ``n`` and broadcastable array arguments. tolerances : dictionary of floats, optional Tolerances and numerical thresholds used by the algorithm. Currently, valid keys of the dictionary are: - ``eps`` - the convergence threshold of Lentz' algorithm - ``tiny`` - not strictly a "tolerance", but a very small positive number used to avoid division by zero The default `eps` is the precision of the appropriate dtype, and the default `tiny` is the precision squared. [1]_ advises that ``eps`` is "as small as you like", but for most purposes, it should not be set smaller than the default because it may prevent convergence of the algorithm. [1]_ also advises that ``tiny`` should be less than typical values of ``eps * b(n)``, so the default is a good choice unless the :math:`b_n` are very small. See [1]_ for details. maxiter : int, default: 100 The maximum number of iterations of the algorithm to perform. log : bool, default: False If True, `a` and `b` return the (natural) logarithm of the terms, `tolerances` contains the logarithm of the tolerances, and the result object reports the logarithm of the convergent. Returns ------- res : _RichResult An object similar to an instance of `scipy.optimize.OptimizeResult` with the following attributes. The descriptions are written as though the values will be scalars; however, if `f` returns an array, the outputs will be arrays of the same shape. success : bool array ``True`` where the algorithm terminated successfully (status ``0``); ``False`` otherwise. status : int array An integer representing the exit status of the algorithm. - ``0`` : The algorithm converged to the specified tolerances. - ``-2`` : The maximum number of iterations was reached. - ``-3`` : A non-finite value was encountered. f : float array The convergent which satisfied a termination criterion. nit : int array The number of iterations of the algorithm that were performed. nfev : int array The number of terms that were evaluated. Notes ----- A generalized continued fraction is an expression of the form .. math:: b_0 + \frac{a_1}{b_1 + \frac{a_2}{b_2 + \frac{a_3}{b_3 + \cdots}}} Successive "convergents" approximate the infinitely recursive continued fraction with a finite number of terms :math:`a_n` and :math:`b_n`, which are provided by callables `a` and `b`, respectively. This implementation follows the modified Lentz algorithm ([1]_, [2]_) to evaluate successive convergents until a termination condition is satisfied. References ---------- .. [1] Press, William H., and Saul A. Teukolsky. "Evaluating continued fractions and computing exponential integrals." Computers in Physics 2.5 (1988): 88-89. .. [2] Lentz's algorithm. Wikipedia. https://en.wikipedia.org/wiki/Lentz%27s_algorithm .. [3] Continued fraction. Wikipedia. https://en.wikipedia.org/wiki/Continued_fraction .. [4] Generalized continued fraction. Wikipedia. https://en.wikipedia.org/wiki/Generalized_continued_fraction Examples -------- The "simple continued fraction" of :math:`\pi` is given in [3]_ as .. math:: 3 + \frac{1}{7 + \frac{1}{15 + \frac{1}{1 + \cdots}}} where the :math:`b_n` terms follow no obvious pattern: >>> b = [3, 7, 15, 1, 292, 1, 1, 1, 2, 1, 3, 1] and the :math:`a_n` terms are all :math:`1`. In this case, all the terms have been precomputed, so we call `_continued_fraction` with simple callables which simply return the precomputed coefficients: >>> import numpy as np >>> from scipy.special._continued_fraction import _continued_fraction >>> res = _continued_fraction(a=lambda n: 1, b=lambda n: b[n], maxiter=len(b) - 1) >>> (res.f - np.pi) / np.pi np.float64(7.067899292141148e-15) A generalized continued fraction for :math:`\pi` is given by: .. math:: 3 + \frac{1^2}{6 + \frac{3^2}{6 + \frac{5^2}{6 + \cdots}}} We define the coefficient callables as: >>> def a(n): ... return (2*n - 1)**2 >>> >>> def b(n): ... if n == 0: ... return 3 ... else: ... return 6 Then the continued fraction can be evaluated as: >>> res = _continued_fraction(a, b) >>> res success: False status: -2 f: 3.1415924109719846 nit: 100 nfev: 101 Note that the requested tolerance was not reached within the (default) maximum number of iterations because it converges very slowly. An expression that converges more rapidly is expressed as the difference between two continued fractions. We will compute both of them in one vectorized call to `_continued_fraction`. >>> u, v = 5, 239 >>> >>> def a(n, a1, _): ... # The shape of the output must be the shape of the arguments ... shape = a1.shape ... if n == 0: ... return np.zeros(shape) ... elif n == 1: ... return a1 ... else: ... return np.full(shape, (n-1)**2) >>> >>> def b(n, _, uv): ... shape = uv.shape ... if n == 0: ... return np.zeros(shape) ... return np.full(shape, (2*n - 1)*uv) >>> >>> res = _continued_fraction(a, b, args=([16, 4], [u, v])) >>> res success: [ True True] status: [0 0] f: [ 3.158e+00 1.674e-02] nit: [10 4] nfev: [11 5] Note that the second term converged in fewer than half the number of iterations as the first. The approximation of :math:`\pi` is the difference between the two: >>> pi = res.f[0] - res.f[1] >>> (pi - np.pi) / np.pi np.float64(2.8271597168564594e-16) If it is more efficient to compute the :math:`a_n` and :math:`b_n` terms together, consider instantiating a class with a method that computes both terms and stores the results in an attribute. Separate methods of the class retrieve the coefficients, and these methods are passed to `_continued_fraction` as arguments `a` and `b`. Similarly,if the coefficients can be computed recursively in terms of previous coefficients, use a class to maintain state between callable evaluations. """ res = _continued_fraction_iv(a, b, args, tolerances, maxiter, log) a, b, args, eps, tiny, maxiter, log, a0, b0, shape, dtype, xp = res callback = None # don't want to test it, but easy to add later # The EIM framework was designed for the case in where there would # be only one callable, and all arguments of the callable would be # arrays. We're going a bit beyond that here, since we have two callables, # and the first argument is an integer (the number of the term). Rather # than complicate the framework, we wrap the user-provided callables to # make this problem fit within the existing framework. def a(n, *args, a=a): n = int(xp.real(xp_ravel(n))[0]) return a(n, *args) def b(n, *args, b=b): n = int(xp.real(xp_ravel(n))[0]) return b(n, *args) def func(n, *args): return xp.stack((a(n, *args), b(n, *args)), axis=-1) status = xp.full_like(a0, eim._EINPROGRESS, dtype=xp.int32) # in progress nit, nfev = 0, 1 # one function evaluation (per function) performed above maxiter = 100 if maxiter is None else maxiter # Quotations describing the algorithm are from [1]_ # "... as small as you like, say eps" if eps is None: eps = xp.finfo(dtype).eps if not log else np.log(xp.finfo(dtype).eps) # "The parameter tiny should be less than typical values of eps |b_n|" if tiny is None: tiny = xp.finfo(dtype).eps**2 if not log else 2*np.log(xp.finfo(dtype).eps) # "Set f0 and C0 to the value b0 or to tiny if b0=0. Set D0 = 0. zero = -xp.inf if log else 0 fn = xp.where(b0 == zero, tiny, b0) Cnm1 = xp_copy(fn) Dnm1 = xp.full_like(fn, zero) CnDn = xp.full_like(fn, xp.inf) work = _RichResult(n=0, fn=fn, Cnm1=Cnm1, Dnm1=Dnm1, CnDn=CnDn, eps=eps, tiny=tiny, nit=nit, nfev=nfev, status=status) res_work_pairs = [('status', 'status'), ('f', 'fn'), ('nit', 'nit'), ('nfev', 'nfev')] def pre_func_eval(work): work.n = xp.reshape(xp.asarray(work.n + 1), (-1,)) return work.n def post_func_eval(n, ab, work): an, bn = ab[..., 0], ab[..., 1] zero = 0 if not log else -xp.inf # "Set D_n = 1/(b_n + a_n D_{n-1}) or 1/tiny, if the denominator vanishes" denominator = (bn + an*work.Dnm1 if not log else _logaddexp(bn, an + work.Dnm1, xp=xp)) denominator[denominator == zero] = tiny Dn = (1/denominator if not log else -denominator) # "Set C_n = b_n + a_n / C_{n-1} (or =tiny, if the expression vanishes)" Cn = (bn + an / work.Cnm1 if not log else _logaddexp(bn, an - work.Cnm1, xp=xp)) Cn[Cn == zero] = tiny # "and set f_n = f_{n-1} C_n D_n" work.CnDn = (Cn * Dn if not log else Cn + Dn) work.fn = (work.fn * work.CnDn if not log else work.fn + work.CnDn) work.Cnm1, work.Dnm1 = Cn, Dn def check_termination(work): # Check for all terminal conditions and record statuses. stop = xp.zeros_like(work.CnDn, dtype=xp.bool) # "You quit when |D_n C_n - 1| is as small as you like, say eps" pij = xp.full_like(work.CnDn, xp.pi*1j) if log else None residual = (xp.abs(work.CnDn - 1) if not log else xp.real(_logaddexp(work.CnDn, pij, xp=xp))) i = residual < work.eps work.status[i] = eim._ECONVERGED stop[i] = True # If function value is NaN, report failure. i = (~xp.isfinite(work.fn) if not log else ~(xp.isfinite(work.fn) | (work.fn == -xp.inf))) work.status[i] = eim._EVALUEERR stop[i] = True return stop def post_termination_check(work): pass def customize_result(res, shape): # Only needed pre-NEP 50 res['f'] = xp.asarray(res['f'], dtype=dtype) res['f'] = res['f'][()] if res['f'].ndim == 0 else res['f'] return shape return eim._loop(work, callback, shape, maxiter, func, args, dtype, pre_func_eval, post_func_eval, check_termination, post_termination_check, customize_result, res_work_pairs, xp=xp)