Merge pull request #3521 from CliMA/aj/velocity_madness

WIP on updating velocities for dycore paper

config/model_configs/sphere_aquaplanet_rhoe_equilmoist_allsky_gw_res.yml

@@ -2,7 +2,7 @@ z_max: 60000.0
 z_elem: 31
 dz_bottom: 50.0
 rayleigh_sponge: true
-dt: "400secs"
+dt: "360secs"
 t_end: "1days"
 dt_save_state_to_disk: "24hours"
 vert_diff: "FriersonDiffusion"

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-197
+198
 
 # **README**
 #
@@ -20,6 +20,10 @@
 
 
 #=
+198
+- Added terms to the implicit solver that result in changes in the
+aquaplanet (ρe_tot) equil allsky monin_obukhov varying insol gravity wave (gfdl_restart) high top 1-moment
+
 197
 - Added single column hydrostatic balance reproducibility test
 

src/cache/precipitation_precomputed_quantities.jl

@@ -40,3 +40,23 @@ function set_precipitation_precomputed_quantities!(Y, p, t)
     )
     return nothing
 end
+
+
+"""
+    set_sedimentation_precomputed_quantities!(Y, p, t)
+Updates the sedimentation terminal velocity stored in `p`
+for the non-equilibrium microphysics scheme
+"""
+function set_sedimentation_precomputed_quantities!(Y, p, t)
+    @assert (p.atmos.moisture_model isa NonEquilMoistModel)
+
+    (; ᶜwₗ, ᶜwᵢ) = p.precomputed
+
+    FT = eltype(Y)
+
+    # compute the precipitation terminal velocity [m/s]
+    # TODO - the actual parameterization will be added in the next PR
+    @. ᶜwₗ = FT(0)
+    @. ᶜwᵢ = FT(0)
+    return nothing
+end

src/cache/precomputed_quantities.jl

@@ -48,6 +48,8 @@ function precomputed_quantities(Y, atmos)
         ᶜspecific = specific_gs.(Y.c),
         ᶜu = similar(Y.c, C123{FT}),
         ᶠu³ = similar(Y.f, CT3{FT}),
+        ᶜwₜqₜ = similar(Y.c, Geometry.WVector{FT}),
+        ᶜwₕhₜ = similar(Y.c, Geometry.WVector{FT}),
         ᶜK = similar(Y.c, FT),
         ᶜts = similar(Y.c, TST),
         ᶜp = similar(Y.c, FT),
@@ -59,6 +61,9 @@ function precomputed_quantities(Y, atmos)
     )
     cloud_diagnostics_tuple =
         similar(Y.c, @NamedTuple{cf::FT, q_liq::FT, q_ice::FT})
+    sedimentation_quantities =
+        atmos.moisture_model isa NonEquilMoistModel ?
+        (; ᶜwₗ = similar(Y.c, FT), ᶜwᵢ = similar(Y.c, FT)) : (;)
     precipitation_sgs_quantities =
         atmos.precip_model isa Microphysics0Moment ?
         (; ᶜSqₜᵖʲs = similar(Y.c, NTuple{n, FT}), ᶜSqₜᵖ⁰ = similar(Y.c, FT)) :
@@ -174,6 +179,7 @@ function precomputed_quantities(Y, atmos)
         advective_sgs_quantities...,
         diagnostic_sgs_quantities...,
         vert_diff_quantities...,
+        sedimentation_quantities...,
         precipitation_quantities...,
         cloud_diagnostics_tuple,
         smagorinsky_lilly_quantities...,
@@ -517,9 +523,30 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     #     (; ᶜmixing_length) = p.precomputed
     #     compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     # end
