Merge pull request #4165 from CliMA/sa/implicit_water_advection

sajjadazimi · web-flow · commit 0116312013ff · 2025-12-15T22:07:10.000Z
implicit water advection
diff --git a/.buildkite/ci_driver.jl b/.buildkite/ci_driver.jl
@@ -184,8 +184,8 @@ if config.parsed_args["check_conservation"]
     @info "    Net water change / total water: $water_conservation"
 
     @test energy_conservation ≈ 0 atol = 100 * eps(FT)
-    @test mass_conservation ≈ 0 atol = 100 * eps(FT)
-    @test water_conservation ≈ 0 atol = 100 * eps(FT)
+    @test mass_conservation ≈ 0 atol = 250 * eps(FT)
+    @test water_conservation ≈ 0 atol = 250 * eps(FT)
 end
 
 # Visualize the solution
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -194,8 +194,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         @. Yₜ.c.ρb_rim -= ᶜprecipdivᵥ(ᶠρ * ᶠright_bias(- ᶜwᵢ * specific(ρb_rim, ρ)))
     end
 
-    # TODO - decide if this needs to be explicit or implicit
-    #vertical_advection_of_water_tendency!(Yₜ, Y, p, t)
+    vertical_advection_of_water_tendency!(Yₜ, Y, p, t)
 
     # This is equivalent to grad_v(Φ) + grad_v(p) / ρ
     ᶜΦ_r = @. lazy(phi_r(thermo_params, ᶜts))
diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl
@@ -196,6 +196,10 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
                 name -> (@name(c.ρe_tot), name) => similar(Y.c, TridiagonalRow),
                 available_condensate_mass_names,
             )...,
+            MatrixFields.unrolled_map(
+                name -> (@name(c.ρq_tot), name) => similar(Y.c, TridiagonalRow),
+                available_condensate_mass_names,
+            )...,
             (@name(c.uₕ), @name(c.uₕ)) =>
                 !isnothing(atmos.turbconv_model) ||
                     !disable_momentum_vertical_diffusion(
@@ -213,6 +217,14 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
                 name -> (name, name) => similar(Y.c, TridiagonalRow),
                 diffused_scalar_names,
             )...,
+            MatrixFields.unrolled_map(
+                name -> (@name(c.ρe_tot), name) => similar(Y.c, TridiagonalRow),
+                available_condensate_mass_names,
+            )...,
+            MatrixFields.unrolled_map(
+                name -> (@name(c.ρq_tot), name) => similar(Y.c, TridiagonalRow),
+                available_condensate_mass_names,
+            )...,
             (@name(c.ρe_tot), @name(c.ρq_tot)) =>
                 similar(Y.c, TridiagonalRow),
             MatrixFields.unrolled_map(
@@ -347,8 +359,10 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
            !(atmos.moisture_model isa DryModel)
             gs_scalar_subalg = if !(atmos.moisture_model isa DryModel)
                 MatrixFields.BlockLowerTriangularSolve(
-                    @name(c.ρq_tot),
                     available_condensate_mass_names...,
+                    alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                        @name(c.ρq_tot),
+                    ),
                 )
             else
                 MatrixFields.BlockDiagonalSolve()
@@ -588,51 +602,21 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         (@name(c.ρq_rim), @name(ᶜwᵢ), FT(1)),
         (@name(c.ρb_rim), @name(ᶜwᵢ), FT(1)),
     )
+    internal_energy_func(name) =
+        (name == @name(c.ρq_liq) || name == @name(c.ρq_rai)) ? TD.internal_energy_liquid :
+        (name == @name(c.ρq_ice) || name == @name(c.ρq_sno)) ? TD.internal_energy_ice :
+        nothing
     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)
