Try out columnwise

.buildkite/Manifest-v1.11.toml

@@ -361,7 +361,7 @@ version = "0.5.17"
 deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
-version = "0.30.1"
+version = "0.30.2"
 
 [[deps.ClimaComms]]
 deps = ["Adapt", "Logging", "LoggingExtras"]
@@ -376,7 +376,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 = "cd1b32a3b8f3eab3e27e08180923a5c43c727f04"
+git-tree-sha1 = "c2d5b7090d6973c816f327431d830d78ac89fd5c"
+repo-rev = "ck/working_columnwise"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.14.33"
 weakdeps = ["CUDA", "Krylov"]

.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

config/default_configs/default_config.yml

@@ -128,7 +128,7 @@ t_start:
   value: "0secs"
 t_end:
   help: "Simulation end time. Examples: [`1200days`, `40secs`]"
-  value: "10days"
+  value: "1200secs"
 output_dir:
   help: "Output directory"
   value: ~

examples/hybrid/driver.jl

@@ -10,14 +10,35 @@ import ClimaAtmos as CA
 import Random
 Random.seed!(1234)
 
-if !(@isdefined config)
+import ClimaCore
+ClimaCore.Fields.local_geometry_field(bc::Base.AbstractBroadcasted) =
+    ClimaCore.Fields.local_geometry_field(axes(bc))
+# 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
+
+ClimaCore.Operators.fd_shmem_is_supported(bc::Base.Broadcast.Broadcasted) = false
+ClimaCore.Operators.use_fd_shmem() = false
+# The existing implementation limits our ability to apply
+# the same expressions from within kernels
+ClimaComms.device(topology::ClimaCore.Topologies.DeviceIntervalTopology) =
+    ClimaComms.CUDADevice()
+ClimaCore.Fields.error_mismatched_spaces(::Type, ::Type) = nothing # causes unsupported dynamic function invocation
+
+
+# if !(@isdefined config)
     (; config_file, job_id) = CA.commandline_kwargs()
     config = CA.AtmosConfig(config_file; job_id)
-end
+# end
 simulation = CA.AtmosSimulation(config)
 sol_res = CA.solve_atmos!(simulation)
 
-include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl"))
-if ClimaComms.iamroot(config.comms_ctx)
-    make_plots(Val(Symbol(simulation.job_id)), simulation.output_dir)
-end
+# include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl"))
+# if ClimaComms.iamroot(config.comms_ctx)
+#     make_plots(Val(Symbol(simulation.job_id)), simulation.output_dir)
+# end

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ₕ -=