-
+    (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed
+    @. ᶜwₜqₜ = Geometry.WVector(0)
+    @. ᶜwₕhₜ = Geometry.WVector(0)
+    #
+    # TODO - uncomment in the next PR. Right now for the purpose of testing
+    # we want to merge with 0 sedimentation and precipitation
+    #
+    if moisture_model isa NonEquilMoistModel
+        set_sedimentation_precomputed_quantities!(Y, p, t)
+        #    (; ᶜwₗ, ᶜwᵢ) = p.precomputed
+        #    @. ᶜwₜqₜ += Geometry.WVector(ᶜwₗ * Y.c.ρq_liq + ᶜwᵢ * Y.c.ρq_ice) / Y.c.ρ
+        #    @. ᶜwₕhₜ += Geometry.WVector(
+        #        ᶜwₗ * Y.c.ρq_liq * (TD.internal_energy_liquid(thermo_params, ᶜts) + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -ᶜwₗ) + Geometry.UVWVector(ᶜu))/2) +
+        #        ᶜwᵢ * Y.c.ρq_ice * (TD.internal_energy_ice(thermo_params, ᶜts)    + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -ᶜwᵢ) + Geometry.UVWVector(ᶜu))/2)
+        #    ) / Y.c.ρ
+    end
     if precip_model isa Microphysics1Moment
         set_precipitation_precomputed_quantities!(Y, p, t)
+        #    (; ᶜwᵣ, ᶜwₛ) = p.precomputed
+        #    @. ᶜwₜqₜ += Geometry.WVector(ᶜwᵣ * Y.c.ρq_rai + ᶜwₛ * Y.c.ρq_sno) / Y.c.ρ
+        #    @. ᶜwₕhₜ += Geometry.WVector(
+        #        ᶜwᵣ * Y.c.ρq_rai * (TD.internal_energy_liquid(thermo_params, ᶜts) + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -ᶜwᵣ) + Geometry.UVWVector(ᶜu))/2) +
+        #        ᶜwₛ * Y.c.ρq_sno * (TD.internal_energy_ice(thermo_params, ᶜts)    + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -ᶜwₛ) + Geometry.UVWVector(ᶜu))/2)
+        #    ) / Y.c.ρ
     end
 
     if turbconv_model isa PrognosticEDMFX

src/prognostic_equations/implicit/implicit_solver.jl

@@ -194,17 +194,16 @@ function ImplicitEquationJacobian(
         (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
     )
 
-    diffused_scalar_names =
-        (@name(c.ρe_tot), available_tracer_names..., ρatke_if_available...)
+    diffused_scalar_names = (@name(c.ρe_tot), available_tracer_names...)
     diffusion_blocks = if use_derivative(diffusion_flag)
         (
             MatrixFields.unrolled_map(
                 name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow),
-                diffused_scalar_names,
+                (diffused_scalar_names..., ρatke_if_available...),
             )...,
             MatrixFields.unrolled_map(
                 name -> (name, name) => similar(Y.c, TridiagonalRow),
-                diffused_scalar_names,
+                (diffused_scalar_names..., ρatke_if_available...),
             )...,
             (
                 is_in_Y(@name(c.ρq_tot)) ?
@@ -218,11 +217,25 @@ function ImplicitEquationJacobian(
                     diffuse_momentum(atmos.vert_diff) ?
                 similar(Y.c, TridiagonalRow) : FT(-1) * I,
         )
-    else
+    elseif atmos.moisture_model isa DryModel
         MatrixFields.unrolled_map(
             name -> (name, name) => FT(-1) * I,
-            (diffused_scalar_names..., @name(c.uₕ)),
+            (diffused_scalar_names..., ρatke_if_available..., @name(c.uₕ)),
+        )
+    else
+        (
+            MatrixFields.unrolled_map(
+                name -> (name, name) => similar(Y.c, TridiagonalRow),
+                diffused_scalar_names,
+            )...,
+            (@name(c.ρe_tot), @name(c.ρq_tot)) =>
+                similar(Y.c, TridiagonalRow),
+            MatrixFields.unrolled_map(
+                name -> (name, name) => FT(-1) * I,
+                (ρatke_if_available..., @name(c.uₕ)),
+            )...,
         )
+
     end
 
     sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX
@@ -299,7 +312,9 @@ function ImplicitEquationJacobian(
 
     alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ))
     alg =
-        if use_derivative(diffusion_flag) || use_derivative(sgs_advection_flag)
+        if use_derivative(diffusion_flag) ||
+           use_derivative(sgs_advection_flag) ||
+           !(atmos.moisture_model isa DryModel)
             alg₁_subalg₂ =
                 if atmos.turbconv_model isa PrognosticEDMFX &&
                    use_derivative(sgs_advection_flag)
@@ -402,6 +417,12 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
         p.precomputed.ᶜK,
         p.precomputed.ᶜts,
         p.precomputed.ᶜp,
+        p.precomputed.ᶜwₜqₜ,
+        p.precomputed.ᶜwₕhₜ,
+        (
+            p.atmos.moisture_model isa NonEquilMoistModel ?
+            (; p.precomputed.ᶜwₗ, p.precomputed.ᶜwᵢ) : (;)
+        )...,
         (
             p.atmos.precip_model isa Microphysics1Moment ?
             (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;)
@@ -463,7 +484,7 @@ end
 
 function update_implicit_equation_jacobian!(A, Y, p, dtγ)
     (; matrix, diffusion_flag, sgs_advection_flag, topography_flag) = A
-    (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ) = p
+    (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜh_tot) = p
     (;
         ᶜtemp_C3,
         ∂ᶜK_∂ᶜuₕ,
@@ -582,6 +603,66 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
             ∂ᶜK_∂ᶠu₃ - (I_u₃,)
     end
 
+
+    tracer_info = (
+        (@name(c.ρq_liq), @name(q_liq), @name(ᶜwₗ)),
+        (@name(c.ρq_ice), @name(q_ice), @name(ᶜwᵢ)),
+        (@name(c.ρq_rai), @name(q_rai), @name(ᶜwᵣ)),
+        (@name(c.ρq_sno), @name(q_sno), @name(ᶜwₛ)),
+    )
+    if !(p.atmos.moisture_model isa DryModel) || use_derivative(diffusion_flag)
+        ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)]
+        @. ∂ᶜρe_tot_err_∂ᶜρe_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρe_tot)) - (I,)
+    end
+
+    if !(p.atmos.moisture_model isa DryModel)
+        @. ∂ᶜρe_tot_err_∂ᶜρe_tot +=
+            dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
+            DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ᶠright_bias_matrix() ⋅
+            DiagonalMatrixRow(
+                -(1 + ᶜkappa_m) / ᶜρ * ifelse(
+                    ᶜh_tot == 0,
+                    (Geometry.WVector(FT(0)),),
+                    p.ᶜwₕhₜ / ᶜh_tot,
+                ),
+            )
+
+        ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)]
+        @. ∂ᶜρe_tot_err_∂ᶜρq_tot =
+            dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
+            DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ᶠright_bias_matrix() ⋅
+            DiagonalMatrixRow(
+                -(ᶜkappa_m) * ∂e_int_∂q_tot / ᶜρ * ifelse(
+                    ᶜh_tot == 0,
+                    (Geometry.WVector(FT(0)),),
+                    p.ᶜwₕhₜ / ᶜh_tot,
+                ),
+            )
+
+        ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)]
+        @. ∂ᶜρq_tot_err_∂ᶜρq_tot =
+            dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
+            DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ᶠright_bias_matrix() ⋅
+            DiagonalMatrixRow(
+                -1 / ᶜρ * ifelse(
+                    ᶜspecific.q_tot == 0,
+                    (Geometry.WVector(FT(0)),),
+                    p.ᶜwₜqₜ / ᶜspecific.q_tot,
+                ),
+            ) - (I,)
+
+        MatrixFields.unrolled_foreach(tracer_info) do (ρqₚ_name, _, wₚ_name)
+            MatrixFields.has_field(Y, ρqₚ_name) || return
+            ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name]
+            ᶜwₚ = MatrixFields.get_field(p, wₚ_name)
+            @. ∂ᶜρqₚ_err_∂ᶜρqₚ =
+                dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
+                DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ᶠright_bias_matrix() ⋅
+                DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,)
+        end
+
+    end
+
     if use_derivative(diffusion_flag)
         (; ᶜK_h, ᶜK_u) = p
         @. ᶜdiffusion_h_matrix =
