Root finding using a loop

I have one equation defined in the function

def fun(x, y, z, v, b):
       Y = (z*(np.sign(x) * (np.abs(x))**(y-1))) - (v*np.sign(b) * (np.abs(b))**(v-1))/(1-b**v)
       return Y.flatten()

that I want to solve for the value of x, given the values of Z0, SS (year 1: Z0=1.2, SS=2, ...) and different combinations of alpha and kappa, for which I am creating a grid.

Z0 = [1.2, 5, 3, 2.5, 4.2]
SS = [2, 3, 2.2, 3.5, 5]

ngrid = 10
kv = np.linspace(0.05, 2, ngrid)
av = np.linspace(1.5, 4, ngrid)

q0 = []                    
for z in range(len(Z0)):
    zz = Z0[z]
    ss = SS[z]
    for i in range(ngrid):
        for j in range(ngrid):
            kappa = kv[i]
            alpha = av[j]
            res0 = root(lambda x: fun(x, alpha, zz, kappa, ss), x0=np.ones(range(ngrid)))
            q0 = res0.x
            print(q0)

where y = alpha; v=kappa, z = Z0; b = S.

I am getting all [], [], ....

Not sure what is going on. Thanks for your help

Solution

Before you attempt to use res0.x, check res0.success. In this case, you'll find that it is False in each case. When res0.success is False, take a look at res0.message for information about why root failed.

During development and debugging, you might also consider getting the solver working for just one set of parameter values before you embed root in three nested loops. For example, here are a few lines from an ipython session (variables were defined in previous lines, not shown):

In [37]: res0 = root(lambda x: fun(x, av[0], Z0[0], kv[0], SS[0]), x0=np.ones(range(ngrid)))

In [38]: res0.success
Out[38]: False

In [39]: res0.message
Out[39]: 'Improper input parameters were entered.'

The message suggests that something is wrong with the input parameters. You call root like this:

        res0 = root(lambda x: fun(x, alpha, zz, kappa, ss), x0=np.ones(range(ngrid)))

A close look at that line shows the problem: the initial guess is np.ones(range(ngrid)):

In [41]: np.ones(range(ngrid))
Out[41]: array([], shape=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), dtype=float64)

That's not what you want! The use of range looks like a simple typo (or "thinko"). The initial guess should be

x0=np.ones(ngrid)

In ipython, we get:

In [50]: res0 = root(lambda x: fun(x, av[0], Z0[0], kv[0], SS[0]), x0=np.ones(ngrid))

In [51]: res0.success
Out[51]: True

In [52]: res0.x
Out[52]: 
array([-0.37405428, -0.37405428, -0.37405428, -0.37405428, -0.37405428,
       -0.37405428, -0.37405428, -0.37405428, -0.37405428, -0.37405428])

All the return values are the same (and this happens for other parameters values), which suggests that you are solving a scalar equation. A closer look at fun shows that you only use x in element-wise operations, so you are in fact solving just a scalar equation. In that case, you can use x0=1:

In [65]: res0 = root(lambda x: fun(x, av[0], Z0[0], kv[0], SS[0]), x0=1)

In [66]: res0.success
Out[66]: True

In [67]: res0.x
Out[67]: array([-0.37405428])

You could also consider using root_scalar instead of root.