[Python] Find correct root of parametrized function given solution for one set of parameters

Let's say I have a function foo(x, a, b) and I want to find a specific one of its (potentially multiple) roots x0, i.e. a value x0 such that foo(x0, a, b) == 0. I know that for (a, b) == (0, 0) the root I want is x0 == 0 and that the function changes continuously with a and b, so I can "follow" the root from (0, 0) to the desired (a, b).

Here's an example function.

def foo(x, a, b):
return (1 + a) * np.sin(a + b - x) - x




For (a, b) == (0, 0) I want to the root at 0, for (2, 0) I want the one near 1.5 and for (2, 1) I want the one near 2.2.

Now, this problem seems like one that may be common enough to have a prepared, fast solver in scipy or another standard package (or tools to easily and efficiently construct one). However, I don't know what terms to search for to find it (or verify that it doesn't exist). Is there a ready-made tool for this? What is this kind of problem called?


Clarifications:

• Depending on (a, b), the "correct" root may disappear (e.g. for (1, 3) in the example). When this happens, returning nan is the preferred behavior, though this is not super important.
• By "fast" I mostly mean that many parameter sets (a, b) can be quickly solved, not just a single one. I will go on calculating the root for a lot of different parameters, e.g. for plotting and integrating over them.

Here's a quickly put together reference implementation that does pretty much what I want for the example above (and creates the plot). It's not very fast, of course, which somewhat limits me in my actual application.

import functools
import numpy as np
from scipy.optimize import root_scalar
from matplotlib import pyplot as plt

def foo(x, a, b):
return (1 + a) * np.sin(a + b - x) - x

fig, ax = plt.subplots()
ax.grid()
ax.set_xlabel("x")
ax.set_ylabel("foo")
x = np.linspace(-np.pi, np.pi, 201)
ax.plot(x, foo(x, 0, 0), label="(a, b) = (0, 0)")
ax.plot(x, foo(x, 2, 0), label="(a, b) = (2, 0)")
ax.plot(x, foo(x, 2, 1), label="(a, b) = (2, 1)")
ax.legend()
plt.show()

# Semi-bodged solver for reference:

def _dfoo(x, a, b):
return -(1 + a) * np.cos(a + b - x) - 1

def _solve_fooroot(guess, a, b):
if np.isnan(guess):
return np.nan
# Determine limits for finding the root.
# Allow for slightly larger limits to account for numerical imprecision.
maxlim = 1.1 * (1 + a)
y0 = foo(guess, a, b)
if y0 == 0:
return guess
dy0 = _dfoo(guess, a, b)
estimate = -y0 / dy0
# Too small estimates are no good.
if np.abs(estimate) < 1e-2 * maxlim:
estimate = np.sign(estimate) * 1e-2 * maxlim
for lim in np.arange(guess, guess + 10 * estimate, 1e-1 * estimate):
if np.sign(foo(lim, a, b)) != np.sign(y0):
bracket = sorted([guess, lim])
break
else:
return np.nan
sol = root_scalar(foo, (a, b), bracket=bracket)
return sol.root

@functools.cache
def _fooroot(an, astep, bn, bstep):
if an == 0:
if bn == 0:
return 0
guessan, guessbn = an, bn - 1
else:
guessan, guessbn = an - 1, bn
# Avoid reaching maximum recursion depth.
for thisbn in range(0, guessbn, 100):
_fooroot(0, astep, thisbn, bstep)
for thisan in range(0, guessan, 100):
_fooroot(thisan, astep, guessbn, bstep)
guess = _fooroot(guessan, astep, guessbn, bstep)
return _solve_fooroot(guess, an * astep, bn * bstep)

@np.vectorize
@functools.cache
def fooroot(a, b):
astep = (-1 if a < 0 else 1) * 1e-2
bstep = (-1 if b < 0 else 1) * 1e-2
guess = _fooroot(int(a / astep), astep, int(b / bstep), bstep)
return _solve_fooroot(guess, a, b)

print(fooroot(0, 0))
print(fooroot(2, 0))
print(fooroot(2, 1))

fig, ax = plt.subplots()
ax.grid()
ax.set_xlabel("b")
ax.set_ylabel("fooroot(a, b)")
b = np.linspace(-3, 3, 201)
for a in [0, 0.2, 0.5]:
ax.plot(b, fooroot(a, b), label=f"a = {a}")
ax.legend()
plt.show()




Continue reading...