I'm deep in the thick of some Monte Carlo analysis using a root-finding algorithm for each sample. For performance reasons I would like to use batch-processing techniques, but scipy.optimize.brentq operates only on scalars, not vectors.

Problem statement:

• Given some function f of at least one argument, which can operate on vectors for at least its first argument, in which case it returns a vector of the same size as the first argument

• Create a function rootsolver (or other appropriate name) which can be called:

  rootsolver(f,a,b,xtol,rtol,maxiter,args,...)
  

where a,b,xtol,rtol,maxiter,args have the same meaning as in scipy.optimize.brentq but all of them except maxiter and args can be vectors (1-D arrays) as long as they are of identical length, and in this case rootsolver proceeds in a vectorized manner:

• choosing initial vector guesses for x
• calling y=f((x,)+args) repeatedly
• choosing updated vector guesses for x on each iteration, based on current and recent values y
• terminating when all of the values x are within sufficient tolerance, or when maxiter is reached
• if a count of iterations is returned (e.g. as in RootResults) it should be a vector of such counts

This requires if-statements to be turned into vectorized equivalents (e.g. conditional assignment based on boolean arrays)

Simple test case:

import numpy as np
from scipy.optimize import brentq
def f(x,a):
    return a-x*x

brentq(f,0,3,args=(2,))   # returns 1.4142135623731364
k=np.arange(1,8)
brentq(f,0,3,args=(k,))   # should return an array which is approximately sqrt(k)

Ideally this would "play well with numba" but I see that the scipy.optimize algorithms are already implemented in C so maybe that isn't a reasonable request.