From 91d743e58057f880fc62e9c521e16de8761339f7 Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Tue, 21 Nov 2023 15:50:20 -0800 Subject: [PATCH] Add implicit diffusion solver --- .buildkite/pipeline.yml | 14 ++ Project.toml | 2 +- config/default_configs/default_config.yml | 3 + ...roclinic_wave_rhoe_equilmoist_expvdiff.yml | 6 +- ...roclinic_wave_rhoe_equilmoist_impvdiff.yml | 13 ++ docs/Manifest.toml | 16 +- examples/Manifest.toml | 32 ++- perf/Manifest.toml | 32 ++- .../implicit/implicit_solver.jl | 192 +++++++++++++----- .../implicit/implicit_tendency.jl | 10 + .../remaining_tendency.jl | 19 +- src/solver/type_getters.jl | 17 +- src/solver/types.jl | 6 + 13 files changed, 271 insertions(+), 91 deletions(-) create mode 100644 config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_impvdiff.yml diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 4c3bc4b317..ce75093009 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -266,6 +266,20 @@ steps: --nc_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff --fig_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff --case_name aquaplanet artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff/*" + + - label: ":computer: no lim ARS baroclinic wave (ρe) equilmoist implicit vertdiff" + command: > + julia --color=yes --project=examples examples/hybrid/driver.jl + --config_file $CONFIG_PATH/sphere_baroclinic_wave_rhoe_equilmoist_impvdiff.yml + + julia --color=yes --project=examples post_processing/remap/remap_pipeline.jl + --data_dir sphere_baroclinic_wave_rhoe_equilmoist_impvdiff + --out_dir sphere_baroclinic_wave_rhoe_equilmoist_impvdiff + + julia --color=yes --project=examples post_processing/plot/plot_pipeline.jl + --nc_dir sphere_baroclinic_wave_rhoe_equilmoist_impvdiff + --fig_dir sphere_baroclinic_wave_rhoe_equilmoist_impvdiff --case_name aquaplanet + artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist_impvdiff/*" # Add this back when we figure out what to do with SSP and zalesak # - label: ":computer: SSP zalesak tracer & energy upwind baroclinic wave (ρe_tot) equilmoist" diff --git a/Project.toml b/Project.toml index e96710d249..6f7218a948 100644 --- a/Project.toml +++ b/Project.toml @@ -50,7 +50,7 @@ Artifacts = "1" AtmosphericProfilesLibrary = "0.1" CLIMAParameters = "0.7.25" ClimaComms = "0.5.6" -ClimaCore = "0.11.0" +ClimaCore = "0.11.1" ClimaTimeSteppers = "0.7.14" CloudMicrophysics = "0.15.0" Colors = "0.12" diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index c678117a67..79e3826583 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -243,6 +243,9 @@ prognostic_tke: prognostic_surface: help: "Determines if surface temperature is prognostic [`false` (default), , `true`, `PrognosticSurfaceTemperature`, `PrescribedSurfaceTemperature`]" value: "false" +implicit_diffusion: + help: "Whether to treat the vertical diffusion tendency implicitly [`false` (default), `true`]" + value: false override_τ_precip: help: "If true, sets τ_precip to dt. Otherwise, τ_precip is set to the value in the toml dictionary" value: true diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml index 45ca42e3b9..bc5e5e4c91 100644 --- a/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml +++ b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml @@ -2,10 +2,10 @@ precip_model: "0M" vert_diff: "true" z_elem: 20 dz_bottom: 100 -dt_save_to_disk: "12hours" +dt_save_to_disk: "6hours" initial_condition: "MoistBaroclinicWave" -dt: "40secs" -t_end: "12hours" +dt: "20secs" +t_end: "6hours" job_id: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff" moist: "equil" toml: [toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml] diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_impvdiff.yml b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_impvdiff.yml new file mode 100644 index 0000000000..8eee5cc789 --- /dev/null +++ b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_impvdiff.yml @@ -0,0 +1,13 @@ +precip_model: "0M" +vert_diff: "true" +z_elem: 20 +dz_bottom: 100 +dt_save_to_disk: "6days" +initial_condition: "MoistBaroclinicWave" +implicit_diffusion: true +max_newton_iters_ode: 2 +dt: "400secs" +t_end: "6days" +job_id: "sphere_baroclinic_wave_rhoe_equilmoist_impvdiff" +moist: "equil" +toml: [toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml] diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 61bf8522bb..93108c31ab 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -269,10 +269,14 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.6" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] -git-tree-sha1 = "5c3809c5da88a3b01b8fe18558f50470fd3c33ed" +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] +git-tree-sha1 = "3d6cd2e89f0d99664b23549da9081ade4e7d8fdb" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" -version = "0.11.0" +version = "0.11.1" +weakdeps = ["Krylov"] + + [deps.ClimaCore.extensions] + KrylovExt = "Krylov" [[deps.ClimaTimeSteppers]] deps = ["CUDA", "ClimaComms", "Colors", "DataStructures", "DiffEqBase", "DiffEqCallbacks", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"] @@ -856,6 +860,12 @@ git-tree-sha1 = "17e462054b42dcdda73e9a9ba0c67754170c88ae" uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" version = "0.9.4" +[[deps.KrylovKit]] +deps = ["ChainRulesCore", "GPUArraysCore", "LinearAlgebra", "Printf"] +git-tree-sha1 = "1a5e1d9941c783b0119897d29f2eb665d876ecf3" +uuid = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" +version = "0.6.0" + [[deps.LDLFactorizations]] deps = ["AMD", "LinearAlgebra", "SparseArrays", "Test"] git-tree-sha1 = "70f582b446a1c3ad82cf87e62b878668beef9d13" diff --git a/examples/Manifest.toml b/examples/Manifest.toml index 099bb587e4..89763e4617 100644 --- a/examples/Manifest.toml +++ b/examples/Manifest.toml @@ -314,10 +314,14 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.6" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] -git-tree-sha1 = "5c3809c5da88a3b01b8fe18558f50470fd3c33ed" +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] +git-tree-sha1 = "3d6cd2e89f0d99664b23549da9081ade4e7d8fdb" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" -version = "0.11.0" +version = "0.11.1" +weakdeps = ["Krylov"] + + [deps.ClimaCore.extensions] + KrylovExt = "Krylov" [[deps.ClimaCoreMakie]] deps = ["ClimaCore", "Makie"] @@ -1236,6 +1240,12 @@ git-tree-sha1 = "17e462054b42dcdda73e9a9ba0c67754170c88ae" uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" version = "0.9.4" +[[deps.KrylovKit]] +deps = ["ChainRulesCore", "GPUArraysCore", "LinearAlgebra", "Printf"] +git-tree-sha1 = "1a5e1d9941c783b0119897d29f2eb665d876ecf3" +uuid = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" +version = "0.6.0" + [[deps.LAME_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "f6250b16881adf048549549fba48b1161acdac8c" @@ -1421,9 +1431,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [[deps.LinearAlgebraX]] deps = ["LinearAlgebra", "Mods", "Permutations", "Primes", "SimplePolynomials"] -git-tree-sha1 = "3d93e38b11993bbdbfac967a46dcb36dd9dd733f" +git-tree-sha1 = "8889b8aa6821c0ee73828a2139314b4d596e7dbc" uuid = "9b3f67b0-2d00-526e-9884-9e4938f8fb88" -version = "0.2.1" +version = "0.2.2" [[deps.LinearOperators]] deps = ["FastClosures", "LDLFactorizations", "LinearAlgebra", "Printf", "SparseArrays", "TimerOutputs"] @@ -1757,9 +1767,9 @@ version = "0.11.30" [[deps.PNGFiles]] deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"] -git-tree-sha1 = "5ded86ccaf0647349231ed6c0822c10886d4a1ee" +git-tree-sha1 = "eed372b0fa15624273a9cdb188b1b88476e6a233" uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883" -version = "0.4.1" +version = "0.4.2" [[deps.Packing]] deps = ["GeometryBasics"] @@ -2228,9 +2238,9 @@ version = "0.3.1" [[deps.SimplePolynomials]] deps = ["Mods", "Multisets", "Polynomials", "Primes"] -git-tree-sha1 = "d537c31cf9995236166e3e9afc424a5a1c59ff9d" +git-tree-sha1 = "e14e1a7063179a90c2981faf3c8cbf12aacccdaf" uuid = "cc47b68c-3164-5771-a705-2bc0097375a0" -version = "0.2.14" +version = "0.2.16" [[deps.SimpleRandom]] deps = ["Distributions", "LinearAlgebra", "Random"] @@ -2531,9 +2541,9 @@ version = "0.4.1" [[deps.Unitful]] deps = ["Dates", "LinearAlgebra", "Random"] -git-tree-sha1 = "242982d62ff0d1671e9029b52743062739255c7e" +git-tree-sha1 = "3c793be6df9dd77a0cf49d80984ef9ff996948fa" uuid = "1986cc42-f94f-5a68-af5c-568840ba703d" -version = "1.18.0" +version = "1.19.0" weakdeps = ["ConstructionBase", "InverseFunctions"] [deps.Unitful.extensions] diff --git a/perf/Manifest.toml b/perf/Manifest.toml index cebe7bad1d..a2a51c0e5d 100644 --- a/perf/Manifest.toml +++ b/perf/Manifest.toml @@ -325,10 +325,14 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.6" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] -git-tree-sha1 = "5c3809c5da88a3b01b8fe18558f50470fd3c33ed" +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "WeakValueDicts"] +git-tree-sha1 = "3d6cd2e89f0d99664b23549da9081ade4e7d8fdb" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" -version = "0.11.0" +version = "0.11.1" +weakdeps = ["Krylov"] + + [deps.ClimaCore.extensions] + KrylovExt = "Krylov" [[deps.ClimaCoreMakie]] deps = ["ClimaCore", "Makie"] @@ -1288,6 +1292,12 @@ git-tree-sha1 = "17e462054b42dcdda73e9a9ba0c67754170c88ae" uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" version = "0.9.4" +[[deps.KrylovKit]] +deps = ["ChainRulesCore", "GPUArraysCore", "LinearAlgebra", "Printf"] +git-tree-sha1 = "1a5e1d9941c783b0119897d29f2eb665d876ecf3" +uuid = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" +version = "0.6.0" + [[deps.LAME_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "f6250b16881adf048549549fba48b1161acdac8c" @@ -1479,9 +1489,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [[deps.LinearAlgebraX]] deps = ["LinearAlgebra", "Mods", "Permutations", "Primes", "SimplePolynomials"] -git-tree-sha1 = "3d93e38b11993bbdbfac967a46dcb36dd9dd733f" +git-tree-sha1 = "8889b8aa6821c0ee73828a2139314b4d596e7dbc" uuid = "9b3f67b0-2d00-526e-9884-9e4938f8fb88" -version = "0.2.1" +version = "0.2.2" [[deps.LinearOperators]] deps = ["FastClosures", "LDLFactorizations", "LinearAlgebra", "Printf", "SparseArrays", "TimerOutputs"] @@ -1809,9 +1819,9 @@ version = "0.11.30" [[deps.PNGFiles]] deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"] -git-tree-sha1 = "5ded86ccaf0647349231ed6c0822c10886d4a1ee" +git-tree-sha1 = "eed372b0fa15624273a9cdb188b1b88476e6a233" uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883" -version = "0.4.1" +version = "0.4.2" [[deps.PProf]] deps = ["AbstractTrees", "EnumX", "FlameGraphs", "Libdl", "OrderedCollections", "Profile", "ProgressMeter", "ProtoBuf", "pprof_jll"] @@ -2304,9 +2314,9 @@ version = "0.3.1" [[deps.SimplePolynomials]] deps = ["Mods", "Multisets", "Polynomials", "Primes"] -git-tree-sha1 = "d537c31cf9995236166e3e9afc424a5a1c59ff9d" +git-tree-sha1 = "e14e1a7063179a90c2981faf3c8cbf12aacccdaf" uuid = "cc47b68c-3164-5771-a705-2bc0097375a0" -version = "0.2.14" +version = "0.2.16" [[deps.SimpleRandom]] deps = ["Distributions", "LinearAlgebra", "Random"] @@ -2625,9 +2635,9 @@ version = "0.4.1" [[deps.Unitful]] deps = ["Dates", "LinearAlgebra", "Random"] -git-tree-sha1 = "242982d62ff0d1671e9029b52743062739255c7e" +git-tree-sha1 = "3c793be6df9dd77a0cf49d80984ef9ff996948fa" uuid = "1986cc42-f94f-5a68-af5c-568840ba703d" -version = "1.18.0" +version = "1.19.0" weakdeps = ["ConstructionBase", "InverseFunctions"] [deps.Unitful.extensions] diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index e8e3eb1332..939ed7cb06 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -2,25 +2,23 @@ import LinearAlgebra: I, Adjoint, ldiv! import ClimaCore.MatrixFields: @name using ClimaCore.MatrixFields +abstract type DiffusionDerivativeFlag end +struct UseDiffusionDerivative <: DiffusionDerivativeFlag end +struct IgnoreDiffusionDerivative <: DiffusionDerivativeFlag end +use_derivative(flag::DiffusionDerivativeFlag) = flag == UseDiffusionDerivative() + abstract type EnthalpyDerivativeFlag end -struct IgnoreEnthalpyDerivative <: EnthalpyDerivativeFlag end struct UseEnthalpyDerivative <: EnthalpyDerivativeFlag end -use_enthalpy_derivative(::IgnoreEnthalpyDerivative, _) = false -use_enthalpy_derivative(::UseEnthalpyDerivative, name) = name == @name(c.ρe_tot) +struct IgnoreEnthalpyDerivative <: EnthalpyDerivativeFlag end +use_derivative(flag::EnthalpyDerivativeFlag) = flag == UseEnthalpyDerivative() """ - ImplicitEquationJacobian(Y; [enthalpy_flag], [transform]) + ImplicitEquationJacobian(Y, atmos, diffusion_flag, enthalpy_flag, transform_flag) A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the -implicit step with the state ``Y`` (see below for more details on this). The -`enthalpy_flag`, which can be set to either `UseEnthalpyDerivative()` or -`IgnoreEnthalpyDerivative()`, specifies whether the derivative of total enthalpy -with respect to vertical velocity should be computed or approximated as 0. The -`transform` flag, which can be set to `true` or `false`, specifies whether or -not the error should be transformed from ``E(Y)`` to ``E(Y)/Δt``, which is -required for non-Rosenbrock timestepping schemes from OrdinaryDiffEq.jl. The -default values for these flags are `IgnoreEnthalpyDerivative()` and `false`, -respectively. +implicit step with the state ``Y``. + +# Background When we use an implicit or split implicit-explicit (IMEX) timestepping scheme, we end up with a nonlinear equation of the form ``E(Y) = 0``, where @@ -72,10 +70,27 @@ steps is repeated, i.e., ``ΔY[1]`` is computed and used to update ``Y`` to when the error ``E(Y)`` is sufficiently close to 0 (according to the convergence condition passed to Newton's method), or when the maximum number of iterations is reached. + +# Arguments + +- `Y::FieldVector`: the state of the simulation +- `atmos::AtmosModel`: the model configuration +- `diffusion_flag::DiffusionDerivativeFlag`: whether the derivative of the + diffusion tendency with respect to the quantities being diffused should be + computed or approximated as 0; must be either `UseDiffusionDerivative()` or + `IgnoreDiffusionDerivative()` instead of a `Bool` to ensure type-stability +- `enthalpy_flag::EnthalpyDerivativeFlag`: whether the derivative of total + enthalpy with respect to vertical velocity should be computed or approximated + as 0; must be either `UseEnthalpyDerivative()` or `IgnoreEnthalpyDerivative()` + instead of a `Bool` to ensure type-stability +- `transform_flag::Bool`: whether the error should be transformed from ``E(Y)`` + to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from + OrdinaryDiffEq.jl """ struct ImplicitEquationJacobian{ M <: MatrixFields.FieldMatrix, S <: MatrixFields.FieldMatrixSolver, + D <: DiffusionDerivativeFlag, E <: EnthalpyDerivativeFlag, T <: Fields.FieldVector, R <: Base.RefValue, @@ -86,7 +101,8 @@ struct ImplicitEquationJacobian{ # solves the linear equation E'(Y) * ΔY = E(Y) for ΔY solver::S - # whether to compute ∂(ᶜh_tot)/∂(ᶠu₃) or to approximate it as 0 + # flags that determine how E'(Y) is approximated + diffusion_flag::D enthalpy_flag::E # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs @@ -94,25 +110,27 @@ struct ImplicitEquationJacobian{ temp_x::T # required by OrdinaryDiffEq.jl to run non-Rosenbrock timestepping schemes - transform::Bool + transform_flag::Bool dtγ_ref::R end function ImplicitEquationJacobian( - Y; - enthalpy_flag = IgnoreEnthalpyDerivative(), - transform = false, + Y, + atmos, + diffusion_flag, + enthalpy_flag, + transform_flag, ) FT = Spaces.undertype(axes(Y.c)) one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' - Bidiagonal_C3 = BidiagonalMatrixRow{C3{FT}} - Bidiagonal_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - Quaddiagonal_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - Tridiagonal_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)} + TridiagonalRow = TridiagonalMatrixRow{FT} + BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} + BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} + QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}} + TridiagonalRow_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)} is_in_Y(name) = MatrixFields.has_field(Y, name) - energy_name = @name(c.ρe_tot) tracer_names = ( @name(c.ρq_tot), @name(c.ρq_liq), @@ -135,33 +153,51 @@ function ImplicitEquationJacobian( # which means that multiplying inv(-1) by a Float32 will yield a Float64. matrix = MatrixFields.FieldMatrix( (@name(c.ρ), @name(c.ρ)) => FT(-1) * I, - (@name(c.uₕ), @name(c.uₕ)) => FT(-1) * I, + MatrixFields.unrolled_map( + name -> + (name, name) => + use_derivative(diffusion_flag) ? + similar(Y.c, TridiagonalRow) : FT(-1) * I, + (@name(c.ρe_tot), available_tracer_names...), + )..., + (@name(c.uₕ), @name(c.uₕ)) => + use_derivative(diffusion_flag) && + diffuse_momentum(atmos.vert_diff) ? + similar(Y.c, TridiagonalRow) : FT(-1) * I, MatrixFields.unrolled_map( name -> (name, name) => FT(-1) * I, - (energy_name, available_tracer_names..., other_available_names...), + other_available_names, )..., - (@name(c.ρ), @name(f.u₃)) => similar(Y.c, Bidiagonal_ACT3), - (energy_name, @name(f.u₃)) => - use_enthalpy_derivative(enthalpy_flag, energy_name) ? - similar(Y.c, Quaddiagonal_ACT3) : similar(Y.c, Bidiagonal_ACT3), + (@name(c.ρ), @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + (@name(c.ρe_tot), @name(f.u₃)) => + use_derivative(enthalpy_flag) ? + similar(Y.c, QuaddiagonalRow_ACT3) : + similar(Y.c, BidiagonalRow_ACT3), MatrixFields.unrolled_map( - name -> (name, @name(f.u₃)) => similar(Y.c, Bidiagonal_ACT3), + name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), available_tracer_names, )..., - (@name(f.u₃), @name(c.ρ)) => similar(Y.f, Bidiagonal_C3), - (@name(f.u₃), energy_name) => similar(Y.f, Bidiagonal_C3), - (@name(f.u₃), @name(f.u₃)) => similar(Y.f, Tridiagonal_C3xACT3), + (@name(f.u₃), @name(c.ρ)) => similar(Y.f, BidiagonalRow_C3), + (@name(f.u₃), @name(c.ρe_tot)) => similar(Y.f, BidiagonalRow_C3), + (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), ) - alg = MatrixFields.SchurComplementSolve(@name(f)) + alg = + use_derivative(diffusion_flag) ? + MatrixFields.ApproximateBlockArrowheadIterativeSolve(@name(c)) : + MatrixFields.BlockArrowheadSolve(@name(c)) + + # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration + # and approximates the A[c, c] block with a MainDiagonalPreconditioner. return ImplicitEquationJacobian( matrix, MatrixFields.FieldMatrixSolver(alg, matrix, Y), + diffusion_flag, enthalpy_flag, similar(Y), similar(Y), - transform, + transform_flag, Ref{FT}(), ) end @@ -178,7 +214,7 @@ function ldiv!( b::Fields.FieldVector, ) MatrixFields.field_matrix_solve!(A.solver, x, A.matrix, b) - if A.transform + if A.transform_flag @. x *= A.dtγ_ref[] end end @@ -217,6 +253,7 @@ function Wfact!(A, Y, p, dtγ, t) p.core.ᶜρ_ref, p.core.ᶜp_ref, p.scratch.ᶜtemp_scalar, + p.scratch.ᶠtemp_scalar, p.params, p.atmos, p.precomputed.ᶜh_tot, @@ -238,7 +275,7 @@ function Wfact!(A, Y, p, dtγ, t) end function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) - (; matrix, enthalpy_flag) = A + (; matrix, diffusion_flag, enthalpy_flag) = A (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶠu₃) = p (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p (; energy_upwinding, density_upwinding) = p.atmos.numerics @@ -250,8 +287,6 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) one_CT3xACT3 = CT3(FT(1)) * CT3(FT(1))' R_d = FT(CAP.R_d(params)) cv_d = FT(CAP.cv_d(params)) - κ_d = FT(CAP.kappa_d(params)) - p_ref_theta = FT(CAP.p_ref_theta(params)) T_tri = FT(CAP.T_triple(params)) ᶜρ = Y.c.ρ @@ -280,9 +315,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # ∂(ᶜp)/∂(ᶜK) ≈ DiagonalMatrixRow(-R_d * ᶜρ / cv_d), and # ∂(ᶜp)/∂(ᶜρe_tot) ≈ DiagonalMatrixRow(R_d / cv_d). - # We can express the implicit tendency for tracers (and density) as either + # We can express the implicit advection tendency for scalars as either # ᶜρχₜ = -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³ * ᶠinterp(ᶜχ))) or # ᶜρχₜ = -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind(ᶠu³, ᶜχ))). + # The implicit advection tendency for density is computed with ᶜχ = 1. # This means that either # ∂(ᶜρχₜ)/∂(ᶠu₃) = # -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ( @@ -313,7 +349,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # ∂((ᶜρe_tot + ᶜp) / ᶜρ)/∂(ᶜK) ⋅ ∂(ᶜK)/∂(ᶠu₃) = # ∂(ᶜp)/∂(ᶜK) * DiagonalMatrixRow(1 / ᶜρ) ⋅ ∂(ᶜK)/∂(ᶠu₃). u³_sign(u³) = sign(u³.u³) # sign of the scalar value stored inside u³ - tracer_info = ( + scalar_info = ( (@name(c.ρ), @name(ᶜ1), density_upwinding), (@name(c.ρe_tot), @name(ᶜh_tot), energy_upwinding), (@name(c.ρq_tot), @name(ᶜspecific.q_tot), tracer_upwinding), @@ -324,13 +360,8 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) (@name(c.ρq_rai), @name(ᶜspecific.q_rai), precip_upwinding), (@name(c.ρq_sno), @name(ᶜspecific.q_sno), precip_upwinding), ) - available_tracer_info = - MatrixFields.unrolled_filter(tracer_info) do (ρχ_name, _, _) - MatrixFields.has_field(Y, ρχ_name) - end - MatrixFields.unrolled_foreach( - available_tracer_info, - ) do (ρχ_name, χ_name, upwinding) + MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, χ_name, upwinding) + MatrixFields.has_field(Y, ρχ_name) || return ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] ᶜχ = if χ_name == @name(ᶜ1) @@ -341,7 +372,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) end if upwinding == Val(:none) - if use_enthalpy_derivative(enthalpy_flag, ρχ_name) + if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( @@ -378,7 +409,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), ) # Need to wrap ᶠupwind_matrix in this for the same reason. - if use_enthalpy_derivative(enthalpy_flag, ρχ_name) + if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( @@ -479,4 +510,65 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,) end + + # We can express the implicit diffusion tendency for scalars and horizontal + # velocity as + # ᶜρχₜ = ᶜadvdivᵥ(ᶠρK_E * ᶠgradᵥ(ᶜχ)) and + # ᶜuₕₜ = ᶜadvdivᵥ(ᶠρK_E * ᶠgradᵥ(ᶜuₕ)) / ᶜρ. + # The implicit diffusion tendency for density is the sum of the ᶜρχₜ's, but + # we approximate the derivative of this sum with respect to ᶜρ as 0. + # This means that + # ∂(ᶜρχₜ)/∂(ᶜρχ) = + # ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠρK_E) ⋅ ᶠgradᵥ_matrix() ⋅ + # ∂(ᶜχ)/∂(ᶜρχ) and + # ∂(ᶜuₕₜ)/∂(ᶜuₕ) = + # DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜadvdivᵥ_matrix() ⋅ + # DiagonalMatrixRow(ᶠρK_E) ⋅ ᶠgradᵥ_matrix(). + # In general, ∂(ᶜχ)/∂(ᶜρχ) = DiagonalMatrixRow(1 / ᶜρ), except for the case + # ∂(ᶜh_tot)/∂(ᶜρe_tot) = + # ∂((ᶜρe_tot + ᶜp) / ᶜρ)/∂(ᶜρe_tot) = + # (I + ∂(ᶜp)/∂(ᶜρe_tot)) ⋅ DiagonalMatrixRow(1 / ᶜρ) ≈ + # DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ). + if use_derivative(diffusion_flag) + (; C_E) = p.atmos.vert_diff + ᶠp = ᶠρK_E = p.ᶠtemp_scalar + interior_uₕ = Fields.level(Y.c.uₕ, 1) + ᶜΔz_surface = Fields.Δz_field(interior_uₕ) + @. ᶠp[colidx] = ᶠinterp(ᶜp[colidx]) + @. ᶠρK_E[colidx] = + ᶠinterp(Y.c.ρ[colidx]) * eddy_diffusivity_coefficient( + C_E, + norm(interior_uₕ[colidx]), + ᶜΔz_surface[colidx] / 2, + ᶠp[colidx], + ) + + ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅ + ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ[colidx]) - + (I,) + + tracer_names = ( + @name(c.ρq_tot), + @name(c.ρq_liq), + @name(c.ρq_ice), + @name(c.ρq_rai), + @name(c.ρq_sno), + ) + MatrixFields.unrolled_foreach(tracer_names) do ρχ_name + MatrixFields.has_field(Y, ρχ_name) || return + ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] + @. ∂ᶜρχ_err_∂ᶜρχ[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅ + ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,) + end + + if diffuse_momentum(p.atmos.vert_diff) + ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)] + @. ∂ᶜuₕ_err_∂ᶜuₕ[colidx] = + dtγ * DiagonalMatrixRow(1 / ᶜρ[colidx]) ⋅ ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅ ᶠgradᵥ_matrix() - (I,) + end + end end diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 086992c508..d3f1617525 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -9,6 +9,16 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t) Yₜ .= zero(eltype(Yₜ)) Fields.bycolumn(axes(Y.c)) do colidx implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx) + if p.atmos.diff_mode == Implicit() + vertical_diffusion_boundary_layer_tendency!( + Yₜ, + Y, + p, + t, + colidx, + p.atmos.vert_diff, + ) + end # NOTE: All ρa tendencies should be applied before calling this function pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 366071981d..5dc0bd581a 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -23,15 +23,16 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) edmf_coriolis_tendency!(Yₜ, Y, p, t, colidx, p.atmos.edmf_coriolis) large_scale_advection_tendency!(Yₜ, Y, p, t, colidx, p.atmos.ls_adv) - (; vert_diff) = p.atmos - vertical_diffusion_boundary_layer_tendency!( - Yₜ, - Y, - p, - t, - colidx, - vert_diff, - ) + if p.atmos.diff_mode == Explicit() + vertical_diffusion_boundary_layer_tendency!( + Yₜ, + Y, + p, + t, + colidx, + p.atmos.vert_diff, + ) + end radiation_tendency!(Yₜ, Y, p, t, colidx, p.atmos.radiation_mode) edmfx_entr_detr_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model) diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 49c66c8cb1..5c2e19265d 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -37,6 +37,9 @@ function get_atmos(config::AtmosConfig, params) edmfx_nh_pressure = parsed_args["edmfx_nh_pressure"] @assert edmfx_nh_pressure in (false, true) + implicit_diffusion = parsed_args["implicit_diffusion"] + @assert implicit_diffusion in (true, false) + model_config = get_model_config(parsed_args) vert_diff = get_vertical_diffusion_model(diffuse_momentum, parsed_args, params, FT) @@ -68,6 +71,7 @@ function get_atmos(config::AtmosConfig, params) ), hyperdiff = get_hyperdiffusion_model(parsed_args, FT), vert_diff, + diff_mode = implicit_diffusion ? Implicit() : Explicit(), viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT), rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT), sfc_temperature = get_sfc_temperature_form(parsed_args), @@ -371,9 +375,16 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false is_cts_algo(::CTS.DistributedODEAlgorithm) = true -function jac_kwargs(ode_algo, Y) +function jac_kwargs(ode_algo, Y, atmos) if is_implicit(ode_algo) - A = ImplicitEquationJacobian(Y; transform = use_transform(ode_algo)) + A = ImplicitEquationJacobian( + Y, + atmos, + atmos.diff_mode == Implicit() ? UseDiffusionDerivative() : + IgnoreDiffusionDerivative(), + IgnoreEnthalpyDerivative(), + use_transform(ode_algo), + ) if use_transform(ode_algo) return (; jac_prototype = A, Wfact_t = Wfact!) else @@ -635,7 +646,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) func = if parsed_args["split_ode"] implicit_func = SciMLBase.ODEFunction( implicit_tendency!; - jac_kwargs(ode_algo, Y)..., + jac_kwargs(ode_algo, Y, atmos)..., tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), ) if is_cts_algo(ode_algo) diff --git a/src/solver/types.jl b/src/solver/types.jl index 442388434f..5933a18516 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -262,6 +262,10 @@ struct TestDycoreConsistency end Base.broadcastable(x::AbstractPerformanceMode) = tuple(x) +abstract type AbstractTimesteppingMode end +struct Explicit <: AbstractTimesteppingMode end +struct Implicit <: AbstractTimesteppingMode end + Base.@kwdef struct AtmosNumerics{ EN_UP, TR_UP, @@ -337,6 +341,7 @@ Base.@kwdef struct AtmosModel{ OGW, HD, VD, + DM, VS, RS, ST, @@ -363,6 +368,7 @@ Base.@kwdef struct AtmosModel{ orographic_gravity_wave::OGW = nothing hyperdiff::HD = nothing vert_diff::VD = nothing + diff_mode::DM = nothing viscous_sponge::VS = nothing rayleigh_sponge::RS = nothing sfc_temperature::ST = nothing