ENH: optimize._chandrupatla: add array API support

[mdhaber]

Reference issue

gh-7242

What does this implement/fix?

This adds array API support to scipy.optimize._chandrupatla and paves the way toward adding array API support to the other elementwise iterative methods (e.g. _differentiate, _tanhsinh, _nsum, _chandrupatla_minimize, and the bracket finders). I'll propose making the rootfinder, minimizer, and bracket finders public shortly.

Additional information

The performance of this function compared to the existing bracketing rootfinders is worst when function evaluations are inexpensive, so I compared the performance between

• brentq
• _chandrupatla (NumPy)
• _chandrupatla (PyTorch, CPU)
• _chandrupatla (CuPy)

when finding the root of xp.cos(x) - p for many values of p. The bracket is always $[0, \pi]$. brentq's default absolute x-tolerance is 2e-12, so I set _chandrupatla to 1e-12 to be more than fair, and I confirmed that both solvers meet their respective tolerance. _chandrupatla finds the roots in a single vectorized call whereas brentq loops with a list comprehension.

# import time
# import numpy as np
# import cupy as cp
# import torch
# import matplotlib.pyplot as plt
#
# rng = np.random.default_rng(1638083107694713882823079058616272161)
# from scipy import optimize
# from scipy.optimize._chandrupatla import _chandrupatla
#
# xp = torch
# n = 10
# n = int(n)
# a = xp.asarray(0)
# b = xp.asarray(np.pi)
# ps = xp.asarray(rng.random(size=n) * 2 - 1, dtype=xp.float64)
#
# def f(x, p):
#     return xp.cos(x) - p
#
# tic = time.perf_counter()
# res = _chandrupatla(f, a, b, args=(ps,), xatol=1e-12)
# toc = time.perf_counter()

import time
import numpy as np
import cupy as cp
import torch
import matplotlib.pyplot as plt

rng = np.random.default_rng(1638083107694713882823079058616272161)
from scipy import optimize
from scipy.optimize._chandrupatla import _chandrupatla

ns = np.logspace(0, 6, 30)

times = {np: [], cp: [], torch: [],  'brentq':[]}

for xp in [np, cp, torch, 'brentq']:
    brentq = False
    if xp == 'brentq':
        xp = np
        brentq = True

    for n in ns:
        n = int(n)
        a = xp.asarray(0)
        b = xp.asarray(np.pi)
        ps = xp.asarray(rng.random(size=n) * 2 - 1, dtype=xp.float64)

        def f(x, p):
            return xp.cos(x) - p

        if brentq:
            tic = time.perf_counter()
            ref = [optimize.brentq(f, a, b, args=(p,)) for p in ps]
            toc = time.perf_counter()
            times['brentq'].append(toc - tic)
            np.testing.assert_allclose(ref, np.arccos(ps), atol=2e-12)
            continue

        tic = time.perf_counter()
        res = _chandrupatla(f, a, b, args=(ps,), xatol=1e-12)
        toc = time.perf_counter()
        times[xp].append(toc - tic)
        np.testing.assert_allclose(cp.asnumpy(res.x), np.arccos(cp.asnumpy(ps)), atol=1e-12)

plt.loglog(ns, times[np], label='np')
plt.loglog(ns, times[cp], label='cp')
plt.loglog(ns, times[torch], label='torch')
plt.loglog(ns, times['brentq'], label='brentq')
plt.xlabel('number of roots')
plt.ylabel('execution time (s)')
plt.title('Root of `xp.cos(x) - p`')
plt.legend()
plt.show()

The function has a lot of overhead due to bells and whistles (e.g. input validation with nice error messages, rich result object, callback function support, etc.). But for solving a lot of equations, the overhead of function calls eventually becomes problematic for brentq.

In this case, you could probably get better performance with cython_optimize, but for more expensive functions with overhead, like finding the argus(1) distribution ppf (#17719 (comment)), the advantage is much more pronounced.

import time
import numpy as np
import cupy as cp
import torch
import matplotlib.pyplot as plt

rng = np.random.default_rng(1638083107694713882823079058616272161)
from scipy import optimize, stats
from scipy.optimize._chandrupatla import _chandrupatla

ns = np.logspace(0, 4, 30)

times = {np: [], cp: [], torch: [], 'brentq':[]}

for xp in [np, 'brentq']:
    brentq = False
    if xp == 'brentq':
        xp = np
        brentq = True

    for n in ns:
        n = int(n)
        a = 0.001
        b = 0.999
        ps = np.linspace(0.005, 0.995, n)

        dist = stats.argus(1)
        def f(x, p):
            return dist.cdf(x) - p

        if brentq:
            tic = time.perf_counter()
            ref = [optimize.brentq(f, a, b, args=(p,)) for p in ps]
            toc = time.perf_counter()
            times['brentq'].append(toc - tic)
            continue

        tic = time.perf_counter()
        res = _chandrupatla(f, a, b, args=(ps,), xatol=1e-12)
        toc = time.perf_counter()
        times[xp].append(toc - tic)

plt.loglog(ns, times[np], label='np')
# plt.loglog(ns, times[cp], label='cp')
plt.loglog(ns, times['brentq'], label='brentq')
plt.xlabel('number of probabilities')
plt.ylabel('execution time (s)')
plt.title('Root of `argus(1).cdf(x) - p`')
plt.legend()
plt.show()

newton is definitely faster than this function, but it should be - the algorithm converges faster. (The advantage of bracketing methods is that convergence is guaranteed if the bracket is valid.) We can add that as another method with the same framework.

The tests are quite strict; I'm ironing out a few failures with alternative backends. Done if CI looks good.

The function currently uses fancy indexing assignment. I can investigate working around that for strict array API support later.

[tupui]

Nice, I just have a few questions, the rest is pretty straightforward and LGTM 👍

[tupui]

Last changes LGTM, letting the CI run and then I think it's good to merge.

[mdhaber]

Ok! Remaining failures seem unrelated - a package install conflict and slow tests, hopefully one of which is a temporary glitch.

[tupui]

Ark you have conflicts now