Cc/test lazy level precomputed2

.buildkite/pipeline.yml

@@ -29,6 +29,9 @@ steps:
       - echo "--- Instantiate .buildkite"
       - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true); using CUDA; CUDA.precompile_runtime(); Pkg.status()'"
 
+      - echo "--- dev package"
+      - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"dy/lazy_field_levels\"))'"
+
     agents:
       slurm_cpus_per_task: 8
       slurm_gpus: 1

src/ClimaAtmos.jl

@@ -8,6 +8,8 @@ import LazyBroadcast
 import LazyBroadcast: lazy
 import Thermodynamics as TD
 import Thermodynamics
+import ClimaCore.MatrixFields: @name
+
 
 include("compat.jl")
 include(joinpath("parameters", "Parameters.jl"))

src/cache/cloud_fraction.jl

@@ -62,15 +62,10 @@ NVTX.@annotate function set_cloud_fraction!(
             TD.PhasePartition(thermo_params, ᶜts).ice,
         )
     else
-        @. cloud_diagnostics_tuple = make_named_tuple(
-            ifelse(
-                specific(Y.c.ρq_liq, Y.c.ρ) + specific(Y.c.ρq_ice, Y.c.ρ) > 0,
-                1,
-                0,
-            ),
-            specific(Y.c.ρq_liq, Y.c.ρ),
-            specific(Y.c.ρq_ice, Y.c.ρ),
-        )
+        q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
+        q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
+        @. cloud_diagnostics_tuple =
+            make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice)
     end
 end
 
@@ -193,8 +188,9 @@ NVTX.@annotate function set_cloud_fraction!(
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
     (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
-    (; ᶜρʲs, ᶜtsʲs, ᶜρa⁰, ᶜρ⁰) = p.precomputed
+    (; ᶜρʲs, ᶜtsʲs) = p.precomputed
     (; turbconv_model) = p.atmos
+    ᶜρa⁰_vals = ᶜρa⁰(Y, p)
 
     # TODO - we should make this default when using diagnostic edmf
     # environment
@@ -213,9 +209,9 @@ NVTX.@annotate function set_cloud_fraction!(
     # weight cloud diagnostics by environmental area
     @. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}(
         tuple(
-            draft_area(ᶜρa⁰, ᶜρ⁰),
-            draft_area(ᶜρa⁰, ᶜρ⁰),
-            draft_area(ᶜρa⁰, ᶜρ⁰),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
         ),
     )
     # updrafts
@@ -308,6 +304,7 @@ function quad_loop(
         FT = eltype(x1_hat)
         @assert(x1_hat >= FT(0))
         @assert(x2_hat >= FT(0))
+        # note: ᶜthermo_state is used as a pointwise function here
         _ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
         hc = TD.has_condensate(thermo_params, _ts)
 

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -61,7 +61,7 @@ NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!(
     local_geometry_halflevel,
     turbconv_model,
 )
-    @. u³⁰_halflevel = divide_by_ρa(
+    @. u³⁰_halflevel = specific(
         ρ_level * u³_halflevel -
         unrolled_dotproduct(ρaʲs_level, u³ʲs_halflevel),
         ρ_level,
@@ -93,23 +93,30 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     FT = eltype(Y)
     n = n_mass_flux_subdomains(turbconv_model)
     (; ᶜΦ) = p.core
-    (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed
-    (; q_tot) = p.precomputed.ᶜspecific
+    (; ᶜp, ᶠu³, ᶜK) = p.precomputed
     (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) =
         p.precomputed.sfc_conditions
     (; ᶜρaʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
     (; ᶠu³⁰, ᶜK⁰) = p.precomputed
 
     (; params) = p
+
     thermo_params = CAP.thermodynamics_params(params)
     turbconv_params = CAP.turbconv_params(params)
+    ᶜts = p.precomputed.ᶜts   #TODO replace
+
+    q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
+    ᶜe_tot = ᶜspecific(Y.c.ρe_tot, Y.c.ρ)
+    ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot))
 
     ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1))
     uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1))
     u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half))
