From ef410927258687862bfa4aac1a8e90e0e123a660 Mon Sep 17 00:00:00 2001
From: Gabriele Bozzola <gbozzola@caltech.edu>
Date: Fri, 27 Oct 2023 10:21:26 -0700
Subject: [PATCH] Cache: define AtmosCache

---
 src/cache/cache.jl                            | 85 ++++++++++++++++++-
 src/callbacks/callbacks.jl                    |  1 +
 .../sponge/viscous_sponge.jl                  |  2 +-
 test/coupler_compatibility.jl                 | 27 +++++-
 4 files changed, 111 insertions(+), 4 deletions(-)

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
index 9fd9fec934a..a410d0ef2f7 100644
--- a/src/cache/cache.jl
+++ b/src/cache/cache.jl
@@ -1,3 +1,84 @@
+Base.@kwdef struct AtmosCache{
+    FT <: Real,
+    SIM,
+    AM,
+    SFC,
+    COR,
+    CAP,
+    GHOST,
+    ENV,
+    PREC,
+    SCRA,
+    HYPE,
+    RS,
+    VS,
+    PR,
+    SUB,
+    LSAD,
+    EDMFCOR,
+    FOR,
+    NONGW,
+    ORGW,
+    RAD,
+    DSS,
+    NETFLUXTOA,
+    NETFLUXSFC,
+}
+    """Timestep of the simulation (in seconds). This is also used by callbacks and tendencies"""
+    dt::FT
+
+    """NamedTuple with job_id, output_dir (used in post), and start_date (used for insolation)."""
+    # TODO: Remove simulation
+    simulation::SIM
+
+    """AtmosModel"""
+    atmos::AM
+
+    """ClimaAtmosParameters that have to be used"""
+    params::CAP
+
+    """Variables that are used generally, such as ᶜρ_ref, ᶜΦ"""
+    core::COR
+
+    """Used by update_surface_conditions! in set_precomputed_quantities! and coupler"""
+    sfc_setup::SFC
+
+    """Center and face ghost buffers used by DSS"""
+    ghost_buffer::GHOST
+
+    env_thermo_quad::ENV
+
+    """Quantities that are updated with set_precomputed_quantities!"""
+    precomputed::PREC
+
+    """Pre-allocated areas of memory to store temporary values"""
+    scratch::SCRA
+
+    """Hyperdiffision quantities for grid and subgrid scale quantities, potentially with
+       ghost buffers for DSS"""
+    hyperdiff::HYPE
+
+    """Additional parameters used by the various tendencies"""
+    rayleigh_sponge::RS
+    viscous_sponge::VS
+    precipitation::PR
+    subsidence::SUB
+    large_scale_advection::LSAD
+    edmf_coriolis::EDMFCOR
+    forcing::FOR
+    non_orographic_gravity_wave::NONGW
+    orographic_gravity_wave::ORGW
+    radiation::RAD
+
+    do_dss::DSS
+
+    """Net energy flux coming through top of atmosphere and surface"""
+    net_energy_flux_toa::NETFLUXTOA
+    net_energy_flux_sfc::NETFLUXSFC
+end
+
+Base.broadcastable(x::AtmosCache) = tuple(x)
+
 # Functions on which the model depends:
 # CAP.R_d(params)         # dry specific gas constant
 # CAP.kappa_d(params)     # dry adiabatic exponent
@@ -86,7 +167,7 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
     orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos)
     radiation = radiation_model_cache(Y, atmos, radiation_args...)
 
-    p = (;
+    p = AtmosCache(;
         simulation,
         atmos,
         core,
@@ -108,7 +189,7 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
         radiation,
         env_thermo_quad,
         ghost_buffer,
-        dt = simulation.dt,
+        simulation.dt,
         net_energy_flux_toa,
         net_energy_flux_sfc,
     )
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
index a819bb4080f..5806d564764 100644
--- a/src/callbacks/callbacks.jl
+++ b/src/callbacks/callbacks.jl
@@ -43,6 +43,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
     t = integrator.t
 
     (; ᶜts, sfc_conditions) = p.precomputed
+    (; params) = p
     (; idealized_insolation, idealized_h2o, idealized_clouds) = p.radiation
     (; insolation_tuple, ᶠradiation_flux, radiation_model) = p.radiation
 
diff --git a/src/parameterized_tendencies/sponge/viscous_sponge.jl b/src/parameterized_tendencies/sponge/viscous_sponge.jl
index 783826eee6e..97279acfce3 100644
--- a/src/parameterized_tendencies/sponge/viscous_sponge.jl
+++ b/src/parameterized_tendencies/sponge/viscous_sponge.jl
@@ -30,7 +30,7 @@ add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t) =
 
 function add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t, ::TotalEnergy)
     (; ᶜβ_viscous) = p.viscous_sponge
-    (; ᶜh_tot) = p
+    (; ᶜh_tot) = p.precomputed
     ᶜρ = Y.c.ρ
     @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
 end
diff --git a/test/coupler_compatibility.jl b/test/coupler_compatibility.jl
index 56836196f02..8a22818f258 100644
--- a/test/coupler_compatibility.jl
+++ b/test/coupler_compatibility.jl
@@ -55,7 +55,32 @@ const T2 = 290
     )
     sfc_setup = similar(Spaces.level(Y.f, half), typeof(surface_state))
     @. sfc_setup = (surface_state,)
-    p_overwritten = merge(p, (; sfc_setup))
+    p_overwritten = CA.AtmosCache(;
+        p.simulation,
+        p.atmos,
+        p.core,
+        sfc_setup,
+        p.params,
+        p.do_dss,
+        p.precomputed,
+        p.scratch,
+        p.hyperdiff,
+        p.rayleigh_sponge,
+        p.viscous_sponge,
+        p.precipitation,
+        p.subsidence,
+        p.large_scale_advection,
+        p.edmf_coriolis,
+        p.forcing,
+        p.non_orographic_gravity_wave,
+        p.orographic_gravity_wave,
+        p.radiation,
+        p.env_thermo_quad,
+        p.ghost_buffer,
+        p.simulation.dt,
+        p.net_energy_flux_toa,
+        p.net_energy_flux_sfc,
+    )
 
     # Test that set_precomputed_quantities! can be used to update the surface
     # temperature to T1 and then to T2.