Start using columnwise

.buildkite/Manifest-v1.11.toml

@@ -381,7 +381,9 @@ weakdeps = ["CUDA", "MPI"]
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
-git-tree-sha1 = "c6ab151ea66f3756566abc039c76ae767e490446"
+git-tree-sha1 = "4085d9d00635b3780533539e12d2e2af5ffc6cb2"
+repo-rev = "ck/fixes"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.14.34"
 weakdeps = ["CUDA", "Krylov"]

.buildkite/ci_driver.jl

@@ -11,6 +11,8 @@ redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false)))
 # and run `using PrecompileTools; PrecompileTools.verbose[] = true; include(".buildkite/PrecompileCI/src/PrecompileCI.jl")`
 using PrecompileCI
 import ClimaComms
+import ClimaCore
+ClimaCore.Fields.error_mismatched_spaces(::Type, ::Type) = nothing # causes unsupported dynamic function invocation
 ClimaComms.@import_required_backends
 import ClimaAtmos as CA
 import Random

.buildkite/pipeline.yml

@@ -328,6 +328,8 @@ steps:
         artifact_paths: "baroclinic_wave_equil/output_active/*"
         agents:
           slurm_constraint: icelake|cascadelake|skylake|epyc
+          # slurm_mem: 20GB # columnwise! insufficient mem
+          slurm_mem: 60GB # columnwise!
 
       - label: ":computer: held suarez"
         key: held_suarez
@@ -904,7 +906,7 @@ steps:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_gpus: 1
-          slurm_mem: 16GB
+          slurm_mem: 60GB
 
       - label: "GPU: compare BW with CPU"
         command: >

src/cache/precipitation_precomputed_quantities.jl

@@ -29,6 +29,76 @@ function Kin(ᶜw_precip, ᶜu_air)
     )
 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ᵣ = @. lazy(CM1.terminal_velocity(
+            cmp.pr,
+            cmp.tv.rain,
+            ᶜY.ρ,
+            max(zero(ᶜY.ρ), ᶜY.ρq_rai / ᶜY.ρ),
+        ))
+    ᶜwₛ = @. lazy(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ₗ = @. lazy(CMNe.terminal_velocity(
+        cmc.liquid,
+        cmc.Ch2022.rain,
+        ᶜY.ρ,
+        max(zero(ᶜY.ρ), ᶜY.ρq_liq / ᶜY.ρ),
+    ))
+    ᶜwᵢ = @. lazy(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(eltype(ᶜ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 @. lazy(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(eltype(ᶜ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_velocities!(Y, p, moisture_model, precip_model)
 

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -14,7 +14,13 @@ import ClimaCore.Fields as Fields
 
 function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s)
     s isa Nothing && return NullBroadcasted()
-    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
+    ᶜz = Fields.coordinate_field(Spaces.center_space(axes(ᶜuₕ))).z
+    ᶠz = Fields.coordinate_field(Spaces.face_space(axes(ᶜuₕ))).z
     zmax = z_max(axes(ᶠz))
+    return rayleigh_sponge_tendency_uₕ(ᶜuₕ, s, ᶜz, ᶠz, zmax)
+end
+
+function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s, ᶜz, ᶠz, zmax)
+    s isa Nothing && return NullBroadcasted()
     return @. lazy(-β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ)
 end

src/prognostic_equations/advection.jl

@@ -139,19 +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