-    h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1))
+    h_tot_int_level =
+        Fields.field_values(Fields.level(ᶜh_tot, 1))
     K_int_level = Fields.field_values(Fields.level(ᶜK, 1))
-    q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
+    q_tot_int_level =
+        Fields.field_values(Fields.level(q_tot, 1))
 
     p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
     Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
@@ -305,8 +312,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     (; dt) = p
     dt = float(dt)
     (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core
-    (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
-    (; q_tot) = p.precomputed.ᶜspecific
+    (; ᶜp, ᶠu³, ᶜts, ᶜK) = p.precomputed
     (;
         ᶜρaʲs,
         ᶠu³ʲs,
@@ -321,13 +327,13 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         ᶠnh_pressure³_buoyʲs,
         ᶠnh_pressure³_dragʲs,
     ) = p.precomputed
-    (; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed
+    (; ᶠu³⁰, ᶜK⁰) = p.precomputed
+
+
 
     if microphysics_model isa Microphysics1Moment
-        ᶜq_liqʲs = p.precomputed.ᶜq_liqʲs
-        ᶜq_iceʲs = p.precomputed.ᶜq_iceʲs
-        q_rai = p.precomputed.ᶜqᵣ
-        q_sno = p.precomputed.ᶜqₛ
+        q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
+        q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
     end
 
     thermo_params = CAP.thermodynamics_params(params)
@@ -347,13 +353,25 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
 
     # integral
+    q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
+    ᶜh_tot = @. lazy(
+        TD.total_specific_enthalpy(
+            thermo_params,
+            ᶜts,
+            specific(Y.c.ρe_tot, Y.c.ρ),
+        ),
+    )
+    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+
     for i in 2:Spaces.nlevels(axes(Y.c))
         ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))
         uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i))
         u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half))
         K_level = Fields.field_values(Fields.level(ᶜK, i))
-        h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
-        q_tot_level = Fields.field_values(Fields.level(q_tot, i))
+        h_tot_level =
+            Fields.field_values(Fields.level(ᶜh_tot, i))
+        q_tot_level =
+            Fields.field_values(Fields.level(q_tot, i))
         p_level = Fields.field_values(Fields.level(ᶜp, i))
         Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
         local_geometry_level = Fields.field_values(
@@ -377,8 +395,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
             Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
         u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1
         K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1))
-        h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1))
-        q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1))
+        h_tot_prev_level =
+            Fields.field_values(Fields.level(ᶜh_tot, i - 1))
+        q_tot_prev_level =
+            Fields.field_values(Fields.level(q_tot, i - 1))
         ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1))
         p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1))
         z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1))
@@ -478,7 +498,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                     Fields.field_values(Fields.level(ᶜq_iceʲ, i - 1))
             end
 
