Skip to content

Commit 44261d8

Browse files
costachriscostachris
committed
remove some ᶜfunctions
1 parent b7d9f53 commit 44261d8

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
@@ -62,8 +62,8 @@ NVTX.@annotate function set_cloud_fraction!(
6262
TD.PhasePartition(thermo_params, ᶜts).ice,
6363
)
6464
else
65-
q_liq = ᶜspecific(Y.c.ρq_liq, Y.c.ρ)
66-
q_ice = ᶜspecific(Y.c.ρq_ice, Y.c.ρ)
65+
q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
66+
q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
6767
@. cloud_diagnostics_tuple =
6868
make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice)
6969
end
@@ -304,7 +304,6 @@ function quad_loop(
304304
FT = eltype(x1_hat)
305305
@assert(x1_hat >= FT(0))
306306
@assert(x2_hat >= FT(0))
307-
# note: ᶜthermo_state is used as a pointwise function here
308307
_ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
309308
hc = TD.has_condensate(thermo_params, _ts)
310309

‎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,
@@ -1106,7 +1098,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
11061098
(; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
11071099

11081100
# Environment precipitation sources (to be applied to grid mean)
1109-
ᶜq_tot = ᶜspecific(Y.c.ρq_tot, Y.c.ρ)
1101+
ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
11101102
@. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
11111103
thermo_params,
11121104
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
@@ -374,162 +374,32 @@ function thermo_state(
374374
return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
375375
end
376376

377-
const FieldOrValue = Union{Fields.Field, Base.AbstractBroadcasted, Real}
378-
function ᶜthermo_state(
379-
thermo_params;
380-
ρ = nothing,
381-
p = nothing,
382-
θ = nothing,
383-
e_int = nothing,
384-
q_tot = nothing,
385-
q_pt = nothing,
386-
)
387-
388-
get_ts(
389-
ρ::T,
390-
::Nothing,
391-
θ::T,
392-
::Nothing,
393-
::Nothing,
394-
::Nothing,
395-
) where {T <: FieldOrValue} = TD.PhaseDry_ρθ(thermo_params, ρ, θ)
396-
get_ts(
397-
ρ::T,
398-
::Nothing,
399-
θ::T,
400-
::Nothing,
401-
q_tot::T,
402-
::Nothing,
403-
) where {T <: FieldOrValue} = TD.PhaseEquil_ρθq(thermo_params, ρ, θ, q_tot)
404-
get_ts(
405-
ρ::T,
406-
::Nothing,
407-
θ::T,
408-
::Nothing,
409-
::Nothing,
410-
q_pt,
411-
) where {T <: FieldOrValue} =
412-
TD.PhaseNonEquil_ρθq(thermo_params, ρ, θ, q_pt)
413-
get_ts(
414-
ρ::T,
415-
::Nothing,
416-
::Nothing,
417-
e_int::T,
418-
::Nothing,
419-
::Nothing,
420-
) where {T <: FieldOrValue} = TD.PhaseDry_ρe(thermo_params, ρ, e_int)
421-
get_ts(
422-
ρ::T,
423-
::Nothing,
424-
::Nothing,
425-
e_int::T,
426-
q_tot::T,
427-
::Nothing,
428-
) where {T <: FieldOrValue} = TD.PhaseEquil_ρeq(
429-
thermo_params,
430-
ρ,
431-
e_int,
432-
q_tot,
433-
3,
434-
eltype(thermo_params)(0.003),
435-
)
436-
get_ts(
437-
ρ::T,
438-
::Nothing,
439-
::Nothing,
440-
e_int::T,
441-
::Nothing,
442-
q_pt,
443-
) where {T <: FieldOrValue} =
444-
TD.PhaseNonEquil(thermo_params, e_int, ρ, q_pt)
445-
get_ts(
446-
::Nothing,
447-
p::T,
448-
θ::T,
449-
::Nothing,
450-
::Nothing,
451-
::Nothing,
452-
) where {T <: FieldOrValue} = TD.PhaseDry_pθ(thermo_params, p, θ)
453-
get_ts(
454-
::Nothing,
455-
p::T,
456-
θ::T,
457-
::Nothing,
458-
q_tot::T,
459-
::Nothing,
460-
) where {T <: FieldOrValue} = TD.PhaseEquil_pθq(thermo_params, p, θ, q_tot)
461-
get_ts(
462-
::Nothing,
463-
p::T,
464-
θ::T,
465-
::Nothing,
466-
::Nothing,
467-
q_pt,
468-
) where {T <: FieldOrValue} =
469-
TD.PhaseNonEquil_pθq(thermo_params, p, θ, q_pt)
470-
get_ts(
471-
::Nothing,
472-
p::T,
473-
::Nothing,
474-
e_int::T,
475-
::Nothing,
476-
::Nothing,
477-
) where {T <: FieldOrValue} = TD.PhaseDry_pe(thermo_params, p, e_int)
478-
get_ts(
479-
::Nothing,
480-
p::T,
481-
::Nothing,
482-
e_int::T,
483-
q_tot::T,
484-
::Nothing,
485-
) where {T <: FieldOrValue} =
486-
TD.PhaseEquil_peq(thermo_params, p, e_int, q_tot)
487-
get_ts(
488-
::Nothing,
489-
p::T,
490-
::Nothing,
491-
e_int::T,
492-
::Nothing,
493-
q_pt,
494-
) where {T <: FieldOrValue} =
495-
TD.PhaseNonEquil_peq(thermo_params, p, e_int, q_pt)
496-
497-
return @. lazy(get_ts(ρ, p, θ, e_int, q_tot, q_pt))
498-
end
499-
500377
function thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)
501-
e_tot = ᶜspecific(ᶜY.ρe_tot, ᶜY.ρ)
502-
energy_var = (; e_int = @. lazy(e_tot - K - Φ))
503-
378+
energy_var = (; e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ)
504379
moisture_var = if moisture_model isa DryModel
505380
(;)
506381
elseif moisture_model isa EquilMoistModel
507-
q_tot = ᶜspecific(ᶜY.ρq_tot, ᶜY.ρ)
508-
(; q_tot)
382+
(; q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ))
509383
elseif moisture_model isa NonEquilMoistModel
510-
q_tot = ᶜspecific(ᶜY.ρq_tot, ᶜY.ρ)
511-
q_liq = ᶜspecific(ᶜY.ρq_liq, ᶜY.ρ)
512-
q_ice = ᶜspecific(ᶜY.ρq_ice, ᶜY.ρ)
513-
q_rai = ᶜspecific(ᶜY.ρq_rai, ᶜY.ρ)
514-
q_sno = ᶜspecific(ᶜY.ρq_sno, ᶜY.ρ)
515-
(;
516-
q_pt = @. lazy(
517-
TD.PhasePartition(q_tot, q_liq + q_rai, q_ice + q_sno),
518-
)
384+
q_pt_args = (;
385+
q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ),
386+
q_liq = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ),
387+
q_ice = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ),
519388
)
389+
(; q_pt = TD.PhasePartition(q_pt_args...))
520390
end
521391
return (; energy_var..., moisture_var...)
522392
end
523393

524-
ᶜts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
525-
ᶜthermo_state(
394+
ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
395+
thermo_state(
526396
thermo_params;
527397
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
528398
ρ,
529399
)
530400

531-
ᶜts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
532-
ᶜthermo_state(
401+
ts_sgs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, p) =
402+
thermo_state(
533403
thermo_params;
534404
thermo_vars(moisture_model, microphysics_model, ᶜY, K, Φ)...,
535405
p,
@@ -594,7 +464,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
594464
# @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
595465
# TODO: We should think more about these increments before we use them.
596466
end
597-
ᶜts .= ᶜts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
467+
@. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
598468
@. ᶜp = TD.air_pressure(thermo_params, ᶜts)
599469

600470
if turbconv_model isa PrognosticEDMFX

0 commit comments

Comments
 (0)