Use OpenMP parallelism

.github/workflows/testing.yml

@@ -198,15 +198,15 @@ jobs:
             python -m pip install autopep8
       - name: Check for pep changes
         run: |
-            if [[ $(autopep8 -r --diff pygyro/) ]]
+            if [[ $(autopep8 --ignore=E265,E226,E24,W50,W690 -r --diff pygyro/) ]]
             then
-              echo "Please run 'autopep8 -ir pygyro/' to clean up formatting"
-              autopep8 -r --diff pygyro/
+              echo "Please run 'autopep8 --ignore=E265,E226,E24,W50,W690 -ir pygyro/' to clean up formatting"
+              autopep8 --ignore=E265,E226,E24,W50,W690 -r --diff pygyro/
               exit 1
             fi
-            if [[ $(autopep8 -r --diff plotting/) ]]
+            if [[ $(autopep8 --ignore=E265,E226,E24,W50,W690 -r --diff plotting/) ]]
             then
-              echo "Please run 'autopep8 -ir plotting/' to clean up formatting"
-              autopep8 -r --diff plotting/
+              echo "Please run 'autopep8 --ignore=E265,E226,E24,W50,W690 -ir plotting/' to clean up formatting"
+              autopep8 --ignore=E265,E226,E24,W50,W690 -r --diff plotting/
               exit 1
             fi

Makefile

