Made it a little more pythonic

Author: harpolea

compressible/interface.py

@@ -108,8 +108,6 @@ def states(idir, qx, qy, ng, dx, dt,
     dtdx = dt / dx
     dtdx4 = 0.25 * dtdx
 
-    dq = np.zeros(nvar)
-    q = np.zeros(nvar)
     lvec = np.zeros((nvar, nvar))
     rvec = np.zeros((nvar, nvar))
     e_val = np.zeros(nvar)
@@ -120,8 +118,8 @@ def states(idir, qx, qy, ng, dx, dt,
     for j in range(jlo - 2, jhi + 2):
         for i in range(ilo - 2, ihi + 2):
 
-            dq[:] = dqv[i, j, :]
-            q[:] = qv[i, j, :]
+            dq = dqv[i, j, :]
+            q = qv[i, j, :]
 
             cs = np.sqrt(gamma * q[ip] / q[irho])
 
@@ -133,18 +131,18 @@ def states(idir, qx, qy, ng, dx, dt,
             if (idir == 1):
                 e_val[:] = np.array([q[iu] - cs, q[iu], q[iu], q[iu] + cs])
 
-                lvec[0, 0:ns] = [0.0, -0.5 *
+                lvec[0, :ns] = [0.0, -0.5 *
                                  q[irho] / cs, 0.0, 0.5 / (cs * cs)]
-                lvec[1, 0:ns] = [1.0, 0.0,
+                lvec[1, :ns] = [1.0, 0.0,
                                  0.0, -1.0 / (cs * cs)]
-                lvec[2, 0:ns] = [0.0, 0.0,             1.0, 0.0]
-                lvec[3, 0:ns] = [0.0, 0.5 *
+                lvec[2, :ns] = [0.0, 0.0,             1.0, 0.0]
+                lvec[3, :ns] = [0.0, 0.5 *
                                  q[irho] / cs,  0.0, 0.5 / (cs * cs)]
 
-                rvec[0, 0:ns] = [1.0, -cs / q[irho], 0.0, cs * cs]
-                rvec[1, 0:ns] = [1.0, 0.0,       0.0, 0.0]
-                rvec[2, 0:ns] = [0.0, 0.0,       1.0, 0.0]
-                rvec[3, 0:ns] = [1.0, cs / q[irho],  0.0, cs * cs]
+                rvec[0, :ns] = [1.0, -cs / q[irho], 0.0, cs * cs]
+                rvec[1, :ns] = [1.0, 0.0,       0.0, 0.0]
+                rvec[2, :ns] = [0.0, 0.0,       1.0, 0.0]
+                rvec[3, :ns] = [1.0, cs / q[irho],  0.0, cs * cs]
 
                 # now the species -- they only have a 1 in their corresponding slot
                 e_val[ns:] = q[iu]
@@ -153,20 +151,20 @@ def states(idir, qx, qy, ng, dx, dt,
                     rvec[n, n] = 1.0
 
             else:
-                e_val = np.array([q[iv] - cs, q[iv], q[iv], q[iv] + cs])
+                e_val[:] = np.array([q[iv] - cs, q[iv], q[iv], q[iv] + cs])
 
-                lvec[0, 0:ns] = [0.0, 0.0, -0.5 *
+                lvec[0, :ns] = [0.0, 0.0, -0.5 *
                                  q[irho] / cs, 0.5 / (cs * cs)]
-                lvec[1, 0:ns] = [1.0, 0.0,
+                lvec[1, :ns] = [1.0, 0.0,
                                  0.0,             -1.0 / (cs * cs)]
-                lvec[2, 0:ns] = [0.0, 1.0, 0.0,             0.0]
-                lvec[3, 0:ns] = [0.0, 0.0, 0.5 *
+                lvec[2, :ns] = [0.0, 1.0, 0.0,             0.0]
+                lvec[3, :ns] = [0.0, 0.0, 0.5 *
                                  q[irho] / cs,  0.5 / (cs * cs)]
 
-                rvec[0, 0:ns] = [1.0, 0.0, -cs / q[irho], cs * cs]
-                rvec[1, 0:ns] = [1.0, 0.0, 0.0,       0.0]
-                rvec[2, 0:ns] = [0.0, 1.0, 0.0,       0.0]
-                rvec[3, 0:ns] = [1.0, 0.0, cs / q[irho],  cs * cs]
+                rvec[0, :ns] = [1.0, 0.0, -cs / q[irho], cs * cs]
+                rvec[1, :ns] = [1.0, 0.0, 0.0,       0.0]
+                rvec[2, :ns] = [0.0, 1.0, 0.0,       0.0]
+                rvec[3, :ns] = [1.0, 0.0, cs / q[irho],  cs * cs]
 
                 # now the species -- they only have a 1 in their corresponding slot
                 e_val[ns:] = q[iv]
@@ -179,20 +177,20 @@ def states(idir, qx, qy, ng, dx, dt,
                 # this is one the right face of the current zone,
                 # so the fastest moving eigenvalue is e_val[3] = u + c
                 factor = 0.5 * (1.0 - dtdx * max(e_val[3], 0.0))
-                q_l[i + 1, j, :] = q[:] + factor * dq[:]
+                q_l[i + 1, j, :] = q + factor * dq
 
                 # left face of the current zone, so the fastest moving
                 # eigenvalue is e_val[3] = u - c
                 factor = 0.5 * (1.0 + dtdx * min(e_val[0], 0.0))
-                q_r[i,  j, :] = q[:] - factor * dq[:]
+                q_r[i,  j, :] = q - factor * dq
 
             else:
 
                 factor = 0.5 * (1.0 - dtdx * max(e_val[3], 0.0))
-                q_l[i, j + 1, :] = q[:] + factor * dq[:]
+                q_l[i, j + 1, :] = q + factor * dq
 
                 factor = 0.5 * (1.0 + dtdx * min(e_val[0], 0.0))
-                q_r[i, j, :] = q[:] - factor * dq[:]
+                q_r[i, j, :] = q - factor * dq
 
             # compute the Vhat functions
             for m in range(nvar):
@@ -285,8 +283,6 @@ def riemann_cgf(idir, qx, qy, ng,
     smallrho = 1.e-10
     smallp = 1.e-10
 
-    xn = np.zeros(nspec)
-
     nx = qx - 2 * ng
     ny = qy - 2 * ng
     ilo = ng
@@ -483,12 +479,12 @@ def riemann_cgf(idir, qx, qy, ng,
             # species now
             if (nspec > 0):
                 if (ustar > 0.0):
-                    xn[:] = U_l[i, j, irhoX:irhoX + nspec] / U_l[i, j, idens]
+                    xn = U_l[i, j, irhoX:irhoX + nspec] / U_l[i, j, idens]
 
                 elif (ustar < 0.0):
-                    xn[:] = U_r[i, j, irhoX:irhoX + nspec] / U_r[i, j, idens]
+                    xn = U_r[i, j, irhoX:irhoX + nspec] / U_r[i, j, idens]
                 else:
-                    xn[:] = 0.5 * (U_l[i, j, irhoX:irhoX + nspec] / U_l[i, j, idens] +
+                    xn = 0.5 * (U_l[i, j, irhoX:irhoX + nspec] / U_l[i, j, idens] +
                                    U_r[i, j, irhoX:irhoX + nspec] / U_r[i, j, idens])
 
             # are we on a solid boundary?
@@ -521,7 +517,7 @@ def riemann_cgf(idir, qx, qy, ng,
                 p_state * un_state
 
             if (nspec > 0):
-                F[i, j, irhoX:irhoX + nspec] = xn[:] * rho_state * un_state
+                F[i, j, irhoX:irhoX + nspec] = xn * rho_state * un_state
 
     return F
 
@@ -596,8 +592,6 @@ def riemann_prim(idir, qx, qy, ng,
     smallrho = 1.e-10
     smallp = 1.e-10
 
-    xn = np.zeros(nspec)
-
     nx = qx - 2 * ng
     ny = qy - 2 * ng
     ilo = ng
@@ -774,12 +768,12 @@ def riemann_prim(idir, qx, qy, ng,
             # species now
             if (nspec > 0):
                 if (ustar > 0.0):
-                    xn[:] = q_l[i, j, iX:iX + nspec]
+                    xn = q_l[i, j, iX:iX + nspec]
 
                 elif (ustar < 0.0):
-                    xn[:] = q_r[i, j, iX:iX + nspec]
+                    xn = q_r[i, j, iX:iX + nspec]
                 else:
-                    xn[:] = 0.5 * (q_l[i, j, iX:iX + nspec] +
+                    xn = 0.5 * (q_l[i, j, iX:iX + nspec] +
                                    q_r[i, j, iX:iX + nspec])
 
             # are we on a solid boundary?
@@ -808,7 +802,7 @@ def riemann_prim(idir, qx, qy, ng,
             q_int[i, j, ip] = p_state
 
             if (nspec > 0):
-                q_int[i, j, iX:iX + nspec] = xn[:]
+                q_int[i, j, iX:iX + nspec] = xn
 
     return q_int
 

incompressible/incomp_interface.py

@@ -36,9 +36,6 @@ def mac_vels(qx, qy, ng, dx, dy, dt,
         MAC velocities in the x and y directions
     """
 
-    u_MAC = np.zeros((qx, qy))
-    v_MAC = np.zeros((qx, qy))
-
     # get the full u and v left and right states (including transverse
     # terms) on both the x- and y-interfaces
     u_xl, u_xr, u_yl, u_yr, v_xl, v_xr, v_yl, v_yr = get_interface_states(qx, qy, ng, dx, dy, dt,
@@ -50,8 +47,8 @@ def mac_vels(qx, qy, ng, dx, dy, dt,
     # Riemann problem -- this follows Burger's equation.  We don't use
     # any input velocity for the upwinding.  Also, we only care about
     # the normal states here (u on x and v on y)
-    riemann_and_upwind(qx, qy, ng, u_xl, u_xr, u_MAC)
-    riemann_and_upwind(qx, qy, ng, v_yl, v_yr, v_MAC)
+    u_MAC = riemann_and_upwind(qx, qy, ng, u_xl, u_xr)
+    v_MAC = riemann_and_upwind(qx, qy, ng, v_yl, v_yr)
 
     return u_MAC, v_MAC
 
@@ -96,11 +93,6 @@ def states(qx, qy, ng, dx, dy, dt,
         x and y velocities predicted to the interfaces
     """
 
-    u_xint = np.zeros((qx, qy))
-    u_yint = np.zeros((qx, qy))
-    v_xint = np.zeros((qx, qy))
-    v_yint = np.zeros((qx, qy))
-
     # get the full u and v left and right states (including transverse
     # terms) on both the x- and y-interfaces
     u_xl, u_xr, u_yl, u_yr, v_xl, v_xr, v_yl, v_yr = get_interface_states(qx, qy, ng, dx, dy, dt,
@@ -111,10 +103,10 @@ def states(qx, qy, ng, dx, dy, dt,
 
     # upwind using the MAC velocity to determine which state exists on
     # the interface
-    upwind(qx, qy, ng, u_xl, u_xr, u_MAC, u_xint)
-    upwind(qx, qy, ng, v_xl, v_xr, u_MAC, v_xint)
-    upwind(qx, qy, ng, u_yl, u_yr, v_MAC, u_yint)
-    upwind(qx, qy, ng, v_yl, v_yr, v_MAC, v_yint)
+    u_xint = upwind(qx, qy, ng, u_xl, u_xr, u_MAC)
+    v_xint = upwind(qx, qy, ng, v_xl, v_xr, u_MAC)
+    u_yint = upwind(qx, qy, ng, u_yl, u_yr, v_MAC)
+    v_yint = upwind(qx, qy, ng, v_yl, v_yr, v_MAC)
 
     return u_xint, u_yint, v_xint, v_yint
 
@@ -166,14 +158,6 @@ def get_interface_states(qx, qy, ng, dx, dy, dt,
     v_yl = np.zeros((qx, qy))
     v_yr = np.zeros((qx, qy))
 
-    uhat_adv = np.zeros((qx, qy))
-    vhat_adv = np.zeros((qx, qy))
-
-    u_xint = np.zeros((qx, qy))
-    u_yint = np.zeros((qx, qy))
-    v_xint = np.zeros((qx, qy))
-    v_yint = np.zeros((qx, qy))
-
     nx = qx - 2 * ng
     ny = qy - 2 * ng
     ilo = ng
@@ -216,19 +200,19 @@ def get_interface_states(qx, qy, ng, dx, dy, dt,
 
     # now get the normal advective velocities on the interfaces by solving
     # the Riemann problem.
-    riemann(qx, qy, ng, u_xl, u_xr, uhat_adv)
-    riemann(qx, qy, ng, v_yl, v_yr, vhat_adv)
+    uhat_adv = riemann(qx, qy, ng, u_xl, u_xr)
+    vhat_adv = riemann(qx, qy, ng, v_yl, v_yr)
 
     # now that we have the advective velocities, upwind the left and right
     # states using the appropriate advective velocity.
 
     # on the x-interfaces, we upwind based on uhat_adv
-    upwind(qx, qy, ng, u_xl, u_xr, uhat_adv, u_xint)
-    upwind(qx, qy, ng, v_xl, v_xr, uhat_adv, v_xint)
+    u_xint = upwind(qx, qy, ng, u_xl, u_xr, uhat_adv)
+    v_xint = upwind(qx, qy, ng, v_xl, v_xr, uhat_adv)
 
     # on the y-interfaces, we upwind based on vhat_adv
-    upwind(qx, qy, ng, u_yl, u_yr, vhat_adv, u_yint)
-    upwind(qx, qy, ng, v_yl, v_yr, vhat_adv, v_yint)
+    u_yint = upwind(qx, qy, ng, u_yl, u_yr, vhat_adv)
+    v_yint = upwind(qx, qy, ng, v_yl, v_yr, vhat_adv)
 
     # at this point, these states are the `hat' states -- they only
     # considered the normal to the interface portion of the predictor.
@@ -277,7 +261,7 @@ def get_interface_states(qx, qy, ng, dx, dy, dt,
 
 # xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 @njit(cache=True)
-def upwind(qx, qy, ng, q_l, q_r, s, q_int):
+def upwind(qx, qy, ng, q_l, q_r, s):
     r"""
     upwind the left and right states based on the specified input
     velocity, s.  The resulting interface state is q_int
@@ -292,7 +276,10 @@ def upwind(qx, qy, ng, q_l, q_r, s, q_int):
         left and right states
     s : ndarray
         velocity
-    q_int : ndarray
+
+    Returns
+    -------
+    out : ndarray
         Upwinded state
     """
 
@@ -303,6 +290,8 @@ def upwind(qx, qy, ng, q_l, q_r, s, q_int):
     jlo = ng
     jhi = ng + ny
 
+    q_int = np.zeros((qx, qy))
+
     for j in range(jlo - 1, jhi + 1):
         for i in range(ilo - 1, ihi + 1):
 
@@ -313,10 +302,12 @@ def upwind(qx, qy, ng, q_l, q_r, s, q_int):
             else:
                 q_int[i, j] = q_r[i, j]
 
+    return q_int
+
 
 # xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 @njit(cache=True)
-def riemann(qx, qy, ng, q_l, q_r, s):
+def riemann(qx, qy, ng, q_l, q_r):
     """
     Solve the Burger's Riemann problem given the input left and right
     states and return the state on the interface.
@@ -331,7 +322,10 @@ def riemann(qx, qy, ng, q_l, q_r, s):
         The number of ghost cells
     q_l, q_r : ndarray
         left and right states
-    s : ndarray
+
+    Returns
+    -------
+    out : ndarray
         Interface state
     """
 
@@ -342,6 +336,8 @@ def riemann(qx, qy, ng, q_l, q_r, s):
     jlo = ng
     jhi = ng + ny
 
+    s = np.zeros((qx, qy))
+
     for j in range(jlo - 1, jhi + 1):
         for i in range(ilo - 1, ihi + 1):
 
@@ -352,10 +348,12 @@ def riemann(qx, qy, ng, q_l, q_r, s):
             else:
                 s[i, j] = q_r[i, j]
 
+    return s
+
 
 # xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 @njit(cache=True)
-def riemann_and_upwind(qx, qy, ng, q_l, q_r, q_int):
+def riemann_and_upwind(qx, qy, ng, q_l, q_r):
     r"""
     First solve the Riemann problem given q_l and q_r to give the
     velocity on the interface and: use this velocity to upwind to
@@ -372,11 +370,12 @@ def riemann_and_upwind(qx, qy, ng, q_l, q_r, q_int):
         The number of ghost cells
     q_l, q_r : ndarray
         left and right states
-    q_int : ndarray
+
+    Returns
+    -------
+    out : ndarray
         Upwinded state
     """
 
-    s = np.zeros((qx, qy))
-
-    riemann(qx, qy, ng, q_l, q_r, s)
-    upwind(qx, qy, ng, q_l, q_r, s, q_int)
+    s = riemann(qx, qy, ng, q_l, q_r)
+    return upwind(qx, qy, ng, q_l, q_r, s)