Initial testing for implicit diagnostic EDMFX

config/default_configs/default_config.yml

@@ -59,6 +59,9 @@ split_ode:
 use_krylov_method:
   help: "Whether to use a Krylov method to solve the linear system in Newton's method (only for ClimaTimeSteppers.jl)"
   value: false
+use_full_jacobian:
+  help: "Whether to use the full Jacobian (approximated using finite differences) in the Krylov method (only for ClimaTimeSteppers.jl)"
+  value: true
 use_dynamic_krylov_rtol:
   help: "Whether to use Eisenstat-Walker forcing instead of a constant relative tolerance in the Krylov method (only for ClimaTimeSteppers.jl)"
   value: false
@@ -71,11 +74,17 @@ ode_algo:
 krylov_rtol:
   help: "Relative tolerance of the Krylov method (only for ClimaTimeSteppers.jl; only used if `use_krylov_method` is `true`)"
   value: 0.1
+eisenstat_walker_forcing_init_rtol:
+  help: "Initial relative tolerance of the Krylov method for Eisenstat-Walker forcing (only for ClimaTimeSteppers.jl; only used if `use_krylov_method` and `use_dynamic_krylov_rtol` are `true`)"
+  value: 0.5
+eisenstat_walker_forcing_max_rtol:
+  help: "Maximum relative tolerance of the Krylov method for Eisenstat-Walker forcing (only for ClimaTimeSteppers.jl; only used if `use_krylov_method` and `use_dynamic_krylov_rtol` are `true`)"
+  value: 0.9
 eisenstat_walker_forcing_alpha:
   help: "Value of alpha to use for Eisenstat-Walker forcing (only for ClimaTimeSteppers.jl; only used if `use_krylov_method` and `use_dynamic_krylov_rtol` are `true`)"
   value: 2.0
 jvp_step_adjustment:
-  help: "Amount by which the step size of the forward difference approximation of the Jacobian-vector product in the Krylov method should be scaled (only used if `use_krylov_method` is `true`)"
+  help: "Amount by which the step size of the forward difference approximation of the Jacobian-vector product in the Krylov method should be scaled (only used if `use_krylov_method` and `use_full_jacobian` are `true`)"
   value: 1.0
 # Radiation
 rad:

config/model_configs/diagnostic_edmfx_dycoms_rf01_box_implicit.yml

@@ -0,0 +1,30 @@
+job_id: "diagnostic_edmfx_dycoms_rf01_box_implicitdif_exactjac_acousticgmpre_gmres_0.99999rtol_5newtoniters_dt2mins_withtkeblock"
+initial_condition: DYCOMS_RF01 
+subsidence: DYCOMS 
+edmf_coriolis: DYCOMS_RF01 
+rad: DYCOMS_RF01 
+surface_setup: DYCOMS_RF01
+turbconv: diagnostic_edmfx 
+prognostic_tke: true
+edmfx_upwinding: first_order 
+edmfx_entr_detr: true 
+edmfx_nh_pressure: true 
+edmfx_sgs_flux: true
+moist: equil
+config: box 
+hyperdiff: "true" 
+kappa_4: 1e21
+x_max: 1e8 
+y_max: 1e8 
+x_elem: 2 
+y_elem: 2 
+z_elem: 30 
+z_max: 1500 
+z_stretch: false
+use_krylov_method: true
+use_full_jacobian: true
+krylov_rtol: 0.99999
+max_newton_iters_ode: 5
+dt: "2mins" 
+t_end: "4hours" 
+dt_save_to_disk: "1mins"

examples/Manifest.toml

@@ -286,7 +286,7 @@ uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 version = "1.16.0"
 
 [[deps.ClimaAtmos]]
