From 64cd8f881999f6b7bceb686fd7d934d352c63685 Mon Sep 17 00:00:00 2001 From: Michael Zingale Date: Fri, 3 Mar 2023 19:43:45 -0500 Subject: [PATCH 1/2] clean up the compressible interface construction this changes things to pass in the grid object and variables object directly, making the code more consistent with the rest of pyro --- pyro/compressible/interface.py | 356 ++++++++++++---------------- pyro/compressible/unsplit_fluxes.py | 24 +- 2 files changed, 163 insertions(+), 217 deletions(-) diff --git a/pyro/compressible/interface.py b/pyro/compressible/interface.py index cfa5d9995..eaa40e57b 100644 --- a/pyro/compressible/interface.py +++ b/pyro/compressible/interface.py @@ -3,8 +3,7 @@ @njit(cache=True) -def states(idir, ng, dx, dt, - irho, iu, iv, ip, ix, nspec, +def states(idir, grid, dt, ivars, gamma, qv, dqv): r""" predict the cell-centered state to the edges in one-dimension @@ -65,13 +64,11 @@ def states(idir, ng, dx, dt, ---------- idir : int Are we predicting to the edges in the x-direction (1) or y-direction (2)? - ng : int - The number of ghost cells - dx : float - The cell spacing + grid : Grid2d + The grid object dt : float The timestep - irho, iu, iv, ip, ix : int + ivars: Indices of the density, x-velocity, y-velocity, pressure and species in the state vector nspec : int @@ -89,37 +86,32 @@ def states(idir, ng, dx, dt, State vector predicted to the left and right edges """ - qx, qy, nvar = qv.shape + q_l = grid.scratch_array(nvar=ivars.nq) + q_r = grid.scratch_array(nvar=ivars.nq) - q_l = np.zeros_like(qv) - q_r = np.zeros_like(qv) + ns = ivars.nq - ivars.naux - nx = qx - 2 * ng - ny = qy - 2 * ng - ilo = ng - ihi = ng + nx - jlo = ng - jhi = ng + ny - - ns = nvar - nspec + if idir == 1: + dtdx = dt / grid.dx + else: + dtdx = dt / grid.dy - dtdx = dt / dx dtdx4 = 0.25 * dtdx - lvec = np.zeros((nvar, nvar)) - rvec = np.zeros((nvar, nvar)) - e_val = np.zeros(nvar) - betal = np.zeros(nvar) - betar = np.zeros(nvar) + lvec = np.zeros((ivars.nq, ivars.nq)) + rvec = np.zeros((ivars.nq, ivars.nq)) + e_val = np.zeros(ivars.nq) + betal = np.zeros(ivars.nq) + betar = np.zeros(ivars.nq) # this is the loop over zones. For zone i, we see q_l[i+1] and q_r[i] - for i in range(ilo - 2, ihi + 2): - for j in range(jlo - 2, jhi + 2): + for i in range(grid.ilo - 2, grid.ihi + 2): + for j in range(grid.jlo - 2, grid.jhi + 2): dq = dqv[i, j, :] q = qv[i, j, :] - cs = np.sqrt(gamma * q[ip] / q[irho]) + cs = np.sqrt(gamma * q[ivars.ip] / q[ivars.irho]) lvec[:, :] = 0.0 rvec[:, :] = 0.0 @@ -127,46 +119,46 @@ def states(idir, ng, dx, dt, # compute the eigenvalues and eigenvectors if idir == 1: - e_val[:] = np.array([q[iu] - cs, q[iu], q[iu], q[iu] + cs]) + e_val[:] = np.array([q[ivars.iu] - cs, q[ivars.iu], q[ivars.iu], q[ivars.iu] + cs]) lvec[0, :ns] = [0.0, -0.5 * - q[irho] / cs, 0.0, 0.5 / (cs * cs)] + q[ivars.irho] / cs, 0.0, 0.5 / (cs * cs)] lvec[1, :ns] = [1.0, 0.0, 0.0, -1.0 / (cs * cs)] 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)] + q[ivars.irho] / cs, 0.0, 0.5 / (cs * cs)] - rvec[0, :ns] = [1.0, -cs / q[irho], 0.0, cs * cs] + rvec[0, :ns] = [1.0, -cs / q[ivars.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] + rvec[3, :ns] = [1.0, cs / q[ivars.irho], 0.0, cs * cs] # now the species -- they only have a 1 in their corresponding slot - e_val[ns:] = q[iu] - for n in range(ix, ix + nspec): + e_val[ns:] = q[ivars.iu] + for n in range(ivars.ix, ivars.ix + ivars.naux): lvec[n, n] = 1.0 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[ivars.iv] - cs, q[ivars.iv], q[ivars.iv], q[ivars.iv] + cs]) lvec[0, :ns] = [0.0, 0.0, -0.5 * - q[irho] / cs, 0.5 / (cs * cs)] + q[ivars.irho] / cs, 0.5 / (cs * cs)] lvec[1, :ns] = [1.0, 0.0, 0.0, -1.0 / (cs * cs)] 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)] + q[ivars.irho] / cs, 0.5 / (cs * cs)] - rvec[0, :ns] = [1.0, 0.0, -cs / q[irho], cs * cs] + rvec[0, :ns] = [1.0, 0.0, -cs / q[ivars.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] + rvec[3, :ns] = [1.0, 0.0, cs / q[ivars.irho], cs * cs] # now the species -- they only have a 1 in their corresponding slot - e_val[ns:] = q[iv] - for n in range(ix, ix + nspec): + e_val[ns:] = q[ivars.iv] + for n in range(ivars.ix, ivars.ix + ivars.naux): lvec[n, n] = 1.0 rvec[n, n] = 1.0 @@ -191,7 +183,7 @@ def states(idir, ng, dx, dt, q_r[i, j, :] = q - factor * dq # compute the Vhat functions - for m in range(nvar): + for m in range(ivars.nq): asum = np.dot(lvec[m, :], dq) betal[m] = dtdx4 * (e_val[3] - e_val[m]) * \ @@ -200,7 +192,7 @@ def states(idir, ng, dx, dt, (1.0 - np.copysign(1.0, e_val[m])) * asum # construct the states - for m in range(nvar): + for m in range(ivars.nq): sum_l = np.dot(betal, rvec[:, m]) sum_r = np.dot(betar, rvec[:, m]) @@ -215,8 +207,7 @@ def states(idir, ng, dx, dt, @njit(cache=True) -def riemann_cgf(idir, ng, - idens, ixmom, iymom, iener, irhoX, nspec, +def riemann_cgf(idir, grid, ivars, lower_solid, upper_solid, gamma, U_l, U_r): r""" @@ -271,50 +262,41 @@ def riemann_cgf(idir, ng, Conserved flux """ - qx, qy, nvar = U_l.shape - - F = np.zeros((qx, qy, nvar)) + F = grid.scratch_array(nvar=ivars.nvar) smallc = 1.e-10 smallrho = 1.e-10 smallp = 1.e-10 - nx = qx - 2 * ng - ny = qy - 2 * ng - ilo = ng - ihi = ng + nx - jlo = ng - jhi = ng + ny - - for i in range(ilo - 1, ihi + 1): - for j in range(jlo - 1, jhi + 1): + for i in range(grid.ilo - 1, grid.ihi + 1): + for j in range(grid.jlo - 1, grid.jhi + 1): # primitive variable states - rho_l = U_l[i, j, idens] + rho_l = U_l[i, j, ivars.idens] # un = normal velocity; ut = transverse velocity if idir == 1: - un_l = U_l[i, j, ixmom] / rho_l - ut_l = U_l[i, j, iymom] / rho_l + un_l = U_l[i, j, ivars.ixmom] / rho_l + ut_l = U_l[i, j, ivars.iymom] / rho_l else: - un_l = U_l[i, j, iymom] / rho_l - ut_l = U_l[i, j, ixmom] / rho_l + un_l = U_l[i, j, ivars.iymom] / rho_l + ut_l = U_l[i, j, ivars.ixmom] / rho_l - rhoe_l = U_l[i, j, iener] - 0.5 * rho_l * (un_l**2 + ut_l**2) + rhoe_l = U_l[i, j, ivars.iener] - 0.5 * rho_l * (un_l**2 + ut_l**2) p_l = rhoe_l * (gamma - 1.0) p_l = max(p_l, smallp) - rho_r = U_r[i, j, idens] + rho_r = U_r[i, j, ivars.idens] if idir == 1: - un_r = U_r[i, j, ixmom] / rho_r - ut_r = U_r[i, j, iymom] / rho_r + un_r = U_r[i, j, ivars.ixmom] / rho_r + ut_r = U_r[i, j, ivars.iymom] / rho_r else: - un_r = U_r[i, j, iymom] / rho_r - ut_r = U_r[i, j, ixmom] / rho_r + un_r = U_r[i, j, ivars.iymom] / rho_r + ut_r = U_r[i, j, ivars.ixmom] / rho_r - rhoe_r = U_r[i, j, iener] - 0.5 * rho_r * (un_r**2 + ut_r**2) + rhoe_r = U_r[i, j, ivars.iener] - 0.5 * rho_r * (un_r**2 + ut_r**2) p_r = rhoe_r * (gamma - 1.0) p_r = max(p_r, smallp) @@ -473,54 +455,53 @@ def riemann_cgf(idir, ng, rhoe_state = 0.5 * (rhoestar_l + rhoestar_r) # species now - if nspec > 0: + if ivars.naux > 0: if ustar > 0.0: - xn = U_l[i, j, irhoX:irhoX + nspec] / U_l[i, j, idens] + xn = U_l[i, j, ivars.irhoX:ivars.irhoX + ivars.naux] / U_l[i, j, ivars.idens] elif ustar < 0.0: - xn = U_r[i, j, irhoX:irhoX + nspec] / U_r[i, j, idens] + xn = U_r[i, j, ivars.irhoX:ivars.irhoX + ivars.naux] / U_r[i, j, ivars.idens] else: - 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]) + xn = 0.5 * (U_l[i, j, ivars.irhoX:ivars.irhoX + ivars.naux] / U_l[i, j, ivars.idens] + + U_r[i, j, ivars.irhoX:ivars.irhoX + ivars.nspec] / U_r[i, j, ivars.idens]) # are we on a solid boundary? if idir == 1: - if i == ilo and lower_solid == 1: + if i == grid.ilo and lower_solid == 1: un_state = 0.0 - if i == ihi + 1 and upper_solid == 1: + if i == grid.ihi + 1 and upper_solid == 1: un_state = 0.0 elif idir == 2: - if j == jlo and lower_solid == 1: + if j == grid.jlo and lower_solid == 1: un_state = 0.0 - if j == jhi + 1 and upper_solid == 1: + if j == grid.jhi + 1 and upper_solid == 1: un_state = 0.0 # compute the fluxes - F[i, j, idens] = rho_state * un_state + F[i, j, ivars.idens] = rho_state * un_state if idir == 1: - F[i, j, ixmom] = rho_state * un_state**2 + p_state - F[i, j, iymom] = rho_state * ut_state * un_state + F[i, j, ivars.ixmom] = rho_state * un_state**2 + p_state + F[i, j, ivars.iymom] = rho_state * ut_state * un_state else: - F[i, j, ixmom] = rho_state * ut_state * un_state - F[i, j, iymom] = rho_state * un_state**2 + p_state + F[i, j, ivars.ixmom] = rho_state * ut_state * un_state + F[i, j, ivars.iymom] = rho_state * un_state**2 + p_state - F[i, j, iener] = rhoe_state * un_state + \ + F[i, j, ivars.iener] = rhoe_state * un_state + \ 0.5 * rho_state * (un_state**2 + ut_state**2) * un_state + \ p_state * un_state - if nspec > 0: - F[i, j, irhoX:irhoX + nspec] = xn * rho_state * un_state + if ivars.naux > 0: + F[i, j, ivars.irhoX:ivars.irhoX + ivars.naux] = xn * rho_state * un_state return F @njit(cache=True) -def riemann_prim(idir, ng, - irho, iu, iv, ip, iX, nspec, +def riemann_prim(idir, grid, ivars, lower_solid, upper_solid, gamma, q_l, q_r): r""" @@ -578,48 +559,39 @@ def riemann_prim(idir, ng, Primitive flux """ - qx, qy, nvar = q_l.shape - - q_int = np.zeros((qx, qy, nvar)) + q_int = grid.scratch_array(nvar=ivars.nq) smallc = 1.e-10 smallrho = 1.e-10 smallp = 1.e-10 - nx = qx - 2 * ng - ny = qy - 2 * ng - ilo = ng - ihi = ng + nx - jlo = ng - jhi = ng + ny - - for i in range(ilo - 1, ihi + 1): - for j in range(jlo - 1, jhi + 1): + for i in range(grid.ilo - 1, grid.ihi + 1): + for j in range(grid.jlo - 1, grid.jhi + 1): # primitive variable states - rho_l = q_l[i, j, irho] + rho_l = q_l[i, j, ivars.irho] # un = normal velocity; ut = transverse velocity if idir == 1: - un_l = q_l[i, j, iu] - ut_l = q_l[i, j, iv] + un_l = q_l[i, j, ivars.iu] + ut_l = q_l[i, j, ivars.iv] else: - un_l = q_l[i, j, iv] - ut_l = q_l[i, j, iu] + un_l = q_l[i, j, ivars.iv] + ut_l = q_l[i, j, ivars.iu] - p_l = q_l[i, j, ip] + p_l = q_l[i, j, ivars.ip] p_l = max(p_l, smallp) - rho_r = q_r[i, j, irho] + rho_r = q_r[i, j, ivars.irho] if idir == 1: - un_r = q_r[i, j, iu] - ut_r = q_r[i, j, iv] + un_r = q_r[i, j, ivars.iu] + ut_r = q_r[i, j, ivars.iv] else: - un_r = q_r[i, j, iv] - ut_r = q_r[i, j, iu] + un_r = q_r[i, j, ivars.iv] + ut_r = q_r[i, j, ivars.iu] - p_r = q_r[i, j, ip] + p_r = q_r[i, j, ivars.ip] p_r = max(p_r, smallp) # define the Lagrangian sound speed @@ -760,50 +732,49 @@ def riemann_prim(idir, ng, p_state = pstar # species now - if nspec > 0: + if ivars.naux > 0: if ustar > 0.0: - xn = q_l[i, j, iX:iX + nspec] + xn = q_l[i, j, ivars.iX:ivars.iX + ivars.naux] elif ustar < 0.0: - xn = q_r[i, j, iX:iX + nspec] + xn = q_r[i, j, ivars.iX:ivars.iX + ivars.naux] else: - xn = 0.5 * (q_l[i, j, iX:iX + nspec] + - q_r[i, j, iX:iX + nspec]) + xn = 0.5 * (q_l[i, j, ivars.iX:ivars.iX + ivars.naux] + + q_r[i, j, ivars.iX:ivars.iX + ivars.naux]) # are we on a solid boundary? if idir == 1: - if i == ilo and lower_solid == 1: + if i == grid.ilo and lower_solid == 1: un_state = 0.0 - if i == ihi + 1 and upper_solid == 1: + if i == grid.ihi + 1 and upper_solid == 1: un_state = 0.0 elif idir == 2: - if j == jlo and lower_solid == 1: + if j == grid.jlo and lower_solid == 1: un_state = 0.0 - if j == jhi + 1 and upper_solid == 1: + if j == grid.jhi + 1 and upper_solid == 1: un_state = 0.0 - q_int[i, j, irho] = rho_state + q_int[i, j, ivars.irho] = rho_state if idir == 1: - q_int[i, j, iu] = un_state - q_int[i, j, iv] = ut_state + q_int[i, j, ivars.iu] = un_state + q_int[i, j, ivars.iv] = ut_state else: - q_int[i, j, iu] = ut_state - q_int[i, j, iv] = un_state + q_int[i, j, ivars.iu] = ut_state + q_int[i, j, ivars.iv] = un_state - q_int[i, j, ip] = p_state + q_int[i, j, ivars.ip] = p_state - if nspec > 0: - q_int[i, j, iX:iX + nspec] = xn + if ivars.naux > 0: + q_int[i, j, ivars.iX:ivars.iX + ivars.naux] = xn return q_int @njit(cache=True) -def riemann_hllc(idir, ng, - idens, ixmom, iymom, iener, irhoX, nspec, +def riemann_hllc(idir, grid, ivars, lower_solid, upper_solid, gamma, U_l, U_r): r""" @@ -835,51 +806,42 @@ def riemann_hllc(idir, ng, Conserved flux """ - qx, qy, nvar = U_l.shape - - F = np.zeros((qx, qy, nvar)) + F = grid.scratch_array(nvar=ivars.nvar) smallc = 1.e-10 smallp = 1.e-10 - U_state = np.zeros(nvar) + U_state = np.zeros(ivars.nvar) - nx = qx - 2 * ng - ny = qy - 2 * ng - ilo = ng - ihi = ng + nx - jlo = ng - jhi = ng + ny - - for i in range(ilo - 1, ihi + 1): - for j in range(jlo - 1, jhi + 1): + for i in range(grid.ilo - 1, grid.ihi + 1): + for j in range(grid.jlo - 1, grid.jhi + 1): # primitive variable states - rho_l = U_l[i, j, idens] + rho_l = U_l[i, j, ivars.idens] # un = normal velocity; ut = transverse velocity if idir == 1: - un_l = U_l[i, j, ixmom] / rho_l - ut_l = U_l[i, j, iymom] / rho_l + un_l = U_l[i, j, ivars.ixmom] / rho_l + ut_l = U_l[i, j, ivars.iymom] / rho_l else: - un_l = U_l[i, j, iymom] / rho_l - ut_l = U_l[i, j, ixmom] / rho_l + un_l = U_l[i, j, ivars.iymom] / rho_l + ut_l = U_l[i, j, ivars.ixmom] / rho_l - rhoe_l = U_l[i, j, iener] - 0.5 * rho_l * (un_l**2 + ut_l**2) + rhoe_l = U_l[i, j, ivars.iener] - 0.5 * rho_l * (un_l**2 + ut_l**2) p_l = rhoe_l * (gamma - 1.0) p_l = max(p_l, smallp) - rho_r = U_r[i, j, idens] + rho_r = U_r[i, j, ivars.idens] if idir == 1: - un_r = U_r[i, j, ixmom] / rho_r - ut_r = U_r[i, j, iymom] / rho_r + un_r = U_r[i, j, ivars.ixmom] / rho_r + ut_r = U_r[i, j, ivars.iymom] / rho_r else: - un_r = U_r[i, j, iymom] / rho_r - ut_r = U_r[i, j, ixmom] / rho_r + un_r = U_r[i, j, ivars.iymom] / rho_r + ut_r = U_r[i, j, ivars.ixmom] / rho_r - rhoe_r = U_r[i, j, iener] - 0.5 * rho_r * (un_r**2 + ut_r**2) + rhoe_r = U_r[i, j, ivars.iener] - 0.5 * rho_r * (un_r**2 + ut_r**2) p_r = rhoe_r * (gamma - 1.0) p_r = max(p_r, smallp) @@ -993,33 +955,31 @@ def riemann_hllc(idir, ng, # R region U_state[:] = U_r[i, j, :] - F[i, j, :] = consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, - U_state) + F[i, j, :] = consFlux(idir, gamma, ivars, U_state) elif S_r > 0.0 and S_c <= 0: # R* region HLLCfactor = rho_r * (S_r - un_r) / (S_r - S_c) - U_state[idens] = HLLCfactor + U_state[ivars.idens] = HLLCfactor if idir == 1: - U_state[ixmom] = HLLCfactor * S_c - U_state[iymom] = HLLCfactor * ut_r + U_state[ivars.ixmom] = HLLCfactor * S_c + U_state[ivars.iymom] = HLLCfactor * ut_r else: - U_state[ixmom] = HLLCfactor * ut_r - U_state[iymom] = HLLCfactor * S_c + U_state[ivars.ixmom] = HLLCfactor * ut_r + U_state[ivars.iymom] = HLLCfactor * S_c - U_state[iener] = HLLCfactor * (U_r[i, j, iener] / rho_r + + U_state[ivars.iener] = HLLCfactor * (U_r[i, j, ivars.iener] / rho_r + (S_c - un_r) * (S_c + p_r / (rho_r * (S_r - un_r)))) # species - if nspec > 0: - U_state[irhoX:irhoX + nspec] = HLLCfactor * \ - U_r[i, j, irhoX:irhoX + nspec] / rho_r + if ivars.naux > 0: + U_state[ivars.irhoX:ivars.irhoX + ivars.nspec] = HLLCfactor * \ + U_r[i, j, ivars.irhoX:ivars.irhoX + ivars.nspec] / rho_r # find the flux on the right interface - F[i, j, :] = consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, - U_r[i, j, :]) + F[i, j, :] = consFlux(idir, gamma, ivars, U_r[i, j, :]) # correct the flux F[i, j, :] = F[i, j, :] + S_r * (U_state[:] - U_r[i, j, :]) @@ -1028,26 +988,25 @@ def riemann_hllc(idir, ng, # L* region HLLCfactor = rho_l * (S_l - un_l) / (S_l - S_c) - U_state[idens] = HLLCfactor + U_state[ivars.idens] = HLLCfactor if idir == 1: - U_state[ixmom] = HLLCfactor * S_c - U_state[iymom] = HLLCfactor * ut_l + U_state[ivars.ixmom] = HLLCfactor * S_c + U_state[ivars.iymom] = HLLCfactor * ut_l else: - U_state[ixmom] = HLLCfactor * ut_l - U_state[iymom] = HLLCfactor * S_c + U_state[ivars.ixmom] = HLLCfactor * ut_l + U_state[ivars.iymom] = HLLCfactor * S_c - U_state[iener] = HLLCfactor * (U_l[i, j, iener] / rho_l + + U_state[ivars.iener] = HLLCfactor * (U_l[i, j, ivars.ener] / rho_l + (S_c - un_l) * (S_c + p_l / (rho_l * (S_l - un_l)))) # species - if nspec > 0: - U_state[irhoX:irhoX + nspec] = HLLCfactor * \ - U_l[i, j, irhoX:irhoX + nspec] / rho_l + if ivars.naux > 0: + U_state[ivars.irhoX:ivars.irhoX + ivars.nspec] = HLLCfactor * \ + U_l[i, j, ivars.irhoX:ivars.irhoX + ivars.nspec] / rho_l # find the flux on the left interface - F[i, j, :] = consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, - U_l[i, j, :]) + F[i, j, :] = consFlux(idir, gamma, ivars, U_l[i, j, :]) # correct the flux F[i, j, :] = F[i, j, :] + S_l * (U_state[:] - U_l[i, j, :]) @@ -1056,8 +1015,7 @@ def riemann_hllc(idir, ng, # L region U_state[:] = U_l[i, j, :] - F[i, j, :] = consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, - U_state) + F[i, j, :] = consFlux(idir, gamma, ivars, U_state) # we should deal with solid boundaries somehow here @@ -1065,7 +1023,7 @@ def riemann_hllc(idir, ng, @njit(cache=True) -def consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, U_state): +def consFlux(idir, gamma, ivars, U_state): r""" Calculate the conservative flux. @@ -1075,7 +1033,7 @@ def consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, U_state): Are we predicting to the edges in the x-direction (1) or y-direction (2)? gamma : float Adiabatic index - idens, ixmom, iymom, iener, irhoX : int + ivars The indices of the density, x-momentum, y-momentum, internal energy density and species partial densities in the conserved state vector. nspec : int @@ -1091,28 +1049,28 @@ def consFlux(idir, gamma, idens, ixmom, iymom, iener, irhoX, nspec, U_state): F = np.zeros_like(U_state) - u = U_state[ixmom] / U_state[idens] - v = U_state[iymom] / U_state[idens] + u = U_state[ivars.ixmom] / U_state[ivars.idens] + v = U_state[ivars.iymom] / U_state[ivars.idens] - p = (U_state[iener] - 0.5 * U_state[idens] * (u * u + v * v)) * (gamma - 1.0) + p = (U_state[ivars.iener] - 0.5 * U_state[ivars.idens] * (u * u + v * v)) * (gamma - 1.0) if idir == 1: - F[idens] = U_state[idens] * u - F[ixmom] = U_state[ixmom] * u + p - F[iymom] = U_state[iymom] * u - F[iener] = (U_state[iener] + p) * u + F[ivars.idens] = U_state[ivars.idens] * u + F[ivars.ixmom] = U_state[ivars.ixmom] * u + p + F[ivars.iymom] = U_state[ivars.iymom] * u + F[ivars.iener] = (U_state[ivars.iener] + p) * u - if nspec > 0: - F[irhoX:irhoX + nspec] = U_state[irhoX:irhoX + nspec] * u + if ivars.naux > 0: + F[ivars.irhoX:ivars.irhoX + ivars.nspec] = U_state[ivars.irhoX:ivars.irhoX + ivars.naux] * u else: - F[idens] = U_state[idens] * v - F[ixmom] = U_state[ixmom] * v - F[iymom] = U_state[iymom] * v + p - F[iener] = (U_state[iener] + p) * v + F[ivars.idens] = U_state[ivars.idens] * v + F[ivars.ixmom] = U_state[ivars.ixmom] * v + F[ivars.iymom] = U_state[ivars.iymom] * v + p + F[ivars.iener] = (U_state[ivars.iener] + p) * v - if nspec > 0: - F[irhoX:irhoX + nspec] = U_state[irhoX:irhoX + nspec] * v + if ivars.naux > 0: + F[ivars.irhoX:ivars.irhoX + ivars.naux] = U_state[ivars.irhoX:ivars.irhoX + ivars.naux] * v return F diff --git a/pyro/compressible/unsplit_fluxes.py b/pyro/compressible/unsplit_fluxes.py index 87814a1f1..5f318eb48 100644 --- a/pyro/compressible/unsplit_fluxes.py +++ b/pyro/compressible/unsplit_fluxes.py @@ -217,11 +217,7 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt): tm_states = tc.timer("interfaceStates") tm_states.begin() - V_l, V_r = ifc.states(1, myg.ng, myg.dx, dt, - ivars.irho, ivars.iu, ivars.iv, ivars.ip, ivars.ix, - ivars.naux, - gamma, - q, ldx) + V_l, V_r = ifc.states(1, myg, dt, ivars, gamma, q, ldx) tm_states.end() @@ -236,11 +232,7 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt): # left and right primitive variable states tm_states.begin() - _V_l, _V_r = ifc.states(2, myg.ng, myg.dy, dt, - ivars.irho, ivars.iu, ivars.iv, ivars.ip, ivars.ix, - ivars.naux, - gamma, - q, ldy) + _V_l, _V_r = ifc.states(2, myg, dt, ivars, gamma, q, ldy) V_l = ai.ArrayIndexer(d=_V_l, grid=myg) V_r = ai.ArrayIndexer(d=_V_r, grid=myg) @@ -294,13 +286,11 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt): else: msg.fail("ERROR: Riemann solver undefined") - _fx = riemannFunc(1, myg.ng, - ivars.idens, ivars.ixmom, ivars.iymom, ivars.iener, ivars.irhox, ivars.naux, + _fx = riemannFunc(1, myg, ivars, solid.xl, solid.xr, gamma, U_xl, U_xr) - _fy = riemannFunc(2, myg.ng, - ivars.idens, ivars.ixmom, ivars.iymom, ivars.iener, ivars.irhox, ivars.naux, + _fy = riemannFunc(2, myg, ivars, solid.yl, solid.yr, gamma, U_yl, U_yr) @@ -393,13 +383,11 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt): tm_riem.begin() - _fx = riemannFunc(1, myg.ng, - ivars.idens, ivars.ixmom, ivars.iymom, ivars.iener, ivars.irhox, ivars.naux, + _fx = riemannFunc(1, myg, ivars, solid.xl, solid.xr, gamma, U_xl, U_xr) - _fy = riemannFunc(2, myg.ng, - ivars.idens, ivars.ixmom, ivars.iymom, ivars.iener, ivars.irhox, ivars.naux, + _fy = riemannFunc(2, myg, ivars, solid.yl, solid.yr, gamma, U_yl, U_yr) From e3ab93a5dce6a383216f9914797ec2ce210d8e7d Mon Sep 17 00:00:00 2001 From: Michael Zingale Date: Sun, 5 Mar 2023 10:16:16 -0500 Subject: [PATCH 2/2] try some numba stuff --- pyro/compressible/simulation.py | 2 +- pyro/mesh/patch.py | 56 +++++++++++++++++++++++++-------- 2 files changed, 44 insertions(+), 14 deletions(-) diff --git a/pyro/compressible/simulation.py b/pyro/compressible/simulation.py index 179e397f5..09735a51c 100644 --- a/pyro/compressible/simulation.py +++ b/pyro/compressible/simulation.py @@ -54,7 +54,7 @@ def __init__(self, myd): def cons_to_prim(U, gamma, ivars, myg): """ convert an input vector of conserved variables to primitive variables """ - q = myg.scratch_array(nvar=ivars.nq) + q = myg.scratch_array_n(ivars.nq) q[:, :, ivars.irho] = U[:, :, ivars.idens] q[:, :, ivars.iu] = U[:, :, ivars.ixmom]/U[:, :, ivars.idens] diff --git a/pyro/mesh/patch.py b/pyro/mesh/patch.py index 52b3e85c1..13a7d0754 100644 --- a/pyro/mesh/patch.py +++ b/pyro/mesh/patch.py @@ -34,11 +34,42 @@ import h5py import numpy as np +from numba import int32, float64 +from numba.experimental import jitclass + import pyro.mesh.boundary as bnd from pyro.mesh.array_indexer import ArrayIndexer, ArrayIndexerFC from pyro.util import msg - +grid_spec = [ + ('nx', int32), + ('ny', int32), + ('ng', int32), + ('qx', int32), + ('qy', int32), + ('xmin', float64), + ('xmax', float64), + ('ymin', float64), + ('ymax', float64), + ('ilo', int32), + ('ihi', int32), + ('jlo', int32), + ('jhi', int32), + ('ic', int32), + ('jc', int32), + ('dx', float64), + ('xl', float64[:]), + ('xr', float64[:]), + ('x', float64[:]), + ('dy', float64), + ('yl', float64[:]), + ('yr', float64[:]), + ('y', float64[:]), + ('x2d', float64[:,:]), + ('y2d', float64[:,:]) +] + +@jitclass(grid_spec) class Grid2d: """ the 2-d grid class. The grid object will contain the coordinate @@ -134,24 +165,23 @@ def __init__(self, nx, ny, ng=1, self.y = 0.5*(self.yl + self.yr) # 2-d versions of the zone coordinates (replace with meshgrid?) - x2d = np.repeat(self.x, self.qy) - x2d.shape = (self.qx, self.qy) - self.x2d = x2d + self.x2d = np.repeat(self.x, self.qy).reshape(self.qx, self.qy) + self.y2d = np.transpose(np.repeat(self.y, self.qx).reshape(self.qy, self.qx)) - y2d = np.repeat(self.y, self.qx) - y2d.shape = (self.qy, self.qx) - y2d = np.transpose(y2d) - self.y2d = y2d + def scratch_array(self): + """ + return a standard numpy array dimensioned to have the size + and number of ghostcells as the parent grid + """ + _tmp = np.zeros((self.qx, self.qy), dtype=np.float64) + return ArrayIndexer(d=_tmp, grid=self) - def scratch_array(self, nvar=1): + def scratch_array_n(self, nvar): """ return a standard numpy array dimensioned to have the size and number of ghostcells as the parent grid """ - if nvar == 1: - _tmp = np.zeros((self.qx, self.qy), dtype=np.float64) - else: - _tmp = np.zeros((self.qx, self.qy, nvar), dtype=np.float64) + _tmp = np.zeros((self.qx, self.qy, nvar), dtype=np.float64) return ArrayIndexer(d=_tmp, grid=self) def coarse_like(self, N):