@@ -598,40 +679,38 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
         end
 
         ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)]
-        ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)]
         @. ∂ᶜρe_tot_err_∂ᶜρ =
             dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(
                 (
                     -(1 + ᶜkappa_m) * ᶜspecific.e_tot -
                     ᶜkappa_m * ∂e_int_∂q_tot * ᶜspecific.q_tot
                 ) / ᶜρ,
             )
-        @. ∂ᶜρe_tot_err_∂ᶜρe_tot =
-            dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) -
-            (I,)
+        @. ∂ᶜρe_tot_err_∂ᶜρe_tot +=
+            dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ)
+
         if MatrixFields.has_field(Y, @name(c.ρq_tot))
             ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)]
-            @. ∂ᶜρe_tot_err_∂ᶜρq_tot =
+            ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)]
+            @. ∂ᶜρe_tot_err_∂ᶜρq_tot +=
                 dtγ * ᶜdiffusion_h_matrix ⋅
                 DiagonalMatrixRow(ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ)
+            @. ∂ᶜρq_tot_err_∂ᶜρ =
+                dtγ * ᶜdiffusion_h_matrix ⋅
+                DiagonalMatrixRow(-(ᶜspecific.q_tot) / ᶜρ)
+            @. ∂ᶜρq_tot_err_∂ᶜρq_tot +=
+                dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ)
         end
 
