fix transverse pressure gradient addition

pyro/compressible/unsplit_fluxes.py

@@ -285,9 +285,6 @@ def apply_source_terms(U_xl, U_xr, U_yl, U_yr,
     tm_source.begin()
 
     myg = my_data.grid
-
-    pres = my_data.get_var("pressure")
-
     dens_src = my_aux.get_var("dens_src")
     xmom_src = my_aux.get_var("xmom_src")
     ymom_src = my_aux.get_var("ymom_src")
@@ -301,11 +298,6 @@ def apply_source_terms(U_xl, U_xr, U_yl, U_yr,
     ymom_src[:, :] = U_src[:, :, ivars.iymom]
     E_src[:, :] = U_src[:, :, ivars.iener]
 
-    # apply non-conservative pressure gradient
-    if myg.coord_type == 1:
-        xmom_src.v()[:, :] += -(pres.ip(1) - pres.v()) / myg.Lx.v()
-        ymom_src.v()[:, :] += -(pres.jp(1) - pres.v()) / myg.Ly.v()
-
     my_aux.fill_BC("dens_src")
     my_aux.fill_BC("xmom_src")
     my_aux.fill_BC("ymom_src")
@@ -410,19 +402,42 @@ def apply_transverse_flux(U_xl, U_xr, U_yl, U_yr,
         with source terms added.
     """
 
+    myg = my_data.grid
+
     # Use Riemann Solver to get interface flux using the left and right states
 
-    F_x = riemann.riemann_flux(1, U_xl, U_xr,
-                               my_data, rp, ivars,
-                               solid.xl, solid.xr, tc)
+    if myg.coord_type == 1:
+        # We need pressure from interface state for transverse update for
+        # SphericalPolar geometry. So we need interface conserved states.
+        F_x, U_x = riemann.riemann_flux(1, U_xl, U_xr,
+                                        my_data, rp, ivars,
+                                        solid.xl, solid.xr, tc,
+                                        return_cons=True)
 
-    F_y = riemann.riemann_flux(2, U_yl, U_yr,
-                               my_data, rp, ivars,
-                               solid.yl, solid.yr, tc)
+        F_y, U_y = riemann.riemann_flux(2, U_yl, U_yr,
+                                        my_data, rp, ivars,
+                                        solid.yl, solid.yr, tc,
+                                        return_cons=True)
 
-    # Now we update the conserved states using the transverse fluxes.
+        gamma = rp.get_param("eos.gamma")
 
-    myg = my_data.grid
+        # Find primitive variable since we need pressure in transverse update.
+        qx = comp.cons_to_prim(U_x, gamma, ivars, myg)
+        qy = comp.cons_to_prim(U_y, gamma, ivars, myg)
+
+    else:
+        # Directly calculate the interface flux using Riemann Solver
+        F_x = riemann.riemann_flux(1, U_xl, U_xr,
+                                   my_data, rp, ivars,
+                                   solid.xl, solid.xr, tc,
+                                   return_cons=False)
+
+        F_y = riemann.riemann_flux(2, U_yl, U_yr,
+                                   my_data, rp, ivars,
+                                   solid.yl, solid.yr, tc,
+                                   return_cons=False)
+
+    # Now we update the conserved states using the transverse fluxes.
 
     tm_transverse = tc.timer("transverse flux addition")
     tm_transverse.begin()
@@ -452,6 +467,25 @@ def apply_transverse_flux(U_xl, U_xr, U_yl, U_yr,
             - hdtV.v(buf=b)*(F_x.ip(1, buf=b, n=n)*myg.Ax.ip(1, buf=b) -
                              F_x.v(buf=b, n=n)*myg.Ax.v(buf=b))
 
+    # apply non-conservative pressure gradient for momentum in spherical geometry
+    # Note that we should only apply this pressure gradient
+    # to the momentum corresponding to the transverse direction.
+    # The momentum in the normal direction already updated pressure during reconstruction.
+
+    if myg.coord_type == 1:
+
+        U_xl.v(buf=b, n=ivars.iymom)[:, :] += - (qy.ip_jp(-1, 1, buf=b, n=ivars.ip) -
+                                                 qy.ip(-1, buf=b, n=ivars.ip)) / myg.Ly.v(buf=b)
+
+        U_xr.v(buf=b, n=ivars.iymom)[:, :] += - (qy.jp(1, buf=b, n=ivars.ip) -
+                                                 qy.v(buf=b, n=ivars.ip)) / myg.Ly.v(buf=b)
+
+        U_yl.v(buf=b, n=ivars.ixmom)[:, :] += - (qx.ip_jp(1, -1, buf=b, n=ivars.ip) -
+                                                 qx.jp(-1, buf=b, n=ivars.ip)) / myg.Lx.v(buf=b)
+
+        U_yr.v(buf=b, n=ivars.ixmom)[:, :] += - (qx.ip(1, buf=b, n=ivars.ip) -
+                                                 qx.v(buf=b, n=ivars.ip)) / myg.Lx.v(buf=b)
+
     tm_transverse.end()
 
     return U_xl, U_xr, U_yl, U_yr