Use columnwise for setting tendencies to zero

Author: charleskawczynski

src/prognostic_equations/remaining_tendency.jl

@@ -37,6 +37,103 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
 end
 
+using ClimaCore.RecursiveApply: rzero
+
+#####
+##### Cell center tendencies
+#####
+function ᶜremaining_tendency(ᶜY, ᶠY, p, t)
+    names = propertynames(ᶜY)
+    tends = construct_tendencies(Val(names), ᶜremaining_tendency, ᶜY, ᶠY, p, t)
+    # We cannot broadcast over a NamedTuple, so we need to check that edge case
+    # first.
+    if all(t -> !(t isa Base.Broadcast.Broadcasted), tends)
+        return make_named_tuple(Val(names), tends...)
+    else
+        return lazy.(make_named_tuple.(Val(names), tends...))
+    end
+end
+
+#####
+##### Cell face tendencies
+#####
+function ᶠremaining_tendency(ᶜY, ᶠY, p, t)
+    names = propertynames(ᶠY)
+    tends = construct_tendencies(Val(names), ᶠremaining_tendency, ᶜY, ᶠY, p, t)
+    # We cannot broadcast over a NamedTuple, so we need to check that edge case
+    # first.
+    if all(t -> !(t isa Base.Broadcast.Broadcasted), tends)
+        return make_named_tuple(Val(names), tends...)
+    else
+        return lazy.(make_named_tuple.(Val(names), tends...))
+    end
+end
+
+#####
+##### Individual tendencies
+#####
+
+function ᶜremaining_tendency(::Val{:ρ}, ᶜY, ᶠY, p, t)
+    :ρ in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρ))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:uₕ}, ᶜY, ᶠY, p, t)
+    :uₕ in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.uₕ))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρe_tot}, ᶜY, ᶠY, p, t)
+    :ρe_tot in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρe_tot))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρq_tot}, ᶜY, ᶠY, p, t)
+    :ρq_tot in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_tot))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρq_liq}, ᶜY, ᶠY, p, t)
+    :ρq_liq in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_liq))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρq_ice}, ᶜY, ᶠY, p, t)
+    :ρq_ice in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_ice))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρq_rai}, ᶜY, ᶠY, p, t)
+    :ρq_rai in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_rai))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:ρq_sno}, ᶜY, ᶠY, p, t)
+    :ρq_sno in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_sno))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t)
+    :sgsʲs in propertynames(ᶜY) || return ()
+    ∑tendencies = rzero(eltype(ᶜY.sgsʲs))
+    return ∑tendencies
+end
+function ᶜremaining_tendency(::Val{:sgs⁰}, ᶜY, ᶠY, p, t)
+    :sgs⁰ in propertynames(ᶜY) || return ()
+    ∑tendencies = rzero(eltype(ᶜY.sgs⁰))
+    return ∑tendencies
+end
+function ᶠremaining_tendency(::Val{:u₃}, ᶜY, ᶠY, p, t)
+    :u₃ in propertynames(ᶠY) || return ()
+    ∑tendencies = zero(eltype(ᶠY.u₃))
+    return ∑tendencies
+end
+function ᶠremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t)
+    :sgsʲs in propertynames(ᶠY) || return ()
+    ∑tendencies = rzero(eltype(ᶠY.sgsʲs))
+    return ∑tendencies
+end
+
 """
     remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
 
@@ -64,7 +161,27 @@ Returns:
 """
 NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     Yₜ_lim .= zero(eltype(Yₜ_lim))
-    Yₜ .= zero(eltype(Yₜ))
+    device = ClimaComms.device(axes(Y.c))
+    p_kernel = (;
+        zmax = Spaces.z_max(axes(Y.f)),
+        atmos = p.atmos,
+        params = p.params,
+        dt = p.dt,
+    )
+    if :sfc in propertynames(Y) # columnwise! does not yet handle .sfc
+        parent(Yₜ.sfc) .= zero(Spaces.undertype(axes(Y.c)))
+    end
+    Operators.columnwise!(
+        device,
+        ᶜremaining_tendency,
+        ᶠremaining_tendency,
+        Yₜ.c,
+        Yₜ.f,
+        Y.c,
+        Y.f,
+        p_kernel,
+        t,
+    )
     horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)
     fill_with_nans!(p)  # TODO: would be better to limit this to debug mode (e.g., if p.debug_mode...)
     horizontal_dynamics_tendency!(Yₜ, Y, p, t)

src/utils/utilities.jl

@@ -496,3 +496,69 @@ function issphere(space)
     return Meshes.domain(Spaces.topology(Spaces.horizontal_space(space))) isa
            Domains.SphereDomain
 end
+
+
+"""
+    construct_tendencies(::Val{names}, f, ᶜY, ᶠY, p, t)
+
+Return a tuple of calls to
+
+```
+f(Val(name), ᶜY, ᶠY, p, t)
+```
+for all names in `names`.
+
+For example, `f(Val((:a, :b)), ᶜY, ᶠY, p, t)` will return:
+
+```
+(
+    f(Val(:a), ᶜY, ᶠY, p, t),
+    f(Val(:b), ᶜY, ᶠY, p, t),
+)
+```
+"""
+@generated function construct_tendencies(
+    ::Val{names},
+    f,
+    ᶜY,
+    ᶠY,
+    p,
+    t,
+) where {names}
+    calls = []
+    for name in names
+        push!(calls, :(f(Val($(QuoteNode(name))), ᶜY, ᶠY, p, t)))
+    end
+    return quote
+        ($(calls...),)
+    end
+end
+
+"""
+    make_named_tuple(::Val{names}, vals...) where {names}
+
+Construct a NamedTuple given the names `names` and values `vals`.
+"""
+make_named_tuple(::Val{names}, vals...) where {names} = NamedTuple{names}(vals)
+
+"""
+    add_tend(∑tendencies, tendency)
+
+A helper function which returns `∑tendencies` when `tendency` is a
+`NullBroadcasted` and `lazy.(∑tendencies + tendency)` when `tendency` is not a
+`NullBroadcasted`.
+"""
+function add_tend end
+add_tend(∑tends, t) = lazy.(∑tends .+ t)
+add_tend(∑tends, ::NullBroadcasted) = ∑tends
+
+"""
+    sub_tend(∑tendencies, tendency)
+
+A helper function which returns `∑tendencies` when `tendency` is a
+`NullBroadcasted` and `lazy.(∑tendencies - tendency)` when `tendency` is not a
+`NullBroadcasted`.
+"""
+function sub_tend end
+sub_tend(∑tends, ::NullBroadcasted) = ∑tends
+sub_tend(∑tends, t) = lazy.(∑tends .- t)