Skip to content

Commit bcc9b71

Browse files
costachriscostachris
committed
remove some ᶜfunctions
1 parent ca491b8 commit bcc9b71

12 files changed

+79
-235
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ NVTX.@annotate function set_cloud_fraction!(
5959
TD.PhasePartition(thermo_params, ᶜts).ice,
6060
)
6161
else
62-
q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
63-
q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
62+
q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
63+
q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
6464
@. cloud_diagnostics_tuple =
6565
make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice)
6666
end
@@ -308,7 +308,6 @@ 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
312311
_ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
313312
hc = TD.has_condensate(thermo_params, _ts)
314313

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 13 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -105,18 +105,16 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
105105
turbconv_params = CAP.turbconv_params(params)
106106
ᶜts = p.precomputed.ᶜts #TODO replace
107107

108-
q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
109-
ᶜe_tot = ᶜspecific(Y.c.ρe_tot, Y.c.ρ)
108+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
109+
ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
110110
ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot))
111111

112112
ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1))
113113
uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1))
114114
u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half))
115-
h_tot_int_level =
116-
Fields.field_values(Fields.level(ᶜh_tot, 1))
115+
h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1))
117116
K_int_level = Fields.field_values(Fields.level(ᶜK, 1))
118-
q_tot_int_level =
119-
Fields.field_values(Fields.level(q_tot, 1))
117+
q_tot_int_level = Fields.field_values(Fields.level(ᶜq_tot, 1))
120118

121119
p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
122120
Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
@@ -332,8 +330,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
332330

333331

334332
if microphysics_model isa Microphysics1Moment
335-
q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
336-
q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
333+
q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
334+
q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
337335
end
338336

339337
thermo_params = CAP.thermodynamics_params(params)
@@ -353,7 +351,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
353351
Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
354352

355353
# integral
356-
q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
354+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
357355
ᶜh_tot = @. lazy(
358356
TD.total_specific_enthalpy(
359357
thermo_params,
@@ -368,10 +366,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
368366
uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i))
369367
u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half))
370368
K_level = Fields.field_values(Fields.level(ᶜK, i))
371-
h_tot_level =
372-
Fields.field_values(Fields.level(ᶜh_tot, i))
373-
q_tot_level =
374-
Fields.field_values(Fields.level(q_tot, i))
369+
h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
370+
q_tot_level = Fields.field_values(Fields.level(ᶜq_tot, i))
375371
p_level = Fields.field_values(Fields.level(ᶜp, i))
376372
Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
377373
local_geometry_level = Fields.field_values(
@@ -395,10 +391,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
395391
Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
396392
u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1
397393
K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1))
398-
h_tot_prev_level =
399-
Fields.field_values(Fields.level(ᶜh_tot, i - 1))
400-
q_tot_prev_level =
401-
Fields.field_values(Fields.level(q_tot, i - 1))
394+
h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1))
395+
q_tot_prev_level = Fields.field_values(Fields.level(ᶜq_tot, i - 1))
402396
ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1))
403397
p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1))
404398
z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1))
@@ -498,9 +492,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
498492
Fields.field_values(Fields.level(ᶜq_iceʲ, i - 1))
499493
end
500494

501-
tke_prev_level = Fields.field_values(
502-
Fields.level(ᶜtke⁰, i - 1),
503-
)
495+
tke_prev_level = Fields.field_values(Fields.level(ᶜtke⁰, i - 1))
504496

