In [2]:
import lmfit
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LightSource
from matplotlib import cm
from scipy.optimize import minimize


---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-2-eac01f218db9> in <module>()
----> 1 import lmfit
      2 import numpy as np
      3 from matplotlib import pyplot as plt
      4 from mpl_toolkits.mplot3d import Axes3D
      5 from matplotlib.colors import LightSource

ModuleNotFoundError: No module named 'lmfit'

In [2]:
def rosenbrock(pars):
    x = pars[0]
    y = pars[1]
    tmp1 = y - x * x
    tmp2 = 1 - x
    return 100 * tmp1 * tmp1 + tmp2 * tmp2

In [3]:
npoints = 50
x = np.linspace(-2,2,npoints)
y = np.linspace(-1,3,npoints)
xx, yy = np.meshgrid(x,y)
zz = np.zeros((npoints,npoints))
for i in range(npoints):
    for j in range(npoints):    
        zz[i,j] = rosenbrock([xx[i,j], yy[i,j]])

In [4]:
ls = LightSource(270, 45)
rgb = ls.shade(zz, blend_mode='soft', cmap=cm.coolwarm)
illuminated_surface = ls.shade(zz, cmap=cm.coolwarm)

In [5]:
fig=plt.figure(figsize=(9, 8), dpi= 80, facecolor='w', edgecolor='k')
ax = fig.gca(projection='3d')
surf = ax.plot_surface(xx, yy, zz, rstride=1, cstride=1, linewidth=0, facecolors=illuminated_surface, antialiased=True)



In [104]:
x0 = np.array([-2.0, -1.0])
res = minimize(rosenbrock, x0, method='nelder-mead', options={'xtol': 1e-8, 'disp': True})


Optimization terminated successfully.
         Current function value: 0.000000
         Iterations: 107
         Function evaluations: 201

In [82]:
print(res)


 final_simplex: (array([[1.        , 1.        ],
       [1.        , 1.        ],
       [1.        , 1.00000001]]), array([4.20699532e-18, 6.84024280e-18, 1.01782062e-17]))
           fun: 4.206995321238234e-18
       message: 'Optimization terminated successfully.'
          nfev: 201
           nit: 107
        status: 0
       success: True
             x: array([1., 1.])

In [ ]: