Skip to content

Commit b3f3351

Browse files
Use columnwise for setting tendencies to zero
Update src/utils/utilities.jl Co-authored-by: Gregory L. Wagner <gregory.leclaire.wagner@gmail.com> Update src/utils/utilities.jl Co-authored-by: Gregory L. Wagner <gregory.leclaire.wagner@gmail.com>
1 parent 3443adb commit b3f3351

File tree

2 files changed

+196
-1
lines changed

2 files changed

+196
-1
lines changed

src/prognostic_equations/remaining_tendency.jl

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,115 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
3737
apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
3838
end
3939

40+
using ClimaCore.RecursiveApply: rzero
41+
42+
#####
43+
##### Cell center tendencies
44+
#####
45+
46+
"""
47+
ᶜremaining_tendency(ᶜY, ᶠY, p, t)
48+
49+
Returns a Broadcasted object, for evaluating the cell center remaining
50+
tendency. This method calls `ᶜremaining_tendency(Val(name), ᶜY, ᶠY, p, t)` for
51+
all `propertynames` of `ᶜY`.
52+
"""
53+
function ᶜremaining_tendency(ᶜY, ᶠY, p, t)
54+
names = propertynames(ᶜY)
55+
tends = construct_tendencies(Val(names), ᶜremaining_tendency, ᶜY, ᶠY, p, t)
56+
# We cannot broadcast over a NamedTuple, so we need to check that edge case
57+
# first.
58+
if all(t -> !(t isa Base.Broadcast.Broadcasted), tends)
59+
return make_named_tuple(Val(names), tends...)
60+
else
61+
return lazy.(make_named_tuple.(Val(names), tends...))
62+
end
63+
end
64+
65+
#####
66+
##### Cell face tendencies
67+
#####
68+
69+
"""
70+
ᶠremaining_tendency(ᶜY, ᶠY, p, t)
71+
72+
Returns a Broadcasted object, for evaluating the cell center remaining
73+
tendency. This method calls `ᶠremaining_tendency(Val(name), ᶜY, ᶠY, p, t)` for
74+
all `propertynames` of `ᶠY`.
75+
"""
76+
function ᶠremaining_tendency(ᶜY, ᶠY, p, t)
77+
names = propertynames(ᶠY)
78+
tends = construct_tendencies(Val(names), ᶠremaining_tendency, ᶜY, ᶠY, p, t)
79+
# We cannot broadcast over a NamedTuple, so we need to check that edge case
80+
# first.
81+
if all(t -> !(t isa Base.Broadcast.Broadcasted), tends)
82+
return make_named_tuple(Val(names), tends...)
83+
else
84+
return lazy.(make_named_tuple.(Val(names), tends...))
85+
end
86+
end
87+
88+
#####
89+
##### Individual tendencies
90+
#####
91+
92+
function ᶜremaining_tendency(::Val{:ρ}, ᶜY, ᶠY, p, t)
93+
∑tendencies = zero(eltype(ᶜY.ρ))
94+
return ∑tendencies
95+
end
96+
function ᶜremaining_tendency(::Val{:uₕ}, ᶜY, ᶠY, p, t)
97+
∑tendencies = zero(eltype(ᶜY.uₕ))
98+
return ∑tendencies
99+
end
100+
function ᶜremaining_tendency(::Val{:ρe_tot}, ᶜY, ᶠY, p, t)
101+
∑tendencies = zero(eltype(ᶜY.ρe_tot))
102+
return ∑tendencies
103+
end
104+
function ᶜremaining_tendency(::Val{:ρq_tot}, ᶜY, ᶠY, p, t)
105+
∑tendencies = zero(eltype(ᶜY.ρq_tot))
106+
return ∑tendencies
107+
end
108+
function ᶜremaining_tendency(::Val{:ρq_liq}, ᶜY, ᶠY, p, t)
109+
∑tendencies = zero(eltype(ᶜY.ρq_liq))
110+
return ∑tendencies
111+
end
112+
function ᶜremaining_tendency(::Val{:ρq_ice}, ᶜY, ᶠY, p, t)
113+
∑tendencies = zero(eltype(ᶜY.ρq_ice))
114+
return ∑tendencies
115+
end
116+
function ᶜremaining_tendency(::Val{:ρn_liq}, ᶜY, ᶠY, p, t)
117+
∑tendencies = zero(eltype(ᶜY.ρn_liq))
118+
return ∑tendencies
119+
end
120+
function ᶜremaining_tendency(::Val{:ρn_rai}, ᶜY, ᶠY, p, t)
121+
∑tendencies = zero(eltype(ᶜY.ρn_rai))
122+
return ∑tendencies
123+
end
124+
function ᶜremaining_tendency(::Val{:ρq_rai}, ᶜY, ᶠY, p, t)
125+
∑tendencies = zero(eltype(ᶜY.ρq_rai))
126+
return ∑tendencies
127+
end
128+
function ᶜremaining_tendency(::Val{:ρq_sno}, ᶜY, ᶠY, p, t)
129+
∑tendencies = zero(eltype(ᶜY.ρq_sno))
130+
return ∑tendencies
131+
end
132+
function ᶜremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t)
133+
∑tendencies = rzero(eltype(ᶜY.sgsʲs))
134+
return ∑tendencies
135+
end
136+
function ᶜremaining_tendency(::Val{:sgs⁰}, ᶜY, ᶠY, p, t)
137+
∑tendencies = rzero(eltype(ᶜY.sgs⁰))
138+
return ∑tendencies
139+
end
140+
function ᶠremaining_tendency(::Val{:u₃}, ᶜY, ᶠY, p, t)
141+
∑tendencies = zero(eltype(ᶠY.u₃))
142+
return ∑tendencies
143+
end
144+
function ᶠremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t)
145+
∑tendencies = rzero(eltype(ᶠY.sgsʲs))
146+
return ∑tendencies
147+
end
148+
40149
"""
41150
remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
42151
@@ -64,7 +173,27 @@ Returns:
64173
"""
65174
NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
66175
Yₜ_lim .= zero(eltype(Yₜ_lim))
67-
Yₜ .= zero(eltype(Yₜ))
176+
device = ClimaComms.device(axes(Y.c))
177+
p_kernel = (;
178+
zmax = Spaces.z_max(axes(Y.f)),
179+
atmos = p.atmos,
180+
params = p.params,
181+
dt = p.dt,
182+
)
183+
if :sfc in propertynames(Y) # columnwise! does not yet handle .sfc
184+
parent(Yₜ.sfc) .= zero(Spaces.undertype(axes(Y.c)))
185+
end
186+
Operators.columnwise!(
187+
device,
188+
ᶜremaining_tendency,
189+
ᶠremaining_tendency,
190+
Yₜ.c,
191+
Yₜ.f,
192+
Y.c,
193+
Y.f,
194+
p_kernel,
195+
t,
196+
)
68197
horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)
69198
fill_with_nans!(p) # TODO: would be better to limit this to debug mode (e.g., if p.debug_mode...)
70199
horizontal_dynamics_tendency!(Yₜ, Y, p, t)

