Search code examples
pythonperformancenumba

How do I use numba for the following code?


I am trying to use Numba for the first time in colab and I think I have successfully installed Numba since the @jit is not undefined now but I am getting errors in my code. The following is my attempt:

!apt-get install nvidia-cuda-toolkit
!pip3 install numba
import os
dev_lib_path = !find / -iname 'libdevice'
nvvm_lib_path = !find / -iname 'libnvvm.so'
assert len(dev_lib_path)>0, "Device Lib Missing"
assert len(nvvm_lib_path)>0, "NVVM Missing"
os.environ['NUMBAPRO_LIBDEVICE'] = dev_lib_path[0]
os.environ['NUMBAPRO_NVVM'] = nvvm_lib_path[0]

import numpy as np
import matplotlib.pyplot as plt
import random
from math import *
from random import *
from numba import jit
n=1000
mu=np.random.uniform(0,1,n)
r=[sqrt(-2*log(1-i)) for i in mu]
eta=np.random.uniform(0,1,n)
theta=2*pi*eta;
cuz=[cos(i) for i in theta]
suz=[sin(i) for i in theta]
Zinitial=[a*b for a,b in zip(r,cuz)];
Pinitial=[a*b for a,b in zip(r,suz)];

class Particle:
    def __init__(self, pos, mom, spin):
        self.pos = pos
        self.mom = mom
        self.spin = spin
   
SP = sorted(np.array([Particle(pos = i, mom = j, spin = choice([1, 0])) for i,j in zip(Zinitial,Pinitial)]),key=lambda x:x.pos)
Upi=[];
Downi=[];
count_plot=[];
for j in range(len(SP)):
    if SP[j].spin == 1:
        Upi.append(SP[j].pos)
    else:
        Downi.append(SP[j].pos)
Zavgi=sum(Zinitial)/len(Zinitial)
Zreli=sum(Upi)/len(Upi)-sum(Downi)/len(Downi)
"Observables"


"Time"

iter=10**(5);
dt=1/(2*n);
alf=sqrt(n);

"Dynamics"




@jit(nopython=True)
def f(): 
  counter=0;
  sum1,sum2=0,0;
  Zavg=[Zavgi];
  Zrelm=[Zreli];
  T_plot=[0];
  for i in range(1,iter+1):
        
        t=i*dt;
        T_plot.append(t) 
        Z=[];
        Up=[];
        Down=[];
        c,s=cos(t),sin(t);
        c1,s1=cos(t-dt),sin(t-dt);
        for j in range(n-1):
            collchk=((c*(SP[j].pos)+s*(SP[j].mom))-(c*(SP[j+1].pos)+s*(SP[j+1].mom)))*(c1*(SP[j].pos)+s1*(SP[j].mom)-(c1*(SP[j+1].pos)+s1*(SP[j+1].mom)));

            prel=((c*(SP[j].mom)-s*(SP[j].pos))-(c*(SP[j+1].mom)-s*(SP[j+1].pos)))/2;
               
            rcoeff=1/(1+(prel*alf)**2);
            rand_value=random();
            
            
            if collchk<0:
               
              
               SP[j], SP[j+1]=SP[j+1],SP[j];
               
              
               if rcoeff>rand_value:
                   counter=counter+1
                   SP[j].spin,SP[j+1].spin=SP[j+1].spin,SP[j].spin;
            if SP[j].spin == 1:
                Up.append(c*(SP[j].pos)+s*(SP[j].mom))
            else:
                Down.append(c*(SP[j].pos)+s*(SP[j].mom))
            Z.append(c*(SP[j].pos)+s*(SP[j].mom))

        
        
        Zrel=sum(Up[0:])/len(Up) - sum(Down[0:])/len(Down);
        Zrelm.append(Zrel)
                        
        Zm=sum(Z)/len(Z)
        Zavg.append(Zm)
  return [Zavg, Zrelm, counter]   

Th error that I get with Vandan's code given below:

     Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'sum': cannot determine Numba type of <class 'builtin_function_or_method'>

