Monte carlo simulation, re /r/python thread on numba, cython, etc

mc.pyx

# cython: infer_types=True
# cython: boundscheck=False
# cython: cdvision=True
# distutils: compile_options = ["-O2", "-fopenmp", "-march=native"]
# distutils: link_options = ["-fopenmp"]
cimport cython
from numpy import random as rng
import numpy as np
import numpy.random

from libc cimport math
from libc.stdlib cimport malloc, free
cimport cython.parallel

cdef struct CoordC:
    double x
    double y
    double z

@cython.cdivision(True)
def montecarlo_loop(int ncycle, double[:, :] coords, double beta, double E_T,
                    double boxl, double boxl2):
    numpy.random.seed(0)
    npart = coords.shape[0]
    c_coords = <CoordC*>malloc(coords.shape[0] * sizeof(CoordC))
    for i in range(coords.shape[0]):
        c_coords[i].x = coords[i, 0]
        c_coords[i].y = coords[i, 1]
        c_coords[i].z = coords[i, 2]
    e_diff = 0.0
    nmove = 0
    cdef CoordC disp
    for indx in range(0, ncycle+1):
        nmove = <int>math.floor(npart * rng.random())
        disp.x = get_disp_coord(c_coords[nmove].x, boxl)
        disp.y = get_disp_coord(c_coords[nmove].y, boxl)
        disp.z = get_disp_coord(c_coords[nmove].z, boxl)

        e_diff = shift_ecalc(npart, c_coords, nmove, disp, boxl, boxl2)
        if e_diff < 0.0:
            E_T += e_diff
            c_coords[nmove].x += disp.x
            c_coords[nmove].y += disp.y
            c_coords[nmove].z += disp.z
        else:
            if math.exp(-beta*e_diff) > rng.random():
                E_T += e_diff
                c_coords[nmove].x += disp.x
                c_coords[nmove].y += disp.y
                c_coords[nmove].z += disp.z
    for i in range(coords.shape[0]):
        coords[i, 0] = c_coords[i].x
        coords[i, 1] = c_coords[i].y 
        coords[i, 2] = c_coords[i].z
    print("Culmaltive Energy: ", E_T)
    E_Final = detailed_ecalc(coords.shape[0], coords, boxl, boxl2)
    print("Final Energy: ", E_Final)
    return coords


cdef double get_disp_coord(double coord, double boxl):
    cdef double x = 0.2 * (2.0 * rng.random() - 1.0)
    if math.fabs(x + coord) > boxl:
        x = x - math.copysign(boxl, x)
    return x


@cython.cdivision(True)
cdef double shift_ecalc(int npart, const CoordC* coords, int nmove,
        CoordC disp, double boxl, double boxl2) nogil:
    cdef double energy = 0.0
    cdef int j
    for j in cython.parallel.prange(npart, nogil=True):
        if j != nmove:
            energy += _shift_ecalc_inner(&coords[nmove], &coords[j], &disp, boxl, boxl2)
    return energy


@cython.cdivision(True)
cdef double _shift_ecalc_inner(const CoordC* nmove, const CoordC* coord_j,
    const CoordC* disp, double boxl, double boxl2) nogil:
    cdef double energy = 0
    rx = nmove.x + disp.x - coord_j.x
    ry = nmove.y + disp.y - coord_j.y
    rz = nmove.z + disp.z - coord_j.z
    if math.fabs(rx) > boxl2:
        rx -= math.copysign(boxl, rx)
    if math.fabs(ry) > boxl2:
        ry -= math.copysign(boxl, ry)
    if math.fabs(rz) > boxl2:
        rz -= math.copysign(boxl, rz)

    r_sq = rx*rx + ry*ry + rz*rz
    if r_sq < 49.0:
        LJ = (1.0/r_sq)
        LJ = LJ*LJ*LJ
        LJ = 4.0 * LJ * (LJ-1.0)
        energy += LJ

    rx = nmove.x - coord_j.x
    ry = nmove.y - coord_j.y
    rz = nmove.z - coord_j.z
    if math.fabs(rx) > boxl2:
        rx -= math.copysign(boxl, rx)
    if math.fabs(ry) > boxl2:
        ry -= math.copysign(boxl, ry)
    if math.fabs(rz) > boxl2:
        rz -= math.copysign(boxl, rz)
    r_sq = rx*rx + ry*ry + rz*rz
    if r_sq < 49.0:
        LJ = (1.0/r_sq)
        LJ = LJ*LJ*LJ
        LJ = 4.0*LJ*(LJ-1.0)
        energy -= LJ
    return energy


def detailed_ecalc(npart,coords, boxl, boxl2):
    energy = 0.0
    for i in range(0,npart-1):
        for j in range(i+1,npart):
            rx = coords[i][0] - coords[j][0]
            ry = coords[i][1] - coords[j][1]
            rz = coords[i][2] - coords[j][2]
            if math.fabs(rx) > boxl2:
                rx -= math.copysign(boxl,rx)
            if math.fabs(ry) > boxl2:
                ry -= math.copysign(boxl,ry)
            if math.fabs(rz) > boxl2:
                rz -= math.copysign(boxl,rz)

            r_sq = rx*rx+ ry*ry + rz*rz
            if r_sq < 49.0:
                LJ = (1.0/r_sq)
                LJ = LJ*LJ*LJ
                LJ = 4.0*LJ*(LJ-1.0)
                energy = energy + LJ
    return energy