src/utils/utilities.jl

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,3 +496,69 @@ function issphere(space)
496496
return Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))) isa
497497
Domains.SphereDomain
498498
end
499+
500+
501+
"""
502+
construct_tendencies(::Val{names}, f, ᶜY, ᶠY, p, t)
503+
504+
Return a tuple of calls to
505+
506+
```
507+
f(Val(name), ᶜY, ᶠY, p, t)
508+
```
509+
for all names in `names`.
510+
511+
For example, `f(Val((:a, :b)), ᶜY, ᶠY, p, t)` will return:
512+
513+
```
514+
(
515+
f(Val(:a), ᶜY, ᶠY, p, t),
516+
f(Val(:b), ᶜY, ᶠY, p, t),
517+
)
518+
```
519+
"""
520+
@generated function construct_tendencies(
521+
::Val{names},
522+
f,
523+
ᶜY,
524+
ᶠY,
525+
p,
526+
t,
527+
) where {names}
528+
calls = []
529+
for name in names
530+
push!(calls, :(f(Val($(QuoteNode(name))), ᶜY, ᶠY, p, t)))
531+
end
532+
return quote
533+
($(calls...),)
534+
end
535+
end
536+
537+
"""
538+
make_named_tuple(::Val{names}, vals...) where {names}
539+
540+
Construct a NamedTuple given the names `names` and values `vals`.
541+
"""
542+
make_named_tuple(::Val{names}, vals...) where {names} = NamedTuple{names}(vals)
543+
544+
"""
545+
add_tendency(∑tendencies, tendency)
546+
547+
A helper function which returns `∑tendencies` when `tendency` is a
548+
`NullBroadcasted` and `lazy.(∑tendencies + tendency)` when `tendency` is not a
549+
`NullBroadcasted`.
550+
"""
551+
function add_tendency end
552+
add_tendency(∑tends, t) = lazy.(∑tends .+ t)
553+
add_tendency(∑tends, ::NullBroadcasted) = ∑tends
554+
555+
"""
556+
subtract_tendency(∑tendencies, tendency)
557+
558+
A helper function which returns `∑tendencies` when `tendency` is a
559+
`NullBroadcasted` and `lazy.(∑tendencies - tendency)` when `tendency` is not a
560+
`NullBroadcasted`.
561+
"""
562+
function subtract_tendency end
563+
subtract_tendency(∑tends, ::NullBroadcasted) = ∑tends
564+
subtract_tendency(∑tends, t) = lazy.(∑tends .- t)

0 commit comments

Comments
 (0)