wip

Author: charleskawczynski

.buildkite/ci_driver.jl

@@ -9,6 +9,16 @@ redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false)))
 # To load in the precompiled methods, run `using PrecompileCI` before loading ClimaAtmos. 
 # To see what methods are precompiled, open julia: `julia --project=.buildkite/PrecompileCI`
 # and run `using PrecompileTools; PrecompileTools.verbose[] = true; include(".buildkite/PrecompileCI/src/PrecompileCI.jl")`
+
+# Todo: move this to NullBroadcasts, or is there a better way?
+import NullBroadcasts
+# This makes the following pattern easier:
+# ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge))
+Base.broadcasted(::typeof(+), ::NullBroadcasts.NullBroadcasted, x) = x
+Base.broadcasted(::typeof(+), x, ::NullBroadcasts.NullBroadcasted) = x
+# Base.broadcasted(::typeof(-), ::NullBroadcasts.NullBroadcasted, x) = x
+Base.broadcasted(::typeof(-), x, ::NullBroadcasts.NullBroadcasted) = x
+
 using PrecompileCI
 import ClimaComms
 ClimaComms.@import_required_backends

examples/hybrid/driver.jl

@@ -10,6 +10,15 @@ import ClimaAtmos as CA
 import Random
 Random.seed!(1234)
 
+# Todo: move this to NullBroadcasts, or is there a better way?
+import NullBroadcasts
+# This makes the following pattern easier:
+# ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge))
+Base.broadcasted(::typeof(+), ::NullBroadcasts.NullBroadcasted, x) = x
+Base.broadcasted(::typeof(+), x, ::NullBroadcasts.NullBroadcasted) = x
+# Base.broadcasted(::typeof(-), ::NullBroadcasts.NullBroadcasted, x) = x
+Base.broadcasted(::typeof(-), x, ::NullBroadcasts.NullBroadcasted) = x
+
 if !(@isdefined config)
     (; config_file, job_id) = CA.commandline_kwargs()
     config = CA.AtmosConfig(config_file; job_id)

src/cache/precipitation_precomputed_quantities.jl

@@ -98,6 +98,76 @@ function set_precipitation_velocities!(
     return nothing
 end
 
+function compute_precip_velocities(ᶜY, ᶠY, p, t)
+    cmc = CAP.microphysics_cloud_params(p.params)
+    cmp = CAP.microphysics_1m_params(p.params)
+
+    # compute the precipitation terminal velocity [m/s]
+    ᶜwᵣ = CM1.terminal_velocity(
+        cmp.pr,
+        cmp.tv.rain,
+        ᶜY.ρ,
+        max(zero(ᶜY.ρ), ᶜY.ρq_rai / ᶜY.ρ),
+    )
+    ᶜwₛ = CM1.terminal_velocity(
+        cmp.ps,
+        cmp.tv.snow,
+        ᶜY.ρ,
+        max(zero(ᶜY.ρ), ᶜY.ρq_sno / ᶜY.ρ),
+    )
+    # compute sedimentation velocity for cloud condensate [m/s]
+    ᶜwₗ = CMNe.terminal_velocity(
+        cmc.liquid,
+        cmc.Ch2022.rain,
+        ᶜY.ρ,
+        max(zero(ᶜY.ρ), ᶜY.ρq_liq / ᶜY.ρ),
+    )
+    ᶜwᵢ = CMNe.terminal_velocity(
+        cmc.ice,
+        cmc.Ch2022.small_ice,
+        ᶜY.ρ,
+        max(zero(ᶜY.ρ), ᶜY.ρq_ice / ᶜY.ρ),
+    )
+    return (ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ)
+end
+
+
+compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) =
+    compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t, p.atmos.moisture_model, p.atmos.precip_model)
+compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t, moisture_model, precip_model) = zero(ᶜY.ρ)
+
+function compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t,
+    ::NonEquilMoistModel,
+    ::Microphysics1Moment,
+)
+    (ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ) = compute_precip_velocities(ᶜY, ᶠY, p, t)
+    # compute their contributions to energy and total water advection
+    return Geometry.WVector(
+            ᶜwₗ * ᶜY.ρq_liq +
+            ᶜwᵢ * ᶜY.ρq_ice +
+            ᶜwᵣ * ᶜY.ρq_rai +
+            ᶜwₛ * ᶜY.ρq_sno,
+        ) / ᶜY.ρ
+end
+
+compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) =
+    compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t, p.atmos.moisture_model, p.atmos.precip_model)
+
+compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t, moisture_model, precip_model) = zero(ᶜY.ρ)
+
+function compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t, ::NonEquilMoistModel, ::Microphysics1Moment)
+    thp = CAP.thermodynamics_params(p.params)
+    ᶜts = compute_ᶜts(ᶜY, ᶠY, p, t)
+    (ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ) = compute_precip_velocities(ᶜY, ᶠY, p, t)
+    # compute their contributions to energy and total water advection
+    return @. lazy(Geometry.WVector(
+            ᶜwₗ * ᶜY.ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜY.ᶜu))) +
+            ᶜwᵢ * ᶜY.ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜY.ᶜu))) +
+            ᶜwᵣ * ᶜY.ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜY.ᶜu))) +
+            ᶜwₛ * ᶜY.ρq_sno * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₛ, ᶜY.ᶜu))),
+        ) / ᶜY.ρ)
+end
+
 """
     set_precipitation_cache!(Y, p, precip_model, turbconv_model)
 

src/prognostic_equations/advection.jl

@@ -139,20 +139,6 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     end
     # ... and upwinding correction of energy and total water.
     # (The central advection of energy and total water is done implicitly.)
-    if energy_upwinding != Val(:none)
-        (; ᶜh_tot) = p.precomputed
-        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding)
-        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none))
-        @. Yₜ.c.ρe_tot += vtt - vtt_central
-    end
-
-    if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none)
-        ᶜq_tot = ᶜspecific.q_tot
-        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), tracer_upwinding)
-        vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), Val(:none))
-        @. Yₜ.c.ρq_tot += vtt - vtt_central
-    end
-
     if isnothing(ᶠf¹²)
         # shallow atmosphere
         @. Yₜ.c.uₕ -=