OpenMP on advection

Author: EmilyBourne

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 i in range(nPts_r):
+        for j, q in enumerate(qPts):
+            r = rPts[i]
+            drPhi_0[i, j] = phi_spline.eval(q, r, 0, 1)
+            dthetaPhi_0[i, j] = 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]):
@@ -351,6 +364,7 @@ def poloidal_advection_loop(f: 'float[:,:,:,:]', phi: 'float[:,:,:]', dt: 'float
                             CN0: 'float', kN0: 'float', deltaRN0: 'float', rp: 'float', CTi: 'float', kTi: 'float', deltaRTi: 'float',
                             B0: 'float', tol: 'float', nulBound: 'bool'):
     _, nz, _, _ = f.shape
+    #$omp parallel for firstprivate(phi_spline, pol_spline) private(interp_wt, drPhi_0, dthetaPhi_0, drPhi_k, dthetaPhi_k, endPts_k1_q, endPts_k1_r, endPts_k2_q, endPts_k2_r)
     for j in range(nz):
         solve_2d_system(phi[j], phi_spline, interp_wt, r_bmat, r_l, r_u,
                         r_ipiv, theta_offset, theta_splu)