-            tke_prev_level = Fields.field_values(Fields.level(ᶜtke⁰, i - 1))
+            tke_prev_level = Fields.field_values(
+                Fields.level(ᶜtke⁰, i - 1),
+            )
 
             @. entrʲ_prev_level = entrainment(
                 thermo_params,
@@ -965,7 +987,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     (; dt) = p
     (; ᶜp, ᶜu, ᶜts) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
-    (; ᶜtke⁰) = p.precomputed
     (;
         ᶜlinear_buoygrad,
         ᶜstrain_rate_norm,
@@ -1008,6 +1029,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     @. ᶜprandtl_nvec =
         turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
 
+    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+
     ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0
     # using ᶜu⁰ would be more correct, but this is more consistent with the
@@ -1083,11 +1106,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
     (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
 
     # Environment precipitation sources (to be applied to grid mean)
+    ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
     @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
         thermo_params,
         microphys_0m_params,
         dt,
-        specific(Y.c.ρq_tot, Y.c.ρ),
+        ᶜq_tot,
         ᶜts,
     )
     return nothing

src/cache/precipitation_precomputed_quantities.jl

@@ -112,6 +112,11 @@ function set_precipitation_velocities!(
     cm2p = CAP.microphysics_2m_params(p.params)
     thp = CAP.thermodynamics_params(p.params)
 
+    q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
+    q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
+    q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
+    q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
+
     # compute the precipitation terminal velocity [m/s]
     # TODO sedimentation of snow is based on the 1M scheme
     @. ᶜwnᵣ = getindex(
@@ -253,17 +258,18 @@ function set_precipitation_cache!(
     ::PrognosticEDMFX,
 )
     (; ᶜΦ) = p.core
-    (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs, ᶜρa⁰) = p.precomputed
+    (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs) = p.precomputed
     (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
     (; ᶜts⁰, ᶜtsʲs) = p.precomputed
     thermo_params = CAP.thermodynamics_params(p.params)
+    ᶜρa⁰_vals = ᶜρa⁰(Y, p)
 
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
 
-    @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰
+    @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰_vals
     @. ᶜS_ρe_tot =
         ᶜSqₜᵖ⁰ *
-        ᶜρa⁰ *
+        ᶜρa⁰_vals *
         e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ)
     for j in 1:n
         @. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa
@@ -283,7 +289,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
     (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed
     (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
 
-    (; q_tot, q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific
+    q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
+    q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
+    q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
+    q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
+    q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
 
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
@@ -357,6 +367,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _)
     (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
     (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed
 
+    q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
+    q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
+    n_liq = ᶜspecific(Y.c.ρn_liq, Y.c.ρ)
+    n_rai = ᶜspecific(Y.c.ρn_rai, Y.c.ρ)
+
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
 
@@ -469,14 +484,20 @@ function set_precipitation_surface_fluxes!(
     sfc_ρ = @. lazy(int_ρ * int_J / sfc_J)
 
     # Constant extrapolation to surface, consistent with simple downwinding
-    sfc_wᵣ = sfc_lev(ᶜwᵣ)
-    sfc_wₛ = sfc_lev(ᶜwₛ)
-    sfc_wₗ = sfc_lev(ᶜwₗ)
-    sfc_wᵢ = sfc_lev(ᶜwᵢ)
-    sfc_qᵣ = lazy.(specific.(sfc_lev(Y.c.ρq_rai), sfc_ρ))
-    sfc_qₛ = lazy.(specific.(sfc_lev(Y.c.ρq_sno), sfc_ρ))
-    sfc_qₗ = lazy.(specific.(sfc_lev(Y.c.ρq_liq), sfc_ρ))
-    sfc_qᵢ = lazy.(specific.(sfc_lev(Y.c.ρq_ice), sfc_ρ))
+    ᶜq_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
+    ᶜq_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
+    ᶜq_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
+    ᶜq_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
+    sfc_qᵣ = Fields.field_values(Fields.level(ᶜq_rai, 1))
+    sfc_qₛ = Fields.field_values(Fields.level(ᶜq_sno, 1))
+
+    sfc_qₗ = Fields.field_values(Fields.level(ᶜq_liq, 1))
+    sfc_qᵢ = Fields.field_values(Fields.level(ᶜq_ice, 1))
+
+    sfc_wᵣ = Fields.field_values(Fields.level(ᶜwᵣ, 1))
+    sfc_wₛ = Fields.field_values(Fields.level(ᶜwₛ, 1))
+    sfc_wₗ = Fields.field_values(Fields.level(ᶜwₗ, 1))
+    sfc_wᵢ = Fields.field_values(Fields.level(ᶜwᵢ, 1))
 
     @. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ))
     @. surface_snow_flux = sfc_ρ * (sfc_qₛ * (-sfc_wₛ) + sfc_qᵢ * (-sfc_wᵢ))