vis_sgd.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

N = 100
dt = 1/30 # 30fps
iterations = 600
bounds = [-2, 2, -2, 2]
init_state = [-0.5,-0.5]

class ParticleBox:
	def __init__(self):
		self.state = init_state.copy()

	def step(self, dt):  
		# update positions
		### dt gives it a very small step size
		self.state -= dt * np.array([2*self.state[0],2*self.state[1]])

np.random.seed(0)

box = ParticleBox()

fig = plt.figure()
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
					 xlim=(-3.2, 3.2), ylim=(-2.4, 2.4))

# particles holds the locations of the particles
particles, = ax.plot([], [], 'bo', ms=6)

def f(x,y): # Function to be minimized
	return x*x + y*y

X,Y = np.meshgrid(
	np.linspace(bounds[0],bounds[1],N),
	np.linspace(bounds[2],bounds[3],N))
	
Z = np.zeros(X.shape)
for i in range(N):
	for j in range(N):
		Z[i,j] = f(X[i,j], Y[i,j])
		
ax.pcolor(X,Y,Z)

# rect is the box edge
rect = plt.Rectangle(bounds[::2],
					 bounds[1] - bounds[0],
					 bounds[3] - bounds[2],
					 ec='none', lw=2, fc='none')
ax.add_patch(rect)

def init():
	"""initialize animation"""
	global box, rect
	particles.set_data([], [])
	return particles, rect

def animate(i):
	"""perform animation step"""
	global box, rect, dt
	box.step(dt)
	
	# update pieces of the animation
	particles.set_data(box.state[0], box.state[1])
	return particles, rect

ani = animation.FuncAnimation(fig, animate, frames=iterations,
							  interval=10, blit=True, init_func=init)

plt.show()