Numba implementation for Monte-Carlo simulation

Question

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

Answer 1

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