-        #TODO: tetsing explicit vs implicit
-        #@. ∂ᶜρe_tot_err_∂ᶜρe_tot +=
-        #    dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
-        #    DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅
-        #    DiagonalMatrixRow(
-        #        -(1 + ᶜkappa_m) / ᶜρ * ifelse(
-        #            ᶜh_tot == 0,
-        #            (Geometry.WVector(FT(0)),),
-        #            p.precomputed.ᶜwₕhₜ / ᶜh_tot,
-        #        ),
-        #    )
-
         ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)]
         @. ∂ᶜρe_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρq_tot))
-        #TODO: tetsing explicit vs implicit
-        #@. ∂ᶜρe_tot_err_∂ᶜρq_tot =
-        #    dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
-        #    DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅
-        #    DiagonalMatrixRow(
-        #        -(ᶜkappa_m) * ∂e_int_∂q_tot / ᶜρ * ifelse(
-        #            ᶜh_tot == 0,
-        #            (Geometry.WVector(FT(0)),),
-        #            p.precomputed.ᶜwₕhₜ / ᶜh_tot,
-        #        ),
-        #    )
 
         ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)]
         @. ∂ᶜρq_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρq_tot_err_∂ᶜρq_tot)) - (I,)
-        #TODO: testing explicit vs implicit
-        #@. ∂ᶜρq_tot_err_∂ᶜρq_tot =
-        #    dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅
-        #    DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅
-        #    DiagonalMatrixRow(
-        #        -1 / ᶜρ * ifelse(
-        #           ᶜq_tot == 0,
-        #            (Geometry.WVector(FT(0)),),
-        #            p.precomputed.ᶜwₜqₜ / ᶜq_tot,
-        #        ),
-        #    ) - (I,)
 
         # This scratch variable computation could be skipped if no tracers are present
         @. p.scratch.ᶜbidiagonal_adjoint_matrix_c3 =
@@ -641,6 +625,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
 
         MatrixFields.unrolled_foreach(tracer_info) do (ρχₚ_name, wₚ_name, _)
             MatrixFields.has_field(Y, ρχₚ_name) || return
+
             ∂ᶜρχₚ_err_∂ᶜρχₚ = matrix[ρχₚ_name, ρχₚ_name]
             ᶜwₚ = MatrixFields.get_field(p.precomputed, wₚ_name)
             # TODO: come up with read-able names for the intermediate computations...
@@ -650,6 +635,24 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             @. ∂ᶜρχₚ_err_∂ᶜρχₚ =
                 p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅
                 p.scratch.ᶠband_matrix_wvec - (I,)
+
+            if ρχₚ_name in
+               (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno))
+                ∂ᶜρq_tot_err_∂ᶜρq = matrix[@name(c.ρq_tot), ρχₚ_name]
+                @. ∂ᶜρq_tot_err_∂ᶜρq =
+                    p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅
+                    p.scratch.ᶠband_matrix_wvec
+
+                ∂ᶜρe_tot_err_∂ᶜρq = matrix[@name(c.ρe_tot), ρχₚ_name]
+                e_int_func = internal_energy_func(ρχₚ_name)
+                @. ∂ᶜρe_tot_err_∂ᶜρq =
+                    p.scratch.ᶜbidiagonal_adjoint_matrix_c3 ⋅
+                    p.scratch.ᶠband_matrix_wvec ⋅
+                    DiagonalMatrixRow(
+                        e_int_func(thermo_params, p.precomputed.ᶜts) + ᶜΦ +
+                        $(Kin(ᶜwₚ, p.precomputed.ᶜu)),
+                    )
+            end
         end
 
     end
@@ -719,7 +722,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         for (q_name, e_int_q, ∂cv∂q) in microphysics_tracers
             MatrixFields.has_field(Y, q_name) || continue
             ∂ᶜρe_tot_err_∂ᶜρq = matrix[@name(c.ρe_tot), q_name]
-            @. ∂ᶜρe_tot_err_∂ᶜρq =
+            @. ∂ᶜρe_tot_err_∂ᶜρq +=
                 dtγ * ᶜdiffusion_h_matrix ⋅
                 DiagonalMatrixRow((ᶜkappa_m * (e_int_q - ∂cv∂q * (T - T_0)) - R_v * T) / ᶜρ)
         end
diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl
@@ -69,7 +69,6 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     horizontal_dynamics_tendency!(Yₜ, Y, p, t)
     hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
-    vertical_advection_of_water_tendency!(Yₜ, Y, p, t)
     additional_tendency!(Yₜ, Y, p, t)
     return Yₜ
 end
diff --git a/src/prognostic_equations/water_advection.jl b/src/prognostic_equations/water_advection.jl
@@ -19,20 +19,9 @@ is not a `DryModel`). The tendencies are calculated in a conservative form,
 
 Fluxes at cell faces are reconstructed using a first-order upwind scheme.
 Specifically, face-valued density (`ᶠρ`) is multiplied by a right-biased,
-negated cell-centered specific flux term (e.g., `ᶠright_bias(-ᶜwₜqₜ)` for total water,
-where `ᶜwₜqₜ` represents `w q_t` at cell centers). The resulting face flux
+negated cell-centered specific flux term. The resulting face flux
 is then diverged using the `ᶜprecipdivᵥ` operator.
 
-The precomputed terms `p.precomputed.ᶜwₜqₜ` and `p.precomputed.ᶜwₕhₜ` represent 
-cell-centered specific vertical fluxes (e.g., `w ⋅ q_t` and `w ⋅ h_t`,
-respectively, where `w` is a vertical velocity component).
-
-Additionally, when using prognostic EDMF with non-equilibrium moisture and microphysics,
-this function also applies advective tendencies to the subdomain moist static energy
-(`mse`) and total specific humidity (`q_tot`) fields within each mass flux
-subdomain. These tendencies use sedimentation velocities computed from the precomputed
-fields `ᶜwₕʲs` and `ᶜwₜʲs` for each subdomain j.
-
 Arguments:
 - `Yₜ`: The tendency state vector, modified in place.
 - `Y`: The current state vector.
@@ -42,22 +31,98 @@ Arguments:
 
 Modifies:
 - `Yₜ.c.ρ`, `Yₜ.c.ρe_tot`, and `Yₜ.c.ρq_tot` (always when moisture is present)