@@ -42,7 +42,7 @@ SO_EXT := $(shell $(PYTHON) -c "import sysconfig; print(sysconfig.get_config_var
 
 ifeq ($(ACC), pycc)
 	TOOL := pyccel
-	TOOL_FLAGS := --compiler-family=$(COMP) --flags ' $(FC_FLAGS)' --language=$(LANGUAGE)
+	TOOL_FLAGS := --compiler-family=$(COMP) --flags ' $(FC_FLAGS)' --language=$(LANGUAGE) --openmp
 	NAME_PREFIX := 
 else
 	ifeq ($(ACC), numba)

pygyro/advection/accelerated_advection_steps.py

@@ -41,15 +41,20 @@ def poloidal_advection_step_expl(f: 'float[:,:]',
     multFactor = dt / B0
     multFactor_half = 0.5 * multFactor
 
-    phi_spline.eval_vector(qPts, rPts, drPhi_0, 0, 1)
-    phi_spline.eval_vector(qPts, rPts, dthetaPhi_0, 1, 0)
-
     nPts_r = rPts.shape[0]
     nPts_q = qPts.shape[0]
 
+    #$omp parallel for collapse(2)
+    for j in range(nPts_q):
+        for i, r in enumerate(rPts):
+            q = qPts[j]
+            drPhi_0[j, i] = phi_spline.eval(q, r, 0, 1)
+            dthetaPhi_0[j, i] = phi_spline.eval(q, r, 1, 0)
+
     idx = nPts_r-1
     rMax = rPts[idx]
 
+    #$omp parallel for collapse(2)
     for i in range(nPts_q):
         for j in range(nPts_r):
             # Step one of Heun method
@@ -88,6 +93,7 @@ def poloidal_advection_step_expl(f: 'float[:,:]',
 
     # Find value at the determined point
     if (nulBound):
+        #$omp parallel for collapse(2)
         for i in range(nPts_q):  # theta
             for j in range(nPts_r):  # r
                 if (endPts_k2_r[i, j] < rPts[0]):
@@ -99,6 +105,7 @@ def poloidal_advection_step_expl(f: 'float[:,:]',
                     f[i, j] = pol_spline.eval(
                         endPts_k2_q[i, j], endPts_k2_r[i, j])
     else:
+        #$omp parallel for collapse(2)
         for i in range(nPts_q):  # theta
             for j in range(nPts_r):  # r
                 if (endPts_k2_r[i, j] < rPts[0]):
@@ -152,6 +159,7 @@ def v_parallel_advection_eval_step_loop(f: 'float[:,:,:]', vPts: 'float[:]',
                                            CN0: 'float', kN0: 'float', deltaRN0: 'float', rp: 'float',
                                            CTi: 'float', kTi: 'float', deltaRTi: 'float', bound: 'int'):
     n1, n2, _ = f.shape
+    #$omp parallel for collapse(2) firstprivate(spl, vPts) private(coeffs)
     for j in range(n1):  # z
         for k in range(n2):  # q
             coeffs = spl.coeffs
@@ -197,6 +205,7 @@ def flux_advection_loop(f: 'float[:,:,:,:]', thetaSpline: Spline1D, theta_offset
                         thetaShifts: 'float[:,:,:]', lagrange_coeffs: 'float[:,:,:]'):
     nr, nv, nq, nz = f.shape
 
+    #$omp parallel for collapse(2) firstprivate(thetaSpline)
     for rIdx in range(nr):  # r
         for cIdx in range(nv):  # v
             # find the values of the function at each required point
@@ -252,6 +261,7 @@ def poloidal_advection_step_impl(f: 'float[:,:]', dt: 'float', v: 'float', rPts:
     idx = nPts_r-1
     rMax = rPts[idx]
 
+    #$omp parallel for collapse(2)
     for i in range(nPts_q):
         for j in range(nPts_r):
             # Step one of Heun method
@@ -266,6 +276,7 @@ def poloidal_advection_step_impl(f: 'float[:,:]', dt: 'float', v: 'float', rPts:
     norm = tol+1
     while (norm > tol):
         norm = 0.0
+        #$omp parallel for collapse(2)
         for i in range(nPts_q):
             for j in range(nPts_r):
                 # Handle theta boundary conditions
@@ -315,6 +326,7 @@ def poloidal_advection_step_impl(f: 'float[:,:]', dt: 'float', v: 'float', rPts:
 
     # Find value at the determined point
     if (nulBound):
+        #$omp parallel for collapse(2)
         for i in range(nPts_q):
             for j in range(nPts_r):
                 if (endPts_k2_r[i, j] < rPts[0]):
@@ -326,6 +338,7 @@ def poloidal_advection_step_impl(f: 'float[:,:]', dt: 'float', v: 'float', rPts:
                     f[i, j] = pol_spline.eval(
                         endPts_k2_q[i, j], endPts_k2_r[i, j])
     else:
+        #$omp parallel for collapse(2)
         for i in range(nPts_q):
             for j in range(nPts_r):
                 if (endPts_k2_r[i, j] < rPts[0]):

pygyro/poisson/poisson_tools.py

@@ -25,6 +25,7 @@ def get_perturbed_rho(rho: T, feq: 'float[:,:]', grid: 'float[:,:,:,:]',
 
     nc, = quad_coeffs.shape
 
+    #$omp parallel for collapse(3)
     for i in range(n):
         for j in range(m):
             for k in range(p):
@@ -53,6 +54,7 @@ def get_rho(rho: T, grid: 'float[:,:,:,:]', quad_coeffs: 'float[:]'):
 
     nc, = quad_coeffs.shape
 
+    #$omp parallel for collapse(3)
     for i in range(n):
         for j in range(m):
             for k in range(p):

pygyro/splines/accelerated_spline_interpolators.py

@@ -53,7 +53,7 @@ def solve_system_nonperiodic(ug: 'Final[T[:]]', c: 'T[:]', bmat: 'Final[T[:,:](o
 
 def solve_2d_system(ug: 'float[:,:]', spl: Spline2D, wt: 'float[:,:]',
                     r_bmat: 'float[:,:](order=F)', r_l: np.int32, r_u: np.int32, r_ipiv: 'int32[:]',
-                    theta_offset: int, theta_splu: PeriodicBandedMatrix):
+                    theta_offset: int, theta_splu: Final[PeriodicBandedMatrix]):
     basis1 = spl.basis1
     basis2 = spl.basis2
     n1, n2 = basis1.nbasis, basis2.nbasis
@@ -75,24 +75,26 @@ def solve_2d_system(ug: 'float[:,:]', spl: Spline2D, wt: 'float[:,:]',
     assert sinfo == 0
 
     # Transpose coefficients to self._bwork
+    #$omp parallel for collapse(2)
     for i1 in range(n1):
         for i2 in range(n2):
             wt[i2, i1] = ug[i1, i2]
 
     # Cycle over x2 position and interpolate w along x1 direction.
     # Work on self._bwork
+    #$ omp parallel for default(none) shared(wt, n1, n2, theta_offset) firstprivate(spline1, theta_splu) schedule(static)
     for i2 in range(n2):
         solve_system_periodic(wt[i2, :n1], spline1, theta_offset, theta_splu)
-        # self._interp1.compute_interpolant(wt[i2, :n1], self._spline1)
-        c = spline1.coeffs
-        wt[i2, :] = c
+        wt[i2, :] = spline1.coeffs
 
     # Transpose coefficients to spl.coeffs
+    #$omp parallel for collapse(2)
     for i1 in range(s1):
         for i2 in range(s2):
             w[i1, i2] = wt[i2, i1]
 
     # x1-periodic only: "wrap around" coefficients onto extended array
+    #$omp parallel for collapse(2)
     for i1 in range(p1):
         for i2 in range(s2):
             w[n1 + i1, i2] = w[i1, i2]