update with spherical polar source terms

pyro/compressible/interface.py

@@ -3,7 +3,7 @@
 
 
 @njit(cache=True)
-def states(idir, ng, dx, dt,
+def states(idir, grid, dt,
            irho, iu, iv, ip, ix, nspec,
            gamma, qv, dqv):
     r"""
@@ -65,10 +65,8 @@ 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, Cartesian2d, or SphericalPolar
+        The grid object.
     dt : float
         The timestep
     irho, iu, iv, ip, ix : int
@@ -94,16 +92,13 @@ def states(idir, ng, dx, dt,
     q_l = np.zeros_like(qv)
     q_r = np.zeros_like(qv)
 
-    nx = qx - 2 * ng
-    ny = qy - 2 * ng
-    ilo = ng
-    ihi = ng + nx
-    jlo = ng
-    jhi = ng + ny
-
     ns = nvar - nspec
 
-    dtdx = dt / dx
+    if idir == 1:
+        dtdx = dt / grid.Lx.v()
+    else:
+        dtdx = dt / grid.Ly.v()
+
     dtdx4 = 0.25 * dtdx
 
     lvec = np.zeros((nvar, nvar))
@@ -113,8 +108,8 @@ def states(idir, ng, dx, dt,
     betar = np.zeros(nvar)
 
     # 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, :]

pyro/compressible/simulation.py

@@ -145,6 +145,8 @@ def initialize(self, extra_vars=None, ng=4):
         # some auxiliary data that we'll need to fill GC in, but isn't
         # really part of the main solution
         aux_data = self.data_class(my_grid)
+        aux_data.register_var("dens_src", bc)
+        aux_data.register_var("xmom_src", bc_xodd)
         aux_data.register_var("ymom_src", bc_yodd)
         aux_data.register_var("E_src", bc)
         aux_data.create()
@@ -204,21 +206,29 @@ def evolve(self):
                                             self.ivars, self.solid, self.tc, self.dt)
 
         old_dens = dens.copy()
+        old_xmom = xmom.copy()
         old_ymom = ymom.copy()
 
         # conservative update
-        dtdV = self.dt / g.V.v(n=n)
+        dtdV = self.dt / g.V.v()
 
         for n in range(self.ivars.nvar):
             var = self.cc_data.get_var_by_index(n)
 
             var.v()[:, :] += dtdV * \
-                (Flux_x.v(n=n)*g.Ax_l.v(n=n) - Flux_x.ip(1, n=n)*g.Ax_r.v(n=n) +
-                 Flux_y.v(n=n)*g.Ay_l.v(n=n) - Flux_y.jp(1, n=n)*g.Ay_r.v(n=n))
+                (Flux_x.v(n=n)*g.Ax_l.v() - Flux_x.ip(1, n=n)*g.Ax_r.v() +
+                 Flux_y.v(n=n)*g.Ay_l.v() - Flux_y.jp(1, n=n)*g.Ay_r.v())
 
-        # gravitational source terms
-        ymom[:, :] += 0.5*self.dt*(dens[:, :] + old_dens[:, :])*grav
-        ener[:, :] += 0.5*self.dt*(ymom[:, :] + old_ymom[:, :])*grav
+        # Apply source terms
+
+        if isinstance(g, SphericalPolar):
+            xmom[:, :] += 0.5*self.dt*()
+            ymom[:, :] +=
+            ener[:, :] +=
+        else:
+            # gravitational source terms
+            ymom[:, :] += 0.5*self.dt*(dens[:, :] + old_dens[:, :])*grav
+            ener[:, :] += 0.5*self.dt*(ymom[:, :] + old_ymom[:, :])*grav
 
         if self.particles is not None:
             self.particles.update_particles(self.dt)

pyro/compressible/unsplit_fluxes.py

@@ -174,9 +174,6 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt):
     # =========================================================================
     # Q = (rho, u, v, p, {X})
 
-    dens = my_data.get_var("density")
-    ymom = my_data.get_var("y-momentum")
-
     q = comp.cons_to_prim(my_data.data, gamma, ivars, myg)
 
     # =========================================================================
@@ -217,7 +214,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,
+    V_l, V_r = ifc.states(1, myg, dt,
                           ivars.irho, ivars.iu, ivars.iv, ivars.ip, ivars.ix,
                           ivars.naux,
                           gamma,
@@ -236,7 +233,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,
+    _V_l, _V_r = ifc.states(2, myg, dt,
                             ivars.irho, ivars.iu, ivars.iv, ivars.ip, ivars.ix,
                             ivars.naux,
                             gamma,
@@ -253,29 +250,77 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt):
     # =========================================================================
     # apply source terms
     # =========================================================================
-    grav = rp.get_param("compressible.grav")
 
-    ymom_src = my_aux.get_var("ymom_src")
-    ymom_src.v()[:, :] = dens.v()*grav
-    my_aux.fill_BC("ymom_src")
+    dens = my_data.get_var("density")
+    xmom = my_data.get_var("x-momentum")
+    ymom = my_data.get_var("y-momentum")
+    E = my_data.get_var("energy")
 
+    dens_src = my_aux.get_var("dens_src")
+    xmom_src = my_aux.get_var("xmom_src")
+    ymom_src = my_aux.get_var("ymom_src")
     E_src = my_aux.get_var("E_src")
-    E_src.v()[:, :] = ymom.v()*grav
+
+    grav = rp.get_param("compressible.grav")
+
+    # src = [0.0 for n in range(ivars.nvar)]
+
+    # Calculate external source terms
+    if isinstance(myg, SphericalPolar):
+        # assume gravity points in r-direction in spherical
+        dens_src.v()[:, :] = 0.0
+        xmom_src.v()[:, :] = dens.v()*grav
+        ymom_src.v()[:, :] = 0.0
+        E_src.v()[:, :] = xmom.v()*grav
+
+    else:
+        # assume gravity points in y-direction in cartesian
+        dens_src.v()[:, :] = 0.0
+        xmom_src.v()[:, :] = 0.0
+        ymom_src.v()[:, :] = dens.v()*grav
+        E_src.v()[:, :] = ymom.v()*grav
+
+    # Calculate geometric+ expanded divergence terms for spherical geometry
+    if isinstance(myg, SphericalPolar):
+        dens_src.v()[:, :] += (-2.0*xmom.v()[:, :] -
+                               np.cot(myg.y2d.v()[:, :])*ymom.v()[:, :]) \
+                               / myg.x2d.v()[:, :]
+
+        xmom_src.v()[:, :] += (ymom.v()[:, :]**2 - xmom.v()[:, :]**2 -
+                               np.cot(myg.y2d.v()[:, :])*xmom.v()[:, :]*ymom.v()[:, :]) \
+                               / (dens.v()[:, :] * myg.x2d.v()[:, :])
+
+        ymom_src.v()[:, :] += (-3.0*xmom.v()[:, :]*ymom.v()[:, :] -
+                               np.cot(myg.y2d.v()[:, :])*ymom.v()[:, :]**2) \
+                               / (dens.v()[:, :] * myg.x2d.v()[:, :])
+
+        E_src.v()[:, :] += (-2.0*xmom.v()[:, :] * (E.v()[:, :] + q.v(n=ivars.ip)[:, :])
+                            - np.cot(myg.y2d.v()[:, :])*ymom.v()[:, :] *
+                            (E.v()[:, :] + q.v(n=ivars.ip)[:, :])) \
+                               / (dens.v()[:, :] * myg.x2d.v()[:, :])
+
+    my_aux.fill_BC("dens_src")
+    my_aux.fill_BC("xmom_src")
+    my_aux.fill_BC("ymom_src")
     my_aux.fill_BC("E_src")
 
     # ymom_xl[i,j] += 0.5*dt*dens[i-1,j]*grav
+    U_xl.v(buf=1, n=ivars.ixmom)[:, :] += 0.5*dt*xmom_src.ip(-1, buf=1)
     U_xl.v(buf=1, n=ivars.iymom)[:, :] += 0.5*dt*ymom_src.ip(-1, buf=1)
     U_xl.v(buf=1, n=ivars.iener)[:, :] += 0.5*dt*E_src.ip(-1, buf=1)
 
     # ymom_xr[i,j] += 0.5*dt*dens[i,j]*grav
+    U_xr.v(buf=1, n=ivars.ixmom)[:, :] += 0.5*dt*xmom_src.v(buf=1)
     U_xr.v(buf=1, n=ivars.iymom)[:, :] += 0.5*dt*ymom_src.v(buf=1)
     U_xr.v(buf=1, n=ivars.iener)[:, :] += 0.5*dt*E_src.v(buf=1)
 
     # ymom_yl[i,j] += 0.5*dt*dens[i,j-1]*grav
+    U_yl.v(buf=1, n=ivars.ixmom)[:, :] += 0.5*dt*xmom_src.jp(-1, buf=1)
     U_yl.v(buf=1, n=ivars.iymom)[:, :] += 0.5*dt*ymom_src.jp(-1, buf=1)
     U_yl.v(buf=1, n=ivars.iener)[:, :] += 0.5*dt*E_src.jp(-1, buf=1)
 
     # ymom_yr[i,j] += 0.5*dt*dens[i,j]*grav
+    U_yr.v(buf=1, n=ivars.ixmom)[:, :] += 0.5*dt*xmom_src.v(buf=1)
     U_yr.v(buf=1, n=ivars.iymom)[:, :] += 0.5*dt*ymom_src.v(buf=1)
     U_yr.v(buf=1, n=ivars.iener)[:, :] += 0.5*dt*E_src.v(buf=1)
 
@@ -360,8 +405,8 @@ def unsplit_fluxes(my_data, my_aux, rp, ivars, solid, tc, dt):
     tm_transverse = tc.timer("transverse flux addition")
     tm_transverse.begin()
 
-    dtdx = dt/myg.dx
-    dtdy = dt/myg.dy
+    dtdx = dt/myg.Lx
+    dtdy = dt/myg.Ly
 
     b = (2, 1)