-- `Yₜ.c.sgsʲs.:(j).mse` and `Yₜ.c.sgsʲs.:(j).q_tot` (when using prognostic EDMFX with 
-  non-equilibrium moisture and multi-moment microphysics)
 """
 function vertical_advection_of_water_tendency!(Yₜ, Y, p, t)
 
+    (; params) = p
+    (; ᶜΦ) = p.core
+    (; ᶜu, ᶜts) = p.precomputed
+    thp = CAP.thermodynamics_params(params)
+
     ᶜJ = Fields.local_geometry_field(Y.c).J
     ᶠJ = Fields.local_geometry_field(Y.f).J
-    (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed
-
-    if !(p.atmos.moisture_model isa DryModel)
-        @. Yₜ.c.ρ -=
-            ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₜqₜ)))
-        @. Yₜ.c.ρe_tot -=
-            ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₕhₜ)))
-        @. Yₜ.c.ρq_tot -=
-            ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₜqₜ)))
+
+    microphysics_tracers = (
+        (@name(ρq_liq), @name(ᶜwₗ)),
+        (@name(ρq_ice), @name(ᶜwᵢ)),
+        (@name(ρq_rai), @name(ᶜwᵣ)),
+        (@name(ρq_sno), @name(ᶜwₛ)),
+    )
+    internal_energy_func(name) =
+        (name == @name(ρq_liq) || name == @name(ρq_rai)) ? TD.internal_energy_liquid :
+        (name == @name(ρq_ice) || name == @name(ρq_sno)) ? TD.internal_energy_ice :
+        nothing
+    MatrixFields.unrolled_foreach(microphysics_tracers) do (ρq_name, w_name)
+        MatrixFields.has_field(Y.c, ρq_name) || return
+
+        ᶜρq = MatrixFields.get_field(Y.c, ρq_name)
+        ᶜw = MatrixFields.get_field(p.precomputed, w_name)
+        vtt = @.lazy(
+            -1 * ᶜprecipdivᵥ(
+                ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(
+                    Geometry.WVector(-(ᶜw)) * specific(ᶜρq, Y.c.ρ),
+                ),
+            ),
+        )
+        @. Yₜ.c.ρ += vtt
+        @. Yₜ.c.ρq_tot += vtt
+
+        e_int_func = internal_energy_func(ρq_name)
+        @. Yₜ.c.ρe_tot -= ᶜprecipdivᵥ(
+            ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(
+                Geometry.WVector(-(ᶜw)) * specific(ᶜρq, Y.c.ρ) *
+                (e_int_func(thp, ᶜts) + ᶜΦ + $(Kin(ᶜw, ᶜu))),
+            ),
+        )
+    end
+
+    # For prognostic edmf, augment the energy tendencies with the additional energy contributions
+    # so that the total-grid energy flux remains consistent. Specifically, we enforce that the 
+    # grid-mean energy flux equals the sum of the subdomain (updraft/environment) energy fluxes 
+    # by accounting for the energy carried by sedimenting tracers.
+    if p.atmos.turbconv_model isa PrognosticEDMFX
+        (; ᶜρʲs, ᶜtsʲs, ᶜts⁰) = p.precomputed
+
+        ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰))
+
+        # TODO the following code works for only one updraft 
+        sgs_microphysics_tracers = (
+            (@name(q_liq), @name(ᶜwₗʲs.:(1)), @name(ᶜwₗ)),
+            (@name(q_ice), @name(ᶜwᵢʲs.:(1)), @name(ᶜwᵢ)),
+            (@name(q_rai), @name(ᶜwᵣʲs.:(1)), @name(ᶜwᵣ)),
+            (@name(q_sno), @name(ᶜwₛʲs.:(1)), @name(ᶜwₛ)),
+        )
+        MatrixFields.unrolled_foreach(sgs_microphysics_tracers) do (q_name, wʲ_name, w_name)
+            MatrixFields.has_field(Y.c.sgsʲs.:(1), q_name) || return
+
+            ᶜqʲ = MatrixFields.get_field(Y.c.sgsʲs.:(1), q_name)
+            ᶜwʲ = MatrixFields.get_field(p.precomputed, wʲ_name)
+            ᶜρq = MatrixFields.get_field(Y.c, get_ρχ_name(q_name))
+            ᶜw = MatrixFields.get_field(p.precomputed, w_name)
+
+            e_int_func = internal_energy_func(get_ρχ_name(q_name))
+            # TODO do we need to add kinetic energy of subdomains?
+            @. Yₜ.c.ρe_tot -=
+                ᶜprecipdivᵥ(
+                    ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ * ᶠright_bias(
+                        Geometry.WVector(-(ᶜwʲ)) *
+                        draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)) * ᶜqʲ *
+                        (
+                            e_int_func(thp, ᶜtsʲs.:(1)) - e_int_func(thp, ᶜts) -
+                            $(Kin(ᶜw, ᶜu))
+                        ),
+                    ),
+                )
+
+            ᶜwaq⁰ = @. lazy((ᶜρq * ᶜw - Y.c.sgsʲs.:(1).ρa * ᶜqʲ * ᶜwʲ) / ᶜρ⁰)
+            @. Yₜ.c.ρe_tot -=
+                -1 * ᶜprecipdivᵥ(
+                    ᶠinterp(ᶜρ⁰ * ᶜJ) / ᶠJ * ᶠright_bias(
+                        Geometry.WVector(-(ᶜwaq⁰)) *
+                        (e_int_func(thp, ᶜts⁰) - e_int_func(thp, ᶜts) - $(Kin(ᶜw, ᶜu))),
+                    ),
+                )
+        end
     end
 
     return nothing
