Skip to content

Commit 5b0e533

Browse files
Merge pull request #3540 from CliMA/ck/lazy
Use LazyBroadcast for tendencies
2 parents 989a900 + a50cbd5 commit 5b0e533

File tree

9 files changed

+247
-83
lines changed

9 files changed

+247
-83
lines changed

reproducibility_tests/ref_counter.jl

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
208
1+
209
22

33
# **README**
44
#
@@ -20,6 +20,10 @@
2020

2121

2222
#=
23+
209
24+
- Floating point changes, with minor adjustments to how to sponge tendencies are
25+
added.
26+
2327
208
2428
- Update RRTMGP to v0.19.2, which changes cloud optics slightly
2529

src/ClimaAtmos.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ using NVTX
44
import LazyBroadcast
55
import Thermodynamics as TD
66

7+
include("null_broadcasted.jl")
78
include("compat.jl")
89
include(joinpath("parameters", "Parameters.jl"))
910
import .Parameters as CAP
@@ -112,6 +113,7 @@ include(
112113
)
113114
include(joinpath("parameterized_tendencies", "sponge", "rayleigh_sponge.jl"))
114115
include(joinpath("parameterized_tendencies", "sponge", "viscous_sponge.jl"))
116+
include(joinpath("parameterized_tendencies", "sponge", "sponge_tendencies.jl"))
115117
include(
116118
joinpath(
117119
"parameterized_tendencies",

src/null_broadcasted.jl

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# TODO: use https://github.com/CliMA/NullBroadcasts.jl when released
2+
3+
"""
4+
NullBroadcasted()
5+
6+
A `Base.AbstractBroadcasted` that represents arithmetic object.
7+
8+
An `NullBroadcasted()` can be added to, subtracted from, or multiplied by any
9+
value in a broadcast expression without incurring a runtime performance
10+
penalty.
11+
12+
For example, the following rules hold when broadcasting instances of
13+
`NullBroadcasted`:
14+
15+
```
16+
1 + NullBroadcasted() == 1
17+
NullBroadcasted() + 1 == 1
18+
1 - NullBroadcasted() == 1
19+
1 * NullBroadcasted() == NullBroadcasted()
20+
1 / NullBroadcasted() == NullBroadcasted()
21+
```
22+
"""
23+
struct NullBroadcasted <: Base.AbstractBroadcasted end
24+
Base.broadcastable(x::NullBroadcasted) = x
25+
26+
struct NullBroadcastedStyle <: Base.BroadcastStyle end
27+
Base.BroadcastStyle(::Type{<:NullBroadcasted}) = NullBroadcasted()
28+
29+
# Specialize on AbstractArrayStyle to avoid ambiguities with AbstractBroadcasted.
30+
Base.BroadcastStyle(::NullBroadcasted, ::Base.Broadcast.AbstractArrayStyle) =
31+
NullBroadcasted()
32+
Base.BroadcastStyle(::Base.Broadcast.AbstractArrayStyle, ::NullBroadcasted) =
33+
NullBroadcasted()
34+
35+
# Add another method to avoid ambiguity between the previous two.
36+
Base.BroadcastStyle(::NullBroadcasted, ::NullBroadcasted) = NullBroadcasted()
37+
38+
broadcasted_sum(args) =
39+
if isempty(args)
40+
NullBroadcasted()
41+
elseif length(args) == 1
42+
args[1]
43+
else
44+
Base.broadcasted(+, args...)
45+
end
46+
Base.broadcasted(::NullBroadcasted, ::typeof(+), args...) =
47+
broadcasted_sum(filter(arg -> !(arg isa NullBroadcasted), args))
48+
49+
Base.broadcasted(op::typeof(-), ::NullBroadcasted, arg) =
50+
Base.broadcasted(op, arg)
51+
Base.broadcasted(op::typeof(-), arg, ::NullBroadcasted) =
52+
Base.broadcasted(Base.identity, arg)
53+
Base.broadcasted(op::typeof(-), a::NullBroadcasted) = NullBroadcasted()
54+
Base.broadcasted(op::typeof(-), a::NullBroadcasted, ::NullBroadcasted) =
55+
Base.broadcasted(op, a)
56+
57+
Base.broadcasted(op::typeof(+), ::NullBroadcasted, args...) =
58+
Base.broadcasted(op, args...)
59+
Base.broadcasted(op::typeof(+), arg, ::NullBroadcasted, args...) =
60+
Base.broadcasted(op, arg, args...)
61+
Base.broadcasted(
62+
op::typeof(+),
63+
a::NullBroadcasted,
64+
::NullBroadcasted,
65+
args...,
66+
) = Base.broadcasted(op, a, args...)
67+
68+
Base.broadcasted(op::typeof(*), ::NullBroadcasted, args...) = NullBroadcasted()
69+
Base.broadcasted(op::typeof(*), arg, ::NullBroadcasted) = NullBroadcasted()
70+
Base.broadcasted(op::typeof(*), ::NullBroadcasted, ::NullBroadcasted) =
71+
NullBroadcasted()
72+
Base.broadcasted(op::typeof(/), ::NullBroadcasted, args...) = NullBroadcasted()
73+
Base.broadcasted(op::typeof(/), arg, ::NullBroadcasted) = NullBroadcasted()
74+
Base.broadcasted(op::typeof(/), ::NullBroadcasted, ::NullBroadcasted) =
75+
NullBroadcasted()
76+
77+
Base.broadcasted(op::typeof(identity), a::NullBroadcasted) = a
78+
79+
function skip_materialize(dest, bc::Base.Broadcast.Broadcasted)
80+
if typeof(bc.f) <: typeof(+) || typeof(bc.f) <: typeof(-)
81+
if length(bc.args) == 2 &&
82+
bc.args[1] === dest &&
83+
bc.args[2] === Base.Broadcast.Broadcasted(NullBroadcasted, ())
84+
return true
85+
else
86+
return false
87+
end
88+
else
89+
return false
90+
end
91+
end
92+
93+
Base.Broadcast.instantiate(
94+
bc::Base.Broadcast.Broadcasted{NullBroadcastedStyle},
95+
) = x
96+
97+
Base.Broadcast.materialize!(dest, x::NullBroadcasted) =
98+
error("NullBroadcasted objects cannot be materialized.")
99+
Base.Broadcast.materialize(dest, x::NullBroadcasted) =
100+
error("NullBroadcasted objects cannot be materialized.")

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

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

5+
import LazyBroadcast: @lazy
56
import ClimaCore.Fields as Fields
67

7-
rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
8-
98
αₘ(s::RayleighSponge{FT}, z, α) where {FT} = ifelse(z > s.zd, α, FT(0))
109
ζ_rayleigh(s::RayleighSponge{FT}, z, zmax) where {FT} =
1110
sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
@@ -14,8 +13,9 @@ rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
1413
β_rayleigh_w(s::RayleighSponge{FT}, z, zmax) where {FT} =
1514
αₘ(s, z, s.α_w) * ζ_rayleigh(s, z, zmax)
1615

17-
function rayleigh_sponge_tendency!(Yₜ, Y, p, t, s::RayleighSponge)
18-
ᶜz = Fields.coordinate_field(Y.c).z
19-
zmax = z_max(axes(Y.f))
20-
@. Yₜ.c.uₕ -= β_rayleigh_uₕ(s, ᶜz, zmax) * Y.c.uₕ
16+
function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s)
17+
s isa Nothing && return NullBroadcasted()
18+
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
19+
zmax = z_max(axes(ᶠz))
20+
return @lazy @. -β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ
2121
end
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#####
2+
##### Sponge tendencies
3+
#####
4+
5+
import LazyBroadcast: @lazy
6+
import ClimaCore.Fields as Fields
7+
import ClimaCore.Geometry as Geometry
8+
import ClimaCore.Spaces as Spaces
9+
10+
function z_coordinate_fields(space::Spaces.AbstractSpace)
11+
ᶜz = Fields.coordinate_field(Spaces.center_space(space)).z
12+
ᶠz = Fields.coordinate_field(Spaces.face_space(space)).z
13+
return (; ᶜz, ᶠz)
14+
end
15+
16+
17+
function sponge_tendencies!(Yₜ, Y, p, t)
18+
rs, vs = p.atmos.rayleigh_sponge, p.atmos.viscous_sponge
19+
(; ᶜh_tot, ᶜspecific) = p.precomputed
20+
ᶜuₕ = Y.c.uₕ
21+
ᶠu₃ = Yₜ.f.u₃
22+
ᶜρ = Y.c.ρ
23+
vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, vs)
24+
vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, vs)
25+
vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, vs)
26+
rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rs)
27+
28+
# TODO: fuse, once we fix
29+
# https://github.com/CliMA/ClimaCore.jl/issues/2165
30+
@. Yₜ.c.uₕ += vst_uₕ
31+
@. Yₜ.c.uₕ += rst_uₕ
32+
@. Yₜ.f.u₃.components.data.:1 += vst_u₃
33+
@. Yₜ.c.ρe_tot += vst_ρe_tot
34+
35+
# TODO: can we write this out explicitly?
36+
if vs isa ViscousSponge
37+
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
38+
χ_name == :e_tot && continue
39+
vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, vs)
40+
@. ᶜρχₜ += vst_tracer
41+
@. Yₜ.c.ρ += vst_tracer
42+
end
43+
end
44+
end

src/parameterized_tendencies/sponge/viscous_sponge.jl

Lines changed: 29 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,39 +2,46 @@
22
##### Viscous sponge
33
#####
44

5+
import LazyBroadcast: @lazy
56
import ClimaCore.Fields as Fields
67
import ClimaCore.Geometry as Geometry
78
import ClimaCore.Spaces as Spaces
89

9-
10-
viscous_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
11-
1210
αₘ(s::ViscousSponge{FT}, z) where {FT} = ifelse(z > s.zd, s.κ₂, FT(0))
1311
ζ_viscous(s::ViscousSponge{FT}, z, zmax) where {FT} =
1412
sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
1513
β_viscous(s::ViscousSponge{FT}, z, zmax) where {FT} =
1614
αₘ(s, z) * ζ_viscous(s, z, zmax)
1715

18-
function viscous_sponge_tendency!(Yₜ, Y, p, t, s::ViscousSponge)
19-
(; ᶜh_tot, ᶜspecific) = p.precomputed
20-
ᶜuₕ = Y.c.uₕ
21-
ᶜz = Fields.coordinate_field(Y.c).z
22-
ᶠz = Fields.coordinate_field(Y.f).z
16+
function viscous_sponge_tendency_uₕ(ᶜuₕ, s)
17+
s isa Nothing && return NullBroadcasted()
18+
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
2319
zmax = z_max(axes(ᶠz))
24-
@. Yₜ.c.uₕ +=
25-
β_viscous(s, ᶜz, zmax) * (
26-
wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
27-
Geometry.Covariant12Axis(),
28-
wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
29-
)
20+
return @lazy @. β_viscous(s, ᶜz, zmax) * (
21+
wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
22+
Geometry.Covariant12Axis(),
23+
wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
3024
)
31-
@. Yₜ.f.u₃.components.data.:1 +=
32-
β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(Y.f.u₃.components.data.:1))
25+
)
26+
end
27+
28+
function viscous_sponge_tendency_u₃(u₃, s)
29+
s isa Nothing && return NullBroadcasted()
30+
(; ᶠz) = z_coordinate_fields(axes(u₃))
31+
zmax = z_max(axes(ᶠz))
32+
return @lazy @. β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(u₃.components.data.:1))
33+
end
3334

