Skip to content

Commit 5379b0a

Browse files
costachriscostachris
committed
Simplify function args
1 parent 495625f commit 5379b0a

13 files changed

+147
-172
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ NVTX.@annotate function set_cloud_fraction!(
191191
(; ᶜts⁰, cloud_diagnostics_tuple) = p.precomputed
192192
(; ᶜρʲs, ᶜtsʲs) = p.precomputed
193193
(; turbconv_model) = p.atmos
194-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
194+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
195195

196196
# TODO - we should make this default when using diagnostic edmf
197197
# environment
@@ -308,6 +308,7 @@ function quad_loop(
308308
FT = eltype(x1_hat)
309309
@assert(x1_hat >= FT(0))
310310
@assert(x2_hat >= FT(0))
311+
# note: ᶜthermo_state is used as a pointwise function here
311312
_ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
312313
hc = TD.has_condensate(thermo_params, _ts)
313314

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
361361
specific(Y.c.ρe_tot, Y.c.ρ),
362362
),
363363
)
364-
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
364+
ᶜtke⁰ = ᶜspecific_tke(Y, p)
365365

366366
for i in 2:Spaces.nlevels(axes(Y.c))
367367
ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))

‎src/cache/precipitation_precomputed_quantities.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -262,7 +262,7 @@ function set_precipitation_cache!(
262262
(; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
263263
(; ᶜts⁰, ᶜtsʲs) = p.precomputed
264264
thermo_params = CAP.thermodynamics_params(p.params)
265-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
265+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
266266

267267
n = n_mass_flux_subdomains(p.atmos.turbconv_model)
268268

‎src/cache/precomputed_quantities.jl

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ function implicit_precomputed_quantities(Y, atmos)
4343
ᶜts = similar(Y.c, TST),
4444
ᶜp = similar(Y.c, FT),
4545
)
46-
sgs_quantities = turbconv_model isa AbstractEDMF ? (;) : (;)
46+
sgs_quantities = (;)
4747
prognostic_sgs_quantities =
4848
turbconv_model isa PrognosticEDMFX ?
4949
(;
@@ -354,6 +354,7 @@ function thermo_state(
354354
return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
355355
end
356356

357+
const FieldOrValue = Union{Fields.Field, Base.AbstractBroadcasted, Real}
357358
function ᶜthermo_state(
358359
thermo_params;
359360
ρ = nothing,
@@ -364,15 +365,15 @@ function ᶜthermo_state(
364365
q_pt = nothing,
365366
)
366367

367-
get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, ::Nothing) =
368+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
368369
TD.PhaseDry_ρθ(thermo_params, ρ, θ)
369-
get_ts(ρ, ::Nothing, θ, ::Nothing, q_tot, ::Nothing) =
370+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
370371
TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot)
371-
get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, q_pt) =
372+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} =
372373
TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt)
373-
get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, ::Nothing) =
374+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
374375
TD.PhaseDry_ρe(thermo_params, ρ, e_int)
375-
get_ts(ρ, ::Nothing, ::Nothing, e_int, q_tot, ::Nothing) =
376+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
376377
TD.PhaseEquil_ρeq(
377378
thermo_params,
378379
ρ,
@@ -381,19 +382,19 @@ function ᶜthermo_state(
381382
3,
382383
eltype(thermo_params)(0.003),
383384
)
384-
get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, q_pt) =
385+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} =
385386
TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt)
386-
get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, ::Nothing) =
387+
get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
387388
TD.PhaseDry_pθ(thermo_params, p, θ)
388-
get_ts(::Nothing, p, θ, ::Nothing, q_tot, ::Nothing) =
389+
get_ts(::Nothing, p::T, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
389390
TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot)
390-
get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, q_pt) =
391+
get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} =
391392
TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt)
392-
get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, ::Nothing) =
393+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
393394
TD.PhaseDry_pe(thermo_params, p, e_int)
394-
get_ts(::Nothing, p, ::Nothing, e_int, q_tot, ::Nothing) =
395+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
395396
TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot)
396-
get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, q_pt) =
397+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} =
397398
TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt)
398399

399400
return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt))
@@ -423,22 +424,15 @@ function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ)
423424
return (; energy_var..., moisture_var...)
424425
end
425426

426-
ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
427-
ᶜthermo_state(
428-
thermo_params;
429-
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
430-
ρ,
431-
)
432-
433427
ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
434428
ᶜthermo_state(
435429
thermo_params;
436430
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
437431
ρ,
438432
)
439433

440-
ts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
441-
thermo_state(
434+
ᶜts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
435+
ᶜthermo_state(
442436
thermo_params;
443437
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
444438
p,

‎src/cache/prognostic_edmf_precomputed_quantities.jl

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,22 +24,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
2424
(; ᶜp, ᶜK) = p.precomputed
2525
(; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed
2626

27-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
28-
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
27+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
28+
ᶜtke⁰ = ᶜspecific_tke(Y, p)
2929
set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
3030
set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
3131
# @. ᶜK⁰ += ᶜtke⁰
32-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
32+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
3333

3434
ᶜmse⁰ = p.scratch.ᶜtemp_scalar_2
35-
ᶜmse⁰ .= specific_env_mse(Y.c, p)
35+
ᶜmse⁰ .= ᶜspecific_env_mse(Y, p)
3636

3737
if p.atmos.moisture_model isa NonEquilMoistModel &&
3838
p.atmos.microphysics_model isa Microphysics1Moment
39-
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
40-
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
41-
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
42-
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
39+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p)
40+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p)
41+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p)
42+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p)
4343
@. ᶜts⁰ = TD.PhaseNonEquil_phq(
4444
thermo_params,
4545
ᶜp,
@@ -361,8 +361,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
361361
ᶜdz = Fields.Δz_field(axes(Y.c))
362362
ᶜlg = Fields.local_geometry_field(Y.c)
363363
ᶠlg = Fields.local_geometry_field(Y.f)
364-
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
365-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
364+
ᶜtke⁰ = ᶜspecific_tke(Y, p)
365+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
366366

367367
ᶜvert_div = p.scratch.ᶜtemp_scalar
368368
ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
@@ -447,7 +447,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
447447
(; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed
448448
# First order approximation: Use environmental mean fields.
449449
@. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat
450-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
450+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
451451
@. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat
452452
@. ᶜgradᵥ_θ_liq_ice⁰ =
453453
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat
@@ -524,7 +524,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
524524
)
525525
end
526526
# sources from the environment
527-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
527+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
528528
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰)
529529
return nothing
530530
end
@@ -614,11 +614,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
614614
end
615615