-deps = ["ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "CUDA", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "Distributions", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "ImageFiltering", "Insolation", "Interpolations", "IntervalSets", "JLD2", "LambertW", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "OrdinaryDiffEq", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "TerminalLoggers", "Test", "Thermodynamics", "YAML"]
+deps = ["ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "CUDA", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqCallbacks", "Distributions", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "ImageFiltering", "Insolation", "Interpolations", "IntervalSets", "JLD2", "LambertW", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "OrdinaryDiffEq", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "TerminalLoggers", "Test", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.16.0"
@@ -299,11 +299,11 @@ version = "0.5.3"
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"]
-git-tree-sha1 = "5317c7d3591f991f438ddb9e63da3395a219dda6"
+git-tree-sha1 = "7cc80c84873eed8e8e13b0882fb6c32f7b7fafc0"
 repo-rev = "dy/field_matrix"
 repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.10.48"
+version = "0.10.49"
 
 [[deps.ClimaCoreMakie]]
 deps = ["ClimaCore", "Makie"]
@@ -2095,9 +2095,9 @@ version = "0.4.0+0"
 
 [[deps.RootSolvers]]
 deps = ["DocStringExtensions", "ForwardDiff"]
-git-tree-sha1 = "9fb3462240d2898be5d4acf8925e47f70ec64d07"
+git-tree-sha1 = "bafbed2755bf8b1e42fd21f5bb417ff03e9cf1a1"
 uuid = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
-version = "0.3.5"
+version = "0.4.0"
 
 [[deps.RoundingEmulator]]
 git-tree-sha1 = "40b9edad2e5287e05bd413a38f61a8ff55b9557b"
@@ -2480,9 +2480,9 @@ version = "1.0.1"
 
 [[deps.Thermodynamics]]
 deps = ["DocStringExtensions", "KernelAbstractions", "Random", "RootSolvers"]
-git-tree-sha1 = "ac3b5236da4d028b9a3984ae7b28730d46892dbf"
+git-tree-sha1 = "3f8dfde3a9b02a00509ee1bdabc25045e95101dd"
 uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
-version = "0.11.0"
+version = "0.11.1"
 
 [[deps.ThreadingUtilities]]
 deps = ["ManualMemory"]

src/prognostic_equations/advection.jl

@@ -178,16 +178,16 @@ function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         end
 
         # TODO: Move this to implicit_vertical_advection_tendency!.
-        if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations
-            vertical_transport!(
-                Yₜ.c.sgs⁰.ρatke[colidx],
-                ᶜJ[colidx],
-                ᶜρa⁰[colidx],
-                ᶠu³⁰[colidx],
-                ᶜtke⁰[colidx],
-                dt,
-                edmfx_upwinding,
-            )
-        end
+        # if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations
+        #     vertical_transport!(
+        #         Yₜ.c.sgs⁰.ρatke[colidx],
+        #         ᶜJ[colidx],
+        #         ᶜρa⁰[colidx],
+        #         ᶠu³⁰[colidx],
+        #         ᶜtke⁰[colidx],
+        #         dt,
+        #         edmfx_upwinding,
+        #     )
+        # end
     end
 end

src/prognostic_equations/implicit/implicit_solver.jl

@@ -1,4 +1,4 @@
-import LinearAlgebra: I, Adjoint, ldiv!
+import LinearAlgebra: I, Adjoint, ldiv!, mul!
 using ForwardDiff: Dual
 using ClimaCore.MatrixFields
 
@@ -82,8 +82,7 @@ function ImplicitEquationJacobian(
     end
     for center_var in (
         MatrixFields.@var(c.turbconv),
-        MatrixFields.@var(c.sgs⁰),
-        MatrixFields.@var(c.sgsʲs)
+        MatrixFields.@var(c.sgsʲs),
     )
         MatrixFields.has_var(Y, center_var) || continue
         blocks = (blocks..., (center_var, center_var) => ᶜminus_I)
@@ -93,6 +92,16 @@ function ImplicitEquationJacobian(
         blocks = (blocks..., (face_var, face_var) => ᶠminus_I)
     end
 
+    ρatke_var = MatrixFields.@var(c.sgs⁰.ρatke)
+    if MatrixFields.has_var(Y, ρatke_var)
+        blocks = (
+            blocks...,
+            (ρatke_var, ρatke_var) => similar(Y.c, DiagonalMatrixRow{FT}),
+            (ρatke_var, MatrixFields.@var(f.u₃)) =>
+                similar(Y.c, Bidiagonal_ACT3),
+        )
+    end
+
     sfc_var = MatrixFields.@var(sfc)
     if MatrixFields.has_var(Y, sfc_var)
         sfc_minus_I = similar(Y.sfc, DiagonalMatrixRow{FT})
@@ -137,6 +146,19 @@ function ldiv!(
     ldiv!(A.temp_x, A, A.temp_b)
     x .= A.temp_x
 end
+function mul!(
+    x::AbstractVector,
+    A::ImplicitEquationJacobian,
+    b::AbstractVector,
+    α::Number,
+    β::Number,
+)
+    @assert α == 1
+    @assert β == 0
+    A.temp_b .= b
+    mul!(A.temp_x, A, A.temp_b)
+    x .= A.temp_x
+end
 
 function ldiv!(
     x::Fields.FieldVector,
@@ -148,6 +170,16 @@ function ldiv!(
         @. x *= A.dtγ_ref[]
     end
 end
+function mul!(
+    x::Fields.FieldVector,
+    A::ImplicitEquationJacobian,
+    b::Fields.FieldVector,
+)
+    mul!(x, A.W, b)
+    if A.transform
+        @. x /= A.dtγ_ref[]
+    end
+end
 
 ################################################################################
 
@@ -174,7 +206,8 @@ function _Wfact!(A, Y, p, dtγ, colidx)
     (; W, enthalpy_correction) = A
     (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK∂ᶠu₃) = p
     (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p
-    (; energy_upwinding, tracer_upwinding, density_upwinding) = p
+    (; energy_upwinding, tracer_upwinding, density_upwinding, edmfx_upwinding) =
+        p
 
     FT = Spaces.undertype(axes(Y.c))
     one_ATC3 = CT3(FT(1))'
@@ -183,8 +216,8 @@ function _Wfact!(A, Y, p, dtγ, colidx)
     R_d = FT(CAP.R_d(params))
     cv_d = FT(CAP.cv_d(params))
     κ_d = FT(CAP.kappa_d(params))
-    MSLP = FT(CAP.MSLP(params))
-    T_0 = FT(CAP.T_0(params))
+    p_ref_theta = FT(CAP.p_ref_theta(params))
+    T_tri = FT(CAP.T_triple(params))
 
     ᶜρ = Y.c.ρ
     ᶜuₕ = Y.c.uₕ
@@ -206,17 +239,17 @@ function _Wfact!(A, Y, p, dtγ, colidx)
         DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix()
 
     # We can express the pressure as either
-    # ᶜp = MSLP * (ᶜρθ * R_d / MSLP)^(1 / (1 - κ_d)) or
-    # ᶜp = R_d * (ᶜρe_tot / cv_d + ᶜρ * (T_0 - (ᶜK + ᶜΦ) / cv_d)) + O(ᶜq_tot).
+    # ᶜp = p_ref_theta * (ᶜρθ * R_d / p_ref_theta)^(1 / (1 - κ_d)) or
+    # ᶜp = R_d * (ᶜρe_tot / cv_d + ᶜρ * (T_tri - (ᶜK + ᶜΦ) / cv_d)) + O(ᶜq_tot).
     # In the first case, we find that
     # ∂(ᶜp)/∂(ᶜρ) = 0I,
     # ∂(ᶜp)/∂(ᶜK) = 0I, and
     # ∂(ᶜp)/∂(ᶜρθ) =
     #     DiagonalMatrixRow(
-    #         R_d / (1 - κ_d) * (ᶜρθ * R_d / MSLP)^(κ_d / (1 - κ_d))
+    #         R_d / (1 - κ_d) * (ᶜρθ * R_d / p_ref_theta)^(κ_d / (1 - κ_d))
     #     ).
     # In the second case, we can ignore all O(ᶜq_tot) terms to approximate
-    # ∂(ᶜp)/∂(ᶜρ) ≈ DiagonalMatrixRow(R_d * (T_0 - (ᶜK + ᶜΦ) / cv_d)),
+    # ∂(ᶜp)/∂(ᶜρ) ≈ DiagonalMatrixRow(R_d * (T_tri - (ᶜK + ᶜΦ) / cv_d)),
     # ∂(ᶜp)/∂(ᶜK) ≈ DiagonalMatrixRow(-R_d * ᶜρ / cv_d), and
     # ∂(ᶜp)/∂(ᶜρe_tot) ≈ DiagonalMatrixRow(R_d / cv_d).
 
@@ -262,6 +295,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
         (MatrixFields.@var(c.ρq_ice), MatrixFields.@var(q_ice)),
         (MatrixFields.@var(c.ρq_rai), MatrixFields.@var(q_rai)),
         (MatrixFields.@var(c.ρq_sno), MatrixFields.@var(q_sno)),
+        (MatrixFields.@var(c.sgs⁰.ρatke), nothing),
     )
         MatrixFields.has_var(Y, ρχ_var) || continue
         ∂ᶜρχₜ∂ᶠu₃ = W[ρχ_var, MatrixFields.@var(f.u₃)]
@@ -271,6 +305,8 @@ function _Wfact!(A, Y, p, dtγ, colidx)
             ᶜ1
         elseif ρχ_var == MatrixFields.@var(c.ρe_tot)
             p.ᶜh_tot
+        elseif ρχ_var == MatrixFields.@var(c.sgs⁰.ρatke)
+            p.ᶜtke⁰
         else
             MatrixFields.get_var(ᶜspecific, χ_var)
         end
@@ -279,6 +315,8 @@ function _Wfact!(A, Y, p, dtγ, colidx)
         elseif ρχ_var == MatrixFields.@var(c.ρθ) ||
                ρχ_var == MatrixFields.@var(c.ρe_tot)
             energy_upwinding
+        elseif ρχ_var == MatrixFields.@var(c.sgs⁰.ρatke)
+            edmfx_upwinding
         else
             tracer_upwinding
         end
@@ -343,6 +381,18 @@ function _Wfact!(A, Y, p, dtγ, colidx)
         end
     end
 
+    ρatke_var = MatrixFields.@var(c.sgs⁰.ρatke)
+    if MatrixFields.has_var(Y, ρatke_var)
+        (; ᶜmixing_length, ᶜtke⁰) = p
+        turbconv_params = CAP.turbconv_params(params)
+        c_d = TCP.tke_diss_coeff(turbconv_params)
+        ∂ᶜρatkeₜ∂ᶜρatke = W[ρatke_var, ρatke_var]
+        @. ∂ᶜρatkeₜ∂ᶜρatke[colidx] = DiagonalMatrixRow(
+            -1 + dtγ * 3 / 2 * c_d * sqrt(max(ᶜtke⁰[colidx], 0)) /
+            ᶜmixing_length[colidx],
+        )
+    end
+
     # We can express the implicit tendency for vertical velocity as
     # ᶠu₃ₜ =
     #     -(ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) /
@@ -398,7 +448,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
                 DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅
                 DiagonalMatrixRow(
                     R_d / (1 - κ_d) *
-                    (ᶜρθ[colidx] * R_d / MSLP)^(κ_d / (1 - κ_d)),
+                    (ᶜρθ[colidx] * R_d / p_ref_theta)^(κ_d / (1 - κ_d)),
                 )
             )
     elseif MatrixFields.has_var(Y, MatrixFields.@var(c.ρe_tot))
@@ -407,7 +457,7 @@ function _Wfact!(A, Y, p, dtγ, colidx)
             dtγ * (
                 DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅
                 DiagonalMatrixRow(
-                    R_d * (T_0 - (ᶜK[colidx] + ᶜΦ[colidx]) / cv_d),
+                    R_d * (T_tri - (ᶜK[colidx] + ᶜΦ[colidx]) / cv_d),
                 ) +
                 DiagonalMatrixRow(
                     (

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -13,6 +13,8 @@ function implicit_tendency!(Yₜ, Y, p, t)
         end
         Fields.bycolumn(axes(Y.c)) do colidx
             implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
+            edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
+            edmfx_tke_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
             if p.turbconv_model isa TurbulenceConvection.EDMFModel
                 implicit_sgs_flux_tendency!(
                     Yₜ,
@@ -69,10 +71,15 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
         p
     (; turbconv_model, rayleigh_sponge) = p.atmos
     (; dt) = p.simulation
-    n = n_mass_flux_subdomains(turbconv_model)
+    n = n_prognostic_mass_flux_subdomains(turbconv_model)
     ᶜJ = Fields.local_geometry_field(Y.c).J
     (; ᶜspecific, ᶠu³, ᶜp, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, ᶜtemp_scalar) = p
 
+    advect_tke = use_prognostic_tke(turbconv_model)
+    ᶠu³⁰ = advect_tke ? p.ᶠu³⁰ : nothing
+    ᶜρa⁰ = advect_tke ? (n > 0 ? p.ᶜρa⁰ : Y.c.ρ) : nothing
+    ᶜtke⁰ = advect_tke ? (n > 0 ? p.ᶜspecific⁰.tke : p.ᶜtke⁰) : nothing
+
     ᶜ1 = ᶜtemp_scalar
     @. ᶜ1[colidx] = one(Y.c.ρ[colidx])
     vertical_transport!(
@@ -109,6 +116,17 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
             χ_name == :θ ? energy_upwinding : tracer_upwinding,
         )
     end
+    if advect_tke
+        vertical_transport!(
+            Yₜ.c.sgs⁰.ρatke[colidx],
+            ᶜJ[colidx],
+            ᶜρa⁰[colidx],
+            ᶠu³⁰[colidx],
+            ᶜtke⁰[colidx],
+            dt,
+            edmfx_upwinding,
+        )
+    end
 
     @. Yₜ.f.u₃[colidx] +=
         -(

src/prognostic_equations/remaining_tendency.jl

@@ -49,9 +49,9 @@ function additional_tendency!(Yₜ, Y, p, t)
 
         radiation_tendency!(Yₜ, Y, p, t, colidx, p.radiation_model)
         edmfx_entr_detr_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
-        edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)
+        # edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)
         edmfx_nh_pressure_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
-        edmfx_tke_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
+        # edmfx_tke_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
         explicit_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, p.turbconv_model)
         precipitation_tendency!(Yₜ, Y, p, t, colidx, p.precip_model)
         # NOTE: All ρa tendencies should be applied before calling this function

src/solver/type_getters.jl

@@ -431,14 +431,24 @@ function ode_configuration(::Type{FT}, parsed_args) where {FT}
             max_iters = parsed_args["max_newton_iters_ode"],
             krylov_method = if parsed_args["use_krylov_method"]
                 CTS.KrylovMethod(;
-                    jacobian_free_jvp = CTS.ForwardDiffJVP(;
-                        step_adjustment = FT(
-                            parsed_args["jvp_step_adjustment"],
-                        ),
-                    ),
+                    jacobian_free_jvp = if parsed_args["use_full_jacobian"]
+                        CTS.ForwardDiffJVP(;
+                            step_adjustment = FT(
+                                parsed_args["jvp_step_adjustment"],
+                            ),
+                        )
+                    else
+                        nothing
+                    end,
                     forcing_term = if parsed_args["use_dynamic_krylov_rtol"]
+                        initial_rtol = FT(
+                            parsed_args["eisenstat_walker_forcing_init_rtol"],
+                        )
                         α = FT(parsed_args["eisenstat_walker_forcing_alpha"])
-                        CTS.EisenstatWalkerForcing(; α)
+                        max_rtol = FT(
+                            parsed_args["eisenstat_walker_forcing_max_rtol"],
+                        )
+                        CTS.EisenstatWalkerForcing(; initial_rtol, α, max_rtol)
                     else
                         CTS.ConstantForcing(FT(parsed_args["krylov_rtol"]))
                     end,