Merge pull request #4085 from CliMA/sa/add_implicit_sgs_tracers_massflux

add implicit sgs mass flux for tracers

Author: sajjadazimi

config/model_configs/prognostic_edmfx_rico_implicit_column.yml

@@ -9,6 +9,7 @@ implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_nh_pressure: true
 implicit_sgs_vertdiff: true
+implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_entr_model: "Generalized"
@@ -30,7 +31,7 @@ y_elem: 2
 z_elem: 100
 z_stretch: false
 perturb_initstate: false
-dt: "50secs"
+dt: "100secs"
 t_end: "24hours"
 dt_save_state_to_disk: "60mins"
 toml: [toml/prognostic_edmfx_1M.toml]

config/model_configs/prognostic_edmfx_trmm_implicit_column.yml

@@ -9,8 +9,10 @@ implicit_sgs_advection: true
 implicit_sgs_entr_detr: true
 implicit_nh_pressure: true
 implicit_sgs_vertdiff: true
+implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
+debug_jacobian: true
 edmfx_entr_model: "Generalized"
 edmfx_detr_model: "Generalized"
 edmfx_sgs_mass_flux: true

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -267,18 +267,30 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
         @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1
         if use_derivative(sgs_mass_flux_flag)
             (
+                MatrixFields.unrolled_map(
+                    name ->
+                        (name, get_χʲ_name_from_ρχ_name(name)) =>
+                            similar(Y.c, TridiagonalRow),
+                    available_tracer_names,
+                )...,
+                MatrixFields.unrolled_map(
+                    name ->
+                        (name, @name(c.sgsʲs.:(1).ρa)) =>
+                            similar(Y.c, TridiagonalRow),
+                    available_tracer_names,
+                )...,
+                MatrixFields.unrolled_map(
+                    name ->
+                        (name, @name(f.sgsʲs.:(1).u₃)) =>
+                            similar(Y.c, BidiagonalRow_ACT3),
+                    available_tracer_names,
+                )...,
                 (@name(c.ρe_tot), @name(c.sgsʲs.:(1).mse)) =>
                     similar(Y.c, TridiagonalRow),
-                (@name(c.ρq_tot), @name(c.sgsʲs.:(1).q_tot)) =>
-                    similar(Y.c, TridiagonalRow),
                 (@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)) =>
                     similar(Y.c, BidiagonalRow_ACT3),
-                (@name(c.ρq_tot), @name(f.sgsʲs.:(1).u₃)) =>
-                    similar(Y.c, BidiagonalRow_ACT3),
                 (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)) =>
                     similar(Y.c, TridiagonalRow),
-                (@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)) =>
-                    similar(Y.c, TridiagonalRow),
             )
         else
             ()
@@ -742,6 +754,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
         if use_derivative(sgs_advection_flag)
             (; ᶜgradᵥ_ᶠΦ) = p.core
             (; ᶜρʲs, ᶠu³ʲs, ᶜtsʲs, ᶜKʲs, bdmr_l, bdmr_r, bdmr) = p.precomputed
+
+            # upwinding options for q_tot and mse
             is_third_order =
                 p.atmos.numerics.edmfx_mse_q_tot_upwinding == Val(:third_order)
             ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1
@@ -757,6 +771,24 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                 bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})),
             ) # Need to wrap ᶠupwind_matrix in this for well-defined boundaries.
 