34-
@. Yₜ.c.ρe_tot += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜh_tot))
35-
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
36-
χ_name == :e_tot && continue
37-
@. ᶜρχₜ += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
38-
@. Yₜ.c.ρ += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
39-
end
35+
function viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, s)
36+
s isa Nothing && return NullBroadcasted()
37+
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
38+
zmax = z_max(axes(ᶠz))
39+
return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
40+
end
41+
42+
function viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, s)
43+
s isa Nothing && return NullBroadcasted()
44+
(; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
45+
zmax = z_max(axes(ᶠz))
46+
return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜχ))
4047
end

src/prognostic_equations/remaining_tendency.jl

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,8 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
2323
end
2424

2525
NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
26-
viscous_sponge_tendency!(Yₜ, Y, p, t, p.atmos.viscous_sponge)
27-
26+
sponge_tendencies!(Yₜ, Y, p, t)
2827
# Vertical tendencies
29-
rayleigh_sponge_tendency!(Yₜ, Y, p, t, p.atmos.rayleigh_sponge)
3028
forcing_tendency!(Yₜ, Y, p, t, p.atmos.forcing_type)
3129
subsidence_tendency!(Yₜ, Y, p, t, p.atmos.subsidence)
3230
edmf_coriolis_tendency!(Yₜ, Y, p, t, p.atmos.edmf_coriolis)

