Search code examples
pythonarraysscipy-optimize

How to format the argument of scipy.optimize.fsolve for arrays of data


I'd like to use a solver (scipy.optimize.fsolve) to solve for the root of a function, fun(x,y). In this case, I want the argument (y) to be an array (e.g. an array of data). I would also like to avoid using for-loops to call fsolve for each value in y.

In specifying arg (y) as an array, I am getting an error that the result from the function call is not a proper array of floats. Other errors occur if I make "data" a tuple instead of an array.

Here is an MWE of the problem:

import numpy as np
from scipy.optimize import fsolve

def fun(x, y):
    return x+y

data = np.array([1, 2, 3, 4])
x = fsolve(fun, x0=0, args=data)
print(x)

How can the input to fsolve be corrected so that it solves for the root of fun(x,y) for each value of y in the array (without using a for-loop)?


Solution

  • The function, that is passed to fsolve, takes at least one (possibly vector) argument and returns a value of the same length as mentioned here.

    In your case, you are passing x0=0 and args=np.array([1,2,3,4]) to fsolve. The return value of fun has a different length to x0 (x0 is a scalar and args is an array having shape (4,)).

    The following code solves your problem:

    import numpy as np
    from scipy.optimize import fsolve
    
    def fun(x, y):
        return x+y
    
    data = np.array([1, 2, 3, 4])
    x = fsolve(fun, x0=np.array([0,0,0,0]), args=data)
    print(x)