Interact with python -- f2py or ctypes?

Help
7

@CYehLu
Did you try numba? I have slightly modified your code ( I assume that using np.arrays instead of lists is allowed in “pure” python version)

import numpy as np
from numba import njit

def vorticity_py(u, v, dx, dy):
   ny, nx = u.shape
   vort = np.zeros((ny, nx))
   dx2 = dx * 2
   dy2 = dy * 2 
   for j in range(1, ny-1):
       for i in range(1, nx-1):
           vort[j,i] = (v[j,i+1] - v[j,i-1]) / dx2 - (u[j+1,i] - u[j-1,i]) / dy2       
   return vort

def vorticity_np(u, v, dx, dy):   
   ny, nx = u.shape
   vort = np.zeros((ny, nx))
   dx2 = dx * 2
   dy2 = dy * 2 
   vort[1:-1,1:-1] = (v[1:-1,2:] - v[1:-1,:-2]) / dx2 - (u[2:,1:-1] - u[:-2,1:-1]) / dy2
   return vort

@njit(parallel = True)
def vorticity_jit(u, v, dx, dy):
   ny, nx = u.shape
   vort = np.zeros((ny, nx))
   dx2 = dx * 2
   dy2 = dy * 2
   for j in range(1, ny-1):
       for i in range(1, nx-1):
           vort[j,i] = (v[j,i+1] - v[j,i-1]) / dx2 - (u[j+1,i] - u[j-1,i]) / dy2          
   return vort

# prepare data
dx = 0.5
dy = 0.5
ny = 1000
nx = 1000
u = np.random.rand(ny, nx)
v = np.random.rand(ny, nx)

res = vorticity_py(u,v,dx,dy)
print(res[1,1])
res = vorticity_np(u,v,dx,dy)
print(res[1,1])
res = vorticity_jit(u,v,dx,dy)
print(res[1,1])

Results:
%timeit vorticity_py(u,v,dx,dy)
1.24 s ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit vorticity_np(u,v,dx,dy)
16.5 ms ± 306 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit vorticity_jit(u,v,dx,dy)
9.47 ms ± 175 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

1 Like