adding noneq flags and blocks

Author: Olivia Alcabes

src/prognostic_equations/implicit/approx_jacobian.jl

@@ -17,6 +17,7 @@ use_derivative(::IgnoreDerivative) = false
         sgs_entr_detr_flag,
         sgs_mass_flux_flag,
         sgs_nh_pressure_flag,
+        noneq_cloud_formation_flag,
         approximate_solve_iters,
     )
 
@@ -40,16 +41,19 @@ setting their `DerivativeFlag`s to either `UseDerivative` or `IgnoreDerivative`.
   subgrid-scale mass flux tendency should be computed
 - `sgs_nh_pressure_flag::DerivativeFlag`: whether the derivatives of the
   subgrid-scale non-hydrostatic pressure drag tendency should be computed
+- `noneq_cloud_formation_flag::DerivativeFlag`: whether the derivatives
+  of the nonequilibrium cloud formation sources should be computed
 - `approximate_solve_iters::Int`: number of iterations to take for the
   approximate linear solve required when the `diffusion_flag` is `UseDerivative`
 """
-struct ApproxJacobian{F1, F2, F3, F4, F5, F6} <: JacobianAlgorithm
+struct ApproxJacobian{F1, F2, F3, F4, F5, F6, F7} <: JacobianAlgorithm
     topography_flag::F1
     diffusion_flag::F2
     sgs_advection_flag::F3
     sgs_entr_detr_flag::F4
     sgs_mass_flux_flag::F5
     sgs_nh_pressure_flag::F6
+    noneq_cloud_formation_flag::F7
     approximate_solve_iters::Int
 end
 
@@ -59,6 +63,7 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
         diffusion_flag,
         sgs_advection_flag,
         sgs_mass_flux_flag,
+        noneq_cloud_formation_flag,
     ) = alg
     FT = Spaces.undertype(axes(Y.c))
     CTh = CTh_vector_type(axes(Y.c))
@@ -82,14 +87,12 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
         is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : ()
     sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : ()
 
-    tracer_names = (
-        @name(c.ρq_tot),
-        @name(c.ρq_liq),
-        @name(c.ρq_ice),
-        @name(c.ρq_rai),
-        @name(c.ρq_sno),
-    )
-    available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names)
+    condensate_names =
+        (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno))
+    available_condensate_names =
+        MatrixFields.unrolled_filter(is_in_Y, condensate_names)
+    available_tracer_names =
+        (ρq_tot_if_available..., available_condensate_names...)
 
     sgs_tracer_names = (
         @name(c.sgsʲs.:(1).q_tot),
@@ -139,6 +142,17 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
         (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
     )
 
+    condensate_blocks =
+        if atmos.moisture_model isa NonEquilMoistModel &&
+            use_derivative(noneq_cloud_formation_flag)
+            (
+                (@name(c.ρq_liq), @name(c.ρq_tot)) => similar(Y.c, DiagonalRow),
+                (@name(c.ρq_ice), @name(c.ρq_tot)) => similar(Y.c, DiagonalRow),
+            )
+        else
+            ()
+        end
+
     diffused_scalar_names = (@name(c.ρe_tot), available_tracer_names...)
     diffusion_blocks = if use_derivative(diffusion_flag)
         (
@@ -253,6 +267,7 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
         identity_blocks...,
         sgs_advection_blocks...,
         advection_blocks...,
+        condensate_blocks...,
         diffusion_blocks...,
         sgs_massflux_blocks...,
     )
@@ -318,7 +333,7 @@ function jacobian_cache(alg::ApproxJacobian, Y, atmos)
     temp_matrix_column = similar(first_column(temp_matrix))
 
     return (;
-        matrix = MatrixFields.FieldMatrixWithSolver(matrix, Y, solver_alg),
+        matrix = MatrixFields.FieldMatrixWithSolver(matrix, Y, full_alg),
         temp_matrix,
         temp_matrix_column,
     )
@@ -348,8 +363,9 @@ function update_jacobian!(alg::ApproxJacobian, cache, Y, p, dtγ, t)
         diffusion_flag,
         sgs_advection_flag,
         sgs_entr_detr_flag,
-        sgs_nh_pressure_flag,
         sgs_mass_flux_flag,
+        sgs_nh_pressure_flag,
+        noneq_cloud_formation_flag,
     ) = alg
     (; matrix) = cache
     (; params) = p
@@ -555,6 +571,183 @@ function update_jacobian!(alg::ApproxJacobian, cache, Y, p, dtγ, t)
 
     end
 
+    if p.atmos.moisture_model isa NonEquilMoistModel &&
+        use_derivative(noneq_flag)
+         p_vapₛₗ(tps, ts) = TD.saturation_vapor_pressure(tps, ts, TD.Liquid())
+         p_vapₛᵢ(tps, ts) = TD.saturation_vapor_pressure(tps, ts, TD.Ice())
+ 
+         function ∂p_vapₛₗ_∂T(tps, ts)
+             T = TD.air_temperature(tps, ts)
+             Rᵥ = TD.Parameters.R_v(tps)
+             Lᵥ = TD.latent_heat_vapor(tps, ts)
+             return p_vapₛₗ(tps, ts) * Lᵥ / (Rᵥ * T^2)
+         end
+         function ∂p_vapₛᵢ_∂T(tps, ts)
+             T = TD.air_temperature(tps, ts)
+             Rᵥ = TD.Parameters.R_v(tps)
+             Lₛ = TD.latent_heat_sublim(tps, ts)
+             return p_vapₛᵢ(tps, ts) * Lₛ / (Rᵥ * T^2)
+         end
+ 
+         function ∂qₛₗ_∂T(tps, ts)
+             T = TD.air_temperature(tps, ts)
+             Rᵥ = TD.Parameters.R_v(tps)
+             Lᵥ = TD.latent_heat_vapor(tps, ts)
+             qᵥ_sat_liq = TD.q_vap_saturation_liquid(tps, ts)
+             return qᵥ_sat_liq * (Lᵥ / (Rᵥ * T^2) - 1 / T)
+         end
+         function ∂qₛᵢ_∂T(tps, ts)
+             T = TD.air_temperature(tps, ts)
+             Rᵥ = TD.Parameters.R_v(tps)
+             Lₛ = TD.latent_heat_sublim(tps, ts)
+             qᵥ_sat_ice = TD.q_vap_saturation_ice(tps, ts)
+             return qᵥ_sat_ice * (Lₛ / (Rᵥ * T^2) - 1 / T)
+         end
+ 
+         function Γₗ(tps, ts)
+             cₚ_air = TD.cp_m(tps, ts)
+             Lᵥ = TD.latent_heat_vapor(tps, ts)
+             return 1 + (Lᵥ / cₚ_air) * ∂qₛₗ_∂T(tps, ts)
+         end
+         function Γᵢ(tps, ts)
+             cₚ_air = TD.cp_m(tps, ts)
+             Lₛ = TD.latent_heat_sublim(tps, ts)
+             return 1 + (Lₛ / cₚ_air) * ∂qₛᵢ_∂T(tps, ts)
+         end
+ 
+         cmc = CAP.microphysics_cloud_params(params)
+         τₗ = cmc.liquid.τ_relax
+         τᵢ = cmc.ice.τ_relax
+         function limit(q, dt, n::Int)
+             return q / float(dt) / n
+         end
+     
+         function ∂ρqₗ_err_∂ρqᵪ(tps, ts, cmc, dt, deriv, limit_deriv)
+             FT_inner = eltype(tps)
+             q = TD.PhasePartition(tps, ts)
+             ρ = TD.air_density(tps, ts)
+ 
+             S = CMNe.conv_q_vap_to_q_liq_ice_MM2015(cmc.liquid, thp, q, ρ, Tₐ(tps, ts))
+ 
+             if S > FT_inner(0)
+                 if S <= limit(qᵥ(tps, ts), dt, 2)
+                     if TD.vapor_specific_humidity(q) + TD.liquid_specific_humidity(q) > FT_inner(0)
+                         return deriv
+                     else
+                         return FT_inner(0)
+                     end
+                 else
+                     return limit_deriv
+                 end
+             else
+                 if abs(S) <= limit(qₗ(tps, ts, qₚ(qᵣ)), dt, 2)
+                     if TD.vapor_specific_humidity(q) + TD.liquid_specific_humidity(q) > FT_inner(0)
+                         return -deriv
+                     else
+                         return FT_inner(0)
+                     end
+                 else
+                     return -limit_deriv
+                 end
+             end
+         end
+ 
+         function ∂ρqᵢ_err_∂ρqᵪ(tps, ts, cmc, dt, deriv, limit_deriv)
+             FT_inner = eltype(tps)
+             q = TD.PhasePartition(tps, ts)
+             ρ = TD.air_density(tps, ts)
+ 
+             S = CMNe.conv_q_vap_to_q_liq_ice_MM2015(cmc.ice, thp, q, ρ, Tₐ(tps, ts))
+ 
+             if S > FT_inner(0)
+                 if S <= limit(qᵥ(tps, ts), dt, 2)
+                     if TD.vapor_specific_humidity(q) + TD.ice_specific_humidity(q) > FT_inner(0)
+                         return deriv
+                     else
+                         return FT_inner(0)
+                     end
+                 else
+                     return limit_deriv
+                 end
+             else
+                 if abs(S) <= limit(qᵢ(thp, ts, qₚ(qₛ)), dt, 2)
+                     if TD.vapor_specific_humidity(q) + TD.ice_specific_humidity(q) > FT_inner(0)
+                         return -deriv
+                     else
+                         return FT_inner(0)
+                     end
+                 else
+                     return -limit_deriv
+                 end
+             end
+         end
+ 
+ 
+         ∂ᶜρqₗ_err_∂ᶜρqₗ = matrix[@name(c.ρq_liq), @name(c.ρq_liq)]
+         ∂ᶜρqᵢ_err_∂ᶜρqᵢ = matrix[@name(c.ρq_ice), @name(c.ρq_ice)]
+ 
+         ∂ᶜρqₗ_err_∂ᶜρqₜ = matrix[@name(c.ρq_liq), @name(c.ρq_tot)]
+         ∂ᶜρqᵢ_err_∂ᶜρqₜ = matrix[@name(c.ρq_ice), @name(c.ρq_tot)]
+ 
+         #@. ∂ᶜρqₗ_err_∂ᶜρqₗ -=
+         #    DiagonalMatrixRow(1 / (τₗ * Γₗ(thermo_params, ᶜts)))
+         @. ∂ᶜρqₗ_err_∂ᶜρqₗ +=
+             DiagonalMatrixRow(
+                 ∂ρqₗ_err_∂ρqᵪ(
+                     thermo_params, ᶜts, cmc, dt, (-1 / (τₗ * Γₗ(thermo_params, ᶜts))), (1/(2*float(dt))),
+                 )
+             )
+         
+         #@. ∂ᶜρqᵢ_err_∂ᶜρqᵢ -=
+         #    DiagonalMatrixRow(1 / (τᵢ * Γᵢ(thermo_params, ᶜts)))
+ 
+         @. ∂ᶜρqᵢ_err_∂ᶜρqᵢ +=
+             DiagonalMatrixRow(
+                 ∂ρqᵢ_err_∂ρqᵪ(
+                     thermo_params, ᶜts, cmc, dt, (-1 / (τᵢ * Γᵢ(thermo_params, ᶜts))), (1/(2*float(dt))),
+                     )
+                 )
+ 
+         ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts))
+         ᶜ∂T_∂p = @. lazy(1 / (ᶜρ * TD.gas_constant_air(thermo_params, ᶜts)))
+ 
+         # qₛₗ = p_vapₛₗ / p, qₛᵢ = p_vapₛᵢ / p
+         ᶜ∂qₛₗ_∂p = @. lazy(
+             -p_vapₛₗ(thermo_params, ᶜts) / ᶜp^2 +
+             ∂p_vapₛₗ_∂T(thermo_params, ᶜts) * ᶜ∂T_∂p / ᶜp,
+         )
+         ᶜ∂qₛᵢ_∂p = @. lazy(
+             -p_vapₛᵢ(thermo_params, ᶜts) / ᶜp^2 +
+             ∂p_vapₛᵢ_∂T(thermo_params, ᶜts) * ᶜ∂T_∂p / ᶜp,
+         )
+ 
+         ᶜ∂p_∂ρqₜ = @. lazy(
+             ᶜkappa_m * ∂e_int_∂q_tot +
+             ᶜ∂kappa_m∂q_tot * (
+                 cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ +
+                 ∂e_int_∂q_tot * ᶜspecific.q_tot
+             ),
+         )
+ 
+         #@. ∂ᶜρqₗ_err_∂ᶜρqₜ = DiagonalMatrixRow(
+         #    (1 - ᶜρ * ᶜ∂qₛₗ_∂p * ᶜ∂p_∂ρqₜ) / (τₗ * Γₗ(thermo_params, ᶜts)),
+         #)
+         @. ∂ᶜρqₗ_err_∂ᶜρqₜ = DiagonalMatrixRow(
+             ∂ρqₗ_err_∂ρqᵪ(
+                 thermo_params, ts, cmc, dt, ((1 - ᶜρ * ᶜ∂qₛₗ_∂p * ᶜ∂p_∂ρqₜ) / (τₗ * Γₗ(thermo_params, ᶜts))), FT(0)
+             )
+         )
+ 
+         #@. ∂ᶜρqᵢ_err_∂ᶜρqₜ = DiagonalMatrixRow(
+         #    (1 - ᶜρ * ᶜ∂qₛᵢ_∂p * ᶜ∂p_∂ρqₜ) / (τᵢ * Γᵢ(thermo_params, ᶜts)),
+         #)
+         @. ∂ᶜρqᵢ_err_∂ᶜρqₜ = DiagonalMatrixRow(
+             ∂ρqᵢ_err_∂ρqᵪ(
+                 thermo_params, ts, cmc, dt, ((1 - ᶜρ * ᶜ∂qₛᵢ_∂p * ᶜ∂p_∂ρqₜ) / (τᵢ * Γᵢ(thermo_params, ᶜts))), FT(0)
+             )
+         )
+     end
+
     if use_derivative(diffusion_flag)
         α_vert_diff_tracer = CAP.α_vert_diff_tracer(params)
         (; ᶜK_h, ᶜK_u) = p.precomputed

src/solver/type_getters.jl

@@ -423,6 +423,7 @@ get_jacobian(ode_algo, Y, atmos, parsed_args, output_dir) =
             DerivativeFlag(atmos.sgs_entr_detr_mode),
             DerivativeFlag(atmos.sgs_mf_mode),
             DerivativeFlag(atmos.sgs_nh_pressure_mode),
+            DerivativeFlag(atmos.noneq_cloud_formation_mode),
             parsed_args["approximate_linear_solve_iters"],
         )
         use_exact_jacobian = parsed_args["use_exact_jacobian"]