+            # upwinding options for other tracers
+            is_tracer_upwinding_third_order =
+                p.atmos.numerics.edmfx_tracer_upwinding == Val(:third_order)
+            ᶠtracer_upwind = is_tracer_upwinding_third_order ? ᶠupwind3 : ᶠupwind1
+            ᶠset_tracer_upwind_bcs = Operators.SetBoundaryOperator(;
+                top = Operators.SetValue(zero(CT3{FT})),
+                bottom = Operators.SetValue(zero(CT3{FT})),
+            ) # Need to wrap ᶠtracer_upwind in this for well-defined boundaries.
+            TracerUpwindMatrixRowType =
+                is_tracer_upwinding_third_order ? QuaddiagonalMatrixRow :
+                BidiagonalMatrixRow
+            ᶠtracer_upwind_matrix =
+                is_tracer_upwinding_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix
+            ᶠset_tracer_upwind_matrix_bcs = Operators.SetBoundaryOperator(;
+                top = Operators.SetValue(zero(TracerUpwindMatrixRowType{CT3{FT}})),
+                bottom = Operators.SetValue(zero(TracerUpwindMatrixRowType{CT3{FT}})),
+            ) # Need to wrap ᶠtracer_upwind_matrix in this for well-defined boundaries.
+
             ᶠu³ʲ_data = ᶠu³ʲs.:(1).components.data.:1
 
             ᶜkappa_mʲ = p.scratch.ᶜtemp_scalar
@@ -797,6 +829,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                 p.atmos.microphysics_model isa Microphysics1Moment ||
                 p.atmos.microphysics_model isa Microphysics2Moment
             )
+
                 ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1))))
                 ᶜ∂a∂z =
                     @. lazy(
@@ -832,15 +865,17 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                         dtγ * (
                             DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) -
                             ᶜadvdivᵥ_matrix() ⋅
-                            ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1)))
+                            ᶠset_tracer_upwind_matrix_bcs(
+                                ᶠtracer_upwind_matrix(ᶠu³ʲs.:(1)),
+                            )
                         ) - (I,)
                     ∂ᶜχʲ_err_∂ᶠu₃ʲ =
                         matrix[χʲ_name, @name(f.sgsʲs.:(1).u₃)]
                     @. ∂ᶜχʲ_err_∂ᶠu₃ʲ =
                         dtγ * (
                             -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
-                                ᶠset_upwind_bcs(
-                                    ᶠupwind(CT3(sign(ᶠu³ʲ_data)), ᶜχʲ),
+                                ᶠset_tracer_upwind_bcs(
+                                    ᶠtracer_upwind(CT3(sign(ᶠu³ʲ_data)), ᶜχʲ),
                                 ) * adjoint(C3(sign(ᶠu³ʲ_data))),
                             ) +
                             DiagonalMatrixRow(ᶜχʲ) ⋅ ᶜadvdivᵥ_matrix()
@@ -1267,6 +1302,64 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                         (ᶠu³ʲs.:(1) - ᶠu³) *
                         ᶠinterp((Y.c.sgsʲs.:(1).q_tot - ᶜq_tot)) / ᶠJ,
                     ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ)
