Prepare data structures for adding snow

integration_tests/utils/main.jl

@@ -82,6 +82,8 @@ cent_aux_vars_up(FT) = (;
     area = FT(0),
     q_tot = FT(0),
     θ_liq_ice = FT(0),
+    θ_liq_ice_tendency_precip_formation = FT(0),
+    qt_tendency_precip_formation = FT(0),
 )
 cent_aux_vars_edmf(FT, n_up) = (;
     turbconv = (;
@@ -95,6 +97,8 @@ cent_aux_vars_edmf(FT, n_up) = (;
             q_ice = FT(0),
             T = FT(0),
             cloud_fraction = FT(0),
+            θ_liq_ice_tendency_precip_formation_tot = FT(0),
+            qt_tendency_precip_formation_tot = FT(0),
         ),
         up = ntuple(i -> cent_aux_vars_up(FT), n_up),
         en = (;
@@ -112,18 +116,22 @@ cent_aux_vars_edmf(FT, n_up) = (;
             Hvar = FT(0),
             QTvar = FT(0),
             HQTcov = FT(0),
+            qt_tendency_precip_formation = FT(0),
+            θ_liq_ice_tendency_precip_formation = FT(0),
         ),
-        θ_liq_ice_tendency_rain_evap = FT(0),
-        qt_tendency_rain_evap = FT(0),
-        qr_tendency_rain_evap = FT(0),
+        qr_tendency_evap = FT(0),
+        qs_tendency_melt = FT(0),
+        qs_tendency_dep_sub = FT(0),
         qr_tendency_advection = FT(0),
+        qs_tendency_advection = FT(0),
         en_2m = (;
             tke = cent_aux_vars_en_2m(FT),
             Hvar = cent_aux_vars_en_2m(FT),
             QTvar = cent_aux_vars_en_2m(FT),
             HQTcov = cent_aux_vars_en_2m(FT),
         ),
-        term_vel = FT(0),
+        term_vel_rain = FT(0),
+        term_vel_snow = FT(0),
         KM = FT(0),
         KH = FT(0),
         mixing_length = FT(0),
@@ -184,7 +192,9 @@ cent_prognostic_vars_up(FT) = (; ρarea = FT(0), ρaθ_liq_ice = FT(0), ρaq_tot
 cent_prognostic_vars_en(FT) = (; ρatke = FT(0), ρaHvar = FT(0), ρaQTvar = FT(0), ρaHQTcov = 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), pr = (; qr = FT(0)),
+        en = cent_prognostic_vars_en(FT),
+        up = ntuple(i -> cent_prognostic_vars_up(FT), n_up),
+        pr = (; q_rai = FT(0), q_sno = FT(0)),
     ),
 )
 # cent_prognostic_vars_edmf(FT, n_up) = (;) # could also use this for empty model

src/EDMF_Environment.jl

@@ -1,11 +1,12 @@
-function update_env_precip_tendencies(en_thermo::EnvironmentThermodynamics, state, k, qt_tendency, θ_liq_ice_tendency)
-
+function update_env_precip_tendencies(state, k, qt_tendency, θ_liq_ice_tendency, qr_tendency, qs_tendency)
     aux_en = center_aux_environment(state)
     tendencies_pr = center_tendencies_precipitation(state)
-    en_thermo.qt_tendency_rain_formation[k] = qt_tendency * aux_en.area[k]
-    tendencies_pr.qr[k] += -en_thermo.qt_tendency_rain_formation[k]
-    en_thermo.θ_liq_ice_tendency_rain_formation[k] = θ_liq_ice_tendency * aux_en.area[k]
 
+    aux_en.qt_tendency_precip_formation[k] = qt_tendency * aux_en.area[k]
+    aux_en.θ_liq_ice_tendency_precip_formation[k] = θ_liq_ice_tendency * aux_en.area[k]
+
+    tendencies_pr.q_rai[k] += qr_tendency * aux_en.area[k]
+    tendencies_pr.q_sno[k] += qs_tendency * aux_en.area[k]
     return
 end
 
@@ -44,14 +45,22 @@ function sgs_mean(en_thermo::EnvironmentThermodynamics, grid, state, en, precip,
         mph = precipitation_formation(
             param_set,
             precip.precipitation_model,
-            prog_pr.qr[k],
+            prog_pr.q_rai[k],
+            prog_pr.q_sno[k],
             aux_en.area[k],
             ρ0_c[k],
             dt,
             ts,
         )
         update_sat_unsat(en_thermo, state, k, ts)
-        update_env_precip_tendencies(en_thermo, state, k, mph.qt_tendency, mph.θ_liq_ice_tendency)
+        update_env_precip_tendencies(
+            state,
+            k,
+            mph.qt_tendency,
+            mph.θ_liq_ice_tendency,
+            mph.qr_tendency,
+            mph.qs_tendency,
+        )
     end
     return
 end
@@ -75,7 +84,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
     # arrays for storing quadarature points and ints for labeling items in the arrays
     # a python dict would be nicer, but its 30% slower than this (for python 2.7. It might not be the case for python 3)
     env_len = 8
-    src_len = 6
+    src_len = 8
 
     sqpi_inv = 1.0 / sqrt(π)
     sqrt2 = sqrt(2.0)
@@ -88,7 +97,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
     inner_src = zeros(src_len)
     outer_src = zeros(src_len)
     i_ql, i_qi, i_T, i_cf, i_qt_sat, i_qt_unsat, i_T_sat, i_T_unsat = 1:env_len
-    i_SH_qt, i_Sqt_H, i_SH_H, i_Sqt_qt, i_Sqt, i_SH = 1:src_len
+    i_SH_qt, i_Sqt_H, i_SH_H, i_Sqt_qt, i_Sqt, i_SH, i_Sqr, i_Sqs = 1:src_len
 
     @inbounds for k in real_center_indices(grid)
         if (
@@ -170,7 +179,8 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
                     mph = precipitation_formation(
                         param_set,
                         precip.precipitation_model,
-                        prog_pr.qr[k],
+                        prog_pr.q_rai[k],
+                        prog_pr.q_sno[k],
                         aux_en.area[k],
                         ρ0_c[k],
                         dt,
@@ -192,6 +202,8 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
                     end
                     # products for variance and covariance source terms
                     inner_src[i_Sqt] += mph.qt_tendency * weights[m_h] * sqpi_inv
+                    inner_src[i_Sqr] += mph.qr_tendency * weights[m_h] * sqpi_inv
+                    inner_src[i_Sqs] += mph.qs_tendency * weights[m_h] * sqpi_inv
                     inner_src[i_SH] += mph.θ_liq_ice_tendency * weights[m_h] * sqpi_inv
                     inner_src[i_Sqt_H] += mph.qt_tendency * h_hat * weights[m_h] * sqpi_inv
                     inner_src[i_Sqt_qt] += mph.qt_tendency * qt_hat * weights[m_h] * sqpi_inv
@@ -208,7 +220,14 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
             end
 
             # update environmental variables
-            update_env_precip_tendencies(en_thermo, state, k, outer_src[i_Sqt], outer_src[i_SH])
+            update_env_precip_tendencies(
+                state,
+                k,
+                outer_src[i_Sqt],
+                outer_src[i_SH],
+                outer_src[i_Sqr],
+                outer_src[i_Sqs],
+            )
 
             # update cloudy/dry variables for buoyancy in TKE
             aux_en.cloud_fraction[k] = outer_env[i_cf]
@@ -229,6 +248,7 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
             en_thermo.θ_liq_ice_sat[k] = TD.liquid_ice_pottemp(ts_sat)
 
             # update var/covar rain sources
+            # TODO - check
             en_thermo.Hvar_rain_dt[k] = outer_src[i_SH_H] - outer_src[i_SH] * aux_en.θ_liq_ice[k]
             en_thermo.QTvar_rain_dt[k] = outer_src[i_Sqt_qt] - outer_src[i_Sqt] * aux_en.q_tot[k]
             en_thermo.HQTcov_rain_dt[k] =
@@ -241,13 +261,21 @@ function sgs_quadrature(en_thermo::EnvironmentThermodynamics, grid, state, en, p
             mph = precipitation_formation(
                 param_set,
                 precip.precipitation_model,
-                prog_pr.qr[k],
+                prog_pr.q_rai[k],
+                prog_pr.q_sno[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_env_precip_tendencies(
+                state,
+                k,
+                mph.qt_tendency,
+                mph.θ_liq_ice_tendency,
+                mph.qr_tendency,
+                mph.qs_tendency,
+            )
             update_sat_unsat(en_thermo, state, k, ts)
 
             en_thermo.Hvar_rain_dt[k] = 0.0

src/EDMF_Precipitation.jl

@@ -1,7 +1,7 @@
 """
-Computes the qr advection (down) tendency
+Computes the qr and qs advection (down) tendency
 """
-function compute_rain_advection_tendencies(::PrecipPhysics, grid, state, gm, TS::TimeStepping)
+function compute_precip_advection_tendencies(::PrecipPhysics, grid, state, gm, TS::TimeStepping)
     param_set = parameter_set(gm)
     Δz = grid.Δz
     Δt = TS.dt
@@ -15,42 +15,55 @@ function compute_rain_advection_tendencies(::PrecipPhysics, grid, state, gm, TS:
     # helper to calculate the rain velocity
     # TODO: assuming gm.W = 0
     # TODO: verify translation
-    term_vel = aux_tc.term_vel
+    term_vel_rain = aux_tc.term_vel_rain
+    term_vel_snow = aux_tc.term_vel_snow
+
     @inbounds for k in real_center_indices(grid)
-        term_vel[k] = CM1.terminal_velocity(param_set, rain_type, ρ_0_c[k], prog_pr.qr[k])
+        term_vel_rain[k] = CM1.terminal_velocity(param_set, rain_type, ρ_0_c[k], prog_pr.q_rai[k])
+        term_vel_snow[k] = CM1.terminal_velocity(param_set, snow_type, ρ_0_c[k], prog_pr.q_sno[k])
     end
 
     @inbounds for k in reverse(real_center_indices(grid))
         # check stability criterion
-        CFL_out = Δt / Δz * term_vel[k]
+        CFL_out_rain = Δt / Δz * term_vel_rain[k]
+        CFL_out_snow = Δt / Δz * term_vel_snow[k]
         if is_toa_center(grid, k)
-            CFL_in = 0.0
+            CFL_in_rain = 0.0
+            CFL_in_snow = 0.0
         else
-            CFL_in = Δt / Δz * term_vel[k + 1]
+            CFL_in_rain = Δt / Δz * term_vel_rain[k + 1]
+            CFL_in_snow = Δt / Δz * term_vel_snow[k + 1]
         end
-        if max(CFL_in, CFL_out) > CFL_limit
-            error("Time step is too large for rain fall velocity!")
+        if max(CFL_in_rain, CFL_in_snow, CFL_out_rain, CFL_out_snow) > CFL_limit
+            error("Time step is too large for precipitation fall velocity!")
         end
 
         ρ_0_cut = ccut_downwind(ρ_0_c, grid, k)
-        QR_cut = ccut_downwind(prog_pr.qr, grid, k)
-        w_cut = ccut_downwind(term_vel, grid, k)
-        ρQRw_cut = ρ_0_cut .* QR_cut .* w_cut
+        QR_cut = ccut_downwind(prog_pr.q_rai, grid, k)
+        QS_cut = ccut_downwind(prog_pr.q_sno, grid, k)
+        w_rain_cut = ccut_downwind(term_vel_rain, grid, k)
+        w_snow_cut = ccut_downwind(term_vel_snow, grid, k)
+        ρQRw_cut = ρ_0_cut .* QR_cut .* w_rain_cut
+        ρQSw_cut = ρ_0_cut .* QS_cut .* w_snow_cut
         ∇ρQRw = c∇_downwind(ρQRw_cut, grid, k; bottom = FreeBoundary(), top = SetValue(0))
+        ∇ρQSw = c∇_downwind(ρQSw_cut, grid, k; bottom = FreeBoundary(), top = SetValue(0))
 
         ρ_0_c_k = ρ_0_c[k]
 
         # TODO - some positivity limiters are needed
         aux_tc.qr_tendency_advection[k] = ∇ρQRw / ρ_0_c_k
-        tendencies_pr.qr[k] += aux_tc.qr_tendency_advection[k]
+        aux_tc.qs_tendency_advection[k] = ∇ρQSw / ρ_0_c_k
+        tendencies_pr.q_rai[k] += aux_tc.qr_tendency_advection[k]
+        tendencies_pr.q_sno[k] += aux_tc.qs_tendency_advection[k]
     end
     return
 end
 
 """
-Computes the tendencies to θ_liq_ice, qt and qr due to rain evaporation
+Computes the tendencies to θ_liq_ice, qt and qr due to
+evaporation, sublimation, deposition and melting
 """
-function compute_rain_evap_tendencies(::PrecipPhysics, grid, state, gm, TS::TimeStepping)
+function compute_precip_sink_tendencies(::PrecipPhysics, grid, state, gm, TS::TimeStepping)
     param_set = parameter_set(gm)
     Δt = TS.dt
     p0_c = center_ref_state(state).p0
@@ -60,22 +73,36 @@ function compute_rain_evap_tendencies(::PrecipPhysics, grid, state, gm, TS::Time
     prog_gm = center_prog_grid_mean(state)
     prog_pr = center_prog_precipitation(state)
     tendencies_pr = center_tendencies_precipitation(state)
+    tendencies_gm = center_tendencies_grid_mean(state)
 
     @inbounds for k in real_center_indices(grid)
+        # When we fuse loops, this should hopefully disappear
         q_tot_gm = prog_gm.q_tot[k]
         T_gm = aux_gm.T[k]
-        # When we fuse loops, this should hopefully disappear
-        ts = TD.PhaseEquil_pTq(param_set, p0_c[k], T_gm, q_tot_gm)
+        p0 = p0_c[k]
+        ρ0 = ρ0_c[k]
+        ts = TD.PhaseEquil_pTq(param_set, p0, T_gm, q_tot_gm)
         q = TD.PhasePartition(ts)
+        qv = TD.vapor_specific_humidity(ts)
+        qr = prog_pr.q_rai[k]
+        qs = prog_pr.q_sno[k]
 
+        # TODO - handle the limiters elsewhere
         # TODO - move limiters elsewhere
-        qt_tendency =
-            min(prog_pr.qr[k] / Δt, -CM1.evaporation_sublimation(param_set, rain_type, q, prog_pr.qr[k], ρ0_c[k], T_gm))
+        qt_tendency = min(
+            prog_pr.q_rai[k] / Δt,
+            -CM1.evaporation_sublimation(param_set, rain_type, q, prog_pr.q_rai[k], ρ0_c[k], T_gm),
+        )
+
+        aux_tc.qr_tendency_evap[k] = -qt_tendency
+        aux_tc.qs_tendency_melt[k] = 0.0
+        aux_tc.qs_tendency_dep_sub[k] = 0.0
+
+        tendencies_gm.q_tot[k] += qt_tendency
+        tendencies_pr.q_rai[k] += -qt_tendency
+        tendencies_pr.q_sno[k] += 0.0
 
-        aux_tc.qt_tendency_rain_evap[k] = qt_tendency
-        aux_tc.qr_tendency_rain_evap[k] = -qt_tendency
-        tendencies_pr.qr[k] += aux_tc.qr_tendency_rain_evap[k]
-        aux_tc.θ_liq_ice_tendency_rain_evap[k] = θ_liq_ice_helper(ts, qt_tendency)
+        tendencies_gm.θ_liq_ice[k] += θ_liq_ice_helper(ts, qt_tendency)
     end
     return
 end

src/EDMF_Updrafts.jl

@@ -1,3 +1,5 @@
+# TODO - move to EDMF_Precipitation
+# TODO - do all three precip tendencies computations in one loop over k
 """
 Computes tendencies to qt and θ_liq_ice due to precipitation formation
 """
@@ -13,11 +15,14 @@ function compute_precipitation_formation_tendencies(
     p0_c = center_ref_state(state).p0
     ρ0_c = center_ref_state(state).ρ0
     aux_up = center_aux_updrafts(state)
+    aux_tc = center_aux_turbconv(state)
     prog_pr = center_prog_precipitation(state)
     tendencies_pr = center_tendencies_precipitation(state)
 
-    @inbounds for i in 1:(up.n_updrafts)
-        @inbounds for k in real_center_indices(grid)
+    n_up = up.n_updrafts
+
+    @inbounds for k in real_center_indices(grid)
+        @inbounds for i in 1:(n_up)
             T_up = aux_up[i].T[k]
             q_tot_up = aux_up[i].q_tot[k]
             ts_up = TD.PhaseEquil_pTq(param_set, p0_c[k], T_up, q_tot_up)
@@ -26,19 +31,22 @@ function compute_precipitation_formation_tendencies(
             mph = precipitation_formation(
                 param_set,
                 precip.precipitation_model,
-                prog_pr.qr[k],
+                prog_pr.q_rai[k],
+                prog_pr.q_sno[k],
                 aux_up[i].area[k],
                 ρ0_c[k],
                 dt,
                 ts_up,
             )
-            up_thermo.qt_tendency_rain_formation[i, k] = mph.qt_tendency * aux_up[i].area[k]
-            up_thermo.θ_liq_ice_tendency_rain_formation[i, k] = mph.θ_liq_ice_tendency * aux_up[i].area[k]
+            aux_up[i].θ_liq_ice_tendency_precip_formation[k] = mph.θ_liq_ice_tendency * aux_up[i].area[k]
+            aux_up[i].qt_tendency_precip_formation[k] = mph.qt_tendency * aux_up[i].area[k]
+
+            tendencies_pr.q_rai[k] += mph.qr_tendency * aux_up[i].area[k]
+            tendencies_pr.q_sno[k] += mph.qs_tendency * aux_up[i].area[k]
         end
+        aux_tc.bulk.θ_liq_ice_tendency_precip_formation_tot[k] =
+            sum(i -> aux_up[i].θ_liq_ice_tendency_precip_formation[k], 1:n_up)
+        aux_tc.bulk.qt_tendency_precip_formation_tot[k] = sum(i -> aux_up[i].qt_tendency_precip_formation[k], 1:n_up)
     end
-    # 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_pr.qr) .+= -up_thermo.qt_tendency_rain_formation_tot
     return
 end