How to solve a nonlinear function for a vector of parameters in python

Question

I have a non-linear function which I want to solve under different parameter values.

Here is a MWE:

import numpy  as np 
import tensorflow as tf
import scipy.optimize 

# fixed parameters
kon = 0.01
mu  = 1.5
fi  = 0.5 
kappa = 22
w = 0.63

# varying parameters 
n =100
xs = tf.random.normal(shape=(n,), stddev=0.2)
eps = tf.random.normal(shape=(n,), stddev=0.17)
z = tf.sigmoid(tf.random.normal(shape=(n,), stddev=0.22))

My solution

def get_leisure(z, eps, x0):
    hvec = np.empty((0,))
    # leisure today
    for ze,ei,xs in zip(z, eps, x0):
        ei=np.exp(ei)
        xs=np.exp(xs)
        # define the function for leisure 
        def leisure_function(hi):
            return (mu/fi)*np.log(hi) -(1-mu)*kappa*(hi)**(1+(1/fi))- mu*(np.log(w*ei*xs)-np.log(kon))-np.log(ze)

        htemp = scipy.optimize.newton_krylov(leisure_function, 0.5)
        hvec = np.append(hvec, htemp)
    return hvec

My question:

Because the number of cases I have to loop over to solve for hi unknown, might be large, is there a way a better way to do it? e.g. to avoid the loop?

I am not an experienced python user, and I would appreciate any help.

Answer 1

Just to illustrate the point I wanted to make in my comment.

If I run your code as it is and time it:

import numpy  as np 
import tensorflow as tf
import scipy.optimize 

from timeit import default_timer as timer
# fixed parameters
kon = 0.01
mu  = 1.5
fi  = 0.5 
kappa = 22
w = 0.63

# varying parameters 
n=10000
x0 = tf.random.normal(shape=(n,), stddev=0.2)
eps = tf.random.normal(shape=(n,), stddev=0.17)
z = tf.sigmoid(tf.random.normal(shape=(n,), stddev=0.22))

def get_leisure(z, eps, x0):
    hvec = np.empty((0,))
    # leisure today
    for ze,ei,xs in zip(z, eps, x0):
        ei=np.exp(ei)
        xs=np.exp(xs)
        # define the function for leisure 
        def leisure_function(hi):
            return (mu/fi)*np.log(hi) -(1-mu)*kappa*(hi)**(1+(1/fi))- mu*(np.log(w*ei*xs)-np.log(kon))-np.log(ze)

        htemp = scipy.optimize.newton_krylov(leisure_function, 0.5)
        hvec = np.append(hvec, htemp)
    return hvec

start = timer()
msh855_result = get_leisure(z, eps, x0)
end = timer()
print(f'elapsed time: { end - start} s')

It takes about 30.04823809700065 s in a generic machine.

The same problem in the same machine but with a vectorized approach

start = timer()
e_eps = np.exp(eps)
e_x0 = np.exp(x0)
c_1=mu/fi * np.ones(n)
c_2=(1-mu)*kappa * np.ones(n)
c_3= 1 + (1/fi)
c_4=mu*np.log(w*e_eps*e_x0/kon)+np.log(z)
def fun(x):
    return c_1[0] * np.log(x) - c_2[0] * ((x) ** c_3)-c_4
v_result = scipy.optimize.newton_krylov(fun, 0.5 * np.ones(n))
end = timer()
print(f'elapsed time: { end - start} s')

It takes just 0.051494838000508025 s

And if you're not convinced the result is pretty close :

>>> sum(np.sqrt(((msh855_result - v_result)**2)/n))
2.0087031341897715e-06