From d5c330086ddc49a77ca0fae01d96d09551fafd6c Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Mon, 11 Dec 2023 16:46:37 -0800 Subject: [PATCH] Add missing blocks to the dycore Jacobian --- .buildkite/pipeline.yml | 3 +- perf/flame.jl | 4 +- regression_tests/ref_counter.jl | 2 +- src/cache/cache.jl | 3 + src/parameters/Parameters.jl | 2 +- .../implicit/implicit_solver.jl | 302 +++++++++++++----- src/utils/abbreviations.jl | 1 + src/utils/utilities.jl | 53 +++ 8 files changed, 292 insertions(+), 78 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index c55f2ed2481..084a77cb404 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -768,7 +768,7 @@ steps: agents: slurm_gpus_per_task: 1 slurm_ntasks: 4 - slurm_mem: 32G + slurm_mem: 40G - label: "GPU: GPU moist Held-Suarez" command: @@ -830,6 +830,7 @@ steps: artifact_paths: "gpu_implicit_barowave_moist/*" agents: slurm_gpus: 1 + slurm_mem: 16G - label: "Perf: CPU implicit baro wave" command: > diff --git a/perf/flame.jl b/perf/flame.jl index d3780dbddda..e0de0014203 100644 --- a/perf/flame.jl +++ b/perf/flame.jl @@ -41,12 +41,12 @@ allocs_limit["flame_perf_target"] = 148_256 allocs_limit["flame_perf_target_tracers"] = 180_512 allocs_limit["flame_perf_target_edmfx"] = 7_005_552 allocs_limit["flame_perf_diagnostics"] = 25_356_928 -allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_311_040 +allocs_limit["flame_perf_target_diagnostic_edmfx"] = 1_315_072 allocs_limit["flame_sphere_baroclinic_wave_rhoe_equilmoist_expvdiff"] = 4_018_252_656 allocs_limit["flame_perf_target_threaded"] = 1_276_864 allocs_limit["flame_perf_target_callbacks"] = 37_277_112 -allocs_limit["flame_perf_gw"] = 3_226_429_472 +allocs_limit["flame_perf_gw"] = 3_226_429_728 allocs_limit["flame_perf_target_prognostic_edmfx_aquaplanet"] = 1_258_848 # Ideally, we would like to track all the allocations, but this becomes too diff --git a/regression_tests/ref_counter.jl b/regression_tests/ref_counter.jl index b4f334f2651..aaacbe66291 100644 --- a/regression_tests/ref_counter.jl +++ b/regression_tests/ref_counter.jl @@ -1 +1 @@ -141 +142 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..affca4dae21 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,47 +161,86 @@ 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), ) - alg = - use_derivative(diffusion_flag) ? + 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.BlockLowerTriangularSolve(@name(c.uₕ)) + alg = if use_derivative(diffusion_flag) + alg₁ = MatrixFields.StationaryIterativeSolve(; + P_alg = MatrixFields.BlockDiagonalPreconditioner(), + n_iters = approximate_solve_iters, + ) MatrixFields.ApproximateBlockArrowheadIterativeSolve( - @name(c); + names₁...; + alg₁, + alg₂, + P_alg₁ = MatrixFields.MainDiagonalPreconditioner(), n_iters = approximate_solve_iters, - ) : MatrixFields.BlockArrowheadSolve(@name(c)) - - # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration - # and approximates the A[c, c] block with a MainDiagonalPreconditioner. + ) + else + MatrixFields.BlockArrowheadSolve(names₁...; alg₂) + end return ImplicitEquationJacobian( matrix, @@ -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 = T_tri * Δcv_v - 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])) - + ᶜkappa_m[colidx] * 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]), + ) - + ᶜkappa_m[colidx] * 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])) - + ᶜkappa_m[colidx] * + ᶠ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]), + ) - + ᶜkappa_m[colidx] * ᶠ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,34 @@ 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] * Δ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 +675,61 @@ 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] * Δ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] * Δe_int_v0 / ᶜρ[colidx]) + end + + ρ⁰_name = + p.atmos.turbconv_model isa PrognosticEDMFX ? @name(ᶜρa⁰) : @name(ᶜρ) + scalar_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( + scalar_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 +741,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..5b2d8a4efe5 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -126,6 +126,59 @@ 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()), + Geometry.components(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 + N = length(full_CT_axis) + return Geometry.AxisTensor( + (Geometry.Contravariant3Axis(), CTh_axis), + view(Geometry.components(gⁱʲ), N:N, 1:(N - 1)), + ) +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)