Skip to content

Commit 2eb7859

Browse files
costachriscostachris
committed
Modify implicit solver, format
1 parent 0f54b24 commit 2eb7859

File tree

8 files changed

+45
-51
lines changed

8 files changed

+45
-51
lines changed

‎src/cache/precomputed_quantities.jl

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,7 @@ function precomputed_quantities(Y, atmos)
162162

163163
edonly_quantities =
164164
atmos.turbconv_model isa EDOnlyEDMFX ?
165-
(;
166-
ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
167-
) : (;)
165+
(; ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),) : (;)
168166

169167
sgs_quantities = (;
170168
ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace),
@@ -301,8 +299,7 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
301299
end
302300

303301
function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model)
304-
@. ᶜK +=
305-
ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Y.f.u₃, CT3(ᶠu₃⁰ - Y.f.u₃))) / 2 / Y.c.ρ
302+
@. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Y.f.u₃, CT3(ᶠu₃⁰ - Y.f.u₃))) / 2 / Y.c.ρ
306303
for j in 1:n_mass_flux_subdomains(turbconv_model)
307304
ᶜρaʲ = Y.c.sgsʲs.:($j).ρa
308305
ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃

‎src/diagnostics/edmfx_diagnostics.jl

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1160,7 +1160,9 @@ function compute_tke!(
11601160
end
11611161

11621162
ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model))
1163-
ᶜtke = @. lazy(specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
1163+
ᶜtke = @. lazy(
1164+
specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
1165+
)
11641166
if isnothing(out)
11651167
return Base.materialize(ᶜtke)
11661168
else

‎src/prognostic_equations/advection.jl

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
194194
turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ :
195195
p.precomputed.ᶠu³⁰
196196
) : nothing
197-
ᶜρa⁰ = advect_tke ? (n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ) : nothing
197+
ᶜρa⁰ =
198+
advect_tke ?
199+
(n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ) :
200+
nothing
198201
ᶜρ⁰ = if advect_tke
199202
if n > 0
200203
(; ᶜts⁰) = p.precomputed
@@ -205,7 +208,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
205208
else
206209
nothing
207210
end
208-
ᶜtke⁰ = advect_tke ? (@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) : nothing
211+
ᶜtke⁰ =
212+
advect_tke ?
213+
(@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) :
214+
nothing
209215
ᶜa_scalar = p.scratch.ᶜtemp_scalar
210216
ᶜω³ = p.scratch.ᶜtemp_CT3
211217
ᶠω¹² = p.scratch.ᶠtemp_CT12

