-
Notifications
You must be signed in to change notification settings - Fork 20
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #3540 from CliMA/ck/lazy
Use LazyBroadcast for tendencies
- Loading branch information
Showing
9 changed files
with
247 additions
and
83 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| # TODO: use https://github.com/CliMA/NullBroadcasts.jl when released | ||
|
|
||
| """ | ||
| NullBroadcasted() | ||
| A `Base.AbstractBroadcasted` that represents arithmetic object. | ||
| An `NullBroadcasted()` can be added to, subtracted from, or multiplied by any | ||
| value in a broadcast expression without incurring a runtime performance | ||
| penalty. | ||
| For example, the following rules hold when broadcasting instances of | ||
| `NullBroadcasted`: | ||
| ``` | ||
| 1 + NullBroadcasted() == 1 | ||
| NullBroadcasted() + 1 == 1 | ||
| 1 - NullBroadcasted() == 1 | ||
| 1 * NullBroadcasted() == NullBroadcasted() | ||
| 1 / NullBroadcasted() == NullBroadcasted() | ||
| ``` | ||
| """ | ||
| struct NullBroadcasted <: Base.AbstractBroadcasted end | ||
| Base.broadcastable(x::NullBroadcasted) = x | ||
|
|
||
| struct NullBroadcastedStyle <: Base.BroadcastStyle end | ||
| Base.BroadcastStyle(::Type{<:NullBroadcasted}) = NullBroadcasted() | ||
|
|
||
| # Specialize on AbstractArrayStyle to avoid ambiguities with AbstractBroadcasted. | ||
| Base.BroadcastStyle(::NullBroadcasted, ::Base.Broadcast.AbstractArrayStyle) = | ||
| NullBroadcasted() | ||
| Base.BroadcastStyle(::Base.Broadcast.AbstractArrayStyle, ::NullBroadcasted) = | ||
| NullBroadcasted() | ||
|
|
||
| # Add another method to avoid ambiguity between the previous two. | ||
| Base.BroadcastStyle(::NullBroadcasted, ::NullBroadcasted) = NullBroadcasted() | ||
|
|
||
| broadcasted_sum(args) = | ||
| if isempty(args) | ||
| NullBroadcasted() | ||
| elseif length(args) == 1 | ||
| args[1] | ||
| else | ||
| Base.broadcasted(+, args...) | ||
| end | ||
| Base.broadcasted(::NullBroadcasted, ::typeof(+), args...) = | ||
| broadcasted_sum(filter(arg -> !(arg isa NullBroadcasted), args)) | ||
|
|
||
| Base.broadcasted(op::typeof(-), ::NullBroadcasted, arg) = | ||
| Base.broadcasted(op, arg) | ||
| Base.broadcasted(op::typeof(-), arg, ::NullBroadcasted) = | ||
| Base.broadcasted(Base.identity, arg) | ||
| Base.broadcasted(op::typeof(-), a::NullBroadcasted) = NullBroadcasted() | ||
| Base.broadcasted(op::typeof(-), a::NullBroadcasted, ::NullBroadcasted) = | ||
| Base.broadcasted(op, a) | ||
|
|
||
| Base.broadcasted(op::typeof(+), ::NullBroadcasted, args...) = | ||
| Base.broadcasted(op, args...) | ||
| Base.broadcasted(op::typeof(+), arg, ::NullBroadcasted, args...) = | ||
| Base.broadcasted(op, arg, args...) | ||
| Base.broadcasted( | ||
| op::typeof(+), | ||
| a::NullBroadcasted, | ||
| ::NullBroadcasted, | ||
| args..., | ||
| ) = Base.broadcasted(op, a, args...) | ||
|
|
||
| Base.broadcasted(op::typeof(*), ::NullBroadcasted, args...) = NullBroadcasted() | ||
| Base.broadcasted(op::typeof(*), arg, ::NullBroadcasted) = NullBroadcasted() | ||
| Base.broadcasted(op::typeof(*), ::NullBroadcasted, ::NullBroadcasted) = | ||
| NullBroadcasted() | ||
| Base.broadcasted(op::typeof(/), ::NullBroadcasted, args...) = NullBroadcasted() | ||
| Base.broadcasted(op::typeof(/), arg, ::NullBroadcasted) = NullBroadcasted() | ||
| Base.broadcasted(op::typeof(/), ::NullBroadcasted, ::NullBroadcasted) = | ||
| NullBroadcasted() | ||
|
|
||
| Base.broadcasted(op::typeof(identity), a::NullBroadcasted) = a | ||
|
|
||
| function skip_materialize(dest, bc::Base.Broadcast.Broadcasted) | ||
| if typeof(bc.f) <: typeof(+) || typeof(bc.f) <: typeof(-) | ||
| if length(bc.args) == 2 && | ||
| bc.args[1] === dest && | ||
| bc.args[2] === Base.Broadcast.Broadcasted(NullBroadcasted, ()) | ||
| return true | ||
| else | ||
| return false | ||
| end | ||
| else | ||
| return false | ||
| end | ||
| end | ||
|
|
||
| Base.Broadcast.instantiate( | ||
| bc::Base.Broadcast.Broadcasted{NullBroadcastedStyle}, | ||
| ) = x | ||
|
|
||
| Base.Broadcast.materialize!(dest, x::NullBroadcasted) = | ||
| error("NullBroadcasted objects cannot be materialized.") | ||
| Base.Broadcast.materialize(dest, x::NullBroadcasted) = | ||
| error("NullBroadcasted objects cannot be materialized.") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,44 @@ | ||
| ##### | ||
| ##### Sponge tendencies | ||
| ##### | ||
|
|
||
| import LazyBroadcast: @lazy | ||
| import ClimaCore.Fields as Fields | ||
| import ClimaCore.Geometry as Geometry | ||
| import ClimaCore.Spaces as Spaces | ||
|
|
||
| function z_coordinate_fields(space::Spaces.AbstractSpace) | ||
| ᶜz = Fields.coordinate_field(Spaces.center_space(space)).z | ||
| ᶠz = Fields.coordinate_field(Spaces.face_space(space)).z | ||
| return (; ᶜz, ᶠz) | ||
| end | ||
|
|
||
|
|
||
| function sponge_tendencies!(Yₜ, Y, p, t) | ||
| rs, vs = p.atmos.rayleigh_sponge, p.atmos.viscous_sponge | ||
| (; ᶜh_tot, ᶜspecific) = p.precomputed | ||
| ᶜuₕ = Y.c.uₕ | ||
| ᶠu₃ = Yₜ.f.u₃ | ||
| ᶜρ = Y.c.ρ | ||
| vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, vs) | ||
| vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, vs) | ||
| vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, vs) | ||
| rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rs) | ||
|
|
||
| # TODO: fuse, once we fix | ||
| # https://github.com/CliMA/ClimaCore.jl/issues/2165 | ||
| @. Yₜ.c.uₕ += vst_uₕ | ||
| @. Yₜ.c.uₕ += rst_uₕ | ||
| @. Yₜ.f.u₃.components.data.:1 += vst_u₃ | ||
| @. Yₜ.c.ρe_tot += vst_ρe_tot | ||
|
|
||
| # TODO: can we write this out explicitly? | ||
| if vs isa ViscousSponge | ||
| for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) | ||
| χ_name == :e_tot && continue | ||
| vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, vs) | ||
| @. ᶜρχₜ += vst_tracer | ||
| @. Yₜ.c.ρ += vst_tracer | ||
| end | ||
| end | ||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.