[WIP] fix the sponge source (#336)

I forgot to include the kinetic energy source

Author: zingale

pyro/compressible/simulation.py

@@ -427,9 +427,19 @@ def evolve(self):
         if self.rp.get_param("sponge.do_sponge"):
             kappa_f = get_sponge_factor(self.cc_data.data, self.ivars, self.rp, myg)
 
+            U_old = self.cc_data.data.copy()
+
             self.cc_data.data[:, :, self.ivars.ixmom] /= (1.0 + self.dt * kappa_f)
             self.cc_data.data[:, :, self.ivars.iymom] /= (1.0 + self.dt * kappa_f)
 
+            # for energy, there is no change in density from the sponge, so we
+            # can just apply the change in kinetic energy from the velocity update
+            dke = 0.5 * ((self.cc_data.data[..., self.ivars.ixmom]**2 +
+                          self.cc_data.data[..., self.ivars.iymom]**2) -
+                         (U_old[..., self.ivars.ixmom]**2 +
+                          U_old[..., self.ivars.iymom]**2)) / self.cc_data.data[..., self.ivars.idens]
+            self.cc_data.data[..., self.ivars.iener] += dke
+
         if self.particles is not None:
             self.particles.update_particles(self.dt)
 

pyro/compressible_fv4/simulation.py

@@ -54,9 +54,14 @@ def substep(self, myd):
         if self.rp.get_param("sponge.do_sponge"):
             kappa_f = get_sponge_factor(myd.data, self.ivars, self.rp, myg)
 
+            # momentum
             k.v(n=self.ivars.ixmom)[:, :] -= kappa_f.v() * myd.data.v(n=self.ivars.ixmom)
             k.v(n=self.ivars.iymom)[:, :] -= kappa_f.v() * myd.data.v(n=self.ivars.iymom)
 
+            # total energy
+            k.v(n=self.ivars.iener)[:, :] -= kappa_f.v() * (myd.data.v(n=self.ivars.ixmom)**2 / myd.data.v(n=self.ivars.idens) +
+                                                            myd.data.v(n=self.ivars.iymom)**2 / myd.data.v(n=self.ivars.idens))
+
         return k
 
     def preevolve(self):

pyro/compressible_rk/simulation.py

@@ -39,9 +39,14 @@ def substep(self, myd):
         if self.rp.get_param("sponge.do_sponge"):
             kappa_f = compressible.get_sponge_factor(myd.data, self.ivars, self.rp, myg)
 
+            # momentum
             k.v(n=self.ivars.ixmom)[:, :] -= kappa_f.v() * myd.data.v(n=self.ivars.ixmom)
             k.v(n=self.ivars.iymom)[:, :] -= kappa_f.v() * myd.data.v(n=self.ivars.iymom)
 
+            # total energy
+            k.v(n=self.ivars.iener)[:, :] -= kappa_f.v() * (myd.data.v(n=self.ivars.ixmom)**2 / myd.data.v(n=self.ivars.idens) +
+                                                            myd.data.v(n=self.ivars.iymom)**2 / myd.data.v(n=self.ivars.idens))
+
         return k
 
     def method_compute_timestep(self):