Make prognostic rain a ClimaCore field

integration_tests/utils/main.jl

@@ -116,8 +116,11 @@ cent_prognostic_vars(FT, n_up) = (; cent_prognostic_vars_gm(FT)..., cent_prognos
 cent_prognostic_vars_gm(FT) = (; u = FT(0), v = FT(0), θ_liq_ice = FT(0), q_tot = FT(0))
 cent_prognostic_vars_up(FT) = (; area = FT(0), θ_liq_ice = FT(0), q_tot = FT(0))
 cent_prognostic_vars_en(FT) = (; tke = FT(0), Hvar = FT(0), QTvar = FT(0), HQTcov = FT(0))
-cent_prognostic_vars_edmf(FT, n_up) =
-    (; turbconv = (; en = cent_prognostic_vars_en(FT), up = ntuple(i -> cent_prognostic_vars_up(FT), n_up)))
+cent_prognostic_vars_edmf(FT, n_up) = (;
+    turbconv = (;
+        en = cent_prognostic_vars_en(FT), up = ntuple(i -> cent_prognostic_vars_up(FT), n_up), ra = (; qr = FT(0)),
+    ),
+)
 # cent_prognostic_vars_edmf(FT, n_up) = (;) # could also use this for empty model
 
 # Face only

src/EDMF_Environment.jl

@@ -39,7 +39,9 @@ end
 function update_env_precip_tendencies(en_thermo::EnvironmentThermodynamics, state, k, qt_tendency, θ_liq_ice_tendency)
 
     aux_en = center_aux_environment(state)
+    tendencies_ra = center_tendencies_rain(state)
     en_thermo.qt_tendency_rain_formation[k] = qt_tendency * aux_en.area[k]
+    tendencies_ra.qr[k] += -en_thermo.qt_tendency_rain_formation[k]
     en_thermo.θ_liq_ice_tendency_rain_formation[k] = θ_liq_ice_tendency * aux_en.area[k]
 
     return
@@ -67,13 +69,14 @@ function sgs_mean(en_thermo::EnvironmentThermodynamics, grid, state, en, rain, d
     p0_c = center_ref_state(state).p0
     ρ0_c = center_ref_state(state).ρ0
     aux_en = center_aux_environment(state)
+    prog_ra = center_prog_rain(state)
 
     @inbounds for k in real_center_indices(grid)
         # condensation
         q_tot_en = aux_en.q_tot[k]
         ts = TD.PhaseEquil_pθq(param_set, p0_c[k], aux_en.θ_liq_ice[k], q_tot_en)
         # autoconversion and accretion
-        mph = precipitation_formation(param_set, rain.rain_model, rain.QR.values[k], aux_en.area[k], ρ0_c[k], dt, ts)
+        mph = precipitation_formation(param_set, rain.rain_model, prog_ra.qr[k], aux_en.area[k], ρ0_c[k], dt, ts)
         update_cloud_dry(en_thermo, state, k, ts)
         update_env_precip_tendencies(en_thermo, state, k, mph.qt_tendency, mph.θ_liq_ice_tendency)
     end
@@ -88,6 +91,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, r
     ρ0_c = center_ref_state(state).ρ0
     prog_en = center_prog_environment(state)
     aux_en = center_aux_environment(state)
+    prog_ra = center_prog_rain(state)
 
     #TODO - remember you output source terms multipierd by dt (bec. of instanteneous autoconcv)
     #TODO - add tendencies for gm H, QT and QR due to rain
@@ -193,7 +197,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, r
                     mph = precipitation_formation(
                         param_set,
                         rain.rain_model,
-                        rain.QR.values[k],
+                        prog_ra.qr[k],
                         aux_en.area[k],
                         ρ0_c[k],
                         dt,
@@ -263,8 +267,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, r
         else
             # if variance and covariance are zero do the same as in SA_mean
             ts = TD.PhaseEquil_pθq(param_set, p0_c[k], aux_en.θ_liq_ice[k], aux_en.q_tot[k])
-            mph =
-                precipitation_formation(param_set, rain.rain_model, rain.QR.values[k], aux_en.area[k], ρ0_c[k], dt, ts)
+            mph = precipitation_formation(param_set, rain.rain_model, prog_ra.qr[k], aux_en.area[k], ρ0_c[k], dt, ts)
             update_env_precip_tendencies(en_thermo, state, k, mph.qt_tendency, mph.θ_liq_ice_tendency)
             update_cloud_dry(en_thermo, state, k, ts)
 

src/EDMF_Rain.jl

@@ -1,5 +1,4 @@
 function initialize_io(rain::RainVariables, Stats::NetCDFIO_Stats)
-    add_profile(Stats, "qr_mean")
     add_ts(Stats, "rwp_mean")
     add_ts(Stats, "cutoff_rain_rate")
     return
@@ -14,8 +13,6 @@ function io(
     en_thermo::EnvironmentThermodynamics,
     TS::TimeStepping,
 )
-    write_profile(Stats, "qr_mean", rain.QR.values)
-
     rain_diagnostics(rain, grid, state, up_thermo, en_thermo)
     write_ts(Stats, "rwp_mean", rain.mean_rwp)
 
@@ -32,12 +29,13 @@ function rain_diagnostics(
     en_thermo::EnvironmentThermodynamics,
 )
     ρ0_c = center_ref_state(state).ρ0
+    prog_ra = center_prog_rain(state)
 
     rain.mean_rwp = 0.0
     rain.cutoff_rain_rate = 0.0
 
     @inbounds for k in real_center_indices(grid)
-        rain.mean_rwp += ρ0_c[k] * rain.QR.values[k] * grid.Δz
+        rain.mean_rwp += ρ0_c[k] * prog_ra.qr[k] * grid.Δz
 
         # rain rate from cutoff microphysics scheme defined as a total amount of removed water
         # per timestep per EDMF surface area [mm/h]
@@ -56,20 +54,22 @@ end
 """
 Computes the qr advection (down) tendency
 """
-function compute_rain_advection_tendencies(rain::RainPhysics, grid, state, gm, TS::TimeStepping, QR::RainVariable)
+function compute_rain_advection_tendencies(rain::RainPhysics, grid, state, gm, TS::TimeStepping)
     param_set = parameter_set(gm)
     Δz = grid.Δz
     Δt = TS.dt
     CFL_limit = 0.5
 
     ρ_0_c = center_ref_state(state).ρ0
+    tendencies_ra = center_tendencies_rain(state)
+    prog_ra = center_prog_rain(state)
 
     # helper to calculate the rain velocity
     # TODO: assuming gm.W = 0
     # TODO: verify translation
     term_vel = center_field(grid)
     @inbounds for k in real_center_indices(grid)
-        term_vel[k] = CM1.terminal_velocity(param_set, rain_type, ρ_0_c[k], QR.values[k])
+        term_vel[k] = CM1.terminal_velocity(param_set, rain_type, ρ_0_c[k], prog_ra.qr[k])
     end
 
     @inbounds for k in reverse(real_center_indices(grid))
@@ -85,7 +85,7 @@ function compute_rain_advection_tendencies(rain::RainPhysics, grid, state, gm, T
         end
 
         ρ_0_cut = ccut_downwind(ρ_0_c, grid, k)
-        QR_cut = ccut_downwind(QR.values, grid, k)
+        QR_cut = ccut_downwind(prog_ra.qr, grid, k)
         w_cut = ccut_downwind(term_vel, grid, k)
         ρQRw_cut = ρ_0_cut .* QR_cut .* w_cut
         ∇ρQRw = c∇_downwind(ρQRw_cut, grid, k; bottom = FreeBoundary(), top = SetValue(0))
@@ -94,20 +94,23 @@ function compute_rain_advection_tendencies(rain::RainPhysics, grid, state, gm, T
 
         # TODO - some positivity limiters are needed
         rain.qr_tendency_advection[k] = ∇ρQRw / ρ_0_c_k
+        tendencies_ra.qr[k] += rain.qr_tendency_advection[k]
     end
     return
 end
 
 """
 Computes the tendencies to θ_liq_ice, qt and qr due to rain evaporation
 """
-function compute_rain_evap_tendencies(rain::RainPhysics, grid, state, gm, TS::TimeStepping, QR::RainVariable)
+function compute_rain_evap_tendencies(rain::RainPhysics, grid, state, gm, TS::TimeStepping)
     param_set = parameter_set(gm)
     Δt = TS.dt
     p0_c = center_ref_state(state).p0
     ρ0_c = center_ref_state(state).ρ0
     aux_gm = center_aux_grid_mean(state)
     prog_gm = center_prog_grid_mean(state)
+    prog_ra = center_prog_rain(state)
+    tendencies_ra = center_tendencies_rain(state)
 
     @inbounds for k in real_center_indices(grid)
         q_tot_gm = prog_gm.q_tot[k]
@@ -118,10 +121,11 @@ function compute_rain_evap_tendencies(rain::RainPhysics, grid, state, gm, TS::Ti
 
         # TODO - move limiters elsewhere
         qt_tendency =
-            min(QR.values[k] / Δt, -CM1.evaporation_sublimation(param_set, rain_type, q, QR.values[k], ρ0_c[k], T_gm))
+            min(prog_ra.qr[k] / Δt, -CM1.evaporation_sublimation(param_set, rain_type, q, prog_ra.qr[k], ρ0_c[k], T_gm))
 
         rain.qt_tendency_rain_evap[k] = qt_tendency
         rain.qr_tendency_rain_evap[k] = -qt_tendency
+        tendencies_ra.qr[k] += rain.qr_tendency_rain_evap[k]
         rain.θ_liq_ice_tendency_rain_evap[k] = θ_liq_ice_helper(ts, qt_tendency)
     end
     return
@@ -139,12 +143,10 @@ function update_rain(
     rain_phys::RainPhysics,
     TS::TimeStepping,
 )
+    prog_ra = center_prog_rain(state)
+    tendencies_ra = center_tendencies_rain(state)
     @inbounds for k in real_center_indices(grid)
-        rain_var.QR.values[k] +=
-            (
-                rain_phys.qr_tendency_advection[k] + rain_phys.qr_tendency_rain_evap[k] -
-                en_thermo.qt_tendency_rain_formation[k] - up_thermo.qt_tendency_rain_formation_tot[k]
-            ) * TS.dt
+        prog_ra.qr[k] += tendencies_ra.qr[k] * TS.dt
     end
     return
 end

src/EDMF_Updrafts.jl

@@ -239,6 +239,8 @@ function compute_rain_formation_tendencies(
     ρ0_c = center_ref_state(state).ρ0
     prog_up = center_prog_updrafts(state)
     aux_up = center_aux_updrafts(state)
+    prog_ra = center_prog_rain(state)
+    tendencies_ra = center_tendencies_rain(state)
 
     @inbounds for i in 1:(up.n_updrafts)
         @inbounds for k in real_center_indices(grid)
@@ -250,7 +252,7 @@ function compute_rain_formation_tendencies(
             mph = precipitation_formation(
                 param_set,
                 rain.rain_model,
-                rain.QR.values[k],
+                prog_ra.qr[k],
                 prog_up[i].area[k],
                 ρ0_c[k],
                 dt,
@@ -263,5 +265,6 @@ function compute_rain_formation_tendencies(
     # TODO - to be deleted once we sum all tendencies elsewhere
     up_thermo.θ_liq_ice_tendency_rain_formation_tot .= up_sum(up_thermo.θ_liq_ice_tendency_rain_formation)
     up_thermo.qt_tendency_rain_formation_tot .= up_sum(up_thermo.qt_tendency_rain_formation)
+    parent(tendencies_ra.qr) .+= -up_thermo.qt_tendency_rain_formation_tot
     return
 end

src/Turbulence_PrognosticTKE.jl

@@ -404,7 +404,6 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     en_thermo = edmf.EnvThermo
     n_updrafts = up.n_updrafts
     prog_gm = center_prog_grid_mean(state)
-    tendencies_gm = center_tendencies_grid_mean(state)
     prog_en = center_prog_environment(state)
 
     # Update aux / pre-tendencies filters. TODO: combine these into a function that minimizes traversals
@@ -413,11 +412,19 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     set_updraft_surface_bc(edmf, grid, state, gm, Case)
     update_aux!(edmf, gm, grid, state, Case, param_set, TS)
 
+    tendencies_gm = center_tendencies_grid_mean(state)
+    tendencies_up = center_tendencies_updrafts(state)
+    tendencies_en = center_tendencies_environment(state)
+    tendencies_ra = center_tendencies_rain(state)
+    parent(tendencies_gm) .= 0
+    parent(tendencies_up) .= 0
+    parent(tendencies_en) .= 0
+    parent(tendencies_ra) .= 0
     compute_rain_formation_tendencies(up_thermo, grid, state, edmf.UpdVar, edmf.Rain, TS.dt, param_set) # causes division error in dry bubble first time step
     microphysics(en_thermo, grid, state, edmf.EnvVar, edmf.Rain, TS.dt, param_set) # saturation adjustment + rain creation
     if edmf.Rain.rain_model == "clima_1m"
-        compute_rain_evap_tendencies(edmf.RainPhys, grid, state, gm, TS, edmf.Rain.QR)
-        compute_rain_advection_tendencies(edmf.RainPhys, grid, state, gm, TS, edmf.Rain.QR)
+        compute_rain_evap_tendencies(edmf.RainPhys, grid, state, gm, TS)
+        compute_rain_advection_tendencies(edmf.RainPhys, grid, state, gm, TS)
     end
 
     # compute tendencies

src/diagnostics.jl

@@ -105,6 +105,9 @@ function io_dictionary_aux(state)
         "env_RH" => (; dims = ("zc", "t"), group = "profiles", field = center_aux_environment(state).RH),
         "env_thetal" => (; dims = ("zc", "t"), group = "profiles", field = center_aux_environment(state).θ_liq_ice),
         "env_cloud_fraction" => (; dims = ("zc", "t"), group = "profiles", field = center_aux_environment(state).cloud_fraction),
+
+        "qr_mean" => (; dims = ("zc", "t"), group = "profiles", field = center_prog_rain(state).qr),
+
     )
     return io_dict
 end

src/dycore_api.jl

@@ -46,6 +46,7 @@ prognostic_tc(state, fl) = prognostic(state, fl).turbconv
 center_prog_updrafts(state) = prognostic_tc(state, CentField()).up
 face_prog_updrafts(state) = prognostic_tc(state, FaceField()).up
 center_prog_environment(state) = prognostic_tc(state, CentField()).en
+center_prog_rain(state) = prognostic_tc(state, CentField()).ra
 face_prog_environment(state) = prognostic_tc(state, FaceField()).en
 
 #= Auxiliary fields for TurbulenceConvection =#
@@ -63,4 +64,6 @@ tendencies_tc(state, fl) = tendencies(state, fl).turbconv
 center_tendencies_tc(state) = tendencies_tc(state, CentField())
 face_tendencies_tc(state) = tendencies_tc(state, FaceField())
 center_tendencies_updrafts(state) = tendencies_tc(state, CentField()).up
+center_tendencies_environment(state) = tendencies_tc(state, CentField()).en
+center_tendencies_rain(state) = tendencies_tc(state, CentField()).ra
 face_tendencies_updrafts(state) = tendencies_tc(state, FaceField()).up

src/types.jl

@@ -186,26 +186,13 @@ Base.@kwdef struct Tan2018{FT}
     alpha0::FT
 end
 
-struct RainVariable{T}
-    name::String
-    units::String
-    values::T
-    function RainVariable(grid, name, units)
-        values = center_field(grid)
-        return new{typeof(values)}(name, units, values)
-    end
-end
-
 Base.@kwdef mutable struct RainVariables
     rain_model::String = "default_rain_model"
     mean_rwp::Float64 = 0
     cutoff_rain_rate::Float64 = 0
-    QR::RainVariable
 end
 function RainVariables(namelist, grid::Grid)
 
-    QR = RainVariable(grid, "qr_mean", "kg/kg")
-
     rain_model = parse_namelist(
         namelist,
         "microphysics",
@@ -218,7 +205,7 @@ function RainVariables(namelist, grid::Grid)
         error("rain model not recognized")
     end
 
-    return RainVariables(; rain_model, QR)
+    return RainVariables(; rain_model)
 end
 
 struct RainPhysics{T}
@@ -240,18 +227,6 @@ struct RainPhysics{T}
     end
 end
 
-struct VariableDiagnostic{T}
-    values::T
-    name::String
-    units::String
-    function VariableDiagnostic(grid, loc, name, units)
-        # Value at the current timestep
-        values = field(grid, loc)
-        # Placement on staggered grid
-        return new{typeof(values)}(values, name, units)
-    end
-end
-
 struct UpdraftVariable{A}
     new::A # TODO: deprecate this!
     name::String