Skip to content

Commit 36e7a7e

Browse files
costachriscostachris
committed
remove more c funcs
1 parent 44261d8 commit 36e7a7e

13 files changed

+98
-70
lines changed

‎src/cache/cloud_fraction.jl

Lines changed: 4 additions & 4 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, p)
193+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
194194

195195
# TODO - we should make this default when using diagnostic edmf
196196
# environment
@@ -209,9 +209,9 @@ NVTX.@annotate function set_cloud_fraction!(
209209
# weight cloud diagnostics by environmental area
210210
@. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}(
211211
tuple(
212-
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
213-
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
214-
draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
212+
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
213+
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
214+
draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
215215
),
216216
)
217217
# updrafts

‎src/cache/precipitation_precomputed_quantities.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -257,14 +257,14 @@ function set_precipitation_cache!(
257257
(; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
258258
(; ᶜts⁰, ᶜtsʲs) = p.precomputed
259259
thermo_params = CAP.thermodynamics_params(p.params)
260-
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
260+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, p.atmos.turbconv_model))
261261

262262
n = n_mass_flux_subdomains(p.atmos.turbconv_model)
263263

264-
@. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰_vals
264+
@. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰
265265
@. ᶜS_ρe_tot =
266266
ᶜSqₜᵖ⁰ *
267-
ᶜρa⁰_vals *
267+
ᶜρa⁰ *
268268
e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ)
269269
for j in 1:n
270270
@. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa

‎src/cache/precomputed_quantities.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -319,9 +319,9 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
319319
return nothing
320320
end
321321

322-
function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰_vals, ᶠu₃⁰, turbconv_model)
322+
function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model)
323323
@. ᶜK +=
324-
ᶜρa⁰_vals * ᶜinterp(dot(ᶠu₃⁰ - Y.f.u₃, CT3(ᶠu₃⁰ - Y.f.u₃))) / 2 / Y.c.ρ
324+
ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Y.f.u₃, CT3(ᶠu₃⁰ - Y.f.u₃))) / 2 / Y.c.ρ
325325
for j in 1:n_mass_flux_subdomains(turbconv_model)
326326
ᶜρaʲ = Y.c.sgsʲs.:($j).ρa
327327
ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃

‎src/cache/prognostic_edmf_precomputed_quantities.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
362362
ᶜlg = Fields.local_geometry_field(Y.c)
363363
ᶠlg = Fields.local_geometry_field(Y.f)
364364
ᶜtke⁰ = ᶜspecific_tke(Y, p)
365-
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
365+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
366366

367367
ᶜvert_div = p.scratch.ᶜtemp_scalar
368368
ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
@@ -480,7 +480,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
480480
for j in 1:n
481481
ᶠu³ʲ = ᶠu³ʲs.:($j)
482482
@. ᶜtke_exch +=
483-
Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰_vals *
483+
Y.c.sgsʲs.:($$j).ρa * ᶜdetrʲs.:($$j) / ᶜρa⁰ *
484484
(1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
485485
end
486486

@@ -507,7 +507,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
507507
@. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
508508

509509
ρatke_flux_values = Fields.field_values(ρatke_flux)
510-
ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰_vals, 1)) # TODO: replace by surface value
510+
ρa_sfc_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) # TODO: replace by surface value
511511
ustar_values = Fields.field_values(ustar)
512512
sfc_local_geometry_values = Fields.field_values(
513513
Fields.level(Fields.local_geometry_field(Y.f), half),

‎src/diagnostics/Diagnostics.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ import ..SlabOceanSST
5858
import ..draft_area
5959
import ..compute_gm_mixing_length!
6060
import ..horizontal_integral_at_boundary
61-
import ..ᶜρa⁰
61+
import ..ρa⁰
6262
import ..ᶜspecific_tke
6363
import ..ᶜspecific_env_value
6464

‎src/diagnostics/edmfx_diagnostics.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -633,16 +633,16 @@ 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, cache)
636+
ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model))
637637
if isnothing(out)
638638
return draft_area.(
639-
ᶜρa⁰_vals,
639+
ᶜρa⁰,
640640
TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰),
641641
)
642642
else
643643
out .=
644644
draft_area.(
645-
ᶜρa⁰_vals,
645+
ᶜρa⁰,
646646
TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰),
647647
)
648648
end