616616
# Precipitation sources and sinks from the environment
617-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
618-
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
619-
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
620-
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
621-
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
617+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
618+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p)
619+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p)
620+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p)
621+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p)
622622
ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰))
623623
compute_precipitation_sources!(
624624
ᶜSᵖ,

‎src/diagnostics/edmfx_diagnostics.jl

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -633,7 +633,7 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} =
633633

634634
function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX)
635635
thermo_params = CAP.thermodynamics_params(cache.params)
636-
ᶜρa⁰_vals = ᶜρa⁰(state.c, cache)
636+
ᶜρa⁰_vals = ᶜρa⁰(state, cache)
637637
if isnothing(out)
638638
return draft_area.(
639639
ᶜρa⁰_vals,
@@ -953,7 +953,7 @@ function compute_clwen!(
953953
moisture_model::NonEquilMoistModel,
954954
turbconv_model::PrognosticEDMFX,
955955
)
956-
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state.c, cache)
956+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state, cache)
957957
if isnothing(out)
958958
return Base.materialize(ᶜq_liq⁰)
959959
else
@@ -1018,7 +1018,7 @@ function compute_clien!(
10181018
moisture_model::NonEquilMoistModel,
10191019
turbconv_model::PrognosticEDMFX,
10201020
)
1021-
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state.c, cache)
1021+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state, cache)
10221022
if isnothing(out)
10231023
return Base.materialize(ᶜq_ice⁰)
10241024
else
@@ -1067,7 +1067,7 @@ function compute_husraen!(
10671067
microphysics_model_model::Microphysics1Moment,
10681068
turbconv_model::PrognosticEDMFX,
10691069
)
1070-
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state.c, cache)
1070+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state, cache)
10711071
if isnothing(out)
10721072
return Base.materialize(ᶜq_rai⁰)
10731073
else
@@ -1116,7 +1116,7 @@ function compute_hussnen!(
11161116
microphysics_model_model::Microphysics1Moment,
11171117
turbconv_model::PrognosticEDMFX,
11181118
)
1119-
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state.c, cache)
1119+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state, cache)
11201120
if isnothing(out)
11211121
return Base.materialize(ᶜq_sno⁰)
11221122
else
@@ -1150,7 +1150,7 @@ function compute_tke!(
11501150
time,
11511151
turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX},
11521152
)
1153-
ᶜtke = ᶜspecific_tke(state.c.sgs⁰, state.c, cache)
1153+
ᶜtke = ᶜspecific_tke(state, cache)
11541154
if isnothing(out)
11551155
return Base.materialize(ᶜtke)
11561156
else

‎src/parameterized_tendencies/microphysics/precipitation.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ function precipitation_tendency!(
150150
# Source terms from EDMFX environment
151151
(; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed
152152

153-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
153+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
154154

155155
# Update from environment precipitation and cloud formation sources/sinks
156156
@. Yₜ.c.ρq_liq += ᶜρa⁰_vals * ᶜSqₗᵖ⁰

‎src/prognostic_equations/advection.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ 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⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y.c, p)) : @. lazy(Y.c.ρ)) : nothing
197+
ᶜρa⁰_vals = advect_tke ? (n > 0 ? (ᶜρa⁰(Y, p)) : @. lazy(Y.c.ρ)) : nothing
198198
ᶜρ⁰ = if advect_tke
199199
if n > 0
200200
(; ᶜts⁰) = p.precomputed
@@ -205,7 +205,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
205205
else
206206
nothing
207207
end
208-
ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)) : nothing
208+
ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y, p)) : nothing
209209
ᶜa_scalar = p.scratch.ᶜtemp_scalar
210210
ᶜω³ = p.scratch.ᶜtemp_CT3
211211
ᶠω¹² = p.scratch.ᶠtemp_CT12

‎src/prognostic_equations/edmfx_entr_detr.jl

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -531,13 +531,13 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF
531531
(; ᶠu₃⁰) = p.precomputed
532532

533533
ᶜmse⁰ = p.scratch.ᶜtemp_scalar
534-
ᶜmse⁰ .= specific_env_mse(Y.c, p)
534+
ᶜmse⁰ .= ᶜspecific_env_mse(Y, p)
535535
if p.atmos.moisture_model isa NonEquilMoistModel &&
536536
p.atmos.microphysics_model isa Microphysics1Moment
537-
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
538-
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
539-
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
540-
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
537+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p)
538+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p)
539+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p)
540+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p)
541541
end
542542

543543
for j in 1:n
@@ -547,7 +547,7 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF
547547
ᶜmseʲ = Y.c.sgsʲs.:($j).mse
548548
ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
549549

550-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
550+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
551551

552552
@. Yₜ.c.sgsʲs.:($$j).ρa += Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲ - ᶜdetrʲ)
553553

0 commit comments

Comments
 (0)