Skip to content

Commit f0645ec

Browse files
charleskawczynskicharleskawczynski
authored
Merge pull request #3613 from CliMA/ck/lazy_func
Use lazy broadcast function instead of macro
2 parents ab00b3b + aff47d7 commit f0645ec

File tree

11 files changed

+106
-92
lines changed

11 files changed

+106
-92
lines changed

‎src/ClimaAtmos.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module ClimaAtmos
33
using NVTX
44
import NullBroadcasts: NullBroadcasted
55
import LazyBroadcast
6-
import LazyBroadcast: @lazy
6+
import LazyBroadcast: lazy
77
import Thermodynamics as TD
88
import Thermodynamics
99

‎src/cache/eddy_diffusivity_coefficient.jl

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,8 @@ function compute_eddy_diffusivity_coefficient(
44
)
55
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
66
ᶠz_sfc = Fields.level(ᶠz, Fields.half)
7-
return @lazy @. eddy_diffusivity_coefficient_H(
8-
vert_diff.D₀,
9-
vert_diff.H,
10-
ᶠz_sfc,
11-
ᶜz,
7+
return @. lazy(
8+
eddy_diffusivity_coefficient_H(vert_diff.D₀, vert_diff.H, ᶠz_sfc, ᶜz),
129
)
1310
end
1411

@@ -19,10 +16,12 @@ function compute_eddy_diffusivity_coefficient(
1916
)
2017
interior_uₕ = Fields.level(ᶜuₕ, 1)
2118
ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
22-
return @lazy @. eddy_diffusivity_coefficient(
23-
vert_diff.C_E,
24-
norm(interior_uₕ),
25-
ᶜΔz_surface / 2,
26-
ᶜp,
19+
return @. lazy(
20+
eddy_diffusivity_coefficient(
21+
vert_diff.C_E,
22+
norm(interior_uₕ),
23+
ᶜΔz_surface / 2,
24+
ᶜp,
25+
),
2726
)
2827
end

‎src/cache/precomputed_quantities.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ end
184184
# Interpolates the third contravariant component of Y.c.uₕ to cell faces.
185185
function compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
186186
ᶜJ = Fields.local_geometry_field(ᶜρ).J
187-
return @lazy @. ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ))
187+
return @. lazy(ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ)))
188188
end
189189

190190
"""
@@ -212,7 +212,7 @@ function surface_velocity(ᶠu₃, ᶠuₕ³)
212212
sfc_u₃ = Fields.level(ᶠu₃.components.data.:1, half)
213213
sfc_uₕ³ = Fields.level(ᶠuₕ³.components.data.:1, half)
214214
sfc_g³³ = g³³_field(sfc_u₃)
215-
return @lazy @. -sfc_uₕ³ / sfc_g³³ # u³ = uₕ³ + w³ = uₕ³ + w₃ * g³³
215+
return @. lazy(-sfc_uₕ³ / sfc_g³³) # u³ = uₕ³ + w³ = uₕ³ + w₃ * g³³
216216
end
217217

218218
"""

‎src/parameterized_tendencies/radiation/held_suarez.jl

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -142,16 +142,17 @@ function held_suarez_forcing_tendency_ρe_tot(
142142
MSLP,
143143
)
144144

145-
return @lazy @. -compute_ΔρT(
146-
thermo_params,
147-
ts_surf,
148-
ᶜρ,
149-
ᶜp,
150-
lat,
151-
z_surface,
152-
hs_params,
153-
) * cv_d
154-
return nothing
145+
return @. lazy(
146+
-compute_ΔρT(
147+
thermo_params,
148+
ts_surf,
149+
ᶜρ,
150+
ᶜp,
151+
lat,
152+
z_surface,
153+
hs_params,
154+
) * cv_d,
155+
)
155156
end
156157