File "<ipython-input-1-cddc88c01635>", line 52:
def f(SP, Zavgi, Zreli, alf, dt, n):
    <source elided>

        Zrel = sum(Up[0:]) / len(Up) - sum(Down[0:]) / len(Down);
        ^
        

In the end I want to plot the lists that I have returned. Any help would be appreciated and if there is a way to use Numba even for the class definiton that would be great.

EDIT:

import numpy as np
import matplotlib.pyplot as plt
import random
from math import *
from random import *
from numba import jit




"Dynamics"

@jit(nopython=True)
def f(SP, Zavgi, Zreli, alf, dt, n):
    "Time"
    counter = 0;
    sum1, sum2 = 0, 0;
    Zavg = np.array([Zavgi]);
    Zrelm = np.array([Zreli]);
    T_plot = [0];
    for i in range(1, iter + 1):

        t = i * dt;
        T_plot.append(t)
        Z = [];
        Up = [];
        Down = [];
        c, s = cos(t), sin(t);
        c1, s1 = cos(t - dt), sin(t - dt);
        for j in range(n - 1):
            collchk = ((c * (SP[j][0]) + s * (SP[j][1])) - (c * (SP[j + 1][0]) + s * (SP[j + 1][1]))) * (
                    c1 * (SP[j][0]) + s1 * (SP[j][1]) - (c1 * (SP[j + 1][0]) + s1 * (SP[j + 1][1])));

            prel = ((c * (SP[j][1]) - s * (SP[j][0])) - (c * (SP[j + 1][1]) - s * (SP[j + 1][0]))) / 2;

            rcoeff = 1 / (1 + (prel * alf) ** 2);
            rand_value = random();

            if collchk < 0:

                SP[j], SP[j + 1] = SP[j + 1], SP[j];

                if rcoeff > rand_value:
                    counter = counter + 1
                    SP[j][2], SP[j + 1][2] = SP[j + 1][2], SP[j][2];
            if SP[j][2] == 1:
                Up.append(c * (SP[j][0]) + s * (SP[j][1]))
            else:
                Down.append(c * (SP[j][0]) + s * (SP[j][1]))
            Z.append(c * (SP[j][0]) + s * (SP[j][1]))

        Zrel = sum(Up[0:]) / len(Up) - sum(Down[0:]) / len(Down);
        Zrelm = np.append(Zrelm, Zrel)

        Zm = sum(Z) / len(Z)
        Zavg = np.append(Zavg, Zm)
        
        
    return Zavg, Zrelm, counter,T_plot



if __name__ == '__main__':

    n = 1000
    mu = np.random.uniform(0, 1, n)
    r = [sqrt(-2 * log(1 - i)) for i in mu]
    eta = np.random.uniform(0, 1, n)
    theta = 2 * pi * eta;
    cuz = [cos(i) for i in theta]
    suz = [sin(i) for i in theta]
    Zinitial = [a * b for a, b in zip(r, cuz)];
    Pinitial = [a * b for a, b in zip(r, suz)];

    iter = 10 ** (6);
    dt = 1 / (100 * n);
    alf = sqrt(n);


    SP = np.array(sorted(np.array([  np.array([i,j,choice([0,1])]) for i, j in zip(Zinitial, Pinitial)]),
                key=lambda x: x[0]))
    Upi = [];
    Downi = [];
    count_plot = [];
    for j in range(len(SP)):
        if SP[j][2] == 1:
            Upi.append(SP[j][0])
        else:
            Downi.append(SP[j][0])
    Zavgi = sum(Zinitial) / len(Zinitial)
    Zreli = sum(Upi) / len(Upi) - sum(Downi) / len(Downi)


    Zavg, Zrelm, counter,T_plot = f(SP, Zavgi, Zreli, alf, dt, n)

    plt.plot(T_plot, Zrelm)

