From 957c89189c2dd7b85876766eda9e2dd513620f0e Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Thu, 31 Aug 2023 15:38:20 -0700 Subject: [PATCH] Try out new implicit solver interface --- .buildkite/pipeline.yml | 4 + Project.toml | 1 + examples/Manifest.toml | 16 +- src/ClimaAtmos.jl | 3 +- src/cache/cache.jl | 7 +- .../sgs_flux_tendencies.jl | 1 - .../implicit/implicit_solver.jl | 500 ++++++++++++++++++ .../implicit/schur_complement_W.jl | 385 -------------- src/prognostic_equations/implicit/wfact.jl | 292 ---------- src/solver/type_getters.jl | 17 +- src/utils/abbreviations.jl | 14 +- 11 files changed, 530 insertions(+), 710 deletions(-) create mode 100644 src/prognostic_equations/implicit/implicit_solver.jl delete mode 100644 src/prognostic_equations/implicit/schur_complement_W.jl delete mode 100644 src/prognostic_equations/implicit/wfact.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 876148ff6a0..e5b38aeaeed 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,6 +26,7 @@ steps: - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()' - echo "--- Instantiate project" + - "julia --project -e 'using Pkg; Pkg.add(url=\"https://github.com/CliMA/ClimaCore.jl\", rev=\"dy/field_matrix\")'" - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project -e 'using Pkg; Pkg.precompile()'" - "julia --project -e 'using Pkg; Pkg.status()'" @@ -36,16 +37,19 @@ steps: - echo "--- Instantiate test" - "julia --project=test -e 'using Pkg; Pkg.develop(path = \".\")'" + - "julia --project=test -e 'using Pkg; Pkg.add(url=\"https://github.com/CliMA/ClimaCore.jl\", rev=\"dy/field_matrix\")'" - "julia --project=test -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=test -e 'using Pkg; Pkg.precompile()'" - "julia --project=test -e 'using Pkg; Pkg.status()'" - echo "--- Instantiate examples" + - "julia --project=examples -e 'using Pkg; Pkg.add(url=\"https://github.com/CliMA/ClimaCore.jl\", rev=\"dy/field_matrix\")'" - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=examples -e 'using Pkg; Pkg.precompile()'" - "julia --project=examples -e 'using Pkg; Pkg.status()'" - echo "--- Instantiate perf" + - "julia --project=perf -e 'using Pkg; Pkg.add(url=\"https://github.com/CliMA/ClimaCore.jl\", rev=\"dy/field_matrix\")'" - "julia --project=perf -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=perf -e 'using Pkg; Pkg.precompile()'" - "julia --project=perf -e 'using Pkg; Pkg.status()'" diff --git a/Project.toml b/Project.toml index 3d131d68bdf..8a496c38ae6 100644 --- a/Project.toml +++ b/Project.toml @@ -22,6 +22,7 @@ DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5" Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" diff --git a/examples/Manifest.toml b/examples/Manifest.toml index a9b3a39f844..a5c8cbbff01 100644 --- a/examples/Manifest.toml +++ b/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", "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", "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"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.16.0" @@ -299,9 +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 = "f8c0f248eb34fb222455830e35a17f4d1bef9e1c" +git-tree-sha1 = "5317c7d3591f991f438ddb9e63da3395a219dda6" +repo-rev = "dy/field_matrix" +repo-url = "https://github.com/CliMA/ClimaCore.jl.git" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" -version = "0.10.49" +version = "0.10.48" [[deps.ClimaCoreMakie]] deps = ["ClimaCore", "Makie"] @@ -2133,9 +2135,9 @@ version = "0.4.0+0" [[deps.RootSolvers]] deps = ["DocStringExtensions", "ForwardDiff"] -git-tree-sha1 = "bafbed2755bf8b1e42fd21f5bb417ff03e9cf1a1" +git-tree-sha1 = "9fb3462240d2898be5d4acf8925e47f70ec64d07" uuid = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74" -version = "0.4.0" +version = "0.3.5" [[deps.RoundingEmulator]] git-tree-sha1 = "40b9edad2e5287e05bd413a38f61a8ff55b9557b" @@ -2530,9 +2532,9 @@ version = "1.0.1" [[deps.Thermodynamics]] deps = ["DocStringExtensions", "KernelAbstractions", "Random", "RootSolvers"] -git-tree-sha1 = "3f8dfde3a9b02a00509ee1bdabc25045e95101dd" +git-tree-sha1 = "ac3b5236da4d028b9a3984ae7b28730d46892dbf" uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" -version = "0.11.1" +version = "0.11.0" [[deps.ThreadingUtilities]] deps = ["ManualMemory"] diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 8c490a36f49..f715d50f29f 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -43,9 +43,8 @@ include(joinpath("utils", "discrete_hydrostatic_balance.jl")) include(joinpath("prognostic_equations", "pressure_work.jl")) include(joinpath("prognostic_equations", "zero_velocity.jl")) -include(joinpath("prognostic_equations", "implicit", "wfact.jl")) -include(joinpath("prognostic_equations", "implicit", "schur_complement_W.jl")) include(joinpath("prognostic_equations", "implicit", "implicit_tendency.jl")) +include(joinpath("prognostic_equations", "implicit", "implicit_solver.jl")) include(joinpath("prognostic_equations", "remaining_tendency.jl")) include(joinpath("prognostic_equations", "forcing", "large_scale_advection.jl")) # TODO: should this be in tendencies/? diff --git a/src/cache/cache.jl b/src/cache/cache.jl index 387ce26a017..0e4abae3522 100644 --- a/src/cache/cache.jl +++ b/src/cache/cache.jl @@ -1,4 +1,4 @@ -using LinearAlgebra: ×, norm, dot +using LinearAlgebra: ×, norm, dot, Adjoint import ClimaAtmos.Parameters as CAP using ClimaCore: Operators, Fields, Limiters, Geometry, Spaces @@ -96,10 +96,7 @@ function default_cache( ᶜp_ref, ᶜT = similar(Y.c, FT), ᶜf, - ∂ᶜK∂ᶠu₃_data = similar( - Y.c, - Operators.StencilCoefs{-half, half, NTuple{2, FT}}, - ), + ∂ᶜK∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}), params, energy_upwinding, tracer_upwinding, diff --git a/src/dycore_equations_deprecated/sgs_flux_tendencies.jl b/src/dycore_equations_deprecated/sgs_flux_tendencies.jl index 7ce166120cb..dcb99eb0602 100644 --- a/src/dycore_equations_deprecated/sgs_flux_tendencies.jl +++ b/src/dycore_equations_deprecated/sgs_flux_tendencies.jl @@ -1,4 +1,3 @@ -using LinearAlgebra import OrdinaryDiffEq as ODE import Logging import TerminalLoggers diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl new file mode 100644 index 00000000000..f3c54019442 --- /dev/null +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -0,0 +1,500 @@ +import LinearAlgebra: I, Adjoint, ldiv! +using ForwardDiff: Dual +using ClimaCore.MatrixFields + +struct ImplicitEquationJacobian{ + W <: MatrixFields.FieldMatrix, + S <: MatrixFields.FieldMatrixSolver, + FT, + T <: Fields.FieldVector{FT}, +} + W::W + solver::S + + # whether to compute ∂(ᶜh_tot)/∂(ᶠu₃) or to approximate it as a matrix of 0s + enthalpy_correction::Bool + + # required by OrdinaryDiffEq.jl to run non-Rosenbrock integration algorithms + transform::Bool + dtγ_ref::Base.RefValue{FT} + + # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs + temp_b::T + temp_x::T +end + +function ImplicitEquationJacobian( + Y; + enthalpy_correction = false, + transform = false, +) + FT = Spaces.undertype(axes(Y.c)) + one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' + + # FieldMatrix does not support I yet, so we need to define fields that store + # -I. + ᶜminus_I = similar(Y.c, DiagonalMatrixRow{FT}) + ᶠminus_I = similar(Y.f, DiagonalMatrixRow{FT}) + ᶜminus_I .= (-I,) + ᶠminus_I .= (-I,) + + 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)} + blocks = ( + (MatrixFields.@var(c.ρ), MatrixFields.@var(c.ρ)) => ᶜminus_I, + (MatrixFields.@var(c.uₕ), MatrixFields.@var(c.uₕ)) => ᶜminus_I, + (MatrixFields.@var(c.ρ), MatrixFields.@var(f.u₃)) => + similar(Y.c, Bidiagonal_ACT3), + (MatrixFields.@var(f.u₃), MatrixFields.@var(c.ρ)) => + similar(Y.f, Bidiagonal_C3), + (MatrixFields.@var(f.u₃), MatrixFields.@var(f.u₃)) => + similar(Y.f, Tridiagonal_C3xACT3), + ) + for energy_var in (MatrixFields.@var(c.ρθ), MatrixFields.@var(c.ρe_tot)) + MatrixFields.has_var(Y, energy_var) || continue + blocks = ( + blocks..., + (energy_var, energy_var) => ᶜminus_I, + (energy_var, MatrixFields.@var(f.u₃)) => + energy_var == MatrixFields.@var(c.ρe_tot) && + enthalpy_correction ? similar(Y.c, Quaddiagonal_ACT3) : + similar(Y.c, Bidiagonal_ACT3), + (MatrixFields.@var(f.u₃), energy_var) => + similar(Y.f, Bidiagonal_C3), + ) + end + for tracer_var in ( + MatrixFields.@var(c.ρq_tot), + MatrixFields.@var(c.ρq_liq), + MatrixFields.@var(c.ρq_ice), + MatrixFields.@var(c.ρq_rai), + MatrixFields.@var(c.ρq_sno), + ) + MatrixFields.has_var(Y, tracer_var) || continue + blocks = ( + blocks..., + (tracer_var, tracer_var) => ᶜminus_I, + (tracer_var, MatrixFields.@var(f.u₃)) => + similar(Y.c, Bidiagonal_ACT3), + ) + end + for center_var in ( + MatrixFields.@var(c.turbconv), + MatrixFields.@var(c.sgs⁰), + MatrixFields.@var(c.sgsʲs) + ) + MatrixFields.has_var(Y, center_var) || continue + blocks = (blocks..., (center_var, center_var) => ᶜminus_I) + end + for face_var in (MatrixFields.@var(f.turbconv), MatrixFields.@var(f.sgsʲs)) + MatrixFields.has_var(Y, face_var) || continue + blocks = (blocks..., (face_var, face_var) => ᶠminus_I) + end + + sfc_var = MatrixFields.@var(sfc) + if MatrixFields.has_var(Y, sfc_var) + sfc_minus_I = similar(Y.sfc, DiagonalMatrixRow{FT}) + sfc_minus_I .= (-I,) + blocks = (blocks..., (sfc_var, sfc_var) => sfc_minus_I) + end + + W = MatrixFields.FieldMatrix(blocks...) + alg = MatrixFields.SchurComplementSolve(; + vars₁ = VariableSet(MatrixFields.@var(f.u₃)), + ) + solver = MatrixFields.linear_solver(alg, W, Y) + return ImplicitEquationJacobian( + W, + solver, + enthalpy_correction, + transform, + Ref{FT}(), + similar(Y), + similar(Y), + ) +end + +# We only use W, but OrdinaryDiffEq.jl and ClimaTimeSteppers.jl require us to +# pass jac_prototype and then call similar(jac_prototype) to obtain W. This is a +# workaround to avoid unnecessary allocations. +Base.similar(A::ImplicitEquationJacobian) = A + +# This function is used by OrdinaryDiffEq.jl instead of ldiv!. +linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve! +_linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b) + +# This method for ldiv! is called by Krylov.jl from inside ClimaTimeSteppers.jl. +# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a +# related issue that requires the same workaround. +function ldiv!( + x::AbstractVector, + A::ImplicitEquationJacobian, + b::AbstractVector, +) + A.temp_b .= b + ldiv!(A.temp_x, A, A.temp_b) + x .= A.temp_x +end + +function ldiv!( + x::Fields.FieldVector, + A::ImplicitEquationJacobian, + b::Fields.FieldVector, +) + MatrixFields.linear_solve!(A.solver, x, A.W, b) + if A.transform + @. x *= A.dtγ_ref[] + end +end + +################################################################################ + +function Wfact!(A, Y, p, _dtγ, t) + NVTX.@range "Wfact!" color = colorant"green" begin + FT = eltype(Y) + dtγ = FT(_dtγ) # Convert t from a Float64 to an FT. + A.dtγ_ref[] = dtγ + fill_with_nans!(p) + set_precomputed_quantities!(Y, p, t) + Fields.bycolumn(axes(Y.c)) do colidx + _Wfact!(A, Y, p, dtγ, colidx) + # if colidx in ( + # Fields.ColumnIndex{1}((1,), 1), + # Fields.ColumnIndex{2}((1, 1), 1), + # ) + # verify_wfact(A, Y, p, dtγ, t, colidx) + # end + end + end +end + +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 + + FT = Spaces.undertype(axes(Y.c)) + one_ATC3 = CT3(FT(1))' + minus_one_C3xACT3 = -C3(FT(1)) * CT3(FT(1))' + 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)) + MSLP = FT(CAP.MSLP(params)) + T_0 = FT(CAP.T_0(params)) + + ᶜρ = Y.c.ρ + ᶜuₕ = Y.c.uₕ + ᶠu₃ = Y.f.u₃ + ᶜJ = Fields.local_geometry_field(Y.c).J + ᶠg³³ = g³³_field(Y.f) + + # We can express the kinetic energy as + # ᶜK = + # adjoint(CT12(ᶜuₕ)) * ᶜuₕ / 2 + + # ᶜinterp(adjoint(CT3(ᶠu₃)) * ᶠu₃) / 2 + + # adjoint(CT3(ᶜuₕ)) * ᶜinterp(ᶠu₃). + # This means that + # ∂(ᶜK)/∂(ᶠu₃) = + # ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + + # DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix(). + @. ∂ᶜK∂ᶠu₃[colidx] = + ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[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). + # 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)) + # ). + # In the second case, we can ignore all O(ᶜq_tot) terms to approximate + # ∂(ᶜp)/∂(ᶜρ) ≈ DiagonalMatrixRow(R_d * (T_0 - (ᶜK + ᶜΦ) / cv_d)), + # ∂(ᶜp)/∂(ᶜK) ≈ DiagonalMatrixRow(-R_d * ᶜρ / cv_d), and + # ∂(ᶜp)/∂(ᶜρe_tot) ≈ DiagonalMatrixRow(R_d / cv_d). + + # We can express the implicit tendency for scalars as either + # ᶜρχₜ = -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³ * ᶠinterp(ᶜχ))) or + # ᶜρχₜ = -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind(ᶠu³, ᶜχ))). + # This means that either + # ∂(ᶜρχₜ)/∂(ᶠu₃) = + # -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ( + # ∂(ᶠu³)/∂(ᶠu₃) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) + + # DiagonalMatrixRow(ᶠu³) ⋅ ᶠinterp_matrix() ⋅ ∂(ᶜχ)/∂(ᶠu₃) + # ) or + # ∂(ᶜρχₜ)/∂(ᶠu₃) = + # -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ( + # ∂(ᶠupwind(ᶠu³, ᶜχ))/∂(ᶠu³) ⋅ ∂(ᶠu³)/∂(ᶠu₃) + + # ᶠupwind_matrix(ᶠu³) ⋅ ∂(ᶜχ)/∂(ᶠu₃) + # ). + # In addition, we have that + # ∂(ᶠu³)/∂(ᶠu₃) = DiagonalMatrixRow(ᶠg³³ * one_CT3xACT3) and + # ∂(ᶠupwind(ᶠu³, ᶜχ))/∂(ᶠu³) = + # DiagonalMatrixRow( + # u³_sign(ᶠu³) * ᶠupwind(CT3(u³_sign(ᶠu³)), ᶜχ) * (one_AC3,) + # ). + # Since one_AC3 * one_CT3xACT3 = one_ACT3, we can simplify the product of + # these derivatives to + # ∂(ᶠupwind(ᶠu³, ᶜχ))/∂(ᶠu³) ⋅ ∂(ᶠu³)/∂(ᶠu₃) = + # DiagonalMatrixRow( + # u³_sign(ᶠu³) * ᶠupwind(CT3(u³_sign(ᶠu³)), ᶜχ) * ᶠg³³ * (one_ATC3,) + # ). + # In general, ∂(ᶜχ)/∂(ᶠu₃) = 0I, except for the case + # ∂(ᶜh_tot)/∂(ᶠu₃) = + # ∂((ᶜρe_tot + ᶜp) / ᶜρ)/∂(ᶜK) ⋅ ∂(ᶜK)/∂(ᶠu₃) = + # ∂(ᶜp)/∂(ᶜK) * DiagonalMatrixRow(1 / ᶜρ) ⋅ ∂(ᶜK)/∂(ᶠu₃). + # For simplicity, we approximate the value of ∂(ᶜρχₜ)/∂(ᶠu₃) for FCT (both + # Boris-Book and Zalesak) using the value for first-order upwinding. + u³_sign(u³) = sign(u³.u³) # sign of the scalar value stored inside u³ + for (ρχ_var, χ_var) in ( + (MatrixFields.@var(c.ρ), nothing), + (MatrixFields.@var(c.ρθ), MatrixFields.@var(θ)), + (MatrixFields.@var(c.ρe_tot), nothing), + (MatrixFields.@var(c.ρq_tot), MatrixFields.@var(q_tot)), + (MatrixFields.@var(c.ρq_liq), MatrixFields.@var(q_liq)), + (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.has_var(Y, ρχ_var) || continue + ∂ᶜρχₜ∂ᶠu₃ = W[ρχ_var, MatrixFields.@var(f.u₃)] + ᶜχ = if ρχ_var == MatrixFields.@var(c.ρ) + ᶜ1 = p.ᶜtemp_scalar + @. ᶜ1[colidx] = 1 + ᶜ1 + elseif ρχ_var == MatrixFields.@var(c.ρe_tot) + p.ᶜh_tot + else + MatrixFields.get_var(ᶜspecific, χ_var) + end + upwinding = if ρχ_var == MatrixFields.@var(c.ρ) + density_upwinding + elseif ρχ_var == MatrixFields.@var(c.ρθ) || + ρχ_var == MatrixFields.@var(c.ρe_tot) + energy_upwinding + else + tracer_upwinding + end + if upwinding == Val(:none) + if ρχ_var == MatrixFields.@var(c.ρe_tot) && enthalpy_correction + @. ∂ᶜρχₜ∂ᶠu₃[colidx] = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( + DiagonalMatrixRow( + ᶠg³³[colidx] * + (one_CT3xACT3,) * + ᶠinterp(ᶜχ[colidx]), + ) + + DiagonalMatrixRow(ᶠu³[colidx]) ⋅ ᶠinterp_matrix() ⋅ + ∂ᶜK∂ᶠu₃[colidx] * (-R_d / cv_d) + ) + ) + else + @. ∂ᶜρχₜ∂ᶠu₃[colidx] = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * + ᶠg³³[colidx] * + (one_CT3xACT3,) * + ᶠinterp(ᶜχ[colidx]), + ) + ) + end + else + ᶠupwind = upwinding == Val(:third_order) ? ᶠupwind3 : ᶠupwind1 + ᶠupwind_matrix = + upwinding == Val(:third_order) ? ᶠupwind3_matrix : + ᶠupwind1_matrix + if ρχ_var == MatrixFields.@var(c.ρe_tot) && enthalpy_correction + @. ∂ᶜρχₜ∂ᶠu₃[colidx] = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( + DiagonalMatrixRow( + u³_sign(ᶠu³[colidx]) * + ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]) * + ᶠg³³[colidx] * + (one_ATC3,), + ) + + ᶠupwind_matrix(ᶠu³[colidx]) ⋅ ∂ᶜK∂ᶠu₃[colidx] * + (-R_d / cv_d) + ) + ) + else + @. ∂ᶜρχₜ∂ᶠu₃[colidx] = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * + u³_sign(ᶠu³[colidx]) * + ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]) * + ᶠg³³[colidx] * + (one_ATC3,), + ) + ) + end + end + end + + # We can express the implicit tendency for vertical velocity as + # ᶠu₃ₜ = + # -(ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / + # ᶠinterp(ᶜρ). + # The derivative of this expression with respect to density is + # ∂(ᶠu₃ₜ)/∂(ᶜρ) = + # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρ) + + # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ - ᶜρ_ref)) ⋅ ∂(ᶠinterp(ᶜρ - ᶜρ_ref))/∂(ᶜρ) + + # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ)) ⋅ ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ ∂(ᶜp)/∂(ᶜρ) + + # DiagonalMatrixRow(-(ᶠgradᵥ_ᶜΦ) / ᶠinterp(ᶜρ)) ⋅ ᶠinterp_matrix() + + # DiagonalMatrixRow( + # (ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / + # ᶠinterp(ᶜρ)^2 + # ) ⋅ ᶠinterp_matrix() = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ ∂(ᶜp)/∂(ᶜρ) + + # DiagonalMatrixRow( + # (ᶠgradᵥ(ᶜp - ᶜp_ref) - ᶠinterp(ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / + # ᶠinterp(ᶜρ)^2 + # ) ⋅ ᶠinterp_matrix(). + # If the pressure is computed using potential temperature, then + # ∂(ᶠu₃ₜ)/∂(ᶠu₃) = 0I and + # ∂(ᶠu₃ₜ)/∂(ᶜρθ) = + # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρθ) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ ∂(ᶜp)/∂(ᶜρθ). + # If the pressure is computed using total energy, then + # ∂(ᶠu₃ₜ)/∂(ᶠu₃) = + # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶠu₃) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ ∂(ᶜp)/∂(ᶠu₃) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ ∂(ᶜp)/∂(ᶜK) ⋅ + # ∂(ᶜK)/∂(ᶠu₃) and + # ∂(ᶠu₃ₜ)/∂(ᶜρe_tot) = + # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρe_tot) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ + # ∂(ᶜp)/∂(ᶜρe_tot). + ∂ᶠu₃ₜ∂ᶜρ = W[MatrixFields.@var(f.u₃), MatrixFields.@var(c.ρ)] + ∂ᶠu₃ₜ∂ᶠu₃ = W[MatrixFields.@var(f.u₃), MatrixFields.@var(f.u₃)] + if MatrixFields.has_var(Y, MatrixFields.@var(c.ρθ)) + ᶜρθ = Y.c.ρθ + ∂ᶠu₃ₜ∂ᶜρθ = W[MatrixFields.@var(f.u₃), MatrixFields.@var(c.ρθ)] + @. ∂ᶠu₃ₜ∂ᶜρ[colidx] = + dtγ * ( + DiagonalMatrixRow( + ( + ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - + ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] + ) / abs2(ᶠinterp(ᶜρ[colidx])), + ) ⋅ ᶠinterp_matrix() + ) + @. ∂ᶠu₃ₜ∂ᶠu₃[colidx] = DiagonalMatrixRow((minus_one_C3xACT3,)) + @. ∂ᶠu₃ₜ∂ᶜρθ[colidx] = + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow( + R_d / (1 - κ_d) * + (ᶜρθ[colidx] * R_d / MSLP)^(κ_d / (1 - κ_d)), + ) + ) + elseif MatrixFields.has_var(Y, MatrixFields.@var(c.ρe_tot)) + ∂ᶠu₃ₜ∂ᶜρe_tot = W[MatrixFields.@var(f.u₃), MatrixFields.@var(c.ρe_tot)] + @. ∂ᶠu₃ₜ∂ᶜρ[colidx] = + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow( + R_d * (T_0 - (ᶜK[colidx] + ᶜΦ[colidx]) / cv_d), + ) + + DiagonalMatrixRow( + ( + ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - + ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] + ) / abs2(ᶠinterp(ᶜρ[colidx])), + ) ⋅ ᶠinterp_matrix() + ) + @. ∂ᶠu₃ₜ∂ᶠu₃[colidx] = + DiagonalMatrixRow((minus_one_C3xACT3,)) + + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ ∂ᶜK∂ᶠu₃[colidx] + ) + @. ∂ᶠu₃ₜ∂ᶜρe_tot[colidx] = + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() * + R_d / cv_d + ) + end + + # We can express the Rayleigh sponge tendency modification as + # ᶠu₃ₜ -= p.ᶠβ_rayleigh_w * ᶠu₃, + # which translates to + # ∂(ᶠu₃ₜ)/∂(ᶠu₃) += DiagonalMatrixRow(p.ᶠβ_rayleigh_w * minus_one_C3xACT3). + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃ₜ∂ᶠu₃[colidx] += + dtγ * + DiagonalMatrixRow(p.ᶠβ_rayleigh_w[colidx] * (minus_one_C3xACT3,)) + end + + # Note: We need to add parentheses to all broadcast expressions of the form + # dtγ .* (matrix_field .⋅ matrix_field) in order to avoid inference failures + # that cause errors during materialization. +end + +################################################################################ + +function autodiff_wfact_block(Y, p, dtγ, t, Yₜ_var, Y_var, colidx) + Y_field = MatrixFields.get_var(Y, Y_var) + bot_level = Operators.left_idx(axes(Y_field)) + top_level = Operators.right_idx(axes(Y_field)) + partials = ntuple(_ -> 0, top_level - bot_level + 1) + + Yᴰ = Dual.(Y, partials...) + Yᴰ_field = MatrixFields.get_var(Yᴰ, Y_var) + ith_ε(i) = Dual.(0, Base.setindex(partials, 1, i)...) + set_level_εs!(level) = + parent(Spaces.level(Yᴰ_field, level)) .+= ith_ε(level - bot_level + 1) + foreach(set_level_εs!, bot_level:top_level) + + (; atmos) = p + convert_to_duals(fields) = + Fields._values(similar(Fields.FieldVector(; fields...), eltype(Yᴰ))) + dry_atmos_name_pairs = map(propertynames(atmos)) do name + name => name == :moisture_model ? DryModel() : atmos.:($name) + end + dry_atmos = AtmosModel(; dry_atmos_name_pairs...) + pᴰ = (; + p..., + convert_to_duals(temporary_quantities(atmos, axes(Y.c), axes(Y.f)))..., + convert_to_duals(precomputed_quantities(Y, dry_atmos))..., + atmos = dry_atmos, + sfc_setup = nothing, + ) + + Yₜᴰ = similar(Yᴰ) + implicit_tendency!(Yₜᴰ, Yᴰ, pᴰ, t) + Yₜᴰ_field = MatrixFields.get_var(Yₜᴰ, Yₜ_var) + ∂Yₜ∂Y_array_block = + vcat(map(d -> [d.partials.values...]', parent(Yₜᴰ_field[colidx]))...) + return Yₜ_var == Y_var ? -I + dtγ * ∂Yₜ∂Y_array_block : + dtγ * ∂Yₜ∂Y_array_block +end + +function verify_wfact(A, Y, p, dtγ, t, colidx) + for (Yₜ_var, Y_var) in keys(A.W.var_tup_map) + computed_block = map( + x -> x[1], + MatrixFields.column_field2array(A.W[Yₜ_var, Y_var][colidx]), + ) + Yₜ_var == Y_var && computed_block == -I && continue + reference_block = + autodiff_wfact_block(Y, p, dtγ, t, Yₜ_var, Y_var, colidx) + max_error = maximum(abs.(computed_block - reference_block)) + @info t, Yₜ_var, Y_var, maximum(reference_block), max_error + # display(computed_block) + # display(reference_block) + end +end diff --git a/src/prognostic_equations/implicit/schur_complement_W.jl b/src/prognostic_equations/implicit/schur_complement_W.jl deleted file mode 100644 index e2b5dcb27e8..00000000000 --- a/src/prognostic_equations/implicit/schur_complement_W.jl +++ /dev/null @@ -1,385 +0,0 @@ -##### -##### Schur Complement for wfact -##### - -import LinearAlgebra - -import ClimaCore.Spaces as Spaces -import ClimaCore.Fields as Fields -import ClimaCore.Operators as Operators -using ClimaCore.Utilities: half - -struct SchurComplementW{ET, F, FT, J1, J2, J3, J4, J5, J6, J7, S, T} - # whether this struct is used to compute Wfact_t or Wfact - transform::Bool - - # flags for computing the Jacobian - flags::F - - # reference to dtγ, which is specified by the ODE solver - dtγ_ref::FT - - # nonzero blocks of the "dycore Jacobian" - ∂ᶜρₜ∂ᶠ𝕄::J1 - ∂ᶜ𝔼ₜ∂ᶠ𝕄::J2 - ∂ᶠ𝕄ₜ∂ᶜ𝔼::J3 - ∂ᶠ𝕄ₜ∂ᶜρ::J3 - ∂ᶠ𝕄ₜ∂ᶠ𝕄::J4 - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field::J5 - - # nonzero blocks of the "TC Jacobian" - ∂ᶜTCₜ∂ᶜTC::J6 - ∂ᶠTCₜ∂ᶠTC::J7 - - # cache for the Schur complement linear solve - S::S - - # whether to test the Jacobian and linear solver - test::Bool - - # cache that is used to evaluate ldiv! - temp1::T - temp2::T -end - -function tracer_variables(::Type{FT}, ᶜ𝕋_names) where {FT} - (; zip(ᶜ𝕋_names, bidiag_ntuple(FT, Val(length(ᶜ𝕋_names))))...) -end - -function bidiag_ntuple(::Type{FT}, ::Val{N}) where {FT, N} - ntuple( - i -> Operators.StencilCoefs{-half, half, NTuple{2, FT}}((FT(0), FT(0))), - Val(N), - ) -end - -function SchurComplementW(Y, transform, flags, test = false) - @assert length(filter(isequal(:ρ), propertynames(Y.c))) == 1 - @assert length(filter(is_energy_var, propertynames(Y.c))) == 1 - @assert length(filter(is_momentum_var, propertynames(Y.c))) == 1 - @assert length(filter(is_momentum_var, propertynames(Y.f))) == 1 - - FT = Spaces.undertype(axes(Y.c)) - dtγ_ref = Ref(zero(FT)) - - bidiag_type = Operators.StencilCoefs{-half, half, NTuple{2, FT}} - tridiag_type = Operators.StencilCoefs{-1, 1, NTuple{3, FT}} - quaddiag_type = Operators.StencilCoefs{-(1 + half), 1 + half, NTuple{4, FT}} - - ∂ᶜ𝔼ₜ∂ᶠ𝕄_type = - flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact && :ρe_tot in propertynames(Y.c) ? - quaddiag_type : bidiag_type - ∂ᶜρₜ∂ᶠ𝕄 = Fields.Field(bidiag_type, axes(Y.c)) - ∂ᶜ𝔼ₜ∂ᶠ𝕄 = Fields.Field(∂ᶜ𝔼ₜ∂ᶠ𝕄_type, axes(Y.c)) - ∂ᶠ𝕄ₜ∂ᶜ𝔼 = Fields.Field(bidiag_type, axes(Y.f)) - ∂ᶠ𝕄ₜ∂ᶜρ = Fields.Field(bidiag_type, axes(Y.f)) - ∂ᶠ𝕄ₜ∂ᶠ𝕄 = Fields.Field(tridiag_type, axes(Y.f)) - ᶜ𝕋_names = filter(is_tracer_var, propertynames(Y.c)) - - # TODO: can we make this work instead? - # cf = Fields.coordinate_field(axes(Y.c)) - # named_tuple_field(z) = tracer_variables(FT, ᶜ𝕋_names) - # ∂ᶜ𝕋ₜ∂ᶠ𝕄_field = named_tuple_field.(cf) - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field = fill(tracer_variables(FT, ᶜ𝕋_names), axes(Y.c)) - - if :turbconv in propertynames(Y.c) - ᶜTC = Y.c.turbconv - ᶠTC = Y.f.turbconv - - ∂ᶜTCₜ∂ᶜTC_type = - DataLayouts.replace_basetype(FT, tridiag_type, eltype(ᶜTC)) - ∂ᶠTCₜ∂ᶠTC_type = - DataLayouts.replace_basetype(FT, tridiag_type, eltype(ᶠTC)) - - ∂ᶜTCₜ∂ᶜTC = similar(ᶜTC, ∂ᶜTCₜ∂ᶜTC_type) - ∂ᶠTCₜ∂ᶠTC = similar(ᶠTC, ∂ᶠTCₜ∂ᶠTC_type) - - for var_prop_chain in Fields.property_chains(ᶜTC) - ∂ᶜvarₜ∂ᶜvar = - Fields.single_field(∂ᶜTCₜ∂ᶜTC, var_prop_chain, identity) - ∂ᶜvarₜ∂ᶜvar .= tuple(tridiag_type((0, 0, 0))) - end - for var_prop_chain in Fields.property_chains(ᶠTC) - ∂ᶠvarₜ∂ᶠvar = - Fields.single_field(∂ᶠTCₜ∂ᶠTC, var_prop_chain, identity) - ∂ᶠvarₜ∂ᶠvar .= tuple(tridiag_type((0, 0, 0))) - end - else - ∂ᶜTCₜ∂ᶜTC = nothing - ∂ᶠTCₜ∂ᶠTC = nothing - end - - S = Fields.Field(tridiag_type, axes(Y.f)) - N = Spaces.nlevels(axes(Y.f)) - - ᶜ𝕋_names = filter(is_tracer_var, propertynames(Y.c)) - ET = if isempty(ᶜ𝕋_names) - Nothing - else - hspace = Spaces.horizontal_space(axes(∂ᶜ𝕋ₜ∂ᶠ𝕄_field)) - device = ClimaComms.device(hspace) - cid = Fields.ColumnIndex((1, 1), 1) - _∂ᶜ𝕋ₜ∂ᶠ𝕄_field = - device isa ClimaComms.CUDADevice ? ∂ᶜ𝕋ₜ∂ᶠ𝕄_field : - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field[cid] - typeof(getproperty(_∂ᶜ𝕋ₜ∂ᶠ𝕄_field, ᶜ𝕋_names[1])) - end - SchurComplementW{ - ET, - typeof(flags), - typeof(dtγ_ref), - typeof(∂ᶜρₜ∂ᶠ𝕄), - typeof(∂ᶜ𝔼ₜ∂ᶠ𝕄), - typeof(∂ᶠ𝕄ₜ∂ᶜρ), - typeof(∂ᶠ𝕄ₜ∂ᶠ𝕄), - typeof(∂ᶜ𝕋ₜ∂ᶠ𝕄_field), - typeof(∂ᶜTCₜ∂ᶜTC), - typeof(∂ᶠTCₜ∂ᶠTC), - typeof(S), - typeof(Y), - }( - transform, - flags, - dtγ_ref, - ∂ᶜρₜ∂ᶠ𝕄, - ∂ᶜ𝔼ₜ∂ᶠ𝕄, - ∂ᶠ𝕄ₜ∂ᶜ𝔼, - ∂ᶠ𝕄ₜ∂ᶜρ, - ∂ᶠ𝕄ₜ∂ᶠ𝕄, - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field, - ∂ᶜTCₜ∂ᶜTC, - ∂ᶠTCₜ∂ᶠTC, - S, - test, - similar(Y), - similar(Y), - ) -end - -∂ᶜ𝕋ₜ∂ᶠ𝕄_field_eltype(A::T) where {T <: SchurComplementW} = - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field_eltype(T) -∂ᶜ𝕋ₜ∂ᶠ𝕄_field_eltype(::Type{T}) where {ET, T <: SchurComplementW{ET}} = ET - -# We only use Wfact, but the implicit/IMEX solvers require us to pass -# jac_prototype, then call similar(jac_prototype) to obtain J and Wfact. Here -# is a temporary workaround to avoid unnecessary allocations. -Base.similar(w::SchurComplementW) = w - -#= -x = [xᶜρ - xᶜ𝔼 - xᶜ𝕄 - ⋮ - xᶜ𝕋[i] - ⋮ - xᶜTC - xᶠ𝕄 - xᶠTC], -b = [bᶜρ - bᶜ𝔼 - bᶜ𝕄 - ⋮ - bᶜ𝕋[i] - ⋮ - bᶜTC - bᶠ𝕄 - bᶠTC], and -A = -I + dtγ J = - [ -I 0 0 ⋯ 0 ⋯ 0 dtγ ∂ᶜρₜ∂ᶠ𝕄 0 - 0 -I 0 ⋯ 0 ⋯ 0 dtγ ∂ᶜ𝔼ₜ∂ᶠ𝕄 0 - 0 0 -I ⋯ 0 ⋯ 0 0 0 - ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮ ⋮ - 0 0 0 ⋯ -I ⋯ 0 dtγ ∂ᶜ𝕋[i]ₜ∂ᶠ𝕄 0 - ⋮ ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮ - 0 0 0 ⋯ 0 ⋯ dtγ ∂ᶜTCₜ∂ᶜTC - I 0 0 - dtγ ∂ᶠ𝕄ₜ∂ᶜρ dtγ ∂ᶠ𝕄ₜ∂ᶜ𝔼 0 ⋯ 0 ⋯ 0 dtγ ∂ᶠ𝕄ₜ∂ᶠ𝕄 - I 0 - 0 0 0 ⋯ 0 ⋯ 0 0 dtγ ∂ᶠTCₜ∂ᶠTC - I]. - -To simplify our notation, let us denote -A = [-I 0 0 ⋯ 0 ⋯ 0 Aρ𝕄 0 - 0 -I 0 ⋯ 0 ⋯ 0 A𝔼𝕄 0 - 0 0 -I ⋯ 0 ⋯ 0 0 0 - ⋮ ⋮ ⋮ ⋱ ⋮ ⋮ ⋮ ⋮ - 0 0 0 ⋯ -I ⋯ 0 A𝕋𝕄[i] 0 - ⋮ ⋮ ⋮ ⋮ ⋱ 0 ⋮ ⋮ - 0 0 0 ⋯ 0 ⋯ AᶜTC - I 0 0 - A𝕄ρ A𝕄𝔼 0 ⋯ 0 ⋯ 0 A𝕄𝕄 - I 0 - 0 0 0 ⋯ 0 ⋯ 0 0 AᶠTC - I] - -If A x = b, then - -xᶜρ + Aρ𝕄 xᶠ𝕄 = bᶜρ ==> xᶜρ = -bᶜρ + Aρ𝕄 xᶠ𝕄 (1) - -xᶜ𝔼 + A𝔼𝕄 xᶠ𝕄 = bᶜ𝔼 ==> xᶜ𝔼 = -bᶜ𝔼 + A𝔼𝕄 xᶠ𝕄 (2) - -xᶜ𝕄 = bᶜ𝕄 ==> xᶜ𝕄 = -bᶜ𝕄 (3) - -xᶜ𝕋[i] + A𝕋𝕄[i] xᶠ𝕄 = bᶜ𝕋[i] ==> xᶜ𝕋[i] = -bᶜ𝕋[i] + A𝕋𝕄[i] xᶠ𝕄 (4) - (AᶜTC - I) xᶜTC = bᶜTC (5) - A𝕄ρ xᶜρ + A𝕄𝔼 xᶜ𝔼 + (A𝕄𝕄 - I) xᶠ𝕄 = bᶠ𝕄 (6) - (AᶠTC - I) xᶠTC = bᶠTC (7) - -Substituting (1) and (2) into (6) gives us - A𝕄ρ (-bᶜρ + Aρ𝕄 xᶠ𝕄) + A𝕄𝔼 (-bᶜ𝔼 + A𝔼𝕄 xᶠ𝕄) + (A𝕄𝕄 - I) xᶠ𝕄 = bᶠ𝕄 ==> - (A𝕄ρ Aρ𝕄 + A𝕄𝔼 A𝔼𝕄 + A𝕄𝕄 - I) xᶠ𝕄 = bᶠ𝕄 + A𝕄ρ bᶜρ + A𝕄𝔼 bᶜ𝔼 ==> - xᶠ𝕄 = (A𝕄ρ Aρ𝕄 + A𝕄𝔼 A𝔼𝕄 + A𝕄𝕄 - I) \ (bᶠ𝕄 + A𝕄ρ bᶜρ + A𝕄𝔼 bᶜ𝔼) - -Given xᶠ𝕄, we can use (1), (2), (3), and (4) to get xᶜρ, xᶜ𝔼, xᶜ𝕄, and xᶜ𝕋[i]. - -Note: The matrix S = A𝕄ρ Aρ𝕄 + A𝕄𝔼 A𝔼𝕄 + A𝕄𝕄 - I is the "Schur complement" of -the large -I block in A. -=# - -# Function required by OrdinaryDiffEq.jl -linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve! -_linsolve!(x, A, b, update_matrix = false; kwargs...) = - LinearAlgebra.ldiv!(x, A, b) - -# Function required by Krylov.jl (x and b can be AbstractVectors) -# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a -# related issue that requires the same workaround. -function LinearAlgebra.ldiv!(x, A::SchurComplementW, b) - A.temp1 .= b - LinearAlgebra.ldiv!(A.temp2, A, A.temp1) - x .= A.temp2 -end - -function LinearAlgebra.ldiv!( - x::Fields.FieldVector, - A::SchurComplementW, - b::Fields.FieldVector, -) - (; dtγ_ref, S, transform) = A - (; ∂ᶜρₜ∂ᶠ𝕄, ∂ᶜ𝔼ₜ∂ᶠ𝕄, ∂ᶠ𝕄ₜ∂ᶜ𝔼, ∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶠ𝕄ₜ∂ᶠ𝕄, ∂ᶜ𝕋ₜ∂ᶠ𝕄_field) = A - (; ∂ᶜTCₜ∂ᶜTC, ∂ᶠTCₜ∂ᶠTC) = A - dtγ = dtγ_ref[] - cond = Operators.bandwidths(eltype(∂ᶜ𝔼ₜ∂ᶠ𝕄)) != (-half, half) - if cond - str = "The linear solver cannot yet be run with the given ∂ᶜ𝔼ₜ/∂ᶠ𝕄 \ - block, since it has more than 2 diagonals. So, ∂ᶜ𝔼ₜ/∂ᶠ𝕄 will \ - be set to 0 for the Schur complement computation. Consider \ - changing the ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode or the energy variable." - @warn str maxlog = 1 - end - NVTX.@range "linsolve" color = colorant"lime" begin - # Initialize x to -b, which correctly sets all the components of x that - # correspond to variables without implicit tendencies. - @. x = -b - # TODO: Figure out why moving this into _ldiv_serial! results in a lot - # of allocations for EDMFX. - - # Compute Schur complement - Fields.bycolumn(axes(x.c)) do colidx - _ldiv_serial!( - A, - x.c[colidx], - x.f[colidx], - b.c[colidx], - b.f[colidx], - dtγ, - transform, - cond, - ∂ᶜρₜ∂ᶠ𝕄[colidx], - ∂ᶜ𝔼ₜ∂ᶠ𝕄[colidx], - ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx], - ∂ᶠ𝕄ₜ∂ᶜρ[colidx], - ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx], - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field[colidx], - isnothing(∂ᶜTCₜ∂ᶜTC) ? nothing : ∂ᶜTCₜ∂ᶜTC[colidx], - isnothing(∂ᶠTCₜ∂ᶠTC) ? nothing : ∂ᶠTCₜ∂ᶠTC[colidx], - S[colidx], - ) - end - end -end - -function _ldiv_serial!( - A::SchurComplementW, - xc, - xf, - bc, - bf, - dtγ, - transform, - cond, - ∂ᶜρₜ∂ᶠ𝕄, - ∂ᶜ𝔼ₜ∂ᶠ𝕄, - ∂ᶠ𝕄ₜ∂ᶜ𝔼, - ∂ᶠ𝕄ₜ∂ᶜρ, - ∂ᶠ𝕄ₜ∂ᶠ𝕄, - ∂ᶜ𝕋ₜ∂ᶠ𝕄_field, - ∂ᶜTCₜ∂ᶜTC, - ∂ᶠTCₜ∂ᶠTC, - S, -) - dtγ² = dtγ^2 - # TODO: Extend LinearAlgebra.I to work with stencil fields. Allow more - # than 2 diagonals per Jacobian block. - FT = eltype(eltype(S)) - tridiag_type = Operators.StencilCoefs{-1, 1, NTuple{3, FT}} - I = tuple(Operators.StencilCoefs{-1, 1}((zero(FT), one(FT), zero(FT)))) - compose = Operators.ComposeStencils() - apply = Operators.ApplyStencil() - if cond - @. S = dtγ² * compose(∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶜρₜ∂ᶠ𝕄) + dtγ * ∂ᶠ𝕄ₜ∂ᶠ𝕄 - I - else - @. S = - dtγ² * compose(∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶜρₜ∂ᶠ𝕄) + - dtγ² * compose(∂ᶠ𝕄ₜ∂ᶜ𝔼, ∂ᶜ𝔼ₜ∂ᶠ𝕄) + - dtγ * ∂ᶠ𝕄ₜ∂ᶠ𝕄 - I - end - - xᶜρ = xc.ρ - bᶜρ = bc.ρ - ᶜ𝔼_name = filter(is_energy_var, propertynames(xc))[1] - xᶜ𝔼 = getproperty(xc, ᶜ𝔼_name)::typeof(xc.ρ) - bᶜ𝔼 = getproperty(bc, ᶜ𝔼_name)::typeof(xc.ρ) - ᶠ𝕄_name = filter(is_momentum_var, propertynames(xf))[1] - xᶠ𝕄 = getproperty(xf, ᶠ𝕄_name).components.data.:1 - bᶠ𝕄 = getproperty(bf, ᶠ𝕄_name).components.data.:1 - - # Compute xᶠ𝕄. - @. xᶠ𝕄 = bᶠ𝕄 + dtγ * (apply(∂ᶠ𝕄ₜ∂ᶜρ, bᶜρ) + apply(∂ᶠ𝕄ₜ∂ᶜ𝔼, bᶜ𝔼)) - Operators.column_thomas_solve!(S, xᶠ𝕄) - - # Compute the remaining components of x that correspond to variables with - # implicit tendencies. - @. xᶜρ = -bᶜρ + dtγ * apply(∂ᶜρₜ∂ᶠ𝕄, xᶠ𝕄) - @. xᶜ𝔼 = -bᶜ𝔼 + dtγ * apply(∂ᶜ𝔼ₜ∂ᶠ𝕄, xᶠ𝕄) - map(filter(is_tracer_var, propertynames(xc))) do ᶜ𝕋_name - Base.@_inline_meta - xᶜ𝕋 = getproperty(xc, ᶜ𝕋_name)::typeof(xc.ρ) - bᶜ𝕋 = getproperty(bc, ᶜ𝕋_name)::typeof(xc.ρ) - ∂ᶜ𝕋ₜ∂ᶠ𝕄 = getproperty(∂ᶜ𝕋ₜ∂ᶠ𝕄_field, ᶜ𝕋_name)::∂ᶜ𝕋ₜ∂ᶠ𝕄_field_eltype(A) - @. xᶜ𝕋 = -bᶜ𝕋 + dtγ * apply(∂ᶜ𝕋ₜ∂ᶠ𝕄, xᶠ𝕄) - end - if :turbconv in propertynames(xc) - xᶜTC = xc.turbconv - xᶠTC = xf.turbconv - bᶜTC = bc.turbconv - bᶠTC = bf.turbconv - for var_prop_chain in Fields.property_chains(xᶜTC) - xᶜvar = Fields.single_field(xᶜTC, var_prop_chain, identity) - bᶜvar = Fields.single_field(bᶜTC, var_prop_chain, identity) - xᶜvar .= bᶜvar - teye = tuple(tridiag_type((0, 1, 0))) - ∂ᶜvarₜ∂ᶜvar = - Fields.single_field(∂ᶜTCₜ∂ᶜTC, var_prop_chain, identity) - @. ∂ᶜvarₜ∂ᶜvar = ∂ᶜvarₜ∂ᶜvar * dtγ - teye - Operators.column_thomas_solve!(∂ᶜvarₜ∂ᶜvar, xᶜvar) - end - for var_prop_chain in Fields.property_chains(xᶠTC) - xᶠvar = Fields.single_field(xᶠTC, var_prop_chain, identity) - bᶠvar = Fields.single_field(bᶠTC, var_prop_chain, identity) - xᶠvar .= bᶠvar - teye = tuple(tridiag_type((0, 1, 0))) - ∂ᶠvarₜ∂ᶠvar = - Fields.single_field(∂ᶠTCₜ∂ᶠTC, var_prop_chain, identity) - @. ∂ᶠvarₜ∂ᶠvar = ∂ᶠvarₜ∂ᶠvar * dtγ - teye - Operators.column_thomas_solve!(∂ᶠvarₜ∂ᶠvar, xᶠvar) - end - end - - # Apply transform (if needed) - if transform - xc .*= dtγ - xf .*= dtγ - end - return nothing -end diff --git a/src/prognostic_equations/implicit/wfact.jl b/src/prognostic_equations/implicit/wfact.jl deleted file mode 100644 index f418d369205..00000000000 --- a/src/prognostic_equations/implicit/wfact.jl +++ /dev/null @@ -1,292 +0,0 @@ -##### -##### Wfact -##### - -using LinearAlgebra: norm_sqr -import ClimaCore.Spaces as Spaces -import ClimaCore.Operators as Operators -import ClimaCore.Fields as Fields - -# From the chain rule, we know that -# ∂(ᶜρχₜ)/∂(ᶠu₃_data) = ∂(ᶜρχₜ)/∂(ᶠu³_data) * ∂(ᶠu³_data)/∂(ᶠu₃_data), -# where ∂(ᶠu³_data)/∂(ᶠu₃_data) = ᶠg³³. -# If ᶜρχₜ = -ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³ * ᶠinterp(ᶜχ)), then -# ∂(ᶜρχₜ)/∂(ᶠu³_data) = -# -ᶜadvdivᵥ_stencil(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³_unit * ᶠinterp(ᶜχ)) - -# ᶜadvdivᵥ_stencil(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³) * ᶠinterp_stencil(1) * -# ∂(ᶜχ)/∂(ᶠu₃_data). -# If ᶜρχₜ = -ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind(ᶠu³, ᶜχ)), then -# ∂(ᶜρχₜ)/∂(ᶠu₃_data) = -# -ᶜadvdivᵥ_stencil(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind(ᶠu³, ᶜχ) / ᶠu³_data) - -# ᶜadvdivᵥ_stencil(ᶠwinterp(ᶜJ, ᶜρ)) * ᶠupwind_stencil(ᶠu³, 1) * -# ∂(ᶜχ)/∂(ᶠu₃_data). -# Since ᶠu³_data can be 0, we need to modify the last derivative by replacing -# ᶠu³ with CT3(ᶠu³_data + eps(ᶠu³_data)), which lets us avoid divisions by 0. -# Since Operator2Stencil has not yet been extended to upwinding operators, -# ᶠupwind_stencil is not available. -# For simplicity, we approximate the value of ∂(ᶜρχₜ)/∂(ᶠu³_data) for FCT -# (both Boris-Book and Zalesak) using the value for first-order upwinding. -# In the following function, we assume that ∂(ᶜχ)/∂(ᶠu₃_data) = 0; if this is -# not the case, the additional term should be added to this function's result. -get_data_plus_ε(vector) = vector.u³ + eps(vector.u³) -set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, ::Val{:none}) = - @. ∂ᶜρχₜ∂ᶠu₃_data = - -(ᶜadvdivᵥ_stencil(ᶠwinterp(ᶜJ, ᶜρ) * one(ᶠu³) * ᶠinterp(ᶜχ) * ᶠg³³)) -set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, ::Val{:first_order}) = - @. ∂ᶜρχₜ∂ᶠu₃_data = -(ᶜadvdivᵥ_stencil( - ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(CT3(get_data_plus_ε(ᶠu³)), ᶜχ) / - get_data_plus_ε(ᶠu³) * ᶠg³³, - )) -set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, ::Val{:third_order}) = - @. ∂ᶜρχₜ∂ᶠu₃_data = -(ᶜadvdivᵥ_stencil( - ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind3(CT3(get_data_plus_ε(ᶠu³)), ᶜχ) / - get_data_plus_ε(ᶠu³) * ᶠg³³, - )) -set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, ::Val{:boris_book}) = - set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, Val(:first_order)) -set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, ::Val{:zalesak}) = - set_∂ᶜρχₜ∂ᶠu₃!(∂ᶜρχₜ∂ᶠu₃_data, ᶜJ, ᶜρ, ᶠu³, ᶜχ, ᶠg³³, Val(:first_order)) - -function validate_flags!(Y, flags, energy_upwinding) - # TODO: Add Operator2Stencil for UpwindBiasedProductC2F to ClimaCore - # to allow exact Jacobian calculation. - :ρe_tot in propertynames(Y.c) && - energy_upwinding !== Val(:none) && - flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact && - error( - "∂ᶜ𝔼ₜ∂ᶠ𝕄_mode must be :no_∂ᶜp∂ᶜK when using ρe_tot with upwinding", - ) -end - -function Wfact!(W, Y, p, dtγ, t) - NVTX.@range "Wfact!" color = colorant"green" begin - fill_with_nans!(p) - set_precomputed_quantities!(Y, p, t) - Fields.bycolumn(axes(Y.c)) do colidx - Wfact!(W, Y, p, dtγ, t, colidx) - end - end -end - -function Wfact!(W, Y, p, dtγ, t, colidx) - (; flags, dtγ_ref) = W - (; ∂ᶜρₜ∂ᶠ𝕄, ∂ᶜ𝔼ₜ∂ᶠ𝕄, ∂ᶠ𝕄ₜ∂ᶜ𝔼, ∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶠ𝕄ₜ∂ᶠ𝕄, ∂ᶜ𝕋ₜ∂ᶠ𝕄_field) = W - (; ᶜspecific, ᶠu³, ᶜK, ᶜp) = p - (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params, ∂ᶜK∂ᶠu₃_data) = p - (; energy_upwinding, tracer_upwinding, density_upwinding) = p - - ᶜρ = Y.c.ρ - ᶜuₕ = Y.c.uₕ - ᶠu₃ = Y.f.u₃ - ᶜJ = Fields.local_geometry_field(Y.c).J - ᶠg³³ = g³³_field(Y.f) - - validate_flags!(Y, flags, energy_upwinding) - FT = Spaces.undertype(axes(Y.c)) - compose = Operators.ComposeStencils() - - R_d = FT(CAP.R_d(params)) - κ_d = FT(CAP.kappa_d(params)) - cv_d = FT(CAP.cv_d(params)) - T_tri = FT(CAP.T_triple(params)) - MSLP = FT(CAP.MSLP(params)) - - dtγ_ref[] = dtγ - - # We can express the pressure as - # ᶜp = R_d * (ᶜρe_tot / cv_d + ᶜρ * (-(ᶜK + ᶜΦ) / cv_d + T_tri)) + O(ᶜq_tot) - # We will ignore all O(ᶜq_tot) terms when computing derivatives of pressure. - - # ᶜK = - # ( - # dot(C123(ᶜuₕ), CT123(ᶜuₕ)) + - # ᶜinterp(dot(C123(ᶠu₃), CT123(ᶠu₃))) + - # 2 * dot(CT123(ᶜuₕ), ᶜinterp(C123(ᶠu₃))) - # ) / 2 = - # ( - # dot(C123(ᶜuₕ), CT123(ᶜuₕ)) + - # ᶜinterp(ᶠu₃_data^2 * dot(C123(ᶠu₃_unit), CT123(ᶠu₃_unit))) + - # 2 * dot(CT123(ᶜuₕ), ᶜinterp(ᶠu₃_data * C123(ᶠu₃_unit))) - # ) / 2 = - # ∂(ᶜK)/∂(ᶠu₃_data) = - # ( - # ᶜinterp_stencil(2 * ᶠu₃_data * dot(C123(ᶠu₃_unit), CT123(ᶠu₃_unit))) + - # 2 * dot(CT123(ᶜuₕ), ᶜinterp_stencil(C123(ᶠu₃_unit))) - # ) / 2 = - # ᶜinterp_stencil(dot(C123(ᶠu₃_unit), CT123(ᶠu₃))) + - # dot(CT123(ᶜuₕ), ᶜinterp_stencil(C123(ᶠu₃_unit))) - @. ∂ᶜK∂ᶠu₃_data[colidx] = - ᶜinterp_stencil(dot(C123(one(ᶠu₃[colidx])), CT123(ᶠu₃[colidx]))) - @. ∂ᶜK∂ᶠu₃_data.coefs.:1[colidx] += dot( - CT123(ᶜuₕ[colidx]), - getindex(ᶜinterp_stencil(C123(one(ᶠu₃[colidx]))), 1), - ) - @. ∂ᶜK∂ᶠu₃_data.coefs.:2[colidx] += dot( - CT123(ᶜuₕ[colidx]), - getindex(ᶜinterp_stencil(C123(one(ᶠu₃[colidx]))), 2), - ) - # TODO: Figure out why rewriting this as shown below incurs allocations: - # @inline map_dot(vector, vectors) = - # map(vector_coef -> dot(vector, vector_coef), vectors) - # @. ∂ᶜK∂ᶠu₃_data[colidx] = - # ᶜinterp_stencil(dot(C123(one(ᶠu₃[colidx])), CT123(ᶠu₃[colidx]))) + - # map_dot(CT123(ᶜuₕ[colidx]), ᶜinterp_stencil(C123(one(ᶠu₃[colidx])))) - - ᶜ1 = p.ᶜtemp_scalar - @. ᶜ1[colidx] = one(ᶜρ[colidx]) - set_∂ᶜρχₜ∂ᶠu₃!( - ∂ᶜρₜ∂ᶠ𝕄[colidx], - ᶜJ[colidx], - ᶜρ[colidx], - ᶠu³[colidx], - ᶜ1[colidx], - ᶠg³³[colidx], - density_upwinding, - ) - - if :ρθ in propertynames(Y.c) - set_∂ᶜρχₜ∂ᶠu₃!( - ∂ᶜ𝔼ₜ∂ᶠ𝕄[colidx], - ᶜJ[colidx], - ᶜρ[colidx], - ᶠu³[colidx], - ᶜspecific.θ[colidx], - ᶠg³³[colidx], - energy_upwinding, - ) - elseif :ρe_tot in propertynames(Y.c) - (; ᶜh_tot) = p - set_∂ᶜρχₜ∂ᶠu₃!( - ∂ᶜ𝔼ₜ∂ᶠ𝕄[colidx], - ᶜJ[colidx], - ᶜρ[colidx], - ᶠu³[colidx], - ᶜh_tot[colidx], - ᶠg³³[colidx], - energy_upwinding, - ) - if flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact - # ∂(ᶜh_tot)/∂(ᶠu₃_data) = - # ∂(ᶜp / ᶜρ)/∂(ᶠu₃_data) = - # ∂(ᶜp / ᶜρ)/∂(ᶜK) * ∂(ᶜK)/∂(ᶠu₃_data) - # If we ignore the dependence of pressure on moisture, - # ∂(ᶜp / ᶜρ)/∂(ᶜK) = -R_d / cv_d - if energy_upwinding === Val(:none) - @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 -= compose( - ᶜadvdivᵥ_stencil( - ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * ᶠu³[colidx], - ), - compose( - ᶠinterp_stencil(ᶜ1[colidx]), - -R_d / cv_d * ∂ᶜK∂ᶠu₃_data[colidx], - ), - ) - end - end - end - for (∂ᶜ𝕋ₜ∂ᶠ𝕄, ᶜχ, _) in matching_subfields(∂ᶜ𝕋ₜ∂ᶠ𝕄_field, ᶜspecific) - set_∂ᶜρχₜ∂ᶠu₃!( - ∂ᶜ𝕋ₜ∂ᶠ𝕄[colidx], - ᶜJ[colidx], - ᶜρ[colidx], - ᶠu³[colidx], - ᶜχ[colidx], - ᶠg³³[colidx], - tracer_upwinding, - ) - end - - # We use map_get_data to convert ∂(ᶠu₃ₜ)/∂(X) to ∂(ᶠu₃_data)ₜ/∂(X). - @inline map_get_data(vectors) = map(vector -> vector.u₃, vectors) - - # ᶠu₃ₜ = -(ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / ᶠinterp(ᶜρ) - if :ρθ in propertynames(Y.c) - # ∂(ᶠu₃ₜ)/∂(ᶜρθ) = - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) * ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρθ) - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) = -1 / ᶠinterp(ᶜρ) - # If we ignore the dependence of pressure on moisture, - # ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρθ) = - # ᶠgradᵥ_stencil( - # R_d / (1 - κ_d) * (ᶜρθ * R_d / MSLP)^(κ_d / (1 - κ_d)) - # ) - ᶜρθ = Y.c.ρθ - @. ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx] = map_get_data( - -1 / ᶠinterp(ᶜρ[colidx]) * ᶠgradᵥ_stencil( - R_d / (1 - κ_d) * (ᶜρθ[colidx] * R_d / MSLP)^(κ_d / (1 - κ_d)), - ), - ) - - # ∂(ᶠu₃ₜ)/∂(ᶜρ) = - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ - ᶜρ_ref)) * ∂(ᶠinterp(ᶜρ - ᶜρ_ref))/∂(ᶜρ) + - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ)) * ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ - ᶜρ_ref)) = -ᶠgradᵥ_ᶜΦ / ᶠinterp(ᶜρ) - # ∂(ᶠinterp(ᶜρ - ᶜρ_ref))/∂(ᶜρ) = ᶠinterp_stencil(1) - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ)) = - # (ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / ᶠinterp(ᶜρ)^2 - # ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = ᶠinterp_stencil(1) - @. ∂ᶠ𝕄ₜ∂ᶜρ[colidx] = map_get_data( - ( - ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - - ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] - ) / abs2(ᶠinterp(ᶜρ[colidx])) * ᶠinterp_stencil(ᶜ1[colidx]), - ) - - # ∂(ᶠu₃ₜ)/∂(ᶠu₃_data) = 0 - ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx] .= - tuple(Operators.StencilCoefs{-1, 1}((FT(0), FT(0), FT(0)))) - elseif :ρe_tot in propertynames(Y.c) - # ∂(ᶠu₃ₜ)/∂(ᶜρe_tot) = - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) * ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρe_tot) - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) = -1 / ᶠinterp(ᶜρ) - # If we ignore the dependence of pressure on moisture, - # ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρe_tot) = ᶠgradᵥ_stencil(R_d / cv_d) - @. ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx] = map_get_data( - -1 / ᶠinterp(ᶜρ[colidx]) * ᶠgradᵥ_stencil(R_d / cv_d * ᶜ1[colidx]), - ) - - # ∂(ᶠu₃ₜ)/∂(ᶜρ) = - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) * ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρ) + - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ - ᶜρ_ref)) * ∂(ᶠinterp(ᶜρ - ᶜρ_ref))/∂(ᶜρ) + - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ)) * ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) = -1 / ᶠinterp(ᶜρ) - # If we ignore the dependence of pressure on moisture, - # ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρ) = - # ᶠgradᵥ_stencil(R_d * (-(ᶜK + ᶜΦ) / cv_d + T_tri)) - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ - ᶜρ_ref)) = -ᶠgradᵥ_ᶜΦ / ᶠinterp(ᶜρ) - # ∂(ᶠinterp(ᶜρ - ᶜρ_ref))/∂(ᶜρ) = ᶠinterp_stencil(1) - # ∂(ᶠu₃ₜ)/∂(ᶠinterp(ᶜρ)) = - # (ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / ᶠinterp(ᶜρ)^2 - # ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = ᶠinterp_stencil(1) - @. ∂ᶠ𝕄ₜ∂ᶜρ[colidx] = map_get_data( - -1 / ᶠinterp(ᶜρ[colidx]) * ᶠgradᵥ_stencil( - R_d * (-(ᶜK[colidx] + ᶜΦ[colidx]) / cv_d + T_tri), - ) + - ( - ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - - ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] - ) / abs2(ᶠinterp(ᶜρ[colidx])) * ᶠinterp_stencil(ᶜ1[colidx]), - ) - - # ∂(ᶠu₃ₜ)/∂(ᶠu₃_data) = - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) * ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶠu₃_data) = - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) * ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜK) * ∂(ᶜK)/∂(ᶠu₃_data) - # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) = -1 / ᶠinterp(ᶜρ) - # If we ignore the dependence of pressure on moisture, - # ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜK) = ᶠgradᵥ_stencil(-ᶜρ * R_d / cv_d) - @. ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx] = map_get_data( - compose( - -1 / ᶠinterp(ᶜρ[colidx]) * - ᶠgradᵥ_stencil(-(ᶜρ[colidx] * R_d / cv_d)), - ∂ᶜK∂ᶠu₃_data[colidx], - ), - ) - end - - if p.atmos.rayleigh_sponge isa RayleighSponge - # ᶠu₃ₜ -= p.ᶠβ_rayleigh_w * ᶠu₃ - # ∂(ᶠu₃ₜ)/∂(ᶠu₃_data) -= p.ᶠβ_rayleigh_w * ᶠu₃_unit - @. ∂ᶠ𝕄ₜ∂ᶠ𝕄[colidx].coefs.:2 -= p.ᶠβ_rayleigh_w[colidx] - end - - return nothing -end diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index bd1cb42bb3a..9aa1c97dfc9 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -385,20 +385,13 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; is_cts_algo(::DiffEqBase.AbstractODEAlgorithm) = false is_cts_algo(::CTS.DistributedODEAlgorithm) = true -jacobi_flags(::TotalEnergy) = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :no_∂ᶜp∂ᶜK) -jacobi_flags(::PotentialTemperature) = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :exact) - -function jac_kwargs(ode_algo, Y, energy_form) +function jac_kwargs(ode_algo, Y) if is_implicit(ode_algo) - W = SchurComplementW( - Y, - use_transform(ode_algo), - jacobi_flags(energy_form), - ) + A = ImplicitEquationJacobian(Y; transform = use_transform(ode_algo)) if use_transform(ode_algo) - return (; jac_prototype = W, Wfact_t = Wfact!) + return (; jac_prototype = A, Wfact_t = Wfact!) else - return (; jac_prototype = W, Wfact = Wfact!) + return (; jac_prototype = A, Wfact = Wfact!) end else return NamedTuple() @@ -624,7 +617,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) func = if parsed_args["split_ode"] implicit_func = ODE.ODEFunction( implicit_tendency!; - jac_kwargs(ode_algo, Y, atmos.energy_form)..., + jac_kwargs(ode_algo, Y)..., tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), ) if is_cts_algo(ode_algo) diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl index 17acc225a52..33ab6fa10e8 100644 --- a/src/utils/abbreviations.jl +++ b/src/utils/abbreviations.jl @@ -1,4 +1,4 @@ -using ClimaCore: Geometry, Operators +using ClimaCore: Geometry, Operators, MatrixFields # Alternatively, we could use Vec₁₂₃, Vec³, etc., if that is more readable. const C1 = Geometry.Covariant1Vector @@ -69,11 +69,13 @@ const ᶠfct_zalesak = Operators.FCTZalesak( top = Operators.FirstOrderOneSided(), ) -const ᶜinterp_stencil = Operators.Operator2Stencil(ᶜinterp) -const ᶜdivᵥ_stencil = Operators.Operator2Stencil(ᶜdivᵥ) -const ᶜadvdivᵥ_stencil = Operators.Operator2Stencil(ᶜadvdivᵥ) -const ᶠinterp_stencil = Operators.Operator2Stencil(ᶠinterp) -const ᶠgradᵥ_stencil = Operators.Operator2Stencil(ᶠgradᵥ) +const ᶜinterp_matrix = MatrixFields.operator_matrix(ᶜinterp) +const ᶜdivᵥ_matrix = MatrixFields.operator_matrix(ᶜdivᵥ) +const ᶜadvdivᵥ_matrix = MatrixFields.operator_matrix(ᶜadvdivᵥ) +const ᶠinterp_matrix = MatrixFields.operator_matrix(ᶠinterp) +const ᶠgradᵥ_matrix = MatrixFields.operator_matrix(ᶠgradᵥ) +const ᶠupwind1_matrix = MatrixFields.operator_matrix(ᶠupwind1) +const ᶠupwind3_matrix = MatrixFields.operator_matrix(ᶠupwind3) # In order to specify C2F operator boundary conditions with 0 instead of FT(0), # we need to tell ClimaCore how to convert AxisTensor components from Int to FT.