505497
@. entrʲ_prev_level = entrainment(
506498
thermo_params,
@@ -1063,7 +1055,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
10631055
(; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
10641056

10651057
# Environment precipitation sources (to be applied to grid mean)
1066-
ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
1058+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
10671059
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
10681060
thermo_params,
10691061
microphys_0m_params,

‎src/cache/precipitation_precomputed_quantities.jl

Lines changed: 28 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -284,11 +284,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
284284
(; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed
285285
(; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
286286

287-
q_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
288-
q_rai = ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
289-
q_sno = ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
290-
q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
291-
q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
287+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
288+
ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
289+
ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
290+
ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
291+
ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
292292

293293
ᶜSᵖ = p.scratch.ᶜtemp_scalar
294294
ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
@@ -308,11 +308,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
308308
ᶜSqᵣᵖ,
309309
ᶜSqₛᵖ,
310310
Y.c.ρ,
311-
q_tot,
312-
q_liq,
313-
q_ice,
314-
q_rai,
315-
q_sno,
311+
ᶜq_tot,
312+
ᶜq_liq,
313+
ᶜq_ice,
314+
ᶜq_rai,
315+
ᶜq_sno,
316316
ᶜts,
317317
dt,
318318
cmp,
@@ -325,11 +325,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
325325
ᶜSqᵣᵖ,
326326
ᶜSqₛᵖ,
327327
Y.c.ρ,
328-
q_tot,
329-
q_liq,
330-
q_ice,
331-
q_rai,
332-
q_sno,
328+
ᶜq_tot,
329+
ᶜq_liq,
330+
ᶜq_ice,
331+
ᶜq_rai,
332+
ᶜq_sno,
333333
ᶜts,
334334
dt,
335335
cmp,
@@ -474,30 +474,23 @@ function set_precipitation_surface_fluxes!(
474474
sfc_ρ = @. lazy(int_ρ * int_J / sfc_J)
475475

476476
# Constant extrapolation to surface, consistent with simple downwinding
477+
# Temporary scratch variables are used here until CC.field_values supports <lazy> fields
477478
ᶜq_rai = p.scratch.ᶜtemp_scalar
478479
ᶜq_sno = p.scratch.ᶜtemp_scalar_2
479480
ᶜq_liq = p.scratch.ᶜtemp_scalar_3
480481
ᶜq_ice = p.scratch.ᶜtemp_scalar_4
481-
ᶜq_rai .= ᶜspecific(Y.c.ρq_rai, Y.c.ρ)
482-
ᶜq_sno .= ᶜspecific(Y.c.ρq_sno, Y.c.ρ)
483-
ᶜq_liq .= ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
484-
ᶜq_ice .= ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
485-
sfc_qᵣ = Fields.Field(
486-
Fields.field_values(Fields.level(ᶜq_rai, 1)),
487-
sfc_space,
488-
)
489-
sfc_qₛ = Fields.Field(
490-
Fields.field_values(Fields.level(ᶜq_sno, 1)),
491-
sfc_space,
492-
)
493-
sfc_qₗ = Fields.Field(
494-
Fields.field_values(Fields.level(ᶜq_liq, 1)),
495-
sfc_space,
496-
)
497-
sfc_qᵢ = Fields.Field(
498-
Fields.field_values(Fields.level(ᶜq_ice, 1)),
499-
sfc_space,
500-
)
482+
@. ᶜq_rai = specific(Y.c.ρq_rai, Y.c.ρ)
483+
@. ᶜq_sno = specific(Y.c.ρq_sno, Y.c.ρ)
484+
@. ᶜq_liq = specific(Y.c.ρq_liq, Y.c.ρ)
485+
@. ᶜq_ice = specific(Y.c.ρq_ice, Y.c.ρ)
486+
sfc_qᵣ =
487+
Fields.Field(Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space)
488+
sfc_qₛ =
489+
Fields.Field(Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space)
490+
sfc_qₗ =
491+
Fields.Field(Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space)
492+
sfc_qᵢ =
493+
Fields.Field(Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space)
501494
sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space)
502495
sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space)
503496
sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space)

‎src/cache/precomputed_quantities.jl

Lines changed: 12 additions & 142 deletions
Original file line numberDiff line numberDiff line change
@@ -355,162 +355,32 @@ function thermo_state(
355355
return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
356356
end
357357

358-
const FieldOrValue = Union{Fields.Field, Base.AbstractBroadcasted, Real}
359-
function ᶜthermo_state(
360-
thermo_params;
361-
ρ = nothing,
362-
p = nothing,
363-
θ = nothing,
364-
e_int = nothing,
365-
q_tot = nothing,
366-
q_pt = nothing,
367-
)
368-
369-
get_ts(
370-
ρ::T,
371-
::Nothing,
372-
θ::T,
373-
::Nothing,
374-
::Nothing,
375-
::Nothing,
376-
) where {T <: FieldOrValue} = TD.PhaseDry_ρθ(thermo_params, ρ, θ)
377-
get_ts(
378-
ρ::T,
379-
::Nothing,
380-
θ::T,
381-
::Nothing,
382-
q_tot::T,
383-
::Nothing,
384-
) where {T <: FieldOrValue} = TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot)
385-
get_ts(
386-
ρ::T,
387-
::Nothing,
388-
θ::T,
389-
::Nothing,
390-
::Nothing,
391-
q_pt,
392-
) where {T <: FieldOrValue} =
393-
TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt)
394-
get_ts(
395-
ρ::T,
396-
::Nothing,
397-
::Nothing,
398-
e_int::T,
399-
::Nothing,
400-
::Nothing,
401-
) where {T <: FieldOrValue} = TD.PhaseDry_ρe(thermo_params, ρ, e_int)
402-
get_ts(
403-
ρ::T,
404-
::Nothing,
405-
::Nothing,
406-
e_int::T,
407-
q_tot::T,
408-
::Nothing,
409-
) where {T <: FieldOrValue} = TD.PhaseEquil_ρeq(
410-
thermo_params,
411-
ρ,
412-
e_int,
413-
q_tot,
414-
3,
415-
eltype(thermo_params)(0.003),
416-
)
417-
get_ts(
418-
ρ::T,
419-
::Nothing,
420-
::Nothing,
421-
e_int::T,
422-
::Nothing,
423-
q_pt,
424-
) where {T <: FieldOrValue} =
425-
TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt)
426-
get_ts(
427-
::Nothing,
428-
p::T,
429-
θ::T,
430-
::Nothing,
431-
::Nothing,
432-
::Nothing,
433-
) where {T <: FieldOrValue} = TD.PhaseDry_pθ(thermo_params, p, θ)
434-
get_ts(
435-
::Nothing,
436-
p::T,
437-
θ::T,
438-
::Nothing,
439-
q_tot::T,
440-
::Nothing,
441-
) where {T <: FieldOrValue} = TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot)
442-
get_ts(
443-
::Nothing,
444-
p::T,
445-
θ::T,
446-
::Nothing,
447-
::Nothing,
448-
q_pt,
449-
) where {T <: FieldOrValue} =
450-
TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt)
451-
get_ts(
452-
::Nothing,
453-
p::T,
454-
::Nothing,
455-
e_int::T,
456-
::Nothing,
457-
::Nothing,
458-
) where {T <: FieldOrValue} = TD.PhaseDry_pe(thermo_params, p, e_int)
459-
get_ts(
460-
::Nothing,
461-
p::T,
462-
::Nothing,
463-
e_int::T,
464-
q_tot::T,
465-
::Nothing,
466-
) where {T <: FieldOrValue} =
467-
TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot)
468-
get_ts(
469-
::Nothing,
470-
p::T,
471-
::Nothing,
472-
e_int::T,
473-
::Nothing,
474-
q_pt,
475-
) where {T <: FieldOrValue} =
476-
TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt)
477-
478-
return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt))
479-
end
480-
481358
function thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)
482-
e_tot = ᶜspecific(ᶜY.ρe_tot, ᶜY.ρ)
483-
energy_var = (; e_int = @. lazy(e_tot - K - Φ))
484-
359+
energy_var = (; e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ)
485360
moisture_var = if moisture_model isa DryModel
486361
(;)
487362
elseif moisture_model isa EquilMoistModel
488-
q_tot = ᶜspecific(ᶜY.ρq_tot, ᶜY.ρ)
489-
(; q_tot)
363+
(; q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ))
490364
elseif moisture_model isa NonEquilMoistModel
491-
q_tot = ᶜspecific(ᶜY.ρq_tot, ᶜY.ρ)
492-
q_liq = ᶜspecific(ᶜY.ρq_liq, ᶜY.ρ)
493-
q_ice = ᶜspecific(ᶜY.ρq_ice, ᶜY.ρ)
494-
q_rai = ᶜspecific(ᶜY.ρq_rai, ᶜY.ρ)
495-
q_sno = ᶜspecific(ᶜY.ρq_sno, ᶜY.ρ)
496-
(;
497-
q_pt = @. lazy(
498-
TD.PhasePartition(q_tot, q_liq + q_rai, q_ice + q_sno),
499-
)
365+
q_pt_args = (;
366+
q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ),
367+
q_liq = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ),
368+
q_ice = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ),
500369
)
370+
(; q_pt = TD.PhasePartition(q_pt_args...))
501371
end
502372
return (; energy_var..., moisture_var...)
503373
end
504374

505-
ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
506-
ᶜthermo_state(
375+
ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
376+
thermo_state(
507377
thermo_params;
508378
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
509379
ρ,
510380
)
511381

512-
ᶜts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
513-
ᶜthermo_state(
382+
ts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
383+
thermo_state(
514384
thermo_params;
515385
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
516386
p,
@@ -575,7 +445,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
575445
# @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
576446
# TODO: We should think more about these increments before we use them.
577447
end
578-
ᶜts .= ᶜts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
448+
@. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
579449
@. ᶜp = TD.air_pressure(thermo_params, ᶜts)
580450

581451
if turbconv_model isa PrognosticEDMFX

0 commit comments

Comments
 (0)