From cc8e0365c2f49d9fba629261722665b92af0d618 Mon Sep 17 00:00:00 2001 From: Anna Jaruga Date: Thu, 30 Nov 2023 14:41:05 -0800 Subject: [PATCH] Add precipitation to Jacobian --- .../single_column_precipitation_test.yml | 2 + .../implicit/implicit_solver.jl | 72 ++++++++++++++++++- .../implicit/implicit_tendency.jl | 2 +- 3 files changed, 72 insertions(+), 4 deletions(-) diff --git a/config/model_configs/single_column_precipitation_test.yml b/config/model_configs/single_column_precipitation_test.yml index 8caac6f5080..a6d1b4c6bfd 100644 --- a/config/model_configs/single_column_precipitation_test.yml +++ b/config/model_configs/single_column_precipitation_test.yml @@ -11,6 +11,8 @@ moist: "nonequil" precip_model: "1M" precip_upwinding: "first_order" hyperdiff: "false" +implicit_diffusion: true +vert_diff: "true" regression_test: false check_precipitation: true job_id: "single_column_precipitation_test" diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 939ed7cb061..9ea458aeaf0 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -241,12 +241,16 @@ _linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b) function Wfact!(A, Y, p, dtγ, t) NVTX.@range "Wfact!" color = colorant"green" begin # Remove unnecessary values from p to avoid allocations in bycolumn. - (; rayleigh_sponge) = p.atmos + (; rayleigh_sponge, precip_model) = p.atmos p′ = (; p.precomputed.ᶜspecific, p.precomputed.ᶠu³, p.precomputed.ᶜK, p.precomputed.ᶜp, + ( + precip_model isa Microphysics1Moment ? + (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;) + )..., p.core.∂ᶜK_∂ᶠu₃, p.core.ᶜΦ, p.core.ᶠgradᵥ_ᶜΦ, @@ -357,8 +361,6 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) (@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), - (@name(c.ρq_rai), @name(ᶜspecific.q_rai), precip_upwinding), - (@name(c.ρq_sno), @name(ᶜspecific.q_sno), precip_upwinding), ) MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, χ_name, upwinding) MatrixFields.has_field(Y, ρχ_name) || return @@ -571,4 +573,68 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅ ᶠgradᵥ_matrix() - (I,) end end + + # We can express the implicit precipitation tendency for scalars as + # ᶜρχₜ = -(ᶜdivᵥ_ρqₚ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind(ᶠu³ₚ, ᶜχ))). + # This means that + # ∂(ᶜρχₜ)/∂(ᶜρχ) = + # -(ᶜdivᵥ_ρqₚ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) ⋅ ( + # ∂(ᶠupwind(ᶠu³ₚ, ᶜχ))/∂(ᶠu³ₚ) ⋅ ∂(ᶠu³ₚ)/∂(ᶜρχ) + + # ᶠupwind_matrix(ᶠu³ₚ) ⋅ ∂(ᶜχ)/∂(ᶜρχ) + # ). + # In general, ∂(ᶜχ)/∂(ᶜρχ) = DiagonalMatrixRow(1 / ᶜρ), except for the case + # ∂(ᶠu³ₚ)/∂(ᶜρχ) = 0 approximating as zero for now + if use_derivative(diffusion_flag) && + p.atmos.precip_model isa Microphysics1Moment + + #TODO - do we need any of those + # - it's a little bit more work for ρ + # - for energy remember about rescaling by dh/dρe + # - should work out of the box for the rest + + scalar_info = ( + #(@name(c.ρ), density_upwinding, nothing), + #(@name(c.ρe_tot), energy_upwinding, nothing), + #(@name(c.ρq_tot), tracer_upwinding, nothing), + #(@name(c.ρq_liq), tracer_upwinding, nothing), + #(@name(c.ρq_ice), tracer_upwinding, nothing), + #(@name(c.ρq_rai), tracer_upwinding, nothing), + #(@name(c.ρq_sno), tracer_upwinding, nothing), + (@name(c.ρq_rai), precip_upwinding, p.ᶜwᵣ), + (@name(c.ρq_sno), precip_upwinding, p.ᶜwₛ), + ) + + MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, upwinding, ᶜw) + MatrixFields.has_field(Y, ρχ_name) || return + + ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] + + is_third_order = upwinding == Val(:third_order) + ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 + ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix + + if isnothing(ᶜw) + @. ∂ᶜρχ_err_∂ᶜρχ[colidx] -= + dtγ * ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ + ᶠupwind_matrix(ᶠu³[colidx]) ⋅ + DiagonalMatrixRow(1 / ᶜρ[colidx]) + else + ᶜdivᵥ_ρqₚ = Operators.DivergenceF2C( + top = Operators.SetValue(C3(FT(0))), + bottom = Operators.SetDivergence(FT(0)), + ) + ᶜdivᵥ_ρqₚ_matrix = MatrixFields.operator_matrix(ᶜdivᵥ_ρqₚ) + lgf = Fields.local_geometry_field(Y.f) + + @. ∂ᶜρχ_err_∂ᶜρχ[colidx] -= + dtγ * ᶜdivᵥ_ρqₚ_matrix() ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ + ᶠupwind_matrix( + -(ᶠinterp(ᶜw[colidx])) * + CT3(unit_basis_vector_data(CT3, lgf[colidx])), + ) ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) + end + end + end end diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index d3f16175251..1ba5db3d16e 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -134,7 +134,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx) ᶜdivᵥ_ρqₚ = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), - bottom = Operators.Extrapolate(), + bottom = Operators.SetDivergence(FT(0)), ) vertical_transport!(