Skip to content

Commit da36313

Browse files
costachriscostachris
committed
Simplify function args
1 parent ddeeb7b commit da36313

14 files changed

+150
-176
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ NVTX.@annotate function set_cloud_fraction!(
190190
(; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
191191
(; ᶜρʲs, ᶜtsʲs) = p.precomputed
192192
(; turbconv_model) = p.atmos
193-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
193+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
194194

195195
# TODO - we should make this default when using diagnostic edmf
196196
# environment
@@ -304,7 +304,8 @@ function quad_loop(
304304
FT = eltype(x1_hat)
305305
@assert(x1_hat >= FT(0))
306306
@assert(x2_hat >= FT(0))
307-
_ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
307+
# note: ᶜthermo_state is used as a pointwise function here
308+
_ts = ᶜthermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
308309
hc = TD.has_condensate(thermo_params, _ts)
309310

310311
cf = hc ? FT(1) : FT(0) # cloud fraction

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 2 additions & 2 deletions
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))
@@ -1029,7 +1029,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
10291029
@. ᶜprandtl_nvec =
10301030
turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
10311031

1032-
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
1032+
ᶜtke⁰ = ᶜspecific_tke(Y, p)
10331033

10341034
ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
10351035
@. ᶜtke_exch = 0

‎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: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,7 @@ function thermo_state(
373373
return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
374374
end
375375

376+
const FieldOrValue = Union{Fields.Field, Base.AbstractBroadcasted, Real}
376377
function ᶜthermo_state(
377378
thermo_params;
378379
ρ = nothing,
@@ -383,15 +384,15 @@ function ᶜthermo_state(
383384
q_pt = nothing,
384385
)
385386

386-
get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, ::Nothing) =
387+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
387388
TD.PhaseDry_ρθ(thermo_params, ρ, θ)
388-
get_ts(ρ, ::Nothing, θ, ::Nothing, q_tot, ::Nothing) =
389+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
389390
TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot)
390-
get_ts(ρ, ::Nothing, θ, ::Nothing, ::Nothing, q_pt) =
391+
get_ts(ρ::T, ::Nothing, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} =
391392
TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt)
392-
get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, ::Nothing) =
393+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
393394
TD.PhaseDry_ρe(thermo_params, ρ, e_int)
394-
get_ts(ρ, ::Nothing, ::Nothing, e_int, q_tot, ::Nothing) =
395+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
395396
TD.PhaseEquil_ρeq(
396397
thermo_params,
397398
ρ,
@@ -400,19 +401,19 @@ function ᶜthermo_state(
400401
3,
401402
eltype(thermo_params)(0.003),
402403
)
403-
get_ts(ρ, ::Nothing, ::Nothing, e_int, ::Nothing, q_pt) =
404+
get_ts(ρ::T, ::Nothing, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} =
404405
TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt)
405-
get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, ::Nothing) =
406+
get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
406407
TD.PhaseDry_pθ(thermo_params, p, θ)
407-
get_ts(::Nothing, p, θ, ::Nothing, q_tot, ::Nothing) =
408+
get_ts(::Nothing, p::T, θ::T, ::Nothing, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
408409
TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot)
409-
get_ts(::Nothing, p, θ, ::Nothing, ::Nothing, q_pt) =
410+
get_ts(::Nothing, p::T, θ::T, ::Nothing, ::Nothing, q_pt) where {T <: FieldOrValue} =
410411
TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt)
411-
get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, ::Nothing) =
412+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, ::Nothing) where {T <: FieldOrValue} =
412413
TD.PhaseDry_pe(thermo_params, p, e_int)
413-
get_ts(::Nothing, p, ::Nothing, e_int, q_tot, ::Nothing) =
414+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, q_tot::T, ::Nothing) where {T <: FieldOrValue} =
414415
TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot)
415-
get_ts(::Nothing, p, ::Nothing, e_int, ::Nothing, q_pt) =
416+
get_ts(::Nothing, p::T, ::Nothing, e_int::T, ::Nothing, q_pt) where {T <: FieldOrValue} =
416417
TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt)
417418

418419
return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt))
@@ -442,22 +443,15 @@ function thermo_vars(moisture_model, microphysics_model, Y_c, K, Φ)
442443
return (; energy_var..., moisture_var...)
443444
end
444445

445-
ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
446-
ᶜthermo_state(
447-
thermo_params;
448-
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
449-
ρ,
450-
)
451-
452446
ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
453447
ᶜthermo_state(
454448
thermo_params;
455449
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
456450
ρ,
457451
)
458452

459-
ts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
460-
thermo_state(
453+
ᶜts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
454+
ᶜthermo_state(
461455
thermo_params;
462456
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
463457
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,
@@ -369,8 +369,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
369369
ᶜdz = Fields.Δz_field(axes(Y.c))
370370
ᶜlg = Fields.local_geometry_field(Y.c)
371371
ᶠlg = Fields.local_geometry_field(Y.f)
372-
ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
373-
ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
372+
ᶜtke⁰ = ᶜspecific_tke(Y, p)
373+
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
374374

375375
ᶜvert_div = p.scratch.ᶜtemp_scalar
376376
ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
@@ -455,7 +455,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
455455
(; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed
456456
# First order approximation: Use environmental mean fields.
457457
@. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat
458-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
458+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
459459
@. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat
460460
@. ᶜgradᵥ_θ_liq_ice⁰ =
461461
ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat
@@ -567,7 +567,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
567567
)
568568
end
569569
# sources from the environment
570-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
570+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
571571
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰)
572572
return nothing
573573
end
@@ -657,11 +657,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
657657
end
658658

659659
# Precipitation sources and sinks from the environment
660-
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y.c, p)
661-
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y.c, p)
662-
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y.c, p)
663-
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y.c, p)
664-
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y.c, p)
660+
ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p)
661+
ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p)
662+
ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p)
663+
ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p)
664+
ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p)
665665
ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰))
666666
compute_precipitation_sources!(
667667
ᶜ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)