test/parameterized_tendencies/sponge/rayleigh_sponge.jl

Lines changed: 32 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,37 +5,43 @@ using Revise; include("test/parameterized_tendencies/sponge/rayleigh_sponge.jl")
55
using ClimaComms
66
ClimaComms.@import_required_backends
77
import ClimaAtmos as CA
8-
import SurfaceFluxes as SF
9-
import ClimaAtmos.Parameters as CAP
10-
import ClimaCore as CC
8+
using ClimaCore.CommonSpaces
9+
using ClimaCore: Spaces, Fields, Geometry, ClimaCore
1110
using Test
11+
using Base.Broadcast: materialize
12+
13+
pkgversion(ClimaCore) < v"0.14.20" && exit() # CommonSpaces
14+
using ClimaCore.CommonSpaces
1215

13-
include("../../test_helpers.jl")
1416
### Common Objects ###
1517
@testset "Rayleigh-sponge functions" begin
16-
### Boilerplate default integrator objects
17-
config = CA.AtmosConfig(
18-
Dict("initial_condition" => "DryBaroclinicWave");
19-
job_id = "sponge1",
18+
FT = Float64
19+
ᶜspace = ExtrudedCubedSphereSpace(
20+
FT;
21+
z_elem = 10,
22+
z_min = 0,
23+
z_max = 1,
24+
radius = 10,
25+
h_elem = 10,
26+
n_quad_points = 4,
27+
staggering = CellCenter(),
2028
)
21-
(; Y) = generate_test_simulation(config)
22-
zmax = maximum(CC.Fields.coordinate_field(Y.f).z)
23-
z = CC.Fields.coordinate_field(Y.c).z
24-
Y.c.uₕ.components.data.:1 .= ones(axes(Y.c))
25-
Y.c.uₕ.components.data.:2 .= ones(axes(Y.c))
26-
FT = eltype(Y)
27-
ᶜYₜ = zero(Y)
29+
ᶠspace = Spaces.face_space(ᶜspace)
30+
ᶠz = Fields.coordinate_field(ᶠspace).z
31+
ᶜz = Fields.coordinate_field(ᶜspace).z
32+
zmax = maximum(ᶠz)
33+
ᶜuₕ = map(z -> zero(Geometry.Covariant12Vector{eltype(z)}), ᶜz)
34+
@. ᶜuₕ.components.data.:1 = 1
35+
@. ᶜuₕ.components.data.:2 = 1
2836
### Component test begins here
2937
rs = CA.RayleighSponge(; zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
30-
@test CA.β_rayleigh_uₕ.(rs, z, zmax) == @. sin(FT(π) / 2 * z / zmax)^2
31-
CA.rayleigh_sponge_tendency!(ᶜYₜ, Y, nothing, FT(0), rs)
32-
# Test that only required tendencies are affected
33-
for (var_name) in filter(x -> (x != :uₕ), propertynames(Y.c))
34-
@test ᶜYₜ.c.:($var_name) == zeros(axes(Y.c))
35-
end
36-
for (var_name) in propertynames(Y.f)
37-
@test ᶜYₜ.f.:($var_name) == zeros(axes(Y.f))
38-
end
39-
@test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
40-
@test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
38+
expected = @. sin(FT(π) / 2 * ᶜz / zmax)^2
39+
computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, rs)
40+
@test CA.β_rayleigh_uₕ.(rs, ᶜz, zmax) == expected
41+
@test materialize(computed) == .-expected .* ᶜuₕ
42+
43+
# Test when not using a Rayleigh sponge.
44+
computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, nothing)
45+
@test computed isa CA.NullBroadcasted
46+
@. ᶜuₕ += computed # test that it can broadcast
4147
end

0 commit comments

Comments
 (0)