157158
function held_suarez_forcing_tendency_uₕ(
@@ -198,7 +199,10 @@ function held_suarez_forcing_tendency_uₕ(
198199
MSLP,
199200
)
200201

201-
return @lazy @. -(
202-
k_f * height_factor(thermo_params, z_surface, ᶜp, ts_surf, hs_params)
203-
) * ᶜuₕ
202+
return @. lazy(
203+
-(
204+
k_f *
205+
height_factor(thermo_params, z_surface, ᶜp, ts_surf, hs_params)
206+
) * ᶜuₕ,
207+
)
204208
end

‎src/parameterized_tendencies/sponge/rayleigh_sponge.jl

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
##### Rayleigh sponge
33
#####
44

5-
import LazyBroadcast: @lazy
65
import ClimaCore.Fields as Fields
76

87
αₘ(s::RayleighSponge{FT}, z, α) where {FT} = ifelse(z > s.zd, α, FT(0))
@@ -17,5 +16,5 @@ function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s)
1716
s isa Nothing && return NullBroadcasted()
1817
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
1918
zmax = z_max(axes(ᶠz))
20-
return @lazy @. -β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ
19+
return @. lazy(-β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ)
2120
end

‎src/parameterized_tendencies/sponge/viscous_sponge.jl

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
##### Viscous sponge
33
#####
44

5-
import LazyBroadcast: @lazy
65
import ClimaCore.Fields as Fields
76
import ClimaCore.Geometry as Geometry
87
import ClimaCore.Spaces as Spaces
@@ -17,31 +16,33 @@ function viscous_sponge_tendency_uₕ(ᶜuₕ, s)
1716
s isa Nothing && return NullBroadcasted()
1817
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
1918
zmax = z_max(axes(ᶠz))
20-
return @lazy @. β_viscous(s, ᶜz, zmax) * (
21-
wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
22-
Geometry.Covariant12Axis(),
23-
wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
24-
)
19+
return @. lazy(
20+
β_viscous(s, ᶜz, zmax) * (
21+
wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
22+
Geometry.Covariant12Axis(),
23+
wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
24+
)
25+
),
2526
)
2627
end
2728

2829
function viscous_sponge_tendency_u₃(u₃, s)
2930
s isa Nothing && return NullBroadcasted()
3031
(; ᶠz) = z_coordinate_fields(axes(u₃))
3132
zmax = z_max(axes(ᶠz))
32-
return @lazy @. β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(u₃.components.data.:1))
33+
return @. lazy(β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(u₃.components.data.:1)))
3334
end
3435

3536
function viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, s)
3637
s isa Nothing && return NullBroadcasted()
3738
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
3839
zmax = z_max(axes(ᶠz))
39-
return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
40+
return @. lazy(β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜh_tot)))
4041
end
4142

4243
function viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, s)
4344
s isa Nothing && return NullBroadcasted()
4445
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
4546
zmax = z_max(axes(ᶠz))
46-
return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜχ))
47+
return @. lazy(β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜχ)))
4748
end

