numpy/scipy optimize and 3d plotting sample

optimize_test.py

# -*- coding: utf-8 -*-
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
from scipy import optimize

def f(x,y):
    r1 = np.sqrt(x**2+y**2)
    r2 = np.sqrt((x-15)**2+(y-15)**2)
    return np.sin(r1)/np.sqrt(r1) + 1.5*np.exp(-r2)

x = np.arange(0, 30, 0.5)
y = np.arange(0, 30, 0.5)
X, Y = np.meshgrid(x, y)
Z = f(X,Y)

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(X, Y, Z, rstride=3, cstride=3, alpha=0.3)

inits=[[8,8],[10,5],[13,13],[16,16]]
x0s,x1s=[],[]
for x0 in inits:
    opt = optimize.minimize(lambda (x,y): -f(x,y),
                            np.array(x0),method='Nelder-Mead')
    x0s+=x0    
    x1s+=opt.x.tolist()
    
ax.scatter3D(x0s[0::2],x0s[1::2],
    [f(a,b)*1.05 for a,b in zip(x0s[0::2],x0s[1::2])]
    ,c='g',s=100, alpha=1)
ax.scatter3D(x1s[0::2],x1s[1::2],
    [f(a,b)*1.05 for a,b in zip(x1s[0::2],x1s[1::2])]
    ,c='r',s=100, alpha=0.5)

plt.show()