Skip to content

Commit

Permalink
Add precipitation to Jacobian
Browse files Browse the repository at this point in the history
  • Loading branch information
trontrytel committed Dec 2, 2023
1 parent af8e2da commit b09ceff
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 3 deletions.
2 changes: 2 additions & 0 deletions config/model_configs/single_column_precipitation_test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
70 changes: 68 additions & 2 deletions src/prognostic_equations/implicit/implicit_solver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,10 @@ function Wfact!(A, Y, p, dtγ, t)
p.precomputed.ᶠu³,
p.precomputed.ᶜK,
p.precomputed.ᶜp,
(
p.atmos.precip_model isa Microphysics1Moment ?
(; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;)
)...,
p.precomputed.ᶜh_tot,
(
use_derivative(A.diffusion_flag) ?
Expand Down Expand Up @@ -371,8 +375,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
Expand Down Expand Up @@ -584,4 +586,68 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
(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
2 changes: 1 addition & 1 deletion src/prognostic_equations/implicit/implicit_tendency.jl
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,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!(
Expand Down

0 comments on commit b09ceff

Please sign in to comment.