Numba translation of simple Langevin Integrator originally implemented in Cython: https://bitbucket.org/joshua.adelman/stringmethodexamples/src/036ed85f02cd430120bdacc38693f2fbbd6cf025/shared/pylangevin-integrator/cIntegratorSimple.pyx https://bitbucket.org/joshua.adelman/stringmethodexamples/src/036ed85f02cd430120bdacc38693f2fbbd6cf025/shared/p…

numba_langevin.py

from numba import jit, random, int_, float32, void, object_, double, uint32, ulong
import math
import numpy as np

float32_2d = float32[:,::1]


@jit
class Force(object):
    @int_(float32[::1], float32[::1])
    def evaluate(self, coords, force):
        return 0


class Dickson2dRingForce(Force):
    def evaluate(self, coords, force):
        alpha = 3.0
        gamma = 3.0
        chi1 = 2.25
        chi2 = 4.5

        fext = 7.2

        x = coords[0]
        y = coords[1]

        x2y2 = x**2 + y**2
        r = np.sqrt(x2y2)
        invr = 1.0/r
        A = 2.0*alpha*(1.0 - gamma*invr)

        att = np.math.atan2(y,x)
        B1 = chi1*np.sin(2.0*att)*invr*invr 
        B2 = chi2*np.sin(4.0*att)*invr*invr

        fx = x*A - 2.0*y*B1 - 4.0*y*B2
        fy = y*A + 2.0*x*B1 + 4.0*x*B2

        force[0] = -(fx - fext*invr*np.sin(att))
        force[1] = -(fy + fext*invr*np.cos(att))

        return 0


@jit
class SimpleIntegrator(object):

    @void(object_, double, double, double, double, uint32, int_[::1], float32[::1], ulong)
    def __init__(self, forcefield, mass, xi, beta, dt, dims, isperiodic, boxsize, seed):
        self._beta = beta
        self._mass = mass
        self._xi = xi
        self._dt = dt
        self._dims = dims

        self._ff = forcefield
        
        self._sigma = math.sqrt(2.0*dt/(mass*beta*xi))
        self._fp = dt/(mass*xi)
    
        self.boxsize = np.zeros((self._dims,), dtype=np.float32)
        self.isperiodic = np.zeros((self._dims,), dtype=np.int_)

        for k in xrange(self._dims):
            self.boxsize[k] = boxsize[k]
            self.isperiodic[k] = isperiodic[k]

        self.state = random.state_p

    @float32_2d(float32[::1], int_, int_)
    def step_save(self, coords, steps, save_freq):
        F = np.zeros((self._dims,), dtype=np.float32)
        traj = np.zeros((steps/save_freq, self._dims), dtype=np.float32)

        c = 0

        for k in xrange(steps):
            self._ff.evaluate(coords, F)

            for d in xrange(self._dims):
                eta = random.rk_gauss(self.state)
                coords[d] = coords[d] + self._fp*F[d] + self._sigma*eta

                # Wrap the coordinates if periodic in dimension
                if self.isperiodic[d] == 1:
                    coords[d] = coords[d] - math.floor(coords[d]/self.boxsize[d])*self.boxsize[d]

            if k % save_freq == 0:
                for d in xrange(self._dims):
                    traj[c,d] = coords[d]
                c += 1

        return traj

simulate.py

import numpy as np
import os
import time

from numba_langevin import SimpleIntegrator, Dickson2dRingForce


def genrandint():
    """Generates a random integer between 0 and (2^32)-1"""
    x = 0
    for i in range(4):
        x = (x << 8) + ord(os.urandom(1))

    return x


ff = Dickson2dRingForce()
MASS = 1.0
XI = 1.5
BETA = 1.0
NDIMS = 2
DT = 0.005
ISPERIODIC = np.array([0, 0], dtype=np.int)
BOXSIZE = np.array([1.0E8, 1.0E8], dtype=np.float32)

NUM_BLOCKS = int(1e5)
STEPS_PER_BLOCK = 20
BLOCKS_PER_DUMP = 5000

totblocks = NUM_BLOCKS//BLOCKS_PER_DUMP

integrator = SimpleIntegrator(ff, MASS, XI, BETA, DT, NDIMS, ISPERIODIC, BOXSIZE, genrandint())

# Initial coords
x = np.array([-3.0,0.0], dtype=np.float32)

results = np.empty((totblocks*BLOCKS_PER_DUMP, NDIMS), dtype=np.float32)

for dk in xrange(totblocks):
    t1 = time.time()
    ctemp = integrator.step_save(x,BLOCKS_PER_DUMP*STEPS_PER_BLOCK,STEPS_PER_BLOCK)
    x = ctemp[-1,:]

    results[dk*BLOCKS_PER_DUMP:(dk+1)*BLOCKS_PER_DUMP,:] = ctemp

    print 'Block took: {} s'.format(time.time() - t1)

np.save('traj.npy', results)