import matplotlib.pyplot as plt
from numpy import *

N = 10									# Number of spins
T = 10*N								# Number of iterations
set_alpha = array([0.1, 0.01, 0.001])	# Cooling rates

# Initialize spins
S_initial = arange(N)
for i in range(N):
	if random.random() > 0.5:
		S_initial[i] = 1
	else:
		S_initial[i] = -1
# Initialize Gaussian weights
J = arange(N)
for i in range(N):
	J[i] = random.randn()

# Compute minimum energy
E_min = -1*sqrt(J*J).sum()

# Compute the global energy of the chain
def global_energy(S, J):
	S_shift = hstack([S[1:], S[0]])
	energy = -1*dot(J*S, S_shift)
	return energy

# Compute the local energy of the chain
def local_energy(S, J, flip_idx):
	energy_left = J[mod(flip_idx-1,N)]*S[mod(flip_idx-1,N)]*S[flip_idx]
	energy_rght = J[flip_idx]*S[flip_idx]*S[mod(flip_idx+1,N)]
	return energy_left + energy_rght

# Compute initial energy
E_initial = global_energy(S_initial, J)

# Initialize output
output = 0*arange(T*set_alpha.size)
output.resize(T, set_alpha.size)

# Simulate different cooling rates
for i in arange(set_alpha.size):
	# Define cooling rate
	alpha = set_alpha[i]
	# Initialize spins and energy
	S = S_initial
	E = E_initial
	
	# Anneal system
	for t in range(T):
		# Choose spin at random
		flip_idx = int(N*random.random())
		# Compute local energy
		Ei = local_energy(S, J, flip_idx)
		# Flip spin
		S[flip_idx] *= -1
		# Compute new local energy
		Ef = local_energy(S, J, flip_idx)
		# Compute change in energy
		Delta_E = Ef - Ei
		# If energy increases, perhaps accept flip
		if Delta_E > 0:
			# Compute probability to accept
			beta = alpha*t
			p = exp(-beta*Delta_E)
			# If flip not accepted, restore
			if random.random() > p:
				S[flip_idx] *= -1
		# Update and store energy
		E += Delta_E
		output[t, i] = E

# Plot output
fig = plt.figure()
ax = fig.gca()
X = arange(T)
ax.plot(X, output)
flr = 0*arange(T) + E_min
ax.plot(X, flr, 'k--')

# Render plot
plt.show()