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

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)