‎src/parameterized_tendencies/microphysics/precipitation.jl

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

153-
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
153+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
154154

155155
# Update from environment precipitation and cloud formation sources/sinks
156-
@. Yₜ.c.ρq_liq += ᶜρa⁰_vals * ᶜSqₗᵖ⁰
157-
@. Yₜ.c.ρq_ice += ᶜρa⁰_vals * ᶜSqᵢᵖ⁰
158-
@. Yₜ.c.ρq_rai += ᶜρa⁰_vals * ᶜSqᵣᵖ⁰
159-
@. Yₜ.c.ρq_sno += ᶜρa⁰_vals * ᶜSqₛᵖ⁰
156+
@. Yₜ.c.ρq_liq += ᶜρa⁰ * ᶜSqₗᵖ⁰
157+
@. Yₜ.c.ρq_ice += ᶜρa⁰ * ᶜSqᵢᵖ⁰
158+
@. Yₜ.c.ρq_rai += ᶜρa⁰ * ᶜSqᵣᵖ⁰
159+
@. Yₜ.c.ρq_sno += ᶜρa⁰ * ᶜSqₛᵖ⁰
160160

161161
# Update from the updraft precipitation sources
162162
n = n_mass_flux_subdomains(p.atmos.turbconv_model)

‎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, p)) : @. lazy(Y.c.ρ)) : nothing
197+
ᶜρa⁰ = advect_tke ? (n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : @. lazy(Y.c.ρ)) : nothing
198198
ᶜρ⁰ = if advect_tke
199199
if n > 0
200200
(; ᶜts⁰) = p.precomputed
@@ -284,7 +284,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
284284
end
285285

286286
if use_prognostic_tke(turbconv_model) # advect_tke triggers allocations
287-
@. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰_vals, ᶜρ⁰)
287+
@. ᶜa_scalar = ᶜtke⁰ * draft_area(ᶜρa⁰, ᶜρ⁰)
288288
vtt = vertical_transport(ᶜρ⁰, ᶠu³⁰, ᶜa_scalar, dt, edmfx_upwinding)
289289
@. Yₜ.c.sgs⁰.ρatke += vtt
290290
end