‎src/prognostic_equations/edmf_coriolis.jl

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,13 @@ function edmf_coriolis_tendency_uₕ(ᶜuₕ, edmf_coriolis)
1313
(; coriolis_param) = edmf_coriolis
1414
coords = Fields.coordinate_field(ᶜspace)
1515
coriolis_fn(coord) = Geometry.WVector(coriolis_param)
16-
ᶜf_coriolis = @lazy @. Geometry.Contravariant3Vector(coriolis_fn(coords))
17-
ᶜuₕ_g = @lazy @. Geometry.Covariant12Vector(
18-
Geometry.UVVector(prof_ug(ᶜz), prof_vg(ᶜz)),
16+
ᶜf_coriolis = @. lazy(Geometry.Contravariant3Vector(coriolis_fn(coords)))
17+
ᶜuₕ_g = @. lazy(
18+
Geometry.Covariant12Vector(Geometry.UVVector(prof_ug(ᶜz), prof_vg(ᶜz))),
1919
)
2020

2121
# Coriolis
2222
C123 = Geometry.Covariant123Vector
2323
C12 = Geometry.Contravariant12Vector
24-
return @lazy @. - ᶜf_coriolis × (C12(C123(ᶜuₕ)) - C12(C123(ᶜuₕ_g)))
24+
return @. lazy(- ᶜf_coriolis × (C12(C123(ᶜuₕ)) - C12(C123(ᶜuₕ_g))))
2525
end

‎src/prognostic_equations/forcing/large_scale_advection.jl

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ function large_scale_advection_tendency_ρq_tot(
1616
ls_adv isa LargeScaleAdvection || return NullBroadcasted()
1717
(; prof_dTdt, prof_dqtdt) = ls_adv
1818
ᶜz = Fields.coordinate_field(axes(ᶜρ)).z
19-
ᶜdqtdt_hadv = @lazy @. prof_dqtdt(thermo_params, ᶜts, t, ᶜz)
20-
return @lazy @. ᶜρ * ᶜdqtdt_hadv
19+
ᶜdqtdt_hadv = @. lazy(prof_dqtdt(thermo_params, ᶜts, t, ᶜz))
20+
return @. lazy(ᶜρ * ᶜdqtdt_hadv)
2121
end
2222

2323
function large_scale_advection_tendency_ρe_tot(
@@ -30,20 +30,21 @@ function large_scale_advection_tendency_ρe_tot(
3030
ls_adv isa LargeScaleAdvection || return NullBroadcasted()
3131
(; prof_dTdt, prof_dqtdt) = ls_adv
3232
z = Fields.coordinate_field(axes(ᶜρ)).z
33-
ᶜdTdt_hadv = @lazy @. prof_dTdt(thermo_params, ᶜts, t, z)
34-
ᶜdqtdt_hadv = @lazy @. prof_dqtdt(thermo_params, ᶜts, t, z)
33+
ᶜdTdt_hadv = @. lazy(prof_dTdt(thermo_params, ᶜts, t, z))
34+
ᶜdqtdt_hadv = @. lazy(prof_dqtdt(thermo_params, ᶜts, t, z))
3535
T_0 = TD.Parameters.T_0(thermo_params)
3636
Lv_0 = TD.Parameters.LH_v0(thermo_params)
3737
cv_v = TD.Parameters.cv_v(thermo_params)
3838
R_v = TD.Parameters.R_v(thermo_params)
3939
# TODO: should `hv` be a thermo function?
4040
# (hv = cv_v * (ᶜT - T_0) + Lv_0 - R_v * T_0)
41-
return @lazy @. ᶜρ * (
42-
TD.cv_m(thermo_params, ᶜts) * ᶜdTdt_hadv +
43-
(
44-
cv_v * (TD.air_temperature(thermo_params, ᶜts) - T_0) + Lv_0 -
45-
R_v * T_0
46-
) * ᶜdqtdt_hadv
41+
return @. lazy(
42+
ᶜρ * (
43+
TD.cv_m(thermo_params, ᶜts) * ᶜdTdt_hadv +
44+
(
45+
cv_v * (TD.air_temperature(thermo_params, ᶜts) - T_0) + Lv_0 -
46+
R_v * T_0
47+
) * ᶜdqtdt_hadv
48+
),
4749
)
48-
return nothing
4950
end

‎src/prognostic_equations/implicit/implicit_tendency.jl

Lines changed: 32 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
##### Implicit tendencies
33
#####
44

5-
import LazyBroadcast: @lazy
65
import ClimaCore
76
import ClimaCore: Fields, Geometry
87

@@ -50,54 +49,60 @@ end
5049

5150
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:none})
5251
ᶜJ = Fields.local_geometry_field(ᶜρ).J
53-
return @lazy @. -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³ * ᶠinterp(ᶜχ)))
52+
return @. lazy(-(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠu³ * ᶠinterp(ᶜχ))))
5453
end
5554
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:first_order})
5655
ᶜJ = Fields.local_geometry_field(ᶜρ).J
57-
return @lazy @. -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ)))
56+
return @. lazy(-(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ))))
5857
end
5958
@static if pkgversion(ClimaCore) ≥ v"0.14.22"
6059
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:vanleer_limiter})
6160
ᶜJ = Fields.local_geometry_field(ᶜρ).J
62-
return @lazy @. -(ᶜadvdivᵥ(
63-
ᶠwinterp(ᶜJ, ᶜρ) * ᶠlin_vanleer(ᶠu³, ᶜχ, dt),
64-
))
61+
return @. lazy(
62+
-(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠlin_vanleer(ᶠu³, ᶜχ, dt))),
63+
)
6564
end
6665
end
6766
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:third_order})
6867
ᶜJ = Fields.local_geometry_field(ᶜρ).J
69-
return @lazy @. -(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind3(ᶠu³, ᶜχ)))
68+
return @. lazy(-(ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind3(ᶠu³, ᶜχ))))
7069
end
7170
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:boris_book})
7271
ᶜJ = Fields.local_geometry_field(ᶜρ).J
73-
return @lazy @. -(ᶜadvdivᵥ(
74-
ᶠwinterp(ᶜJ, ᶜρ) * (
75-
ᶠupwind1(ᶠu³, ᶜχ) + ᶠfct_boris_book(
76-
ᶠupwind3(ᶠu³, ᶜχ) - ᶠupwind1(ᶠu³, ᶜχ),
77-
ᶜχ / dt - ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ)) / ᶜρ,
78-
)
79-
),
80-
))
72+
return @. lazy(
73+
-(ᶜadvdivᵥ(
74+
ᶠwinterp(ᶜJ, ᶜρ) * (
75+
ᶠupwind1(ᶠu³, ᶜχ) + ᶠfct_boris_book(
76+
ᶠupwind3(ᶠu³, ᶜχ) - ᶠupwind1(ᶠu³, ᶜχ),
77+
ᶜχ / dt -
78+
ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ)) / ᶜρ,
79+
)
80+
),
81+
)),
82+
)
8183
end
8284
function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:zalesak})
8385
ᶜJ = Fields.local_geometry_field(ᶜρ).J
84-
return @lazy @. -(ᶜadvdivᵥ(
85-
ᶠwinterp(ᶜJ, ᶜρ) * (
86-
ᶠupwind1(ᶠu³, ᶜχ) + ᶠfct_zalesak(
87-
ᶠupwind3(ᶠu³, ᶜχ) - ᶠupwind1(ᶠu³, ᶜχ),
88-
ᶜχ / dt,
89-
ᶜχ / dt - ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ)) / ᶜρ,
90-
)
91-
),
92-
))
86+
return @. lazy(
87+
-(ᶜadvdivᵥ(
88+
ᶠwinterp(ᶜJ, ᶜρ) * (
89+
ᶠupwind1(ᶠu³, ᶜχ) + ᶠfct_zalesak(
90+
ᶠupwind3(ᶠu³, ᶜχ) - ᶠupwind1(ᶠu³, ᶜχ),
91+
ᶜχ / dt,
92+
ᶜχ / dt -
93+
ᶜadvdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠupwind1(ᶠu³, ᶜχ)) / ᶜρ,
94+
)
95+
),
96+
)),
97+
)
9398
end
9499

95100
vertical_advection(ᶠu³, ᶜχ, ::Val{:none}) =
96-
@lazy @. -(ᶜadvdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³))
101+
@. lazy(-(ᶜadvdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³)))
97102
vertical_advection(ᶠu³, ᶜχ, ::Val{:first_order}) =
98-
@lazy @. -(ᶜadvdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³))
103+
@. lazy(-(ᶜadvdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³)))
99104
vertical_advection(ᶠu³, ᶜχ, ::Val{:third_order}) =
100-
@lazy @. -(ᶜadvdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³))
105+
@. lazy(-(ᶜadvdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³)))
101106

102107
function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
103108
(; moisture_model, turbconv_model, rayleigh_sponge, precip_model) = p.atmos

‎src/prognostic_equations/remaining_tendency.jl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
2323
return Yₜ
2424
end
2525

26-
import LazyBroadcast: @lazy
2726
import ClimaCore.Fields as Fields
2827
import ClimaCore.Geometry as Geometry
2928
import ClimaCore.Spaces as Spaces

0 commit comments

Comments
 (0)