+
+                # grid-mean tracers
+                if p.atmos.moisture_model isa NonEquilMoistModel && (
+                    p.atmos.microphysics_model isa Microphysics1Moment ||
+                    p.atmos.microphysics_model isa Microphysics2Moment
+                )
+
+                    microphysics_tracers = (
+                        (@name(c.ρq_liq), @name(c.sgsʲs.:(1).q_liq)),
+                        (@name(c.ρq_ice), @name(c.sgsʲs.:(1).q_ice)),
+                        (@name(c.ρq_rai), @name(c.sgsʲs.:(1).q_rai)),
+                        (@name(c.ρq_sno), @name(c.sgsʲs.:(1).q_sno)),
+                        (@name(c.ρn_liq), @name(c.sgsʲs.:(1).n_liq)),
+                        (@name(c.ρn_rai), @name(c.sgsʲs.:(1).n_rai)),
+                    )
+                    MatrixFields.unrolled_foreach(
+                        microphysics_tracers,
+                    ) do (ρχ_name, χʲ_name)
+                        MatrixFields.has_field(Y, ρχ_name) || return
+                        ᶜχʲ = MatrixFields.get_field(Y, χʲ_name)
+
+                        ∂ᶜρχ_err_∂ᶜχʲ =
+                            matrix[ρχ_name, χʲ_name]
+                        @. ∂ᶜρχ_err_∂ᶜχʲ =
+                            dtγ *
+                            -(ᶜadvdivᵥ_matrix()) ⋅
+                            DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅
+                            ᶠset_tracer_upwind_matrix_bcs(
+                                ᶠtracer_upwind_matrix(ᶠu³ʲs.:(1)),
+                            ) ⋅
+                            DiagonalMatrixRow(draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)))
+
+                        ∂ᶜρχ_err_∂ᶜρa =
+                            matrix[ρχ_name, @name(c.sgsʲs.:(1).ρa)]
+                        @. ∂ᶜρχ_err_∂ᶜρa =
+                            dtγ *
+                            -(ᶜadvdivᵥ_matrix()) ⋅
+                            DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅
+                            ᶠset_tracer_upwind_matrix_bcs(
+                                ᶠtracer_upwind_matrix(ᶠu³ʲs.:(1)),
+                            ) ⋅
+                            DiagonalMatrixRow(ᶜχʲ / ᶜρʲs.:(1))
+
+                        ∂ᶜρχ_err_∂ᶠu₃ʲ =
+                            matrix[ρχ_name, @name(f.sgsʲs.:(1).u₃)]
+                        @. ∂ᶜρχ_err_∂ᶠu₃ʲ =
+                            dtγ * (
+                                -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                                    ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ *
+                                    ᶠset_tracer_upwind_bcs(
+                                        ᶠtracer_upwind(CT3(sign(ᶠu³ʲ_data)),
+                                            draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)) * ᶜχʲ,
+                                        ),
+                                    ) * adjoint(C3(sign(ᶠu³ʲ_data))),
+                                )) ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ))
+
+                    end
+                end
             end
         elseif rs isa RayleighSponge
             ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =

src/utils/variable_manipulations.jl

@@ -343,7 +343,7 @@ If `χ_name` has internal structure (e.g. a composite name), the function
 recurses into the child names and prepends `ρ` at the lowest level.
 """
 function get_ρχ_name(χ_name)
-    parent_name = MatrixFields.extract_first(χ_name)
+    parent_name = MatrixFields.FieldName(MatrixFields.extract_first(χ_name))
     child_name = MatrixFields.drop_first(χ_name)
     ρχ_name =
         (child_name == MatrixFields.@name()) ?
@@ -353,6 +353,35 @@ function get_ρχ_name(χ_name)
     return ρχ_name
 end
 
+"""
+    get_χʲ_name_from_ρχ_name(ρχ_name::FieldName)
+
+Construct the `FieldName` corresponding to the specific tracer in the updraft 
+(`χʲ` with j = 1) associated with a given density-weighted tracer on the grid mean (`ρ·χ`).
+
+Given the name of a density-weighted tracer `ρχ_name`, this function returns the 
+corresponding name of the specific tracer in the first subgrid updraft.
+
+The function operates recursively on hierarchical field names:  
+- If `ρχ_name` is a base name (no children), it replaces the `ρ` prefix with the 
+  appropriate updraft-specific prefix (e.g. `sgsʲs.:(1)`) and converts the variable 
+  to its specific form.
+- If `ρχ_name` has internal structure (e.g. a composite name), the function 
+  recurses into the child names and applies the transformation at the lowest level.
+"""
+function get_χʲ_name_from_ρχ_name(ρχ_name)
+    parent_name = MatrixFields.FieldName(MatrixFields.extract_first(ρχ_name))
+    child_name = MatrixFields.drop_first(ρχ_name)
+    χʲ_name =
+        (child_name == MatrixFields.@name()) ?
+        MatrixFields.append_internal_name(
+            @name(sgsʲs.:(1)),
+            specific_tracer_name(ρχ_name),
+        ) :
+        MatrixFields.append_internal_name(parent_name, get_χʲ_name_from_ρχ_name(child_name))
+    return χʲ_name
+end
+
 """
     ρa⁰(ρ, sgsʲs, turbconv_model)