Solution

  • I modified your code a bit and was able to get the function to run. I removed the Particle class and changed all the list instances to numpy arrays.

    This is the output for the Zavg, Zrelm and counter:

    Zavg: [0.07047501 0.06735052 0.06728123 ... 0.39516435 0.3947497  0.39433495] 
    Zrelm: [-0.04179043 -0.04461464 -0.0394889  ... -0.11080628 -0.11087257
         -0.11093883] 
    Counter: 538
    

    Here is the modified code:

    import numpy as np
    import matplotlib.pyplot as plt
    import random
    from math import *
    from random import *
    from numba import jit
    
    
    "Dynamics"
    
    @jit(nopython=True)
    def f(SP, Zavgi, Zreli, alf, dt, n):
        "Time"
        counter = 0;
        sum1, sum2 = 0, 0;
        Zavg = np.array([Zavgi]);
        Zrelm = np.array([Zreli]);
        T_plot = [0];
        for i in range(1, iter + 1):
    
            t = i * dt;
            T_plot.append(t)
            Z = [];
            Up = [];
            Down = [];
            c, s = cos(t), sin(t);
            c1, s1 = cos(t - dt), sin(t - dt);
            for j in range(n - 1):
                collchk = ((c * (SP[j][0]) + s * (SP[j][1])) - (c * (SP[j + 1][0]) + s * (SP[j + 1][1]))) * (
                        c1 * (SP[j][0]) + s1 * (SP[j][1]) - (c1 * (SP[j + 1][0]) + s1 * (SP[j + 1][1])));
    
                prel = ((c * (SP[j][1]) - s * (SP[j][0])) - (c * (SP[j + 1][1]) - s * (SP[j + 1][0]))) / 2;
    
                rcoeff = 1 / (1 + (prel * alf) ** 2);
                rand_value = random();
    
                if collchk < 0:
    
                    SP[j], SP[j + 1] = SP[j + 1], SP[j];
    
                    if rcoeff > rand_value:
                        counter = counter + 1
                        SP[j][2], SP[j + 1][2] = SP[j + 1][2], SP[j][2];
                if SP[j][2] == 1:
                    Up.append(c * (SP[j][0]) + s * (SP[j][1]))
                else:
                    Down.append(c * (SP[j][0]) + s * (SP[j][1]))
                Z.append(c * (SP[j][0]) + s * (SP[j][1]))
    
            Zrel = np.sum(np.array(Up)) / len(Up) - np.sum(np.array(Down)) / len(Down);
            Zrelm = np.append(Zrelm, Zrel)
    
            Zm = np.sum(np.array(Z)) / len(Z)
            Zavg = np.append(Zavg, Zm)
    
    
        return Zavg, Zrelm, counter, T_plot
    
    
    
    if __name__ == '__main__':
    
        n = 1000
        mu = np.random.uniform(0, 1, n)
        r = [sqrt(-2 * log(1 - i)) for i in mu]
        eta = np.random.uniform(0, 1, n)
        theta = 2 * pi * eta;
        cuz = [cos(i) for i in theta]
        suz = [sin(i) for i in theta]
        Zinitial = [a * b for a, b in zip(r, cuz)];
        Pinitial = [a * b for a, b in zip(r, suz)];
    
        iter = 10 ** (5);
        dt = 1 / (2 * n);
        alf = sqrt(n);
    
    
        SP = np.array(sorted(np.array([  np.array([i,j,choice([0,1])]) for i, j in zip(Zinitial, Pinitial)]),
                    key=lambda x: x[0]))
        Upi = [];
        Downi = [];
        count_plot = [];
        for j in range(len(SP)):
            if SP[j][2] == 1:
                Upi.append(SP[j][0])
            else:
                Downi.append(SP[j][0])
        Zavgi = sum(Zinitial) / len(Zinitial)
        Zreli = sum(Upi) / len(Upi) - sum(Downi) / len(Downi)
    
    
        Zavg, Zrelm, counter, T_plot = f(SP, Zavgi, Zreli, alf, dt, n)
    
        print(Zavg, Zrelm, counter)
        plt.plot(T_plot, Zrelm)
        plt.show()
    

    This is how the plot looks like: wave