Start using columnwise

Author: charleskawczynski

src/prognostic_equations/remaining_tendency.jl

@@ -10,9 +10,118 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
 end
 
+prognostic_nt(::Val{names}, tends...) where {names} =
+    NamedTuple{names}(tends)
+
+function ᶜremaining_tendency(ᶜY, ᶠY, p, t)
+    names = propertynames(ᶜY)
+    tends = (
+            ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρq_liq(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρq_ice(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρq_rai(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_ρq_sno(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_sfc(ᶜY, ᶠY, p, t)...,
+            ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)...,
+    )
+    return lazy.(prognostic_nt.(Val(names), tends...))
+end
+function ᶠremaining_tendency(ᶜY, ᶠY, p, t)
+    names = propertynames(ᶠY)
+    tends = (
+            ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)...,
+            ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)...,
+    )
+    return lazy.(prognostic_nt.(Val(names), tends...))
+end
+
+function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t)
+    :ρ in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρ))
+    return (;ρ=∑tendencies)
+end
+function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t)
+    :uₕ in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.uₕ))
+    return (;uₕ=∑tendencies)
+end
+function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t)
+    :ρe_tot in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρe_tot))
+    return (;ρe_tot=∑tendencies)
+end
+function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t)
+    :ρq_tot in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_tot))
+    return (;ρq_tot=∑tendencies)
+end
+function ᶜremaining_tendency_ρq_liq(ᶜY, ᶠY, p, t)
+    :ρq_liq in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_liq))
+    return (;ρq_liq=∑tendencies)
+end
+function ᶜremaining_tendency_ρq_ice(ᶜY, ᶠY, p, t)
+    :ρq_ice in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_ice))
+    return (;ρq_ice=∑tendencies)
+end
+function ᶜremaining_tendency_ρq_rai(ᶜY, ᶠY, p, t)
+    :ρq_rai in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_rai))
+    return (;ρq_rai=∑tendencies)
+end
+function ᶜremaining_tendency_ρq_sno(ᶜY, ᶠY, p, t)
+    :ρq_sno in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.ρq_sno))
+    return (;ρq_sno=∑tendencies)
+end
+function ᶜremaining_tendency_sfc(ᶜY, ᶠY, p, t)
+    :sfc in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.sfc))
+    return (;sfc=∑tendencies)
+end
+function ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)
+    :sgsʲs in propertynames(ᶜY) || return ()
+    ∑tendencies = zero(eltype(ᶜY.sgsʲs))
+    return (;sgsʲs=∑tendencies)
+end
+function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)
+    :u₃ in propertynames(ᶠY) || return ()
+    ∑tendencies = zero(eltype(ᶠY.u₃))
+    return (;u₃=∑tendencies)
+end
+function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)
+    :sgsʲs in propertynames(ᶠY) || return ()
+    ∑tendencies = zero(eltype(ᶠY.sgsʲs))
+    return (;sgsʲs=∑tendencies)
+end
+
+
 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))
+    (localmem_lg, localmem_state) = if device isa ClimaComms.CUDADevice
+        Val(false), Val(true)
+    else
+        Val(false), Val(false)
+    end
+    p_kernel = (;)
+    Operators.columnwise!(
+        device,
+        ᶜremaining_tendency,
+        ᶠremaining_tendency,
+        Yₜ.c,
+        Yₜ.f,
+        Y.c,
+        Y.f,
+        p_kernel,
+        t,
+        localmem_lg,
+        localmem_state
+    )
     horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)
     fill_with_nans!(p)
     horizontal_advection_tendency!(Yₜ, Y, p, t)