Cache: precomputed

post_processing/post_processing_funcs.jl

@@ -228,7 +228,7 @@ function custom_postprocessing(sol, output_dir, p)
     anim = @animate for (Y, t) in zip(sol.u, sol.t)
         CA.set_precomputed_quantities!(Y, p, t) # sets ᶜts
         Plots.plot(
-            vec(TD.air_temperature.(thermo_params, p.ᶜts)),
+            vec(TD.air_temperature.(thermo_params, p.precomputed.ᶜts)),
             vec(Fields.coordinate_field(Y.c).z ./ 1000);
             xlabel = "T [K]",
             ylabel = "z [km]",
@@ -294,23 +294,23 @@ function postprocessing_plane(sol, output_dir, p)
     end
 
     gen_plot_plane(
-        Geometry.UVector.(p.ᶜu),
+        Geometry.UVector.(p.precomputed.ᶜu),
         "horz_velocity.png",
         "Horizontal Velocity",
         "u[m/s]",
         "z[m]",
     )
 
     gen_plot_plane(
-        Geometry.WVector.(p.ᶜu),
+        Geometry.WVector.(p.precomputed.ᶜu),
         "vert_velocity.png",
         "Vertical Velocity",
         "w[m/s]",
         "z[m]",
     )
 
     gen_plot_plane(
-        TD.virtual_pottemp.(thermo_params, p.ᶜts),
+        TD.virtual_pottemp.(thermo_params, p.precomputed.ᶜts),
         "virtual_pottemp.png",
         "Virtual Pottemp",
         "Theta[K]",

src/cache/cache.jl

@@ -52,55 +52,72 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
     net_energy_flux_toa = [Geometry.WVector(FT(0))]
     net_energy_flux_sfc = [Geometry.WVector(FT(0))]
 
-    default_cache = (;
-        is_init = Ref(true),
-        simulation,
-        atmos,
-        sfc_setup = surface_setup(params),
-        limiter,
+    core = (
         ᶜΦ,
         ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
         ᶜρ_ref,
         ᶜp_ref,
         ᶜT = similar(Y.c, FT),
         ᶜf,
         ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
+    )
+
+    sfc_setup = surface_setup(params)
+    scratch = temporary_quantities(Y, atmos)
+    env_thermo_quad = SGSQuadrature(FT)
+
+    precomputed = precomputed_quantities(Y, atmos)
+    precomputing_arguments =
+        (; atmos, core, params, sfc_setup, precomputed, scratch, simulation)
+
+    # Coupler compatibility
+    isnothing(precomputing_arguments.sfc_setup) &&
+        SurfaceConditions.set_dummy_surface_conditions!(precomputing_arguments)
+
+    set_precomputed_quantities!(Y, precomputing_arguments, FT(0))
+
+    radiation_args =
+        atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode ?
+        (params, precomputed.ᶜp) : ()
+
+    hyperdiff = hyperdiffusion_cache(Y, atmos)
+    rayleigh_sponge = rayleigh_sponge_cache(Y, atmos)
+    viscous_sponge = viscous_sponge_cache(Y, atmos)
+    precipitation = precipitation_cache(Y, atmos)
+    subsidence = subsidence_cache(Y, atmos)
+    large_scale_advection = large_scale_advection_cache(Y, atmos)
+    edmf_coriolis = edmf_coriolis_cache(Y, atmos)
+    forcing = forcing_cache(Y, atmos)
+    non_orographic_gravity_wave = non_orographic_gravity_wave_cache(Y, atmos)
+    orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos)
+    radiation = radiation_model_cache(Y, atmos, radiation_args...)
+
+    p = (;
+        simulation,
+        atmos,
+        core,
+        sfc_setup,
         params,
         do_dss,
+        precomputed,
+        scratch,
+        hyperdiff,
+        rayleigh_sponge,
+        viscous_sponge,
+        precipitation,
+        subsidence,
+        large_scale_advection,
+        edmf_coriolis,
+        forcing,
+        non_orographic_gravity_wave,
+        orographic_gravity_wave,
+        radiation,
+        env_thermo_quad,
         ghost_buffer,
+        dt = simulation.dt,
         net_energy_flux_toa,
         net_energy_flux_sfc,
-        env_thermo_quad = SGSQuadrature(FT),
-        precomputed_quantities(Y, atmos)...,
-        scratch = temporary_quantities(Y, atmos),
         numerics,
     )
-    set_precomputed_quantities!(Y, default_cache, FT(0))
-    default_cache.is_init[] = false
-
-    radiation_args =
-        atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode ?
-        (params, default_cache.ᶜp) : ()
-
-    return merge(
-        (;
-            hyperdiff = hyperdiffusion_cache(Y, atmos),
-            rayleigh_sponge = rayleigh_sponge_cache(Y, atmos),
-            viscous_sponge = viscous_sponge_cache(Y, atmos),
-            precipitation = precipitation_cache(Y, atmos),
-            large_scale_advection = large_scale_advection_cache(Y, atmos),
-            subsidence = subsidence_cache(Y, atmos),
-            edmf_coriolis = edmf_coriolis_cache(Y, atmos),
-            forcing = forcing_cache(Y, atmos),
-            non_orographic_gravity_wave = non_orographic_gravity_wave_cache(
-                Y,
-                atmos,
-            ),
-            orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos),
-            radiation = radiation_model_cache(Y, atmos, radiation_args...),
-        ),
-        (; Δt = simulation.dt),
-        (; atmos.turbconv_model),
-        default_cache,
-    )
+    return p
 end

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -105,13 +105,13 @@ function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(Y, p, t)
     (; turbconv_model) = p.atmos
     FT = eltype(Y)
     n = n_mass_flux_subdomains(turbconv_model)
-
-    (; ᶜp, ᶜΦ, ᶠu³, ᶜh_tot, ᶜK) = p
-    (; q_tot) = p.ᶜspecific
+    (; ᶜΦ) = p.core
+    (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed
+    (; q_tot) = p.precomputed.ᶜspecific
     (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) =
-        p.sfc_conditions
-    (; ᶜρaʲs, ᶠu³ʲs, ᶜuʲs, ᶜKʲs, ᶜh_totʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p
-    (; ᶠu³⁰, ᶜK⁰, ᶜh_tot⁰) = p
+        p.precomputed.sfc_conditions
+    (; ᶜρaʲs, ᶠu³ʲs, ᶜKʲs, ᶜh_totʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
+    (; ᶠu³⁰, ᶜK⁰, ᶜh_tot⁰) = p.precomputed
 
     thermo_params = CAP.thermodynamics_params(p.params)
 
@@ -231,9 +231,10 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
     ᶜz = Fields.coordinate_field(Y.c).z
     (; params) = p
     (; dt) = p.simulation
-    (; ᶜp, ᶜΦ, ᶜρ_ref, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p
-    (; q_tot) = p.ᶜspecific
-    (; buoyancy_flux) = p.sfc_conditions
+    (; ᶜΦ, ᶜρ_ref) = p.core
+    (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
+    (; q_tot) = p.precomputed.ᶜspecific
+    (; buoyancy_flux) = p.precomputed.sfc_conditions
     (;
         ᶜρaʲs,
         ᶠu³ʲs,
@@ -247,8 +248,8 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
         ᶜnh_pressureʲs,
         ᶜS_q_totʲs,
         ᶜS_e_totʲs_helper,
-    ) = p
-    (; ᶠu³⁰, ᶜK⁰, ᶜh_tot⁰) = p
+    ) = p.precomputed
+    (; ᶠu³⁰, ᶜK⁰, ᶜh_tot⁰) = p.precomputed
     thermo_params = CAP.thermodynamics_params(params)
     microphys_params = CAP.microphysics_params(params)
 
@@ -618,8 +619,8 @@ Updates the top boundary condition of precomputed quantities stored in `p` for d
 """
 function set_diagnostic_edmf_precomputed_quantities_top_bc!(Y, p, t)
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
-    (; ᶜentrʲs, ᶜdetrʲs, ᶜS_q_totʲs, ᶜS_e_totʲs_helper) = p
-    (; ᶠu³⁰, ᶠu³ʲs, ᶜuʲs) = p
+    (; ᶜentrʲs, ᶜdetrʲs, ᶜS_q_totʲs, ᶜS_e_totʲs_helper) = p.precomputed
+    (; ᶠu³⁰, ᶠu³ʲs, ᶜuʲs) = p.precomputed
 
     # set values for the top level
     i_top = Spaces.nlevels(axes(Y.c))
@@ -665,13 +666,13 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
     ᶜdz = Fields.Δz_field(axes(Y.c))
     (; params) = p
     (; dt) = p.simulation
-    (; ᶜp, ᶜu, ᶜts) = p
-    (; q_tot) = p.ᶜspecific
-    (; ustar, obukhov_length) = p.sfc_conditions
-    (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs) = p
-    (; ᶜtke⁰, ᶠu³⁰, ᶜS_q_tot⁰, ᶜu⁰) = p
-    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜmixing_length) = p
-    (; ᶜK_h, ᶜK_u, ρatke_flux) = p
+    (; ᶜp, ᶜu, ᶜts) = p.precomputed
+    (; q_tot) = p.precomputed.ᶜspecific
+    (; ustar, obukhov_length) = p.precomputed.sfc_conditions
+    (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs) = p.precomputed
+    (; ᶜtke⁰, ᶠu³⁰, ᶜS_q_tot⁰, ᶜu⁰) = p.precomputed
+    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜmixing_length) = p.precomputed
+    (; ᶜK_h, ᶜK_u, ρatke_flux) = p.precomputed
     thermo_params = CAP.thermodynamics_params(params)
     microphys_params = CAP.microphysics_params(params)
     ᶜlg = Fields.local_geometry_field(Y.c)

src/cache/precomputed_quantities.jl

@@ -302,7 +302,8 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     thermo_params = CAP.thermodynamics_params(p.params)
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_args = (thermo_params, energy_form, moisture_model)
-    (; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜΦ) = p
+    (; ᶜΦ) = p.core
+    (; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3
 
     @. ᶜspecific = specific_gs(Y.c)
@@ -333,12 +334,14 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
 
     if energy_form isa TotalEnergy
-        (; ᶜh_tot) = p
+        (; ᶜh_tot) = p.precomputed
         @. ᶜh_tot =
             TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot)
     end
 
-    SurfaceConditions.update_surface_conditions!(Y, p, t)
+    if !isnothing(p.sfc_setup)
+        SurfaceConditions.update_surface_conditions!(Y, p, t)
+    end
 
     if turbconv_model isa PrognosticEDMFX
         set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
@@ -365,10 +368,16 @@ values of the first updraft.
 function output_prognostic_sgs_quantities(Y, p, t)
     (; turbconv_model) = p.atmos
     thermo_params = CAP.thermodynamics_params(p.params)
-    (; ᶜρa⁰, ᶜρ⁰, ᶜtsʲs) = p
+    (; ᶜρa⁰, ᶜρ⁰, ᶜtsʲs) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3
     set_ᶠuₕ³!(ᶠuₕ³, Y)
-    (ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶜK⁺) = similar.((p.ᶠu₃⁰, p.ᶜu⁰, p.ᶠu³⁰, p.ᶜK⁰))
+    (ᶠu₃⁺, ᶜu⁺, ᶠu³⁺, ᶜK⁺) =
+        similar.((
+            p.precomputed.ᶠu₃⁰,
+            p.precomputed.ᶜu⁰,
+            p.precomputed.ᶠu³⁰,
+            p.precomputed.ᶜK⁰,
+        ))
     set_sgs_ᶠu₃!(u₃⁺, ᶠu₃⁺, Y, turbconv_model)
     set_velocity_quantities!(ᶜu⁺, ᶠu³⁺, ᶜK⁺, ᶠu₃⁺, Y.c.uₕ, ᶠuₕ³)
     ᶜts⁺ = ᶜtsʲs.:1
@@ -385,8 +394,8 @@ values of the first updraft.
 """
 function output_diagnostic_sgs_quantities(Y, p, t)
     thermo_params = CAP.thermodynamics_params(p.params)
-    (; ᶜρaʲs, ᶜtsʲs) = p
-    ᶠu³⁺ = p.ᶠu³ʲs.:1
+    (; ᶜρaʲs, ᶜtsʲs) = p.precomputed
+    ᶠu³⁺ = p.precomputed.ᶠu³ʲs.:1
     ᶜu⁺ = @. (C123(Y.c.uₕ) + C123(ᶜinterp(ᶠu³⁺)))
     ᶜts⁺ = @. ᶜtsʲs.:1
     ᶜa⁺ = @. draft_area(ᶜρaʲs.:1, TD.air_density(thermo_params, ᶜts⁺))

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -14,8 +14,10 @@ function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³
 
     thermo_params = CAP.thermodynamics_params(p.params)
     (; turbconv_model) = p.atmos
-    (; ᶜp, ᶜΦ, ᶜh_tot) = p
-    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
+    (; ᶜΦ,) = p.core
+    (; ᶜp, ᶜh_tot) = p.precomputed
+    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰, ᶜq_tot⁰) =
+        p.precomputed
 
     @. ᶜρa⁰ = ρa⁰(Y.c)
     @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
@@ -59,9 +61,10 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
 
-    (; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot) = p
-    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p
-    (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
+    (; ᶜΦ,) = p.core
+    (; ᶜspecific, ᶜp, ᶜh_tot) = p.precomputed
+    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
+    (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
 
     for j in 1:n
         ᶜuʲ = ᶜuʲs.:($j)
@@ -90,7 +93,8 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
         )
         ᶜρ_int_val = Fields.field_values(Fields.level(Y.c.ρ, 1))
         ᶜp_int_val = Fields.field_values(Fields.level(ᶜp, 1))
-        (; ρ_flux_h_tot, ρ_flux_q_tot, ustar, obukhov_length) = p.sfc_conditions
+        (; ρ_flux_h_tot, ρ_flux_q_tot, ustar, obukhov_length) =
+            p.precomputed.sfc_conditions
         buoyancy_flux_val = Fields.field_values(buoyancy_flux)
         ρ_flux_h_tot_val = Fields.field_values(ρ_flux_h_tot)
         ρ_flux_q_tot_val = Fields.field_values(ρ_flux_q_tot)
@@ -168,18 +172,19 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
     FT = eltype(params)
     n = n_mass_flux_subdomains(turbconv_model)
 
-    (; ᶜu, ᶜp, ᶜρ_ref) = p
-    (; ᶜtke⁰, ᶜρa⁰, ᶜu⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) = p
+    (; ᶜρ_ref) = p.core
+    (; ᶜspecific, ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶜu⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) =
+        p.precomputed
     (;
         ᶜmixing_length,
         ᶜlinear_buoygrad,
         ᶜstrain_rate_norm,
         ᶜK_u,
         ᶜK_h,
         ρatke_flux,
-    ) = p
-    (; ᶜuʲs, ᶠu³ʲs, ᶜtsʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p
-    (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
+    ) = p.precomputed
+    (; ᶜuʲs, ᶜtsʲs, ᶠu³ʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p.precomputed
+    (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
 
     ᶜz = Fields.coordinate_field(Y.c).z
     z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)