WIP on updating velocities for dycore paper

CliMA · Jan 14, 2025 · e3f8b25 · e3f8b25
1 parent 8b6f11d
commit e3f8b25
Show file tree

Hide file tree

Showing 4 changed files with 141 additions and 36 deletions.
diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl
@@ -40,3 +40,22 @@ 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]
+    @. ᶜwₗ = FT(0.0004)
+    @. ᶜwᵢ = FT(0.0004)
+    return nothing
+end
diff --git a/src/cache/precomputed_quantities.jl b/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,12 @@ 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 +182,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 +526,31 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     #     (; ᶜmixing_length) = p.precomputed
     #     compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     # end
-
+    @. p.precomputed.ᶜwₜqₜ = Geometry.WVector(0)
+    @. p.precomputed.ᶜwₕhₜ = Geometry.WVector(0)
+    if moisture_model isa NonEquilMoistModel
+        set_sedimentation_precomputed_quantities!(Y, p, t)
+        @. p.precomputed.ᶜwₜqₜ +=
+            Geometry.WVector(
+                p.precomputed.ᶜwₗ * Y.c.ρq_liq + p.precomputed.ᶜwᵢ * Y.c.ρq_ice
+        ) / Y.c.ρ
+        @. p.precomputed.ᶜwₕhₜ +=
+            Geometry.WVector(
+                p.precomputed.ᶜwₗ * Y.c.ρq_liq * (TD.internal_energy_liquid(thermo_params, ᶜts) + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -p.precomputed.ᶜwₗ) + Geometry.UVWVector(ᶜu))/2) +
+                p.precomputed.ᶜwᵢ * Y.c.ρq_ice * (TD.internal_energy_ice(thermo_params, ᶜts)    + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -p.precomputed.ᶜwᵢ) + Geometry.UVWVector(ᶜu))/2)
+        ) / Y.c.ρ
+    end
     if precip_model isa Microphysics1Moment
         set_precipitation_precomputed_quantities!(Y, p, t)
+        @. p.precomputed.ᶜwₜqₜ +=
+            Geometry.WVector(
+                p.precomputed.ᶜwᵣ * Y.c.ρq_rai + p.precomputed.ᶜwₛ * Y.c.ρq_sno
+            ) / Y.c.ρ
+        @. p.precomputed.ᶜwₕhₜ +=
+            Geometry.WVector(
+                p.precomputed.ᶜwᵣ * Y.c.ρq_rai * (TD.internal_energy_liquid(thermo_params, ᶜts) + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -p.precomputed.ᶜwᵣ) + Geometry.UVWVector(ᶜu))/2) +
+                p.precomputed.ᶜwₛ * Y.c.ρq_sno * (TD.internal_energy_ice(thermo_params, ᶜts)    + ᶜΦ + norm_sqr(Geometry.UVWVector(0, 0, -p.precomputed.ᶜwₛ) + Geometry.UVWVector(ᶜu))/2)
+            ) / Y.c.ρ
     end
 
     if turbconv_model isa PrognosticEDMFX

diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl
@@ -194,18 +194,18 @@ 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,10 +218,21 @@ function ImplicitEquationJacobian(
                     diffuse_momentum(atmos.vert_diff) ?
                 similar(Y.c, TridiagonalRow) : FT(-1) * I,
         )
+    elseif atmos.moisture_model isa DryModel
+        MatrixFields.unrolled_map(
+            name -> (name, name) => FT(-1) * I,
+            (diffused_scalar_names..., ρatke_if_available..., @name(c.uₕ)),
+            )
     else
+        (
+        MatrixFields.unrolled_map(
+            name -> (name, name) => similar(Y.c, TridiagonalRow),
+            diffused_scalar_names,
+            )...,
         MatrixFields.unrolled_map(
             name -> (name, name) => FT(-1) * I,
-            (diffused_scalar_names..., @name(c.uₕ)),
+            (ρatke_if_available..., @name(c.uₕ)),
+            )...,
         )
     end
 
@@ -299,7 +310,7 @@ 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 +413,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 +480,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 +599,48 @@ 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 = -1 * (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)
+          )
+
+        ∂ᶜρ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,41 +657,33 @@ 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 =
                 dtγ * ᶜdiffusion_h_matrix ⋅
                 DiagonalMatrixRow(ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ)
         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,)
-        end
+            @. ∂ᶜρq_err_∂ᶜρq +=
+                dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ)
+            end
 
         if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke))
             turbconv_params = CAP.turbconv_params(params)
@@ -678,19 +729,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

diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/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,19 @@ 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