-        tracer_info = (
-            (@name(c.ρq_tot), @name(q_tot)),
-            (@name(c.ρq_liq), @name(q_liq)),
-            (@name(c.ρq_ice), @name(q_ice)),
-            (@name(c.ρq_rai), @name(q_rai)),
-            (@name(c.ρq_sno), @name(q_sno)),
-        )
-        MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name)
+        MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name, _)
             MatrixFields.has_field(Y, ρq_name) || return
             ᶜq = MatrixFields.get_field(ᶜspecific, q_name)
             ∂ᶜρq_err_∂ᶜρ = matrix[ρq_name, @name(c.ρ)]
             ∂ᶜρq_err_∂ᶜρq = matrix[ρq_name, ρq_name]
             @. ∂ᶜρq_err_∂ᶜρ =
                 dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(ᶜq) / ᶜρ)
-            @. ∂ᶜρq_err_∂ᶜρq =
-                dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) - (I,)
+            @. ∂ᶜρq_err_∂ᶜρq +=
+                dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ)
         end
 
         if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke))
@@ -678,19 +757,6 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
                 dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_u_matrix - (I,)
         end
 
-        ᶠlg = Fields.local_geometry_field(Y.f)
-        precip_info =
-            ((@name(c.ρq_rai), @name(ᶜwᵣ)), (@name(c.ρq_sno), @name(ᶜwₛ)))
-        MatrixFields.unrolled_foreach(precip_info) do (ρqₚ_name, wₚ_name)
-            MatrixFields.has_field(Y, ρqₚ_name) || return
-            ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name]
-            ᶜwₚ = MatrixFields.get_field(p, wₚ_name)
-            ᶠtmp = p.ᶠtemp_CT3
-            @. ᶠtmp = CT3(unit_basis_vector_data(CT3, ᶠlg)) * ᶠwinterp(ᶜJ, ᶜρ)
-            @. ∂ᶜρqₚ_err_∂ᶜρqₚ +=
-                dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠtmp) ⋅
-                ᶠright_bias_matrix() ⋅ DiagonalMatrixRow(-(ᶜwₚ) / ᶜρ)
-        end
     end
 
     if p.atmos.turbconv_model isa PrognosticEDMFX

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -144,11 +144,15 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜJ = Fields.local_geometry_field(Y.c).J
     (; ᶠgradᵥ_ᶜΦ) = p.core
     (; ᶜh_tot, ᶜspecific, ᶠu³, ᶜp) = p.precomputed
+    (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed
 
     @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠwinterp(ᶜJ, Y.c.ρ) * ᶠu³)
+    @. Yₜ.c.ρ -= ᶜprecipdivᵥ(ᶠwinterp(ᶜJ, Y.c.ρ) * ᶠright_bias(-(ᶜwₜqₜ)))
 
     # Central advection of active tracers (e_tot and q_tot)
     vertical_transport!(Yₜ.c.ρe_tot, ᶜJ, Y.c.ρ, ᶠu³, ᶜh_tot, dt, Val(:none))
+    @. Yₜ.c.ρe_tot -= ᶜprecipdivᵥ(ᶠwinterp(ᶜJ, Y.c.ρ) * ᶠright_bias(-(ᶜwₕhₜ)))
+
     if !(moisture_model isa DryModel)
         vertical_transport!(
             Yₜ.c.ρq_tot,
@@ -159,6 +163,20 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
             dt,
             Val(:none),
         )
+        @. Yₜ.c.ρq_tot -=
+            ᶜprecipdivᵥ(ᶠwinterp(ᶜJ, Y.c.ρ) * ᶠright_bias(-(ᶜwₜqₜ)))
+    end
+
+    if moisture_model isa NonEquilMoistModel
+        (; ᶜwₗ, ᶜwᵢ) = p.precomputed
+        @. Yₜ.c.ρq_liq -= ᶜprecipdivᵥ(
+            ᶠwinterp(ᶜJ, Y.c.ρ) *
+            ᶠright_bias(Geometry.WVector(-(ᶜwₗ)) * ᶜspecific.q_liq),
+        )
+        @. Yₜ.c.ρq_ice -= ᶜprecipdivᵥ(
+            ᶠwinterp(ᶜJ, Y.c.ρ) *
+            ᶠright_bias(Geometry.WVector(-(ᶜwᵢ)) * ᶜspecific.q_ice),
+        )
     end
 
     if precip_model isa Microphysics1Moment