From 868ebac70878ab80b0500a8e0ca303f95c849afd Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Tue, 24 Oct 2023 16:58:23 -0700
Subject: [PATCH] Cache: viscous_sponge_cache

---
 src/cache/cache.jl                            |  6 ++-
 .../sponge/viscous_sponge.jl                  | 41 ++++++++++++++-----
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
index c16f423c56..13bb6b9fdd 100644
--- a/src/cache/cache.jl
+++ b/src/cache/cache.jl
@@ -127,8 +127,10 @@ function additional_cache(
     end
 
     return merge(
-        (; rayleigh_sponge = rayleigh_sponge_cache(Y, atmos)),
-        viscous_sponge_cache(atmos.viscous_sponge, Y),
+        (;
+            rayleigh_sponge = rayleigh_sponge_cache(Y, atmos),
+            viscous_sponge = viscous_sponge_cache(Y, atmos),
+        ),
         precipitation_cache(Y, precip_model),
         subsidence_cache(Y, atmos.subsidence),
         large_scale_advection_cache(Y, atmos.ls_adv),
diff --git a/src/parameterized_tendencies/sponge/viscous_sponge.jl b/src/parameterized_tendencies/sponge/viscous_sponge.jl
index 804606635c..16c9e2517d 100644
--- a/src/parameterized_tendencies/sponge/viscous_sponge.jl
+++ b/src/parameterized_tendencies/sponge/viscous_sponge.jl
@@ -6,10 +6,13 @@ import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Spaces as Spaces
 
-viscous_sponge_cache(::Nothing, Y) = NamedTuple()
+viscous_sponge_cache(Y, atmos::AtmosModel) =
+    viscous_sponge_cache(Y, atmos.viscous_sponge)
+
+viscous_sponge_cache(Y, ::Nothing) = (;)
 viscous_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
-function viscous_sponge_cache(viscous_sponge::ViscousSponge, Y)
+function viscous_sponge_cache(Y, viscous_sponge::ViscousSponge)
     (; κ₂, zd) = viscous_sponge
     FT = Spaces.undertype(axes(Y.c))
     ᶜz = Fields.coordinate_field(Y.c).z
@@ -22,16 +25,34 @@ function viscous_sponge_cache(viscous_sponge::ViscousSponge, Y)
     return (; ᶜβ_viscous, ᶠβ_viscous)
 end
 
-function viscous_sponge_tendency!(Yₜ, Y, p, t, ::ViscousSponge)
-    (; ᶜβ_viscous, ᶠβ_viscous) = p
+add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t) =
+    add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t, p.atmos.energy_form)
+
+function add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t, ::TotalEnergy)
+    (; ᶜβ_viscous) = p.viscous_sponge
+    (; ᶜh_tot) = p
+    ᶜρ = Y.c.ρ
+    @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
+end
+
+function add_viscous_sponge_energy_tendency!(
+    Yₜ,
+    Y,
+    p,
+    t,
+    ::PotentialTemperature,
+)
+    (; ᶜβ_viscous) = p.viscous_sponge
     ᶜρ = Y.c.ρ
+    @. Yₜ.c.ρθ += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(Y.c.ρθ / ᶜρ))
+end
+
+function viscous_sponge_tendency!(Yₜ, Y, p, t, ::ViscousSponge)
+    (; ᶜβ_viscous, ᶠβ_viscous) = p.viscous_sponge
     ᶜuₕ = Y.c.uₕ
-    if :ρθ in propertynames(Y.c)
-        @. Yₜ.c.ρθ += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(Y.c.ρθ / ᶜρ))
-    elseif :ρe_tot in propertynames(Y.c)
-        (; ᶜh_tot) = p
-        @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
-    end
+
+    add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t)
+
     @. Yₜ.c.uₕ +=
         ᶜβ_viscous * (
             wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(