diff --git a/pyro/compressible_fv4/simulation.py b/pyro/compressible_fv4/simulation.py index b47d19821..5f096be68 100644 --- a/pyro/compressible_fv4/simulation.py +++ b/pyro/compressible_fv4/simulation.py @@ -1,5 +1,6 @@ import pyro.compressible_fv4.fluxes as flx from pyro import compressible_rk +from pyro.compressible import get_external_sources from pyro.mesh import fv, integration @@ -17,24 +18,23 @@ def substep(self, myd): """ myg = myd.grid - grav = self.rp.get_param("compressible.grav") - # compute the source terms -- we need to do these to 4th - # order. Start by evaluating the sources using the - # cell-center quantities (including one ghost cell. - dens_cc = myd.to_centers("density") - ymom_cc = myd.to_centers("y-momentum") + # compute the source terms -- we need to do this first + # using the cell-center data and then convert it back to + # averages - ymom_src = myg.scratch_array() - ymom_src.v(buf=1)[:, :] = dens_cc.v(buf=1)[:, :]*grav + U_cc = myg.scratch_array(nvar=self.ivars.nvar) + for n in range(self.ivars.nvar): + U_cc[:, :, n] = myd.to_centers(myd.names[n]) - E_src = myg.scratch_array() - E_src.v(buf=1)[:, :] = ymom_cc.v(buf=1)[:, :]*grav + # cell-centered sources + S = get_external_sources(myd.t, self.dt, U_cc, + self.ivars, self.rp, myg) - # now bring back to averages -- we only need this in the - # interior (no ghost cells) - ymom_src.v()[:, :] = ymom_src.v()[:, :] - myg.dx**2*ymom_src.lap()/24.0 - E_src.v()[:, :] = E_src.v()[:, :] - myg.dx**2*E_src.lap()/24.0 + # bring the sources back to averages -- we only care about + # the interior (no ghost cells) + for n in range(self.ivars.nvar): + S.v(n=n)[:, :] -= myg.dx**2 * S.lap(n=n) / 24.0 k = myg.scratch_array(nvar=self.ivars.nvar) @@ -43,10 +43,7 @@ def substep(self, myd): for n in range(self.ivars.nvar): k.v(n=n)[:, :] = \ (flux_x.v(n=n) - flux_x.ip(1, n=n))/myg.dx + \ - (flux_y.v(n=n) - flux_y.jp(1, n=n))/myg.dy - - k.v(n=self.ivars.iymom)[:, :] += ymom_src.v()[:, :] - k.v(n=self.ivars.iener)[:, :] += E_src.v()[:, :] + (flux_y.v(n=n) - flux_y.jp(1, n=n))/myg.dy + S.v(n=n) return k diff --git a/pyro/compressible_rk/simulation.py b/pyro/compressible_rk/simulation.py index 56be8d8f8..a4d5aad35 100644 --- a/pyro/compressible_rk/simulation.py +++ b/pyro/compressible_rk/simulation.py @@ -16,17 +16,11 @@ def substep(self, myd): """ myg = myd.grid - grav = self.rp.get_param("compressible.grav") - # compute the source terms - dens = myd.get_var("density") - ymom = myd.get_var("y-momentum") - - ymom_src = myg.scratch_array() - ymom_src.v()[:, :] = dens.v()[:, :]*grav - - E_src = myg.scratch_array() - E_src.v()[:, :] = ymom.v()[:, :]*grav + # source terms -- note: this dt is the entire dt, not the + # stage's dt + S = compressible.get_external_sources(myd.t, self.dt, myd.data, + self.ivars, self.rp, myg) k = myg.scratch_array(nvar=self.ivars.nvar) @@ -36,10 +30,7 @@ def substep(self, myd): for n in range(self.ivars.nvar): k.v(n=n)[:, :] = \ (flux_x.v(n=n) - flux_x.ip(1, n=n))/myg.dx + \ - (flux_y.v(n=n) - flux_y.jp(1, n=n))/myg.dy - - k.v(n=self.ivars.iymom)[:, :] += ymom_src.v()[:, :] - k.v(n=self.ivars.iener)[:, :] += E_src.v()[:, :] + (flux_y.v(n=n) - flux_y.jp(1, n=n))/myg.dy + S.v(n=n) return k