‎src/prognostic_equations/edmfx_sgs_flux.jl

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -111,8 +111,7 @@ function edmfx_sgs_mass_flux_tendency!(
111111
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
112112
@. ᶠu³_diff = ᶠu³⁰ - ᶠu³
113113
@. ᶜa_scalar =
114-
(ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) *
115-
draft_area(ᶜρa⁰, ᶜρ⁰)
114+
(ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
116115
vtt = vertical_transport(
117116
ᶜρ⁰,
118117
ᶠu³_diff,
@@ -186,8 +185,7 @@ function edmfx_sgs_mass_flux_tendency!(
186185
@. ᶠu³_diff = ᶠu³⁰ - ᶠu³
187186

188187
@. ᶜa_scalar =
189-
(ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) *
190-
draft_area(ᶜρa⁰, ᶜρ⁰)
188+
(ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
191189
vtt = vertical_transport(
192190
ᶜρ⁰,
193191
ᶠu³_diff,
@@ -198,8 +196,7 @@ function edmfx_sgs_mass_flux_tendency!(
198196
@. Yₜ.c.ρq_liq += vtt
199197

200198
@. ᶜa_scalar =
201-
(ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) *
202-
draft_area(ᶜρa⁰, ᶜρ⁰)
199+
(ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
203200
vtt = vertical_transport(
204201
ᶜρ⁰,
205202
ᶠu³_diff,
@@ -210,8 +207,7 @@ function edmfx_sgs_mass_flux_tendency!(
210207
@. Yₜ.c.ρq_ice += vtt
211208

212209
@. ᶜa_scalar =
213-
(ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) *
214-
draft_area(ᶜρa⁰, ᶜρ⁰)
210+
(ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
215211
vtt = vertical_transport(
216212
ᶜρ⁰,
217213
ᶠu³_diff,
@@ -222,8 +218,7 @@ function edmfx_sgs_mass_flux_tendency!(
222218
@. Yₜ.c.ρq_rai += vtt
223219

224220
@. ᶜa_scalar =
225-
(ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) *
226-
draft_area(ᶜρa⁰, ᶜρ⁰)
221+
(ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
227222
vtt = vertical_transport(
228223
ᶜρ⁰,
229224
ᶠu³_diff,

‎src/prognostic_equations/edmfx_tke.jl

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,9 @@ function edmfx_tke_tendency!(
5959
FT = eltype(p.params)
6060

6161

62-
ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
62+
ᶜρa⁰ =
63+
turbconv_model isa PrognosticEDMFX ?
64+
(@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
6365
nh_pressure3_buoyʲs =
6466
turbconv_model isa PrognosticEDMFX ?
6567
p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs
@@ -101,7 +103,8 @@ function edmfx_tke_tendency!(
101103
# buoyancy production
102104
@. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad
103105

104-
ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
106+
ᶜtke⁰ =
107+
@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
105108

106109
# entrainment and detraiment
107110
# using ᶜu⁰ and local geometry results in allocation

‎src/prognostic_equations/hyperdiffusion.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,8 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
126126
sgsʲs = nothing
127127
end
128128
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, sgsʲs, turbconv_model))
129-
ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
129+
ᶜtke⁰ =
130+
@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
130131
(; ᶜ∇²tke⁰) = p.hyperdiff
131132
@. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰))
132133
end

‎src/prognostic_equations/implicit/manual_sparse_jacobian.jl

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -634,8 +634,9 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
634634
dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ)
635635
end
636636

637-
foreach_gs_tracer(Y) do ᶜρχ, ρχ_name
637+
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, _)
638638
MatrixFields.has_field(Y, ρχ_name) || return
639+
ᶜρχ = MatrixFields.get_field(Y, ρχ_name)
639640
ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
640641
∂ᶜρχ_err_∂ᶜρ = matrix[ρχ_name, @name(c.ρ)]
641642
∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name]
@@ -655,7 +656,9 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
655656
c_d = CAP.tke_diss_coeff(turbconv_params)
656657
(; dt) = p
657658
turbconv_model = p.atmos.turbconv_model
658-
ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? (@.lazy(ρa⁰(Y.c))) : Y.c.ρ
659+
ᶜρa⁰ =
660+
p.atmos.turbconv_model isa PrognosticEDMFX ?
661+
(@.lazy(ρa⁰(Y.c))) : Y.c.ρ
659662
ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))
660663
ᶜρatke⁰ = Y.c.sgs⁰.ρatke
661664

@@ -989,21 +992,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
989992
Δcv_v * TD.gas_constant_air(thermo_params, ᶜts)
990993
) / abs2(TD.cv_m(thermo_params, ᶜts))
991994

992-
if MatrixFields.has_field(Y, @name(c.ρq_tot))
993-
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
994-
@. ∂ᶜρe_tot_err_∂ᶜρq_tot +=
995-
dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅
996-
DiagonalMatrixRow((
997-
ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ +
998-
ᶜ∂kappa_m∂q_tot * (
999-
cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ +
1000-
∂e_int_∂q_tot * ᶜq_tot
1001-
)
1002-
))
1003-
1004-
else
1005-
ᶜq_tot = lazy(0)
1006-
end
995+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
1007996

1008997
@. ∂ᶜρe_tot_err_∂ᶜρ +=
1009998
dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅
@@ -1014,6 +1003,16 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
10141003
) / ᶜρ,
10151004
)
10161005

1006+
@. ∂ᶜρe_tot_err_∂ᶜρq_tot +=
1007+
dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅
1008+
DiagonalMatrixRow((
1009+
ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ +
1010+
ᶜ∂kappa_m∂q_tot * (
1011+
cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ +
1012+
∂e_int_∂q_tot * ᶜq_tot
1013+
)
1014+
))
1015+
10171016
@. ∂ᶜρe_tot_err_∂ᶜρe_tot +=
10181017
dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅
10191018
DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ)

‎src/utils/variable_manipulations.jl

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -263,20 +263,12 @@ Arguments:
263263
- `gs`: The grid-scale iteration object, which contains the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX) from the state `Y.c`, or `ᶜρaʲs` in the cache for DiagnosticEDMFX.
264264
- `turbconv_model`: The turbulence convection model, used to determine how to access draft data.
265265
"""
266-
function ᶜenv_value(
267-
grid_scale_value,
268-
f_draft,
269-
gs,
270-
)
266+
function ᶜenv_value(grid_scale_value, f_draft, gs)
271267
return @. lazy(grid_scale_value - draft_sum(f_draft, gs))
272268
end
273269

274270

275-
function env_value(
276-
grid_scale_value,
277-
f_draft,
278-
gs,
279-
)
271+
function env_value(grid_scale_value, f_draft, gs)
280272
return grid_scale_value - draft_sum(f_draft, gs)
281273
end
282274

@@ -455,9 +447,8 @@ function ᶜspecific_env_mse(Y, p)
455447
# Numerator: ρa⁰mse⁰ = ρmse - (Σ ρaʲ * mseʲ)
456448

457449
if turbconv_model isa PrognosticEDMFX
458-
ρa⁰mse⁰ =
459-
ᶜenv_value(ᶜρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, Y.c.sgsʲs)
460-
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
450+
ρa⁰mse⁰ = ᶜenv_value(ᶜρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, Y.c.sgsʲs)
451+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
461452
elseif turbconv_model isa DiagnosticEDMFX || turbconv_model isa EDOnlyEDMFX
462453

463454
n = n_mass_flux_subdomains(turbconv_model)

0 commit comments

Comments
 (0)