diff --git a/src/cache/cache.jl b/src/cache/cache.jl index 4d41086148e..4a991930278 100644 --- a/src/cache/cache.jl +++ b/src/cache/cache.jl @@ -103,6 +103,7 @@ end # This is a constant Coriolis frequency that is only used if space is flat function build_cache(Y, atmos, params, surface_setup, dt, t_end, start_date) FT = eltype(params) + CTh = CTh_vector_type(Y.c) ᶜcoord = Fields.local_geometry_field(Y.c).coordinates grav = FT(CAP.grav(params)) @@ -155,6 +156,8 @@ function build_cache(Y, atmos, params, surface_setup, dt, t_end, start_date) ᶜp_ref, ᶜT = similar(Y.c, FT), ᶜf, + ᶜkappa_m = similar(Y.c, FT), + ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CTh{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/parameters/Parameters.jl b/src/parameters/Parameters.jl index 78ed6c7d843..a0f69874b7c 100644 --- a/src/parameters/Parameters.jl +++ b/src/parameters/Parameters.jl @@ -92,7 +92,7 @@ for var in fieldnames(TD.Parameters.ThermodynamicsParameters) @eval $var(ps::ACAP) = TD.Parameters.$var(thermodynamics_params(ps)) end # Thermodynamics derived parameters -for var in [:molmass_ratio, :R_d, :R_v, :cp_d, :cv_v, :cv_l, :cv_d] +for var in [:molmass_ratio, :R_d, :R_v, :e_int_v0, :cp_d, :cv_v, :cv_l, :cv_d] @eval $var(ps::ACAP) = TD.Parameters.$var(thermodynamics_params(ps)) end diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 1fe240212f2..68ec16e3186 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -125,15 +125,22 @@ function ImplicitEquationJacobian( approximate_solve_iters, ) FT = Spaces.undertype(axes(Y.c)) - one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' + CTh = CTh_vector_type(axes(Y.c)) + TridiagonalRow = TridiagonalMatrixRow{FT} BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} + TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - TridiagonalRow_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)} + BidiagonalRow_C3xACTh = + BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} + TridiagonalRow_C3xACT3 = + TridiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} is_in_Y(name) = MatrixFields.has_field(Y, name) + ρq_tot_if_available = is_in_Y(@name(c.ρq_tot)) ? (@name(c.ρq_tot),) : () + tracer_names = ( @name(c.ρq_tot), @name(c.ρq_liq), @@ -154,44 +161,83 @@ function ImplicitEquationJacobian( # 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, + identity_blocks = MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (@name(c.ρ), other_available_names...), + ) + + advection_blocks = ( MatrixFields.unrolled_map( name -> - (name, name) => - use_derivative(diffusion_flag) ? - similar(Y.c, TridiagonalRow) : FT(-1) * I, - (@name(c.ρe_tot), available_tracer_names...), + (name, @name(c.uₕ)) => similar(Y.c, TridiagonalRow_ACTh), + (@name(c.ρ), @name(c.ρe_tot), available_tracer_names...), )..., - (@name(c.uₕ), @name(c.uₕ)) => - use_derivative(diffusion_flag) && ( - !isnothing(atmos.turbconv_model) || - diffuse_momentum(atmos.vert_diff) - ) ? similar(Y.c, TridiagonalRow) : FT(-1) * I, MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - other_available_names, + name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + (@name(c.ρ), available_tracer_names...), )..., - (@name(c.ρ), @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), (@name(c.ρe_tot), @name(f.u₃)) => use_derivative(enthalpy_flag) ? similar(Y.c, QuaddiagonalRow_ACT3) : similar(Y.c, BidiagonalRow_ACT3), MatrixFields.unrolled_map( - name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), - available_tracer_names, + name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), + (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...), )..., - (@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_C3xACTh), (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), ) + diffusion_blocks = if use_derivative(diffusion_flag) + ( + MatrixFields.unrolled_map( + name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow), + (@name(c.ρe_tot), available_tracer_names...), + )..., + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + (@name(c.ρe_tot), available_tracer_names...), + )..., + (@name(c.ρe_tot), @name(c.ρq_tot)) => + similar(Y.c, TridiagonalRow), + (@name(c.uₕ), @name(c.uₕ)) => + !isnothing(atmos.turbconv_model) || + diffuse_momentum(atmos.vert_diff) ? + similar(Y.c, TridiagonalRow) : FT(-1) * I, + ) + else + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (@name(c.ρe_tot), available_tracer_names..., @name(c.uₕ)), + ) + end + + matrix = MatrixFields.FieldMatrix( + identity_blocks..., + advection_blocks..., + diffusion_blocks..., + ) + + names₁ = ( + @name(c.ρ), + @name(c.ρe_tot), + available_tracer_names..., + other_available_names..., + ) # Everything in Y except for ᶜuₕ and ᶠu₃. + + alg₁ = MatrixFields.StationaryIterativeSolve(; + P_alg = MatrixFields.BlockDiagonalPreconditioner(), + n_iters = 5, # TODO: Lower this value. + ) + alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) alg = use_derivative(diffusion_flag) ? MatrixFields.ApproximateBlockArrowheadIterativeSolve( - @name(c); + names₁...; + alg₁, + alg₂, n_iters = approximate_solve_iters, - ) : MatrixFields.BlockArrowheadSolve(@name(c)) + ) : MatrixFields.BlockArrowheadSolve(names₁...; alg₂) # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration # and approximates the A[c, c] block with a MainDiagonalPreconditioner. @@ -251,6 +297,7 @@ function Wfact!(A, Y, p, dtγ, t) p.precomputed.ᶜspecific, p.precomputed.ᶠu³, p.precomputed.ᶜK, + p.precomputed.ᶜts, p.precomputed.ᶜp, ( p.atmos.precip_model isa Microphysics1Moment ? @@ -261,11 +308,18 @@ function Wfact!(A, Y, p, dtγ, t) use_derivative(A.diffusion_flag) ? (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;) )..., + ( + use_derivative(A.diffusion_flag) && + p.atmos.turbconv_model isa AbstractEDMF ? + (; p.precomputed.ᶜtke⁰) : (;) + )..., ( use_derivative(A.diffusion_flag) && p.atmos.turbconv_model isa PrognosticEDMFX ? (; p.precomputed.ᶜρa⁰) : (;) )..., + p.core.ᶜkappa_m, + p.core.∂ᶜK_∂ᶜuₕ, p.core.∂ᶜK_∂ᶠu₃, p.core.ᶜΦ, p.core.ᶠgradᵥ_ᶜΦ, @@ -294,24 +348,33 @@ 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, ᶜts, ᶜp, ᶜkappa_m, ∂ᶜ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))' + CTh = CTh_vector_type(axes(Y.c)) + 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)) + Δcv_v = FT(CAP.cv_v(params)) - cv_d T_tri = FT(CAP.T_triple(params)) + e_int_v0 = FT(CAP.e_int_v0(params)) + thermo_params = CAP.thermodynamics_params(params) ᶜρ = Y.c.ρ ᶜuₕ = Y.c.uₕ ᶠu₃ = Y.f.u₃ ᶜJ = Fields.local_geometry_field(Y.c).J - ᶠg³³ = g³³_field(Y.f) + ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ + ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ + + @. ᶜkappa_m[colidx] = + TD.gas_constant_air(thermo_params, ᶜts[colidx]) / + TD.cv_m(thermo_params, ᶜts[colidx]) # We can express the kinetic energy as # ᶜK = @@ -322,6 +385,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(CTh(ᶜuₕ[colidx])) + + adjoint(ᶜinterp(ᶠu₃[colidx])) * g³ʰ(ᶜgⁱʲ[colidx]), + ) @. ∂ᶜK_∂ᶠu₃[colidx] = ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[colidx]))) + DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix() @@ -378,6 +445,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 +457,35 @@ 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³ʰ(ᶜ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) + ᶠinterp(ᶜχ[colidx]) * g³³(ᶠ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³ʰ(ᶜgⁱʲ[colidx])) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * - ᶠg³³[colidx] * - (one_CT3xACT3,) * - ᶠinterp(ᶜχ[colidx]), + ᶠinterp(ᶜχ[colidx]) * + g³³(ᶠgⁱʲ[colidx]), ) end else @@ -426,6 +505,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³ʰ(ᶜ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 +528,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³³(ᶠ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³ʰ(ᶜgⁱʲ[colidx])) @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * @@ -448,8 +553,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³³(ᶠgⁱʲ[colidx]), ) end end @@ -457,7 +562,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # TODO: Move the vertical advection of ρatke into the implicit tendency. if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) + ∂ᶜρatke_err_∂ᶜuₕ = matrix[@name(c.sgs⁰.ρatke), @name(c.uₕ)] ∂ᶜρatke_err_∂ᶠu₃ = matrix[@name(c.sgs⁰.ρatke), @name(f.u₃)] + + ∂ᶜρatke_err_∂ᶜuₕ[colidx] .= (zero(eltype(∂ᶜρatke_err_∂ᶜuₕ)),) ∂ᶜρatke_err_∂ᶠu₃[colidx] .= (zero(eltype(∂ᶜρatke_err_∂ᶠu₃)),) end @@ -500,14 +608,12 @@ 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_∂ᶠ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() ⋅ DiagonalMatrixRow( - R_d * (T_tri - (ᶜK[colidx] + ᶜΦ[colidx]) / cv_d), + ᶜkappa_m[colidx] * (T_tri * cv_d - ᶜK[colidx] - ᶜΦ[colidx]), ) + DiagonalMatrixRow( ( @@ -517,19 +623,35 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ) ⋅ ᶠinterp_matrix() ) @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] = - dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() * - R_d / cv_d + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶜkappa_m[colidx]) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] + @. ∂ᶠu₃_err_∂ᶜρq_tot[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶜkappa_m[colidx] * (T_tri * Δcv_v - e_int_v0)) + end + + ∂ᶠ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_∂ᶜuₕ[colidx] = + dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ ∂ᶜK_∂ᶜuₕ[colidx] 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(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ + ∂ᶜ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) ⋅ + ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,) end @@ -554,6 +676,64 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) # (I + ∂(ᶜp)/∂(ᶜρe_tot)) ⋅ DiagonalMatrixRow(1 / ᶜρ) ≈ # DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ). if use_derivative(diffusion_flag) + ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] + ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρ[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow( + ( + -(1 + ᶜkappa_m[colidx]) * ᶜspecific.e_tot[colidx] - + ᶜkappa_m[colidx] * + (T_tri * Δcv_v - e_int_v0) * + ᶜspecific.q_tot[colidx] + ) / ᶜρ[colidx], + ) + @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx])) ⋅ + ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow((1 + ᶜkappa_m[colidx]) / ᶜρ[colidx]) - (I,) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρq_tot[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx]), + ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶜkappa_m[colidx] * (T_tri * Δcv_v - e_int_v0) / ᶜρ[colidx], + ) + end + + ρ⁰_name = + p.atmos.turbconv_model isa PrognosticEDMFX ? @name(ᶜρa⁰) : @name(ᶜρ) + tracer_info = ( + (@name(c.ρq_tot), @name(ᶜspecific.q_tot), @name(ᶜK_h), @name(ᶜρ)), + (@name(c.ρq_liq), @name(ᶜspecific.q_liq), @name(ᶜK_h), @name(ᶜρ)), + (@name(c.ρq_ice), @name(ᶜspecific.q_ice), @name(ᶜK_h), @name(ᶜρ)), + (@name(c.ρq_rai), @name(ᶜspecific.q_rai), @name(ᶜK_h), @name(ᶜρ)), + (@name(c.ρq_sno), @name(ᶜspecific.q_sno), @name(ᶜK_h), @name(ᶜρ)), + (@name(c.sgs⁰.ρatke), @name(ᶜtke⁰), @name(ᶜK_u), ρ⁰_name), + ) + MatrixFields.unrolled_foreach( + tracer_info, + ) do (ρχ_name, χ_name, K_name, ρ_name) + MatrixFields.has_field(Y, ρχ_name) || return + ∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)] + ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] + ᶜχ = MatrixFields.get_field(p, χ_name) + ᶜK_χ = MatrixFields.get_field(p, K_name) + ᶜρ_χ = ρ_name == @name(ᶜρ) ? ᶜρ : MatrixFields.get_field(p, ρ_name) + @. ∂ᶜρχ_err_∂ᶜρ[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp(ᶜρ_χ[colidx]) * ᶠinterp(ᶜK_χ[colidx]), + ) ⋅ ᶠgradᵥ_matrix() ⋅ + DiagonalMatrixRow(-(ᶜχ[colidx]) / ᶜρ_χ[colidx]) + @. ∂ᶜρχ_err_∂ᶜρχ[colidx] = + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp(ᶜρ_χ[colidx]) * ᶠinterp(ᶜK_χ[colidx]), + ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ_χ[colidx]) - (I,) + end + if ( !isnothing(p.atmos.turbconv_model) || diffuse_momentum(p.atmos.vert_diff) @@ -565,26 +745,6 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_u[colidx]), ) ⋅ ᶠgradᵥ_matrix() - (I,) end - - ᶜρ⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ - scalar_info = ( - (@name(c.ρe_tot), ᶜρ, p.ᶜK_h, 1 + R_d / cv_d), - (@name(c.ρq_tot), ᶜρ, p.ᶜK_h, 1), - (@name(c.ρq_liq), ᶜρ, p.ᶜK_h, 1), - (@name(c.ρq_ice), ᶜρ, p.ᶜK_h, 1), - (@name(c.ρq_rai), ᶜρ, p.ᶜK_h, 1), - (@name(c.ρq_sno), ᶜρ, p.ᶜK_h, 1), - (@name(c.sgs⁰.ρatke), ᶜρ⁰, p.ᶜK_u, 1), - ) - MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, ᶜρ_χ, ᶜK_χ, s_χ) - MatrixFields.has_field(Y, ρχ_name) || return - ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] - @. ∂ᶜρχ_err_∂ᶜρχ[colidx] = - dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( - ᶠinterp(ᶜρ_χ[colidx]) * ᶠinterp(ᶜK_χ[colidx]), - ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(s_χ / ᶜρ_χ[colidx]) - - (I,) - end end # We can express the implicit precipitation tendency for scalars as 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 2e9c9fb68ef..ab08b19c8e5 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -126,6 +126,58 @@ function g³³_field(field) return g_field.:($end_index) # For both 2D and 3D spaces, g³³ = g[end]. end +""" + g³³(gⁱʲ) + +Extracts the `g³³` sub-tensor from the `gⁱʲ` tensor. +""" +g³³(gⁱʲ) = Geometry.AxisTensor( + (Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()), + gⁱʲ[end], +) + + +""" + g³ʰ(gⁱʲ) + +Extracts the `g³ʰ` sub-tensor from the `gⁱʲ` tensor. +""" +function g³ʰ(gⁱʲ) + full_CT_axis = axes(gⁱʲ)[1] + CTh_axis = if full_CT_axis == Geometry.Contravariant123Axis() + Geometry.Contravariant12Axis() + elseif full_CT_axis == Geometry.Contravariant13Axis() + Geometry.Contravariant1Axis() + elseif full_CT_axis == Geometry.Contravariant23Axis() + Geometry.Contravariant2Axis() + else + error("$full_CT_axis is missing either vertical or horizontal sub-axes") + end + return Geometry.AxisTensor( + (Geometry.Contravariant3Axis(), CTh_axis), + gⁱʲ[end:end, 1:length(CTh_axis)], + ) +end + +""" + CTh_vector_type(space) + +Extracts the (abstract) horizontal contravariant vector type from the given +`AbstractSpace`. +""" +function CTh_vector_type(space) + full_CT_axis = axes(eltype(Fields.local_geometry_field(space).gⁱʲ))[1] + return if full_CT_axis == Geometry.Contravariant123Axis() + Geometry.Contravariant12Vector + elseif full_CT_axis == Geometry.Contravariant13Axis() + Geometry.Contravariant1Vector + elseif full_CT_axis == Geometry.Contravariant23Axis() + Geometry.Contravariant2Vector + else + error("$full_CT_axis is missing either vertical or horizontal sub-axes") + end +end + """ unit_basis_vector_data(type, local_geometry)