CODEWITHSUNDEEP
numpyscipyequation-solving
Jahfet
Jahfet

Reputation: 279

scipy.optimize: faster root finding over 2D grid

I wrote some code using scipy to find the root the following equation:

def equation(x, y):
     return (x / y) * np.log((a * x / b) + 1.0) - 2.0 * c * c

with a, b, and c scalars.

I have values for y on a rectangular grid (say Y, shape 300x200), and need to find the corresponding x solving the equation for each point. I also have a starting estimate for the value of x at each point (X0, shape 300x 200)

At the moment I've managed to solve this by looping through each item in the array Y, and calling:

for index, value in np.ndenumerate(Y):
    result[index] = scipy.optimize.newton(equation, X0[index], args=(value))
    # or other scalar solvers like brentq

This works but too slow to allow me to release my script. Given the fact that the values are organised on a grid and the Y and result array contain "continuous" values e.g. gradually changing from the edge of the array towards the center I am sure there must be a nice array oriented / multi dimensional way to solve that problem, which could also give better performance.

I have tried a number of options but have not succeeded so far. Any idea?

Any help would be appreciated.

Thanks

Upvotes: 4

Views: 2838

Answers (2)

ev-br
ev-br

Reputation: 26040

(expanding a comment) As @askewchan shows in his answer, the runtime here is dominated by actually solving the equations.

Here's what I'd do here: absorb 2*c*c as a multiplicative constant into y, ditto for a/b. What's left is an equation of the form t log (1 + t) = z, with t and z being related to your x and y.
Now tabulate the values of z = t log(1 + t) over the range you need, interpolate t vs z, and you have solutions for your equation. Notice that you only need a 1D interpolation, no matter what are the shapes of the arrays X and Y.

For interpolation, scipy has a range of interpolators. The simplest to use is probably interp1d, or a UnivariateSpline. All of them support vectorized operations, so you can probably vectorize the evaluations. If this is enough for you performance-wise, you're all set.

Depending on the range of the values of x and y you need, you might or might not want to augment the tabulation with the explicit functional behavior at the boundaries [e.g. t log(1 + t) approaches t^2 as t->0 etc]. If you need that, have a look at https://bitbucket.org/burovski/tabulations --- this is ugly as hell, but it works [there's some work in scipy to have a better alternative, but that's in process at the moment].

Upvotes: 2

askewchan
askewchan

Reputation: 46530

It's tough to test without values of Y and X0, but using np.frompyfunc might be a little faster. It takes a scalar function and turns it into a ufunc, which operates on arrays elementwise.

import numpy as np
import scipy.optimize

a, b, c = np.random.rand(3)

def equation(x, y):
    return (x/y)*np.log((a*x/b) + 1.0) - 2.*c*c

def solver(y, x0):
    return scipy.optimize.newton(equation, x0, args=(y,))

f = np.frompyfunc(solver, 2, 1)   # 2 arrays in, 1 array out, apply solver elementwise

Y, X0 = np.random.rand(2, 300, 200)
res = f(Y, X0)

It still doesn't truly vectorize the process, though, and it doesn't speed it up:

In [51]: timeit f(Y, X0)
1 loops, best of 3: 10.4 s per loop

In [52]: timeit OP(Y, X0, a, b, c)
1 loops, best of 3: 10.5 s per loop

But as far as I can tell, the speed issue does not come from the looping, but rather from the fact that you're solving the equation Y.size times. I believe this is as fast as it will get if you truly solve for each value in Y:

In [53]: timeit solver(Y[0,0], X0[0,0])
10000 loops, best of 3: 178 µs per loop

In [54]: Y.size*178e-6
Out[54]: 10.68

Probably you already realized this :P but you have to actually decrease the computation, by making some approximation or assumption. Something like @Zhenya's suggestion is probably necessary, but we'd have to know a bit more about the values of Y, X0, a, b, c.

Upvotes: 1

Related Questions