add du3_dq jacobian for microphysics tracers

Author: szy21

config/model_configs/baroclinic_wave_equil.yml

@@ -1,10 +1,10 @@
-precip_model: "0M"
+precip_model: "1M"
 dt_save_state_to_disk: "2days"
 reproducibility_test: true
 initial_condition: "MoistBaroclinicWave"
 dt: "450secs"
-t_end: "10days"
-moist: "equil"
+t_end: "6days"
+moist: "nonequil"
 disable_surface_flux_tendency: true
 diagnostics:
   - short_name: [pfull, wa, va, rv, hus, ke]

src/parameters/Parameters.jl

@@ -141,7 +141,7 @@ for var in fieldnames(TD.Parameters.ThermodynamicsParameters)
     @eval $var(ps::ACAP) = TD.Parameters.$var(thermodynamics_params(ps))
 end
 # Thermodynamics derived parameters
-for var in [:Rv_over_Rd, :kappa_d, :e_int_v0, :cv_v, :cv_l, :cv_d]
+for var in [:Rv_over_Rd, :kappa_d, :e_int_v0, :e_int_i0, :cv_v, :cv_l, :cv_d]
     @eval $var(ps::ACAP) = TD.Parameters.$var(thermodynamics_params(ps))
 end
 

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -91,11 +91,16 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
         is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : ()
     sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : ()
 
-    condensate_names = (
+    condensate_mass_names = (
         @name(c.ρq_liq),
         @name(c.ρq_ice),
         @name(c.ρq_rai),
         @name(c.ρq_sno),
+    )
+    available_condensate_mass_names =
+        MatrixFields.unrolled_filter(is_in_Y, condensate_mass_names)
+    condensate_names = (
+        condensate_mass_names...,
         @name(c.ρn_liq),
         @name(c.ρn_rai),
         # P3 frozen
@@ -161,6 +166,10 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
             name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),
             active_scalar_names,
         )...,
+        MatrixFields.unrolled_map(
+            name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),
+            available_condensate_mass_names,
+        )...,
         (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACT12),
         (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
     )
@@ -183,6 +192,10 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
                         similar(Y.c, TridiagonalRow),
                 ) : ()
             )...,
+            MatrixFields.unrolled_map(
+                name -> (@name(c.ρe_tot), name) => similar(Y.f, TridiagonalRow),
+                available_condensate_mass_names,
+            )...,
             (@name(c.uₕ), @name(c.uₕ)) =>
                 !isnothing(atmos.turbconv_model) ||
                     !disable_momentum_vertical_diffusion(
@@ -421,10 +434,14 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
     LH_s0 = FT(CAP.LH_s0(params))
     Δcp_l = FT(CAP.cp_l(params) - CAP.cp_v(params))
     Δcp_i = FT(CAP.cp_i(params) - CAP.cp_v(params))
+    Δcv_l = FT(CAP.cp_l(params)) - FT(CAP.cv_v(params))
+    Δcv_i = FT(CAP.cp_i(params)) - FT(CAP.cv_v(params))
+    e_int_v0 = FT(CAP.e_int_v0(params))
+    e_int_s0 = FT(CAP.e_int_i0(params)) + FT(CAP.e_int_v0(params))
     # This term appears a few times in the Jacobian, and is technically
     # minus ∂e_int_∂q_tot
-    thermo_params = CAP.thermodynamics_params(params)
     ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - e_int_v0
+    thermo_params = CAP.thermodynamics_params(params)
     ᶜh_tot = @. lazy(
         TD.total_specific_enthalpy(
             thermo_params,
@@ -532,6 +549,24 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜ∂p∂ρq_tot)
     end
 
+    microphysics_tracers = (
+        (@name(c.ρq_liq), e_int_v0, Δcv_l),
+        (@name(c.ρq_ice), e_int_s0, Δcv_i),
+        (@name(c.ρq_rai), e_int_v0, Δcv_l),
+        (@name(c.ρq_sno), e_int_s0, Δcv_i),
+    )
+    # TODO using unrolled_foreach here allocates! (breaks the flame tests)
+    # MatrixFields.unrolled_foreach(
+    #     microphysics_tracers,
+    # ) do (q_name, e_int_q, ∂cv∂q)
+    for (q_name, e_int_q, ∂cv∂q) in microphysics_tracers
+        MatrixFields.has_field(Y, q_name) || continue
+        ∂ᶠu₃_err_∂ᶜρq = matrix[@name(f.u₃), q_name]
+        @. ∂ᶠu₃_err_∂ᶜρq =
+            dtγ * ᶠp_grad_matrix ⋅
+            DiagonalMatrixRow(ᶜkappa_m * (e_int_q - ∂cv∂q * (T - T_0)) - R_v * T)
+    end
+
     ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)]
     ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)]
     I_u₃ = DiagonalMatrixRow(one_C3xACT3)
@@ -690,6 +725,18 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                 dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ)
         end
 
+        # TODO using unrolled_foreach here allocates! (breaks the flame tests)
+        # MatrixFields.unrolled_foreach(
+        #     microphysics_tracers,
+        # ) do (q_name, e_int_q, ∂cv∂q)
+        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), @name(q_name)]
+            @. ∂ᶜρe_tot_err_∂ᶜρq =
+                dtγ * ᶜdiffusion_h_matrix ⋅
+                DiagonalMatrixRow((ᶜkappa_m * (e_int_q - ∂cv∂q * (T - T_0)) - R_v * T) / ᶜρ)
+        end
+
         MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, _, α)
             MatrixFields.has_field(Y, ρχ_name) || return
             ∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)]
@@ -1218,6 +1265,20 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
                     p.scratch.ᶜtridiagonal_matrix_scalar ⋅
                     DiagonalMatrixRow(ᶜ∂p∂ρq_tot / ᶜρ)
 
+                # TODO using unrolled_foreach here allocates! (breaks the flame tests)
+                # MatrixFields.unrolled_foreach(
+                #     microphysics_tracers,
+                # ) do (q_name, e_int_q, ∂cv∂q)
+                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), @name(q_name)]
+                    @. ∂ᶜρe_tot_err_∂ᶜρq +=
+                        p.scratch.ᶜtridiagonal_matrix_scalar ⋅
+                        DiagonalMatrixRow(
+                            (ᶜkappa_m * (e_int_q - ∂cv∂q * (T - T_0)) - R_v * T) / ᶜρ,
+                        )
+                end
+
                 @. ∂ᶜρe_tot_err_∂ᶜρe_tot +=
                     p.scratch.ᶜtridiagonal_matrix_scalar ⋅
                     DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ)