add mpi support

Author: Hongwei

vmc_jastrow_cp_exact_mpi.py

@@ -0,0 +1,135 @@
+import numpy as np
+from mpi4py import MPI
+import math
+import time
+import numba
+from matplotlib import pyplot
+
+N = 4
+
+@numba.njit
+def coefficient(state, alpha):
+
+	ssum = 0.0
+	for i in range(N):
+		for j in range(i+1, N):
+
+			deno = min(math.fabs(1.0*j - 1.0*i), N*1.0 - math.fabs(1.0*j - 1.0*i) )
+			# print(state[i] * state[j] / deno)
+			ssum += state[i] * state[j] / deno
+
+	return math.exp(-alpha * ssum) 
+
+@numba.njit
+def local_energy(state, coeff, alpha):
+
+	res = 0.0
+	ssum = 0.0 
+
+	for i in range(N):
+		res += state[i] * state[(i+1)%N]
+
+	for i in range(N):
+		if(state[i] * state[(i+1)%N] < 0.0):			
+			state_new = state.copy()
+			# print(state_new)
+			state_new[i] *= -1.0
+			state_new[(i+1)%N] *= -1.0
+
+			ssum += coefficient(state_new, alpha)/coeff
+
+	return res - 0.5 * ssum
+
+
+@numba.njit
+def sampler(alpha, Nsample = 5000, Nskip = 3):
+
+	state = np.ones(N) 
+	state[: N//2] = -1
+
+	state *= 0.5
+	state = state[np.random.permutation(N)]
+
+	ssum = 0.0
+	# coeff_old = coefficient(state, alpha)
+
+	for i in range(Nsample):
+
+		for i in range(Nskip):
+
+			x = np.random.randint(low = 0, high = N)
+			y = x
+
+			while(state[y] * state[x] > 0):
+				y = np.random.randint(low = 0, high = N)
+
+			new_state = state.copy()
+			new_state[x] *= -1.0
+			new_state[y] *= -1.0
+
+			coeff_old = coefficient(state, alpha)
+			coeff_new =	coefficient(new_state, alpha)
+
+			if(np.random.random() < min(1.0, (coeff_new**2)/(coeff_old**2))):
+				state = new_state.copy()
+				coeff_old = coeff_new
+
+		tmp = local_energy(state, coeff_old, alpha)
+		
+
+		ssum += tmp
+
+	return ssum / Nsample 
+
+
+if(__name__ == '__main__'):
+
+	comm = MPI.COMM_WORLD
+	nprocs = comm.Get_size()
+	rank = comm.Get_rank()
+
+	ns = 10000
+	ns = ns // nprocs
+
+
+	if(rank == 0):
+		x, y = [], []
+
+
+	t0 = time.time()
+
+	for i in range(-30, 40):
+
+		alpha = i * 0.1
+
+		# comm.Barrier()
+		mpi_energy = sampler(alpha, ns) / nprocs
+
+		energy = comm.reduce(mpi_energy, root=0)
+
+		if(rank == 0):
+			print("Alpha: %.2f, Energy: %.2f" % (alpha, energy))
+			x.append(alpha)
+			y.append(energy)
+
+	if(rank == 0):
+
+		t1 = time.time()
+		print("Elapsed time: %.2f sec" % (t1 - t0))
+
+		pyplot.xlabel("alpha")
+		pyplot.ylabel("Energy")
+		pyplot.plot(x, y, 'o', label="VMC")
+		pyplot.legend()
+		pyplot.show()
+
+
+
+
+
+
+
+
+
+
+