From d98e7c6b893929164486d5f90efee3cac528076c Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Thu, 31 Aug 2023 15:38:20 -0700 Subject: [PATCH] Update to new implicit solver interface --- docs/src/api.md | 6 + src/ClimaAtmos.jl | 3 +- src/cache/cache.jl | 7 +- .../sgs_flux_tendencies.jl | 1 - .../implicit/implicit_solver.jl | 524 ++++++++++++++++++ .../implicit/schur_complement_W.jl | 383 ------------- src/prognostic_equations/implicit/wfact.jl | 290 ---------- src/solver/type_getters.jl | 17 +- src/utils/abbreviations.jl | 14 +- test/implicit/debugging_tools.jl | 92 +-- test/implicit/ldiv.jl | 49 -- test/implicit/wfact.jl | 42 -- 12 files changed, 593 insertions(+), 835 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 delete mode 100644 test/implicit/ldiv.jl delete mode 100644 test/implicit/wfact.jl diff --git a/docs/src/api.md b/docs/src/api.md index 7f1eb46ccd..469446ed69 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -47,3 +47,9 @@ ClimaAtmos.InitialConditions.LifeCycleTan2018 ClimaAtmos.InitialConditions.Bomex ClimaAtmos.InitialConditions.Soares ``` + +### Implicit Solver + +```@docs +ClimaAtmos.ImplicitEquationJacobian +``` diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 3daa52df33..e75c0b4b62 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -44,9 +44,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 94607c0a2b..c12c1ce77a 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 .Parameters as CAP using ClimaCore: Operators, Fields, Limiters, Geometry, Spaces @@ -94,10 +94,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 b1c4bffc37..e51c9b98e7 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 0000000000..c85d9f335e --- /dev/null +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -0,0 +1,524 @@ +import LinearAlgebra: I, Adjoint, ldiv! +import ClimaCore.MatrixFields: @name +using ClimaCore.MatrixFields + +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) + +""" + ImplicitEquationJacobian(Y; [enthalpy_flag], [transform]) + +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. + +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 +```math + E(Y) = Y_{imp}(Y) - Y = \\hat{Y} + Δt * T_{imp}(Y) - Y. +``` +In this expression, ``Y_{imp}(Y)`` denotes the state at some time ``t + Δt``. +This can be expressed as the sum of ``\\hat{Y}``, the contribution from the +state at time ``t`` (and possibly also at earlier times, depending on the order +of the timestepping scheme), and ``Δt * T_{imp}(Y)``, the contribution from the +implicit tendency ``T_{imp}`` between times ``t`` and ``t + Δt``. The new state +at the end of each implicit step in the timestepping scheme is the value of +``Y`` that solves this equation, i.e., the value of ``Y`` that is consistent +with the state ``Y_{imp}(Y)`` predicted by the implicit step. + +Note: When we use a higher-order timestepping scheme, the full step ``Δt`` is +divided into several sub-steps or "stages", where the duration of stage ``i`` is +``Δt * γ_i`` for some constant ``γ_i`` between 0 and 1. + +In order to solve this equation using Newton's method, we must specify the +derivative ``∂E/∂Y``. Since ``\\hat{Y}`` does not depend on ``Y`` (it is only a +function of the state at or before time ``t``), this derivative is +```math + E'(Y) = Δt * T_{imp}'(Y) - I. +``` +In addition, we must specify how to divide ``E(Y)`` by this derivative, i.e., +how to solve the linear equation +```math + E'(Y) * ΔY = E(Y). +``` + +Note: This equation comes from assuming that there is some ``ΔY`` such that +``E(Y - ΔY) = 0`` and making the first-order approximation +```math + E(Y - ΔY) \\approx E(Y) - E'(Y) * ΔY. +``` + +After initializing ``Y`` to ``Y[0] = \\hat{Y}``, Newton's method executes the +following steps: +- Compute the derivative ``E'(Y[0])``. +- Compute the implicit tendency ``T_{imp}(Y[0])`` and use it to get ``E(Y[0])``. +- Solve the linear equation ``E'(Y[0]) * ΔY[0] = E(Y[0])`` for ``ΔY[0]``. +- Update ``Y`` to ``Y[1] = Y[0] - ΔY[0]``. +If the number of Newton iterations is limited to 1, this new value of ``Y`` is +taken to be the solution of the implicit equation. Otherwise, this sequence of +steps is repeated, i.e., ``ΔY[1]`` is computed and used to update ``Y`` to +``Y[2] = Y[1] - ΔY[1]``, then ``ΔY[2]`` is computed and used to update ``Y`` to +``Y[3] = Y[2] - ΔY[2]``, and so on. The iterative process is terminated either +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. +""" +struct ImplicitEquationJacobian{ + M <: MatrixFields.FieldMatrix, + S <: MatrixFields.FieldMatrixSolver, + E <: EnthalpyDerivativeFlag, + T <: Fields.FieldVector, + R <: Base.RefValue, +} + # stores the matrix E'(Y) = Δt * T_imp'(Y) - I + matrix::M + + # 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 + enthalpy_flag::E + + # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs + temp_b::T + temp_x::T + + # required by OrdinaryDiffEq.jl to run non-Rosenbrock timestepping schemes + transform::Bool + dtγ_ref::R +end + +function ImplicitEquationJacobian( + Y; + enthalpy_flag = IgnoreEnthalpyDerivative(), + transform = false, +) + 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)} + + is_in_Y(name) = MatrixFields.has_field(Y, name) + + energy_name = + MatrixFields.unrolled_findonly(is_in_Y, (@name(c.ρe_tot), @name(c.ρθ))) + tracer_names = ( + @name(c.ρq_tot), + @name(c.ρq_liq), + @name(c.ρq_ice), + @name(c.ρq_rai), + @name(c.ρq_sno), + ) + available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names) + other_names = ( + @name(c.sgs⁰.ρatke), + @name(c.sgsʲs), + @name(f.sgsʲs), + @name(c.turbconv), + @name(f.turbconv), + @name(sfc), + ) + other_available_names = MatrixFields.unrolled_filter(is_in_Y, other_names) + + # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0, + # 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) => FT(-1) * I, + (energy_name, available_tracer_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), + MatrixFields.unrolled_map( + name -> (name, @name(f.u₃)) => similar(Y.c, Bidiagonal_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), + ) + + alg = MatrixFields.SchurComplementSolve(@name(f)) + + return ImplicitEquationJacobian( + matrix, + MatrixFields.FieldMatrixSolver(alg, matrix, Y), + enthalpy_flag, + similar(Y), + similar(Y), + transform, + Ref{FT}(), + ) +end + +# We only use A, but OrdinaryDiffEq.jl and ClimaTimeSteppers.jl require us to +# pass jac_prototype and then call similar(jac_prototype) to obtain A. This is a +# workaround to avoid unnecessary allocations. +Base.similar(A::ImplicitEquationJacobian) = A + +# This method specifies how to solve the equation E'(Y) * ΔY = E(Y) for ΔY. +function ldiv!( + x::Fields.FieldVector, + A::ImplicitEquationJacobian, + b::Fields.FieldVector, +) + MatrixFields.field_matrix_solve!(A.solver, x, A.matrix, b) + if A.transform + @. x *= A.dtγ_ref[] + end +end + +# 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 + +# 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 specifies how to compute E'(Y), which is referred to as "Wfact" in +# OrdinaryDiffEq.jl. +function Wfact!(A, Y, p, dtγ, t) + NVTX.@range "Wfact!" color = colorant"green" begin + # Remove unnecessary values from p to avoid allocations in bycolumn. + (; energy_form, rayleigh_sponge) = p.atmos + p′ = (; + p.ᶜspecific, + p.ᶠu³, + p.ᶜK, + p.ᶜp, + p.∂ᶜK_∂ᶠu₃, + p.ᶜΦ, + p.ᶠgradᵥ_ᶜΦ, + p.ᶜρ_ref, + p.ᶜp_ref, + p.ᶜtemp_scalar, + p.params, + p.energy_upwinding, + p.tracer_upwinding, + p.density_upwinding, + p.atmos, + (energy_form isa TotalEnergy ? (; p.ᶜh_tot) : (;))..., + (rayleigh_sponge isa RayleighSponge ? (; p.ᶠβ_rayleigh_w) : (;))..., + ) + + # Convert dtγ from a Float64 to an FT. + FT = Spaces.undertype(axes(Y.c)) + dtγ′ = FT(dtγ) + + A.dtγ_ref[] = dtγ′ + Fields.bycolumn(axes(Y.c)) do colidx + update_implicit_equation_jacobian!(A, Y, p′, dtγ′, colidx) + end + end +end + +function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) + (; matrix, enthalpy_flag) = 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))' + 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)) + p_ref_theta = FT(CAP.p_ref_theta(params)) + T_tri = FT(CAP.T_triple(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 = 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 / 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_tri - (ᶜ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 tracers (and density) 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₃) + # ). + # For simplicity, we approximate the value of ∂(ᶜρχₜ)/∂(ᶠu₃) for FCT (both + # Boris-Book and Zalesak) using the value for first-order upwinding. + # 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₃). + u³_sign(u³) = sign(u³.u³) # sign of the scalar value stored inside u³ + tracer_info = ( + (@name(c.ρ), @name(ᶜ1), density_upwinding), + (@name(c.ρθ), @name(ᶜspecific.θ), energy_upwinding), + (@name(c.ρe_tot), @name(ᶜh_tot), energy_upwinding), + (@name(c.ρq_tot), @name(ᶜspecific.q_tot), tracer_upwinding), + (@name(c.ρq_liq), @name(ᶜspecific.q_liq), tracer_upwinding), + (@name(c.ρq_ice), @name(ᶜspecific.q_ice), tracer_upwinding), + (@name(c.ρq_rai), @name(ᶜspecific.q_rai), tracer_upwinding), + (@name(c.ρq_sno), @name(ᶜspecific.q_sno), tracer_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) + ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] + + ᶜχ = if χ_name == @name(ᶜ1) + @. p.ᶜtemp_scalar[colidx] = 1 + p.ᶜtemp_scalar + else + MatrixFields.get_field(p, χ_name) + end + + if upwinding == Val(:none) + if use_enthalpy_derivative(enthalpy_flag, ρχ_name) + @. ∂ᶜρχ_err_∂ᶠ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 + @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * + ᶠg³³[colidx] * + (one_CT3xACT3,) * + ᶠinterp(ᶜχ[colidx]), + ) + end + else + is_third_order = upwinding == Val(:third_order) + ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 + ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix + UpwindMatrixRowType = + is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow + + ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(CT3{FT})), + bottom = Operators.SetValue(zero(CT3{FT})), + ) # Need to wrap ᶠupwind in this to give it well-defined boundaries. + ᶠset_upwind_matrix_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), + 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) + @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( + DiagonalMatrixRow( + u³_sign(ᶠu³[colidx]) * + ᶠset_upwind_bcs( + ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]), + ) * + ᶠg³³[colidx] * + (one_ATC3,), + ) + + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³[colidx])) ⋅ + ∂ᶜK_∂ᶠu₃[colidx] * (-R_d / cv_d) + ) + else + @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * + u³_sign(ᶠu³[colidx]) * + ᶠset_upwind_bcs( + ᶠ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₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρθ) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ + # ∂(ᶜp)/∂(ᶜρθ) and + # ∂(ᶠu₃ₜ)/∂(ᶠu₃) = 0I. + # If the pressure is computed using total energy, then + # ∂(ᶠu₃ₜ)/∂(ᶜρe_tot) = + # ∂(ᶠu₃ₜ)/∂(ᶠgradᵥ(ᶜp - ᶜp_ref)) ⋅ ∂(ᶠgradᵥ(ᶜp - ᶜp_ref))/∂(ᶜρe_tot) = + # DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅ + # ∂(ᶜp)/∂(ᶜρe_tot) and + # ∂(ᶠ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₃). + # In addition, we sometimes have the Rayleigh sponge tendency modification + # ᶠu₃ₜ += -p.ᶠβ_rayleigh_w * ᶠu₃, + # which translates to + # ∂(ᶠu₃ₜ)/∂(ᶠu₃) += DiagonalMatrixRow(-p.ᶠβ_rayleigh_w). + # Note: Because ∂(u₃)/∂(u₃) is actually the C3 identity tensor (which we + # denote by one_C3xACT3), rather than the scalar 1, we need to replace I + # with I_u₃ = DiagonalMatrixRow(one_C3xACT3). We cannot use I * one_C3xACT3 + # because UniformScaling only supports Numbers as scale factors. + ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] + ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)] + I_u₃ = DiagonalMatrixRow(one_C3xACT3) + if MatrixFields.has_field(Y, @name(c.ρθ)) + ᶜρθ = Y.c.ρθ + ∂ᶠu₃_err_∂ᶜρθ = matrix[@name(f.u₃), @name(c.ρθ)] + @. ∂ᶠu₃_err_∂ᶜρ[colidx] = + dtγ * DiagonalMatrixRow( + ( + ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - + ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] + ) / abs2(ᶠinterp(ᶜρ[colidx])), + ) ⋅ ᶠinterp_matrix() + @. ∂ᶠu₃_err_∂ᶜρθ[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow( + R_d / (1 - κ_d) * + (ᶜρθ[colidx] * R_d / p_ref_theta)^(κ_d / (1 - κ_d)), + ) + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = + dtγ * + DiagonalMatrixRow(-p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,)) - + (I_u₃,) + else + ∂ᶠu₃_err_∂ᶠu₃[colidx] .= (-I_u₃,) + end + elseif MatrixFields.has_field(Y, @name(c.ρe_tot)) + ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] + @. ∂ᶠu₃_err_∂ᶜρ[colidx] = + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow( + R_d * (T_tri - (ᶜK[colidx] + ᶜΦ[colidx]) / cv_d), + ) + + DiagonalMatrixRow( + ( + ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - + ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] + ) / abs2(ᶠinterp(ᶜρ[colidx])), + ) ⋅ ᶠinterp_matrix() + ) + @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ + ᶠgradᵥ_matrix() * R_d / cv_d + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = + dtγ * ( + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ + ∂ᶜK_∂ᶠu₃[colidx] + DiagonalMatrixRow( + -p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,), + ) + ) - (I_u₃,) + else + @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ + ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,) + end + 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 9ae95ed5bc..0000000000 --- a/src/prognostic_equations/implicit/schur_complement_W.jl +++ /dev/null @@ -1,383 +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 - -NVTX.@annotate 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 - # 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 - -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 e5d9c587c5..0000000000 --- a/src/prognostic_equations/implicit/wfact.jl +++ /dev/null @@ -1,290 +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 - -NVTX.@annotate function Wfact!(W, Y, p, dtγ, t) - fill_with_nans!(p) - Fields.bycolumn(axes(Y.c)) do colidx - Wfact!(W, Y, p, dtγ, t, colidx) - 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)) - p_ref_theta = FT(CAP.p_ref_theta(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 / p_ref_theta)^(κ_d / (1 - κ_d)) - # ) - ᶜρθ = Y.c.ρθ - @. ∂ᶠ𝕄ₜ∂ᶜ𝔼[colidx] = map_get_data( - -1 / ᶠinterp(ᶜρ[colidx]) * ᶠgradᵥ_stencil( - R_d / (1 - κ_d) * - (ᶜρθ[colidx] * R_d / p_ref_theta)^(κ_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 0f82a2b588..19c73b7644 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -382,20 +382,13 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; is_cts_algo(::SciMLBase.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() @@ -712,7 +705,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, 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 9db0375cce..5c9f37ee6f 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 @@ -70,8 +70,10 @@ 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) diff --git a/test/implicit/debugging_tools.jl b/test/implicit/debugging_tools.jl index 2f33b80ece..fadd373fb9 100644 --- a/test/implicit/debugging_tools.jl +++ b/test/implicit/debugging_tools.jl @@ -1,56 +1,58 @@ using ForwardDiff: Dual -using SparseArrays: spdiagm +using ClimaCore: Spaces, Fields, Operators, MatrixFields -using ClimaCore: Spaces, Operators +# TODO: Turn this into a unit test for the implicit solver once the code is more +# modular. + +function autodiff_wfact_block(Y, p, dtγ, t, Yₜ_name, Y_name, colidx) + Y_field = MatrixFields.get_field(Y, Y_name) + bot_level = Operators.left_idx(axes(Y_field)) + top_level = Operators.right_idx(axes(Y_field)) + partials = ntuple(_ -> 0, top_level - bot_level + 1) -get_var(obj, ::Tuple{}) = obj -get_var(obj, tup::Tuple) = get_var(getproperty(obj, tup[1]), Base.tail(tup)) -function exact_column_jacobian_block( - implicit_tendency!, - Y, - p, - t, - i, - j, - h, - Yₜ_name, - Y_name, -) - T = eltype(Y) - Y_var = get_var(Y, Y_name) - Y_var_vert_space = Spaces.column(axes(Y_var), i, j, h) - bot_level = Operators.left_idx(Y_var_vert_space) - top_level = Operators.right_idx(Y_var_vert_space) - partials = ntuple(_ -> zero(T), top_level - bot_level + 1) Yᴰ = Dual.(Y, partials...) - Yᴰ_var = get_var(Yᴰ, Y_name) - ith_ε(i) = Dual.(zero(T), Base.setindex(partials, one(T), i)...) + Yᴰ_field = MatrixFields.get_field(Yᴰ, Y_name) + ith_ε(i) = Dual.(0, Base.setindex(partials, 1, i)...) set_level_εs!(level) = - parent(Spaces.level(Yᴰ_var, level)) .+= ith_ε(level - bot_level + 1) + 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) - col = Spaces.column(get_var(Yₜᴰ, Yₜ_name), i, j, h) - return vcat(map(dual -> [dual.partials.values...]', parent(col))...) + implicit_tendency!(Yₜᴰ, Yᴰ, pᴰ, t) + Yₜᴰ_field = MatrixFields.get_field(Yₜᴰ, Yₜ_name) + ∂Yₜ∂Y_array_block = + vcat(map(d -> [d.partials.values...]', parent(Yₜᴰ_field[colidx]))...) + return Yₜ_name == Y_name ? dtγ * ∂Yₜ∂Y_array_block - I : + dtγ * ∂Yₜ∂Y_array_block end -# Note: These only work for scalar stencils. -vector_column(arg, i, j, h) = parent(Spaces.column(arg, i, j, h)) -function matrix_column(stencil, stencil_input_space, i, j, h) - lbw, ubw = Operators.bandwidths(eltype(stencil)) - coefs_column = Spaces.column(stencil, i, j, h).coefs - row_space = axes(coefs_column) - lrow = Operators.left_idx(row_space) - rrow = Operators.right_idx(row_space) - num_rows = rrow - lrow + 1 - col_space = Spaces.column(stencil_input_space, i, j, h) - lcol = Operators.left_idx(col_space) - rcol = Operators.right_idx(col_space) - num_cols = rcol - lcol + 1 - diag_key_value(diag) = - (diag + lrow - lcol) => view( - parent(getproperty(coefs_column, diag - lbw + 1)), - (max(lrow, lcol - diag):min(rrow, rcol - diag)) .- (lrow - 1), +function verify_wfact(A, Y, p, dtγ, t, colidx) + for (Yₜ_name, Y_name) in keys(A.matrix) + computed_block = map( + x -> x[1], + MatrixFields.column_field2array(A.matrix[Yₜ_name, Y_name][colidx]), ) - return spdiagm(num_rows, num_cols, map(diag_key_value, lbw:ubw)...) + Yₜ_name == Y_name && computed_block == -I && continue + reference_block = + autodiff_wfact_block(Y, p, dtγ, t, Yₜ_name, Y_name, colidx) + max_error = maximum(abs.(computed_block - reference_block)) + @info t, Yₜ_name, Y_name, maximum(reference_block), max_error + # display(computed_block) + # display(reference_block) + end end diff --git a/test/implicit/ldiv.jl b/test/implicit/ldiv.jl deleted file mode 100644 index f376bc7cdb..0000000000 --- a/test/implicit/ldiv.jl +++ /dev/null @@ -1,49 +0,0 @@ - -function verify_matrix(x, A, b, update_matrix = false; kwargs...) - (; dtγ_ref, S, S_column_arrays) = A - (; ∂ᶜρₜ∂ᶠ𝕄, ∂ᶜ𝔼ₜ∂ᶠ𝕄, ∂ᶠ𝕄ₜ∂ᶜ𝔼, ∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶠ𝕄ₜ∂ᶠ𝕄, ∂ᶜ𝕋ₜ∂ᶠ𝕄_field) = A - dtγ = dtγ_ref[] - dtγ² = dtγ^2 - FT = eltype(eltype(S)) - @assert A.test && Operators.bandwidths(eltype(∂ᶜ𝔼ₜ∂ᶠ𝕄)) == (-half, half) - Ni, Nj, _, Nv, Nh = size(Fields.field_values(x.c)) - Nᶜf = DataLayouts.typesize(FT, eltype(x.c)) - J_col = zeros(FT, Nv * Nᶜf + Nv + 1, Nv * Nᶜf + Nv + 1) - for h in 1:Nh, j in 1:Nj, i in 1:Ni - x_col = Fields.FieldVector(; - c = Spaces.column(x.c, i, j, h), - f = Spaces.column(x.f, i, j, h), - ) - b_col = Fields.FieldVector(; - c = Spaces.column(b.c, i, j, h), - f = Spaces.column(b.f, i, j, h), - ) - ᶜρ_position = findfirst(isequal(:ρ), propertynames(x.c)) - ᶜρ_offset = DataLayouts.fieldtypeoffset(FT, eltype(x.c), ᶜρ_position) - ᶜρ_indices = (Nv * ᶜρ_offset + 1):(Nv * (ᶜρ_offset + 1)) - ᶜ𝔼_position = findfirst(is_energy_var, propertynames(x.c)) - ᶜ𝔼_offset = DataLayouts.fieldtypeoffset(FT, eltype(x.c), ᶜ𝔼_position) - ᶜ𝔼_indices = (Nv * ᶜ𝔼_offset + 1):(Nv * (ᶜ𝔼_offset + 1)) - ᶠ𝕄_indices = (Nv * Nᶜf + 1):(Nv * (Nᶜf + 1) + 1) - J_col[ᶜρ_indices, ᶠ𝕄_indices] .= - matrix_column(∂ᶜρₜ∂ᶠ𝕄, axes(x.f), i, j, h) - J_col[ᶜ𝔼_indices, ᶠ𝕄_indices] .= - matrix_column(∂ᶜ𝔼ₜ∂ᶠ𝕄, axes(x.f), i, j, h) - J_col[ᶠ𝕄_indices, ᶜρ_indices] .= - matrix_column(∂ᶠ𝕄ₜ∂ᶜρ, axes(x.c), i, j, h) - J_col[ᶠ𝕄_indices, ᶜ𝔼_indices] .= - matrix_column(∂ᶠ𝕄ₜ∂ᶜ𝔼, axes(x.c), i, j, h) - J_col[ᶠ𝕄_indices, ᶠ𝕄_indices] .= - matrix_column(∂ᶠ𝕄ₜ∂ᶠ𝕄, axes(x.f), i, j, h) - for ᶜ𝕋_position in findall(CA.is_tracer_var, propertynames(x.c)) - ᶜ𝕋_offset = - DataLayouts.fieldtypeoffset(FT, eltype(x.c), ᶜ𝕋_position) - ᶜ𝕋_indices = (Nv * ᶜ𝕋_offset + 1):(Nv * (ᶜ𝕋_offset + 1)) - ᶜ𝕋_name = propertynames(x.c)[ᶜ𝕋_position] - ∂ᶜ𝕋ₜ∂ᶠ𝕄 = getproperty(∂ᶜ𝕋ₜ∂ᶠ𝕄_field, ᶜ𝕋_name) - J_col[ᶜ𝕋_indices, ᶠ𝕄_indices] .= - matrix_column(∂ᶜ𝕋ₜ∂ᶠ𝕄, axes(x.f), i, j, h) - end - @assert (-LinearAlgebra.I + dtγ * J_col) * x_col ≈ b_col - end -end diff --git a/test/implicit/wfact.jl b/test/implicit/wfact.jl deleted file mode 100644 index bb6660473c..0000000000 --- a/test/implicit/wfact.jl +++ /dev/null @@ -1,42 +0,0 @@ - -function verify_wfact_matrix(W, Y, p, dtγ, t) - (; ∂ᶜρₜ∂ᶠ𝕄, ∂ᶜ𝔼ₜ∂ᶠ𝕄, ∂ᶠ𝕄ₜ∂ᶜ𝔼, ∂ᶠ𝕄ₜ∂ᶜρ, ∂ᶠ𝕄ₜ∂ᶠ𝕄, ∂ᶜ𝕋ₜ∂ᶠ𝕄_field) = W - (; ᶜts) = p - - if eltype(ᶜts) <: TD.PhaseEquil - error("This function is incompatible with $(typeof(ᶜts))") - end - - # Checking every column takes too long, so just check one. - i, j, h = 1, 1, 1 - args = (implicit_tendency!, Y, p, t, i, j, h) - ᶜ𝔼_name = filter(CA.is_energy_var, propertynames(Y.c))[1] - - @assert matrix_column(∂ᶜρₜ∂ᶠ𝕄, axes(Y.f), i, j, h) ≈ - exact_column_jacobian_block(args..., (:c, :ρ), (:f, :u₃)) - @assert matrix_column(∂ᶠ𝕄ₜ∂ᶜ𝔼, axes(Y.c), i, j, h) ≈ - exact_column_jacobian_block(args..., (:f, :u₃), (:c, ᶜ𝔼_name)) - @assert matrix_column(∂ᶠ𝕄ₜ∂ᶠ𝕄, axes(Y.f), i, j, h) ≈ - exact_column_jacobian_block(args..., (:f, :u₃), (:f, :u₃)) - for ᶜρc_name in filter(CA.is_tracer_var, propertynames(Y.c)) - ∂ᶜρcₜ∂ᶠ𝕄 = getproperty(∂ᶜ𝕋ₜ∂ᶠ𝕄_field, ᶜρc_name) - ᶜρc_tuple = (:c, ᶜρc_name) - @assert matrix_column(∂ᶜρcₜ∂ᶠ𝕄, axes(Y.f), i, j, h) ≈ - exact_column_jacobian_block(args..., ᶜρc_tuple, (:f, :u₃)) - end - - ∂ᶜ𝔼ₜ∂ᶠ𝕄_approx = matrix_column(∂ᶜ𝔼ₜ∂ᶠ𝕄, axes(Y.f), i, j, h) - ∂ᶜ𝔼ₜ∂ᶠ𝕄_exact = - exact_column_jacobian_block(args..., (:c, ᶜ𝔼_name), (:f, :u₃)) - if flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact - @assert ∂ᶜ𝔼ₜ∂ᶠ𝕄_approx ≈ ∂ᶜ𝔼ₜ∂ᶠ𝕄_exact - else - err = norm(∂ᶜ𝔼ₜ∂ᶠ𝕄_approx .- ∂ᶜ𝔼ₜ∂ᶠ𝕄_exact) / norm(∂ᶜ𝔼ₜ∂ᶠ𝕄_exact) - @assert err < 1e-6 - # Note: the highest value seen so far is ~3e-7 (only applies to ρe_tot) - end - - ∂ᶠ𝕄ₜ∂ᶜρ_approx = matrix_column(∂ᶠ𝕄ₜ∂ᶜρ, axes(Y.c), i, j, h) - ∂ᶠ𝕄ₜ∂ᶜρ_exact = exact_column_jacobian_block(args..., (:f, :u₃), (:c, :ρ)) - @assert ∂ᶠ𝕄ₜ∂ᶜρ_approx ≈ ∂ᶠ𝕄ₜ∂ᶜρ_exact -end