From 89a07d7f41df2bd240be11641f565de62341b8d5 Mon Sep 17 00:00:00 2001
From: Zhi Chen <62574124+zhichen3@users.noreply.github.com>
Date: Sun, 1 Dec 2024 10:42:50 -0500
Subject: [PATCH] Fix transverse pres gradient for spherical geometry. (#296)

during calculation of the interface state, we should only apply pressure gradient to the momentum in the transverse direction since the normal direction momentum is already includes pressure term during reconstruction.
---
 pyro/compressible/unsplit_fluxes.py | 69 ++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 17 deletions(-)

diff --git a/pyro/compressible/unsplit_fluxes.py b/pyro/compressible/unsplit_fluxes.py
index 64c656a9b..637051edb 100644
--- a/pyro/compressible/unsplit_fluxes.py
+++ b/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,25 +402,49 @@ 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()
 
     b = (2, 1)
-    hdtV = 0.5*dt / myg.V
+    hdt = 0.5*dt
+    hdtV = hdt / myg.V
 
     for n in range(ivars.nvar):
 
@@ -452,6 +468,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)[:, :] += - hdt * (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)[:, :] += - hdt * (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)[:, :] += - hdt * (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)[:, :] += - hdt * (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