%matplotlib inline
# %%bash
# pip install mycli
import matplotlib.pyplot as plt
import autograd.numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.colors import LogNorm
from matplotlib import animation
from IPython.display import HTML
from autograd import elementwise_grad, value_and_grad
from scipy.optimize import minimize
from collections import defaultdict
from itertools import zip_longest
from functools import partial
plt.rcParams['lines.linewidth'] = 1
plt.rcParams['lines.color'] = 'r'
# params
fun = 'rosenbrock'
x0 = [0.75, 2.0]
# video
interval = 200
# load path
import h5py
filename="path.jld"
file = h5py.File(filename, 'r')
data= file["pos"][()]
file.close()
julia_path = data[:, :]
print(julia_path.shape)
julia_path
if fun == 'rosenbrock':
# rosenbrock
f = lambda x, y: (1-x)**2+100*(y-x**2)**2
elif fun == 'himmelblau':
# himmelblau
f = lambda x, y: (x**2+y-11)**2+(x+xy**2-7)**2
elif fun == 'both':
# booth
f = lambda x, y: (x+2*y-7)**2+(2*x+x-5)**2
elif fun == 'bohachesvky1':
# bohachevsky1
f = lambda x, y: x**2+2*y**2-0.3*cos(3*pi*x)-0.4*cos(4*pi*y)+0.7
elif fun == 'easom':
# easom
f = lambda x, y: -cos(x)*cos(y)*exp(-(x-pi)**2-(y-pi)**2)
xmin, xmax, xstep = 0, 2, .01
ymin, ymax, ystep = 0., 2.5, .01
x, y = np.meshgrid(np.arange(xmin, xmax + xstep, xstep), np.arange(ymin, ymax + ystep, ystep))
z = f(x, y)
minima = np.array([1., 1.])
f(*minima)
minima_ = minima.reshape(-1, 1)
minima_
f(*minima_)
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1,
edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.plot(*minima_, f(*minima_), 'r*', markersize=10)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z = f(x,y)$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
plt.show()
dz_dx = elementwise_grad(f, argnum=0)(x, y)
dz_dy = elementwise_grad(f, argnum=1)(x, y)
fig, ax = plt.subplots(figsize=(15, 10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.quiver(x, y, x - dz_dx, y - dz_dy, alpha=.5)
ax.plot(*minima_, 'r*', markersize=18)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
plt.show()
path = julia_path
fig, ax = plt.subplots(figsize=(15, 10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.quiver(path[0,:-1], path[1,:-1], path[0,1:]-path[0,:-1], path[1,1:]-path[1,:-1], scale_units='xy', angles='xy', scale=1.1, color='r')
ax.plot(*minima_, 'k*', markersize=25)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
path = julia_path
plt.rcParams['lines.linewidth'] = 2
plt.rcParams['lines.color'] = 'r'
plt.rcParams['lines.markersize'] = 3
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.quiver(julia_path[0,:-1], julia_path[1,:-1], f(*julia_path[::,:-1]),
julia_path[0,1:]-julia_path[0,:-1], julia_path[1,1:]-julia_path[1,:-1], f(*(julia_path[::,1:]-julia_path[::,:-1])),
color='k', length = 0.2)
ax.plot(*minima_, f(*minima_), 'r*', markersize=10)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
plt.rcParams['lines.linewidth'] = 1
fig, ax = plt.subplots(figsize=(15, 10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.plot(*minima_, 'r*', markersize=18)
line, = ax.plot([], [], 'b', label='L-BFGS', lw=2)
point, = ax.plot([], [], 'bo')
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
ax.legend(loc='upper left', prop={'size': 20})
def init():
line.set_data([], [])
point.set_data([], [])
return line, point
def animate(i):
line.set_data(*path[::,:i])
point.set_data(*path[::,i-1:i])
return line, point
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=path.shape[1], interval=interval,
repeat_delay=5, blit=True)
HTML(anim.to_html5_video())
plt.rcParams['lines.linewidth'] = 2.5
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.plot(*minima_, f(*minima_), 'r*', markersize=20)
line, = ax.plot([], [], [], 'b', label='Newton-CG', lw=2)
point, = ax.plot([], [], [], 'bo')
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
def init():
line.set_data([], [])
line.set_3d_properties([])
point.set_data([], [])
point.set_3d_properties([])
return line, point
def animate(i):
line.set_data(path[0,:i], path[1,:i])
line.set_3d_properties(f(*path[::,:i]))
point.set_data(path[0,i-1:i], path[1,i-1:i])
point.set_3d_properties(f(*path[::,i-1:i]))
return line, point
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=path.shape[1], interval=interval,
repeat_delay=5, blit=True)
HTML(anim.to_html5_video())
func = value_and_grad(lambda args: f(*args))
def make_minimize_cb(path=[]):
def minimize_cb(xk):
# note that we make a deep copy of xk
path.append(np.copy(xk))
return minimize_cb
path_ = [x0]
res = minimize(func, x0=x0, method='Newton-CG',
jac=True, tol=1e-20, callback=make_minimize_cb(path_))
dict(res)
path = np.array(path_).T
path.shape
fig, ax = plt.subplots(figsize=(15,10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.quiver(path[0,:-1], path[1,:-1], path[0,1:]-path[0,:-1], path[1,1:]-path[1,:-1], scale_units='xy', angles='xy', scale=1, color='k')
ax.plot(*minima_, 'r*', markersize=18)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.quiver(path[0,:-1], path[1,:-1], f(*path[::,:-1]),
path[0,1:]-path[0,:-1], path[1,1:]-path[1,:-1], f(*(path[::,1:]-path[::,:-1])),
color='k', length = 0.2)
ax.plot(*minima_, f(*minima_), 'r*', markersize=10)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
fig, ax = plt.subplots(figsize=(15, 10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.plot(*minima_, 'r*', markersize=18)
line, = ax.plot([], [], 'b', label='Newton-CG', lw=2)
point, = ax.plot([], [], 'bo')
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
ax.legend(loc='upper left')
def init():
line.set_data([], [])
point.set_data([], [])
return line, point
def animate(i):
line.set_data(*path[::,:i])
point.set_data(*path[::,i-1:i])
return line, point
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=path.shape[1], interval=interval,
repeat_delay=5, blit=True)
HTML(anim.to_html5_video())
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.plot(*minima_, f(*minima_), 'r*', markersize=10)
line, = ax.plot([], [], [], 'b', label='Newton-CG', lw=2)
point, = ax.plot([], [], [], 'bo')
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
def init():
line.set_data([], [])
line.set_3d_properties([])
point.set_data([], [])
point.set_3d_properties([])
return line, point
def animate(i):
line.set_data(path[0,:i], path[1,:i])
line.set_3d_properties(f(*path[::,:i]))
point.set_data(path[0,i-1:i], path[1,i-1:i])
point.set_3d_properties(f(*path[::,i-1:i]))
return line, point
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=path.shape[1], interval=interval,
repeat_delay=5, blit=True)
HTML(anim.to_html5_video())
class TrajectoryAnimation(animation.FuncAnimation):
def __init__(self, *paths, labels=[], fig=None, ax=None, frames=None,
interval=interval, repeat_delay=5, blit=True, **kwargs):
if fig is None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
else:
if ax is None:
ax = fig.gca()
self.fig = fig
self.ax = ax
self.paths = paths
if frames is None:
frames = max(path.shape[1] for path in paths)
self.lines = [ax.plot([], [], label=label, lw=2)[0]
for _, label in zip_longest(paths, labels)]
self.points = [ax.plot([], [], 'o', color=line.get_color())[0]
for line in self.lines]
super(TrajectoryAnimation, self).__init__(fig, self.animate, init_func=self.init_anim,
frames=frames, interval=interval, blit=blit,
repeat_delay=repeat_delay, **kwargs)
def init_anim(self):
for line, point in zip(self.lines, self.points):
line.set_data([], [])
point.set_data([], [])
return self.lines + self.points
def animate(self, i):
for line, point, path in zip(self.lines, self.points, self.paths):
line.set_data(*path[::,:i])
point.set_data(*path[::,i-1:i])
return self.lines + self.points
class TrajectoryAnimation3D(animation.FuncAnimation):
def __init__(self, *paths, zpaths, labels=[], fig=None, ax=None, frames=None,
interval=interval, repeat_delay=5, blit=True, **kwargs):
if fig is None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
else:
if ax is None:
ax = fig.gca()
self.fig = fig
self.ax = ax
self.paths = paths
self.zpaths = zpaths
if frames is None:
frames = max(path.shape[1] for path in paths)
self.lines = [ax.plot([], [], [], label=label, lw=2)[0]
for _, label in zip_longest(paths, labels)]
super(TrajectoryAnimation3D, self).__init__(fig, self.animate, init_func=self.init_anim,
frames=frames, interval=interval, blit=blit,
repeat_delay=repeat_delay, **kwargs)
def init_anim(self):
for line in self.lines:
line.set_data([], [])
line.set_3d_properties([])
return self.lines
def animate(self, i):
for line, path, zpath in zip(self.lines, self.paths, self.zpaths):
line.set_data(*path[::,:i])
line.set_3d_properties(zpath[:i])
return self.lines
methods = [
"CG",
# "BFGS",
"Newton-CG",
"L-BFGS-B",
"TNC",
"SLSQP",
# "dogleg",
# "trust-ncg"
]
minimize_ = partial(minimize, fun=func, x0=x0, jac=True, bounds=[(xmin, xmax), (ymin, ymax)], tol=1e-20)
paths_ = defaultdict(list)
for method in methods:
paths_[method].append(x0)
results = {method: minimize_(method=method, callback=make_minimize_cb(paths_[method])) for method in methods}
paths = [np.array(paths_[method]).T for method in methods]
#append our Julia data
paths.append(julia_path)
methods.append("Our L-BFGS implementation")
# print(methods)
zpaths = [f(*path) for path in paths]
fig, ax = plt.subplots(figsize=(15, 10))
ax.contour(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.jet)
ax.plot(*minima_, 'r*', markersize=10)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
anim = TrajectoryAnimation(*paths, labels=methods, ax=ax)
ax.legend(loc='upper left')
HTML(anim.to_html5_video())
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection='3d', elev=50, azim=50)
ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, edgecolor='none', alpha=.8, cmap=plt.cm.jet)
ax.plot(*minima_, f(*minima_), 'r*', markersize=10)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$z$')
ax.set_xlim((xmin, xmax))
ax.set_ylim((ymin, ymax))
anim = TrajectoryAnimation3D(*paths, zpaths=zpaths, labels=methods, ax=ax)
ax.legend(loc='upper left')
HTML(anim.to_html5_video())