scipy 1.15.3__cp313-cp313t-macosx_14_0_arm64.whl → 1.16.0rc2__cp313-cp313t-macosx_14_0_arm64.whl

scipy/stats/_continued_fraction.py

@@ -0,0 +1,387 @@
+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)