diff --git a/src/cache/cache.jl b/src/cache/cache.jl index 354e1b2813a..82cb746d25e 100644 --- a/src/cache/cache.jl +++ b/src/cache/cache.jl @@ -144,6 +144,7 @@ function build_cache(Y, atmos, params, surface_setup, dt, start_date) ᶜp_ref, ᶜT = similar(Y.c, FT), ᶜf, + ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CT12{FT}}}), ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}), # Used by diagnostics such as hfres, evspblw surface_ct3_unit = CT3.( diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 1fe240212f2..7643f2da14c 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -125,12 +125,15 @@ function ImplicitEquationJacobian( approximate_solve_iters, ) FT = Spaces.undertype(axes(Y.c)) - one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' TridiagonalRow = TridiagonalMatrixRow{FT} BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} + TridiagonalRow_ACT12 = TridiagonalMatrixRow{Adjoint{FT, CT12{FT}}} BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - TridiagonalRow_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)} + BidiagonalRow_C3xACT12 = + BidiagonalMatrixRow{typeof(C3(FT(1)) * CT12(FT(1), FT(1))')} + TridiagonalRow_C3xACT3 = + TridiagonalMatrixRow{typeof(C3(FT(1)) * CT3(FT(1))')} is_in_Y(name) = MatrixFields.has_field(Y, name) @@ -172,26 +175,40 @@ function ImplicitEquationJacobian( name -> (name, name) => FT(-1) * I, other_available_names, )..., - (@name(c.ρ), @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + MatrixFields.unrolled_map( + name -> + (name, @name(c.uₕ)) => similar(Y.c, TridiagonalRow_ACT12), + (@name(c.ρ), @name(c.ρe_tot), available_tracer_names...), + )..., + MatrixFields.unrolled_map( + name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + (@name(c.ρ), available_tracer_names...), + )..., (@name(c.ρe_tot), @name(f.u₃)) => use_derivative(enthalpy_flag) ? similar(Y.c, QuaddiagonalRow_ACT3) : similar(Y.c, BidiagonalRow_ACT3), - MatrixFields.unrolled_map( - name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), - available_tracer_names, - )..., (@name(f.u₃), @name(c.ρ)) => similar(Y.f, BidiagonalRow_C3), (@name(f.u₃), @name(c.ρe_tot)) => similar(Y.f, BidiagonalRow_C3), + (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACT12), (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), ) + scalar_names = ( + @name(c.ρ), + @name(c.ρe_tot), + available_tracer_names..., + other_available_names..., + ) # Everything in Y except for ᶜuₕ and ᶠu₃. + + alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) alg = use_derivative(diffusion_flag) ? MatrixFields.ApproximateBlockArrowheadIterativeSolve( - @name(c); + scalar_names...; + alg₂, n_iters = approximate_solve_iters, - ) : MatrixFields.BlockArrowheadSolve(@name(c)) + ) : MatrixFields.BlockArrowheadSolve(scalar_names...; alg₂) # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration # and approximates the A[c, c] block with a MainDiagonalPreconditioner. @@ -266,6 +283,7 @@ function Wfact!(A, Y, p, dtγ, t) p.atmos.turbconv_model isa PrognosticEDMFX ? (; p.precomputed.ᶜρa⁰) : (;) )..., + p.core.∂ᶜK_∂ᶜuₕ, p.core.∂ᶜK_∂ᶠu₃, p.core.ᶜΦ, p.core.ᶠgradᵥ_ᶜΦ, @@ -294,15 +312,14 @@ end function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) (; matrix, diffusion_flag, enthalpy_flag) = A - (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶠu₃) = p + (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃) = p (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p (; energy_upwinding, density_upwinding) = p.atmos.numerics (; tracer_upwinding, precip_upwinding) = p.atmos.numerics FT = Spaces.undertype(axes(Y.c)) - one_ATC3 = CT3(FT(1))' + one_AC3 = C3(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)) T_tri = FT(CAP.T_triple(params)) @@ -311,7 +328,8 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ᶜuₕ = Y.c.uₕ ᶠu₃ = Y.f.u₃ ᶜJ = Fields.local_geometry_field(Y.c).J - ᶠg³³ = g³³_field(Y.f) + ᶜg³ʰ = g³ʰ_tensor_field(Y.c) + ᶠg³³ = g³³_tensor_field(Y.f) # We can express the kinetic energy as # ᶜK = @@ -322,6 +340,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # ∂(ᶜK)/∂(ᶠu₃) = # ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + # DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix(). + @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow( + adjoint(CT12(ᶜuₕ[colidx])) + + adjoint(ᶜinterp(ᶠu₃[colidx])) * ᶜg³ʰ[colidx], + ) @. ∂ᶜK_∂ᶠu₃[colidx] = ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[colidx]))) + DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix() @@ -378,6 +400,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ) MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, χ_name, upwinding) MatrixFields.has_field(Y, ρχ_name) || return + ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] ᶜχ = if χ_name == @name(ᶜ1) @@ -389,24 +412,33 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) if upwinding == Val(:none) if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot) + @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( + DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅ + ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(ᶜg³ʰ[colidx]) - + (R_d / cv_d) * DiagonalMatrixRow(ᶠu³[colidx]) ⋅ + ᶠinterp_matrix() ⋅ ∂ᶜK_∂ᶜuₕ[colidx] + ) @. ∂ᶜρχ_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) + DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * ᶠg³³[colidx]) - + (R_d / cv_d) * DiagonalMatrixRow(ᶠu³[colidx]) ⋅ + ᶠinterp_matrix() ⋅ ∂ᶜK_∂ᶠu₃[colidx] ) else + @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * ᶠinterp(ᶜχ[colidx]), + ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(ᶜg³ʰ[colidx]) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * - ᶠg³³[colidx] * - (one_CT3xACT3,) * - ᶠinterp(ᶜχ[colidx]), + ᶠinterp(ᶜχ[colidx]) * + ᶠg³³[colidx], ) end else @@ -426,6 +458,21 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ) # Need to wrap ᶠupwind_matrix in this for the same reason. if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot) + @. ∂ᶜρχ_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]), + ) * + (one_AC3,), + ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(ᶜg³ʰ[colidx]) - + (R_d / cv_d) * + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³[colidx])) ⋅ + ∂ᶜK_∂ᶜuₕ[colidx] + ) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ ( @@ -434,13 +481,24 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ᶠset_upwind_bcs( ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]), ) * - ᶠg³³[colidx] * - (one_ATC3,), - ) + + (one_AC3,) * + ᶠg³³[colidx], + ) - + (R_d / cv_d) * ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³[colidx])) ⋅ - ∂ᶜK_∂ᶠu₃[colidx] * (-R_d / cv_d) + ∂ᶜK_∂ᶠu₃[colidx] ) 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]), + ) * + (one_AC3,), + ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(ᶜg³ʰ[colidx]) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * @@ -448,8 +506,8 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ᶠset_upwind_bcs( ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]), ) * - ᶠg³³[colidx] * - (one_ATC3,), + (one_AC3,) * + ᶠg³³[colidx], ) end end @@ -500,9 +558,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # 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_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] + ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)] ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)] I_u₃ = DiagonalMatrixRow(one_C3xACT3) - ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] @. ∂ᶠu₃_err_∂ᶜρ[colidx] = dtγ * ( DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ @@ -519,6 +578,9 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] = dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() * R_d / cv_d + @. ∂ᶠu₃_err_∂ᶜuₕ[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ ∂ᶜK_∂ᶜuₕ[colidx] if p.atmos.rayleigh_sponge isa RayleighSponge @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = dtγ * ( diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl index 46b51bc24ac..ba2151a63d1 100644 --- a/src/utils/abbreviations.jl +++ b/src/utils/abbreviations.jl @@ -74,6 +74,7 @@ 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 ᶠwinterp_matrix = MatrixFields.operator_matrix(ᶠwinterp) const ᶠgradᵥ_matrix = MatrixFields.operator_matrix(ᶠgradᵥ) const ᶠupwind1_matrix = MatrixFields.operator_matrix(ᶠupwind1) const ᶠupwind3_matrix = MatrixFields.operator_matrix(ᶠupwind3) diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index d76e74537a5..00d10e8c619 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -126,6 +126,44 @@ function g³³_field(field) return g_field.:($end_index) # For both 2D and 3D spaces, g³³ = g[end]. end +""" + g³³_tensor_field(field) + +Extracts the `g³³` tensor field from `Fields.local_geometry_field(field)`. +""" +g³³_tensor_field(field) = Base.Broadcast.broadcasted( + Geometry.AxisTensor, + ((Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()),), + g³³_field(field), +) + +""" + g³ʰ_tensor_field(field) + +Extracts the `g³ʰ` tensor field from `Fields.local_geometry_field(field)`. +""" +function g³ʰ_tensor_field(field) + full_axis = axes(eltype(Fields.local_geometry_field(field).gⁱʲ))[1] + all_components = Fields.local_geometry_field(field).gⁱʲ.components + if full_axis == Geometry.Contravariant123Axis() + horizontal_axis = Geometry.Contravariant12Axis() + n_components = 2 + elseif full_axis == Geometry.Contravariant13Axis() + horizontal_axis = Geometry.Contravariant1Axis() + n_components = 1 + elseif full_axis == Geometry.Contravariant23Axis() + horizontal_axis = Geometry.Contravariant2Axis() + n_components = 1 + else + error("Local geometry does not specify a g³ʰ tensor field") + end + g³ʰ_axes = ((Geometry.Contravariant3Axis(), horizontal_axis),) + g³ʰ_field = Base.Broadcast.broadcasted(all_components) do g³ʰ + g³ʰ[end:end, 1:n_components] + end + return Base.Broadcast.broadcasted(Geometry.AxisTensor, g³ʰ_axes, g³ʰ_field) +end + """ unit_basis_vector_data(type, local_geometry)