I'm actually writing a Monte-Carlo code for simple fluid simulation. The code was predicting good results for energy and pressure before Numba implementation, and after its implementation it doesn't predicts good result at all. Since I'm not very familiar with Numba, I don't really know where the error is comming from. Maybe you will have an idea...
Thanks a lot.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Feb 4 13:32:54 2024
@author: tristanmillet
"""
import numpy as np
from numba import njit
# Constants
N = 100
density = 0.8
T = 2.0
beta = 1 / T
L = (N / density) ** (1.0 / 3.0)
cutoff = 2.5
boxVolume = N / density
numEqSteps = 100000
numSampSteps = 100000
numTotalSteps = numEqSteps + numSampSteps
progressFreq = int(numTotalSteps * 0.01)
energy_instant_values = np.empty(numTotalSteps, dtype=np.float64)
virial_instant_values = np.empty(numTotalSteps, dtype=np.float64)
@njit
def latticeDisplace():
positions = np.empty((N, 3), dtype=np.float64)
delta = L / (N ** (1 / 3))
flag = 0
for x in np.arange(delta / 2, L, delta):
for y in np.arange(delta / 2, L, delta):
for z in np.arange(delta / 2, L, delta):
if flag < N:
positions[flag] = [x, y, z]
flag += 1
return positions
@njit
def calc_dist(i, j):
p1 = positions[i]
p2 = positions[j]
dx = p1[0] - p2[0]
dy = p1[1] - p2[1]
dz = p1[2] - p2[2]
if dx > L / 2:
dx -= L
if dx < -L / 2:
dx += L
if dy > L / 2:
dy -= L
if dy < -L / 2:
dy += L
if dz > L / 2:
dz -= L
if dz < -L / 2:
dz += L
return np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
@njit
def calc_LJ_potential(dist):
pot = 4.0 * ((1.0 / dist) ** 12.0 - (1.0 / dist) ** 6.0)
return pot
@njit
def calc_energy_particle(p):
energy_particle = 0.0
for j in range(N):
if j != p:
dist = calc_dist(p, j)
if dist <= cutoff:
energy_particle += calc_LJ_potential(dist)
return energy_particle
@njit
def calc_energy_total():
energy_total = 0.0
for i in range(N):
energy_total += calc_energy_particle(i)
energy_total *= 0.5
# Tail correction
energy_tail_corr = (8.0 / 3.0) * np.pi * density * (1.0 ** 3) * (
(1.0 / 3.0) * ((1.0 / cutoff) ** 9) - ((1.0 / cutoff) ** 3))
energy_total += N * energy_tail_corr
return energy_total
@njit
def calc_virial_particle(p):
virial_particle = 0.0
for j in range(N):
if j != p:
dist = calc_dist(p, j)
if dist <= cutoff:
virial_particle += calc_virial(dist)
return virial_particle
@njit
def calc_virial(dist):
return 48.0 * ((1 / dist) ** 12 - 0.5 * (1 / dist) ** 6)
@njit
def calc_virial_total():
virial_total = 0.0
for i in range(N):
virial_total += calc_virial_particle(i)
return 0.5 * virial_total
@njit
def move_particle(p, displ=0.5):
local_ = positions[p].copy()
for i in range(3):
local_[i] += (np.random.rand() - 0.5) * displ
if local_[i] >= L:
local_[i] -= L
if local_[i] < 0.0:
local_[i] += L
return local_
positions = latticeDisplace()
energy = calc_energy_total()
virial = calc_virial_total()
accept_counter=0
energy_sum = 0.0
virial_sum = 0.0
print(f"Parameters used for simulation: T={T},rho={density}, N={N}")
for step in range(numTotalSteps):
particle_index = np.random.randint(0, N)
prev_particle_energy = calc_energy_particle(particle_index)
prev_particle_virial = calc_virial_particle(particle_index)
prev_particle = np.array(positions[particle_index])
positions[particle_index] = move_particle(particle_index)
new_particle_energy = calc_energy_particle(particle_index)
delta_particle_energy = new_particle_energy - prev_particle_energy
if (delta_particle_energy < 0) or (np.random.rand() < np.exp(-beta * delta_particle_energy)):
energy += delta_particle_energy
new_particle_virial = calc_virial_particle(particle_index)
virial += new_particle_virial - prev_particle_virial
accept_counter += 1
else:
# Restore old configuration
positions[particle_index] = prev_particle
virial_instant_values[step] = virial
energy_instant_values[step] = energy
energy_sum += energy
virial_sum += virial
# Reset sums and counter for sampling
if step == numEqSteps:
energy_sum = 0.0
virial_sum = 0.0
accept_counter = 0
if step % progressFreq == 0:
print(accept_counter)
accept_counter = 0
print(str(int((step * 1.0 / numTotalSteps) * 100)) + "% " + (
"[Equilibration]" if step < numEqSteps else "[Sampling]"))
avgEnergy = energy_sum / numSampSteps / N
pressure_tail_corr = (16.0 / 3.0) * np.pi * (density ** 2) * (1 ** 3) * ((2.0 / 3.0) * ((1 / cutoff) ** 9) - ((1 / cutoff) ** 3))
pressure = (virial_sum / 3.0 / numSampSteps / boxVolume) + density * T + pressure_tail_corr
finalAcceptRate = accept_counter * 1.0 / numSampSteps * 100.0
np.savetxt("instant_energy.txt",energy_instant_values)
np.savetxt("instant_virial.txt",virial_instant_values)
print("Avg. energy = " + str(avgEnergy))
print("Avg. pressure = " + str(pressure))
print("Accept. rate = " + str(finalAcceptRate))
I tried to debug briefly the code by "print" some part of the code, and apparently the Metropolis criteria (which determines if we move or not a particle) is always fulfilled, which is not the case when I don't implement Numba. Here the difference of the result between Numba implementation and without Numba:
without Numba (good prediction)
Avg. energy = -4.789510303401584
Avg. pressure = 5.078151508549905
Accept. rate = 0.469
Density = 0.8
T=2
with Numba:
Avg. energy = 114.5028952818473
Avg. pressure = 407.77783626468135
Accept. rate = 1.999
Density = 0.8
T=2
Don't use global variables. I changed the jitted function to accept positions as first argument:
import numpy as np
from numba import njit
# Constants
N = 100
density = 0.8
T = 2.0
beta = 1 / T
L = (N / density) ** (1.0 / 3.0)
cutoff = 2.5
boxVolume = N / density
numEqSteps = 100000
numSampSteps = 100000
numTotalSteps = numEqSteps + numSampSteps
progressFreq = int(numTotalSteps * 0.01)
energy_instant_values = np.empty(numTotalSteps, dtype=np.float64)
virial_instant_values = np.empty(numTotalSteps, dtype=np.float64)
@njit
def latticeDisplace():
positions = np.empty((N, 3), dtype=np.float64)
delta = L / (N ** (1 / 3))
flag = 0
for x in np.arange(delta / 2, L, delta):
for y in np.arange(delta / 2, L, delta):
for z in np.arange(delta / 2, L, delta):
if flag < N:
positions[flag] = [x, y, z]
flag += 1
return positions
@njit
def calc_dist(positions, i, j):
p1 = positions[i]
p2 = positions[j]
dx = p1[0] - p2[0]
dy = p1[1] - p2[1]
dz = p1[2] - p2[2]
if dx > L / 2:
dx -= L
if dx < -L / 2:
dx += L
if dy > L / 2:
dy -= L
if dy < -L / 2:
dy += L
if dz > L / 2:
dz -= L
if dz < -L / 2:
dz += L
return np.sqrt(dx**2 + dy**2 + dz**2)
@njit
def calc_LJ_potential(dist):
pot = 4.0 * ((1.0 / dist) ** 12.0 - (1.0 / dist) ** 6.0)
return pot
@njit
def calc_energy_particle(positions, p):
energy_particle = 0.0
for j in range(N):
if j != p:
dist = calc_dist(positions, p, j)
if dist <= cutoff:
energy_particle += calc_LJ_potential(dist)
return energy_particle
@njit
def calc_energy_total(positions):
energy_total = 0.0
for i in range(N):
energy_total += calc_energy_particle(positions, i)
energy_total *= 0.5
# Tail correction
energy_tail_corr = (
(8.0 / 3.0)
* np.pi
* density
* (1.0**3)
* ((1.0 / 3.0) * ((1.0 / cutoff) ** 9) - ((1.0 / cutoff) ** 3))
)
energy_total += N * energy_tail_corr
return energy_total
@njit
def calc_virial_particle(positions, p):
virial_particle = 0.0
for j in range(N):
if j != p:
dist = calc_dist(positions, p, j)
if dist <= cutoff:
virial_particle += calc_virial(dist)
return virial_particle
@njit
def calc_virial(dist):
return 48.0 * ((1 / dist) ** 12 - 0.5 * (1 / dist) ** 6)
@njit
def calc_virial_total(positions):
virial_total = 0.0
for i in range(N):
virial_total += calc_virial_particle(positions, i)
return 0.5 * virial_total
@njit
def move_particle(positions, p, displ=0.5):
local_ = positions[p].copy()
for i in range(3):
local_[i] += (np.random.rand() - 0.5) * displ
if local_[i] >= L:
local_[i] -= L
if local_[i] < 0.0:
local_[i] += L
return local_
positions = latticeDisplace()
energy = calc_energy_total(positions)
virial = calc_virial_total(positions)
accept_counter = 0
energy_sum = 0.0
virial_sum = 0.0
print(f"Parameters used for simulation: T={T},rho={density}, N={N}")
for step in range(numTotalSteps):
particle_index = np.random.randint(0, N)
prev_particle_energy = calc_energy_particle(positions, particle_index)
prev_particle_virial = calc_virial_particle(positions, particle_index)
prev_particle = np.array(positions[particle_index])
positions[particle_index] = move_particle(positions, particle_index)
new_particle_energy = calc_energy_particle(positions, particle_index)
delta_particle_energy = new_particle_energy - prev_particle_energy
if (delta_particle_energy < 0) or (
np.random.rand() < np.exp(-beta * delta_particle_energy)
):
energy += delta_particle_energy
new_particle_virial = calc_virial_particle(positions, particle_index)
virial += new_particle_virial - prev_particle_virial
accept_counter += 1
else:
# Restore old configuration
positions[particle_index] = prev_particle
virial_instant_values[step] = virial
energy_instant_values[step] = energy
energy_sum += energy
virial_sum += virial
# Reset sums and counter for sampling
if step == numEqSteps:
energy_sum = 0.0
virial_sum = 0.0
accept_counter = 0
if step % progressFreq == 0:
print(accept_counter)
accept_counter = 0
print(
str(int((step * 1.0 / numTotalSteps) * 100))
+ "% "
+ ("[Equilibration]" if step < numEqSteps else "[Sampling]")
)
avgEnergy = energy_sum / numSampSteps / N
pressure_tail_corr = (
(16.0 / 3.0)
* np.pi
* (density**2)
* (1**3)
* ((2.0 / 3.0) * ((1 / cutoff) ** 9) - ((1 / cutoff) ** 3))
)
pressure = (
(virial_sum / 3.0 / numSampSteps / boxVolume) + density * T + pressure_tail_corr
)
finalAcceptRate = accept_counter * 1.0 / numSampSteps * 100.0
# np.savetxt("instant_energy.txt", energy_instant_values)
# np.savetxt("instant_virial.txt", virial_instant_values)
print("Avg. energy = " + str(avgEnergy))
print("Avg. pressure = " + str(pressure))
print("Accept. rate = " + str(finalAcceptRate))
Prints:
Avg. energy = -4.771041235079697
Avg. pressure = 5.186359541491075
Accept. rate = 0.47800000000000004