‎src/prognostic_equations/edmfx_sgs_flux.jl

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ function edmfx_sgs_mass_flux_tendency!(
4545
(; ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜts) = p.precomputed
4646
thermo_params = CAP.thermodynamics_params(p.params)
4747
ᶜρ⁰ = @. lazy(TD.air_density(thermo_params, ᶜts⁰))
48-
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
48+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
4949
(; dt) = p
5050
ᶜJ = Fields.local_geometry_field(Y.c).J
5151

@@ -81,7 +81,7 @@ function edmfx_sgs_mass_flux_tendency!(
8181
@. ᶠu³_diff = ᶠu³⁰ - ᶠu³
8282

8383
ᶜmse⁰ = ᶜspecific_env_mse(Y, p)
84-
@. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰_vals, ᶜρ⁰)
84+
@. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰)
8585
vtt = vertical_transport(
8686
ᶜρ⁰,
8787
ᶠu³_diff,
@@ -112,7 +112,7 @@ function edmfx_sgs_mass_flux_tendency!(
112112
@. ᶠu³_diff = ᶠu³⁰ - ᶠu³
113113
@. ᶜa_scalar =
114114
(ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) *
115-
draft_area(ᶜρa⁰_vals, ᶜρ⁰)
115+
draft_area(ᶜρa⁰, ᶜρ⁰)
116116
vtt = vertical_transport(
117117
ᶜρ⁰,
118118
ᶠu³_diff,
@@ -187,7 +187,7 @@ function edmfx_sgs_mass_flux_tendency!(
187187

188188
@. ᶜa_scalar =
189189
(ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) *
190-
draft_area(ᶜρa⁰_vals, ᶜρ⁰)
190+
draft_area(ᶜρa⁰, ᶜρ⁰)
191191
vtt = vertical_transport(
192192
ᶜρ⁰,
193193
ᶠu³_diff,
@@ -199,7 +199,7 @@ function edmfx_sgs_mass_flux_tendency!(
199199

200200
@. ᶜa_scalar =
201201
(ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) *
202-
draft_area(ᶜρa⁰_vals, ᶜρ⁰)
202+
draft_area(ᶜρa⁰, ᶜρ⁰)
203203
vtt = vertical_transport(
204204
ᶜρ⁰,
205205
ᶠu³_diff,
@@ -211,7 +211,7 @@ function edmfx_sgs_mass_flux_tendency!(
211211

212212
@. ᶜa_scalar =
213213
(ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) *
214-
draft_area(ᶜρa⁰_vals, ᶜρ⁰)
214+
draft_area(ᶜρa⁰, ᶜρ⁰)
215215
vtt = vertical_transport(
216216
ᶜρ⁰,
217217
ᶠu³_diff,
@@ -223,7 +223,7 @@ function edmfx_sgs_mass_flux_tendency!(
223223

224224
@. ᶜa_scalar =
225225
(ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) *
226-
draft_area(ᶜρa⁰_vals, ᶜρ⁰)
226+
draft_area(ᶜρa⁰, ᶜρ⁰)
227227
vtt = vertical_transport(
228228
ᶜρ⁰,
229229
ᶠu³_diff,
@@ -404,14 +404,14 @@ function edmfx_sgs_diffusive_flux_tendency!(
404404
(; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
405405
(; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
406406
ᶠgradᵥ = Operators.GradientC2F()
407-
ᶜρa⁰_vals = ᶜρa⁰(Y, p)
407+
ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
408408
ᶜtke⁰ = ᶜspecific_tke(Y, p)
409409

410410
if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
411411
ᶠρaK_h = p.scratch.ᶠtemp_scalar
412-
@. ᶠρaK_h = ᶠinterp(ᶜρa⁰_vals) * ᶠinterp(ᶜK_h)
412+
@. ᶠρaK_h = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_h)
413413
ᶠρaK_u = p.scratch.ᶠtemp_scalar
414-
@. ᶠρaK_u = ᶠinterp(ᶜρa⁰_vals) * ᶠinterp(ᶜK_u)
414+
@. ᶠρaK_u = ᶠinterp(ᶜρa⁰) * ᶠinterp(ᶜK_u)
415415

416416
# Total enthalpy diffusion
417417
ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(

‎src/prognostic_equations/edmfx_tke.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ function edmfx_tke_tendency!(
5050
FT = eltype(p.params)
5151

5252

53-
ᶜρa⁰_vals = turbconv_model isa PrognosticEDMFX ? ᶜρa⁰(Y, p) : Y.c.ρ
53+
ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : @. lazy(Y.c.ρ)
5454
nh_pressure3_buoyʲs =
5555
turbconv_model isa PrognosticEDMFX ?
5656
p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs
@@ -73,9 +73,9 @@ function edmfx_tke_tendency!(
7373

7474
if use_prognostic_tke(turbconv_model)
7575
# shear production
76-
@. Yₜ.c.sgs⁰.ρatke += 2 * ᶜρa⁰_vals * ᶜK_u * ᶜstrain_rate_norm
76+
@. Yₜ.c.sgs⁰.ρatke += 2 * ᶜρa⁰ * ᶜK_u * ᶜstrain_rate_norm
7777
# buoyancy production
78-
@. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰_vals * ᶜK_h * ᶜlinear_buoygrad
78+
@. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad
7979

8080
ᶜtke⁰ = ᶜspecific_tke(Y, p)
8181

0 commit comments

Comments
 (0)