From ac53bcca7649c8a89320f1ab287087e043545c4c Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 22 May 2025 09:25:07 -0400 Subject: [PATCH 1/6] Use ck/working_columnwise branch --- .buildkite/Manifest-v1.11.toml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index b711d30feb..3ca3b9bbf3 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.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"] From 9d32277c7979ed029f09be2878c9d4cf622c890b Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 22 May 2025 10:13:53 -0400 Subject: [PATCH 2/6] Try out columnwise --- .../remaining_tendency.jl | 64 ++++++++++++++++++- 1 file changed, 63 insertions(+), 1 deletion(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index a8929739a3..570541ee57 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,9 +10,71 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end +function ᶜremaining_tendency(ᶜY, ᶠY, p, t) + nt = (; + ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t)..., + ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t)..., + ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t)..., + ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., + ) + return nt +end +function ᶠremaining_tendency(ᶜY, ᶠY, p, t) + nt = (; + ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)..., + ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., + ) + return nt +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_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) + device = ClimaComms.device(axes(Y.c)) + p_rt = (;) + Operators.columnwise!( + device, + ᶜremaining_tendency, + ᶠremaining_tendency, + Yₜ.c, + Yₜ.f, + Y.c, + Y.f, + p_rt, + t, + Val(false), + Val(false), + ) Yₜ_lim .= zero(eltype(Yₜ_lim)) - Yₜ .= zero(eltype(Yₜ)) horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t) fill_with_nans!(p) horizontal_advection_tendency!(Yₜ, Y, p, t) From 7aa534c2c4b45b0de415a4bdc0812b7ff5d6d570 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 22 May 2025 10:16:34 -0400 Subject: [PATCH 3/6] wip --- src/prognostic_equations/remaining_tendency.jl | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 570541ee57..8bc684f34d 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -35,6 +35,18 @@ end function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t) :uₕ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.uₕ)) + + ᶜuₕ = ᶜY.c.uₕ + ᶠu₃ = Yₜ.f.u₃ + ᶜρ = Y.c.ρ + (; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos + (; ls_adv, edmf_coriolis) = p.atmos + (; params) = p + thermo_params = CAP.thermodynamics_params(params) + (; ᶜp, sfc_conditions, ᶜts) = p.precomputed + + ∑tendencies += viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) + return (;uₕ=∑tendencies) end function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) From f039fc96039c396fdde1b66dab586e8356455938 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 23 May 2025 09:15:38 -0400 Subject: [PATCH 4/6] wip --- .buildkite/ci_driver.jl | 10 + examples/hybrid/driver.jl | 9 + .../precipitation_precomputed_quantities.jl | 70 ++++++ src/prognostic_equations/advection.jl | 14 -- .../remaining_tendency.jl | 215 ++++++++++++++++-- src/prognostic_equations/water_advection.jl | 16 -- src/utils/utilities.jl | 20 +- 7 files changed, 292 insertions(+), 62 deletions(-) diff --git a/.buildkite/ci_driver.jl b/.buildkite/ci_driver.jl index 3207bcd3a0..cde455fc60 100644 --- a/.buildkite/ci_driver.jl +++ b/.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 diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index 5036e7a100..73fae7553f 100644 --- a/examples/hybrid/driver.jl +++ b/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) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index ca172846d6..2d00026f2a 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/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) diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 304a048869..d57fb5e40c 100644 --- a/src/prognostic_equations/advection.jl +++ b/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ₕ -= diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 8bc684f34d..925715cd20 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,13 +10,37 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end +function type_depth_limit(io::IO, s::String; maxtypedepth::Union{Nothing, Int}) + sz = get(io, :displaysize, displaysize(io))::Tuple{Int, Int} + return Base.type_depth_limit(s, max(sz[2], 120); maxdepth = maxtypedepth) +end + +function assert_concretized(var) + if var isa Base.AbstractBroadcasted + T_lim = type_depth_limit(stdout, string(typeof(var)); maxtypedepth = 3) + error("This should be concretized:\n\n\t$T_lim\n\n") + end + return nothing +end + function ᶜremaining_tendency(ᶜY, ᶠY, p, t) nt = (; ᶜ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_sgsʲs(ᶜY, ᶠY, p, t)..., ) + debug_tendencies = true + if debug_tendencies + val = Operators.getidx(axes(ᶜY), nt, 1, (1,1,1)) + assert_concretized(val) + :ρ in propertynames(ᶜY) && assert_concretized(val.ρ) + :uₕ in propertynames(ᶜY) && assert_concretized(val.uₕ) + :ρe_tot in propertynames(ᶜY) && assert_concretized(val.ρe_tot) + :ρq_tot in propertynames(ᶜY) && assert_concretized(val.ρq_tot) + :sgsʲs in propertynames(ᶜY) && assert_concretized(val.sgsʲs) + end return nt end function ᶠremaining_tendency(ᶜY, ᶠY, p, t) @@ -24,36 +48,182 @@ function ᶠremaining_tendency(ᶜY, ᶠY, p, t) ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)..., ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., ) + debug_tendencies = true + if debug_tendencies + val = Operators.getidx(axes(ᶠY), nt, 1, (1,1,1)) + assert_concretized(val) + :u₃ in propertynames(ᶠY) && assert_concretized(val.u₃) + :sgsʲs in propertynames(ᶠY) && assert_concretized(val.sgsʲs) + end return nt end +water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = + @. lazy(ᶜprecipdivᵥ(ᶠinterp(ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜχ)))) + function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) :ρ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρ)) + ᶜJ = Fields.local_geometry_field(ᶜY).J + ᶠJ = Fields.local_geometry_field(ᶠY).J + + if !(p.atmos.moisture_model isa DryModel) + ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) + cmc = CAP.microphysics_cloud_params(p.params) + cmp = CAP.microphysics_1m_params(p.params) + thp = CAP.thermodynamics_params(p.params) + ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) + end return (;ρ=∑tendencies) end function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t) :uₕ in propertynames(ᶜY) || return () - ∑tendencies = zero(eltype(ᶜY.uₕ)) + ᶜuₕ = ᶜY.uₕ + ∑tendencies = zero(eltype(ᶜuₕ)) - ᶜuₕ = ᶜY.c.uₕ - ᶠu₃ = Yₜ.f.u₃ - ᶜρ = Y.c.ρ - (; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos - (; ls_adv, edmf_coriolis) = p.atmos - (; params) = p - thermo_params = CAP.thermodynamics_params(params) - (; ᶜp, sfc_conditions, ᶜts) = p.precomputed + (; viscous_sponge, rayleigh_sponge) = p.atmos - ∑tendencies += viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) + ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)) + ∑tendencies = lazy.(∑tendencies .+ rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)) return (;uₕ=∑tendencies) end + +function surface_velocity_full(ᶠu₃, ᶠuₕ³) + assert_eltype(ᶠuₕ³, Geometry.Contravariant3Vector) + assert_eltype(ᶠu₃, Geometry.Covariant3Vector) + ᶠlg = Fields.local_geometry_field(axes(ᶠu₃)) + sfc_u₃ = ᶠu₃ # Fields.level(ᶠu₃.components.data.:1, half) + sfc_uₕ³ = ᶠuₕ³ # Fields.level(ᶠuₕ³.components.data.:1, half) + sfc_g³³ = g³³_field(axes(sfc_u₃)) + w₃ = @. lazy(- C3(sfc_uₕ³ / sfc_g³³, ᶠlg)) # u³ = uₕ³ + w³ = uₕ³ + w₃ * g³³ + assert_eltype(w₃, Geometry.Covariant3Vector) + return w₃ +end + +function compute_ᶜts(ᶜY, ᶠY, p, t) + (; moisture_model, precip_model) = p.atmos + thermo_params = CAP.thermodynamics_params(p.params) + thermo_args = (thermo_params, moisture_model, precip_model) + ᶜz = Fields.coordinate_field(ᶜY).z + FT = Spaces.undertype(axes(ᶜY)) + grav = FT(CAP.grav(p.params)) + ᶜΦ = @. lazy(grav * ᶜz) + + ᶜρ = ᶜY.ρ + ᶜuₕ = ᶜY.uₕ + ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) + ᶠu₃ = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) + ᶜK = compute_kinetic(ᶜuₕ, ᶠu₃) + ᶜspecific = @. lazy(specific_gs(ᶜY)) + return @. lazy(ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, ᶜρ)) +end + +assert_eltype(bc::Base.AbstractBroadcasted, ::Type{T}) where {T} = + assert_eltype(eltype(bc), T) +assert_eltype(f::Fields.Field, ::Type{T}) where {T} = + assert_eltype(Fields.field_values(f), T) +assert_eltype(data::DataLayouts.AbstractData, ::Type{T}) where {T} = + assert_eltype(eltype(data), T) +assert_eltype(::Type{S}, ::Type{T}) where {S, T} = + @assert S <: T "Type $S should be a subtype of $T" + +function compute_ᶠu₃_with_bcs(ᶠu₃, ᶠuₕ³) + assert_eltype(ᶠu₃, Geometry.Covariant3Vector) + assert_eltype(ᶠuₕ³, Geometry.Contravariant3Vector) + ᶠz = Fields.coordinate_field(axes(ᶠu₃)).z + sfc_u₃ = surface_velocity_full(ᶠu₃, ᶠuₕ³) + # todo: generalize this with z_min + return @. lazy(ifelse(iszero(ᶠz), sfc_u₃, ᶠu₃)) +end function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) :ρe_tot in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρe_tot)) + + (; moisture_model, viscous_sponge, precip_model) = p.atmos + (; energy_upwinding) = p.atmos.numerics + thermo_params = CAP.thermodynamics_params(p.params) + thermo_args = (thermo_params, moisture_model, precip_model) + ᶜz = Fields.coordinate_field(ᶜY).z + ᶜρ = ᶜY.ρ + ᶜuₕ = ᶜY.uₕ + ᶜρe_tot = ᶜY.ρe_tot + ᶜts = compute_ᶜts(ᶜY, ᶠY, p, t) + ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts)) + ᶜe_tot = @. lazy(ᶜρe_tot / ᶜρ) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot)) + + if !(p.atmos.moisture_model isa DryModel) + ᶜwₕhₜ = compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) + cmc = CAP.microphysics_cloud_params(p.params) + cmp = CAP.microphysics_1m_params(p.params) + thp = CAP.thermodynamics_params(p.params) + ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₕhₜ)) + # @show "test 1" + end + if energy_upwinding != Val(:none) + # @show energy_upwinding + # @show "test 2" + (; dt) = p + ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) + # ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) + # ᶠz = Fields.coordinate_field(axes(ᶠY.u₃)).z + # sfc_u₃ = surface_velocity_full(ᶠY.u₃, ᶠuₕ³) + # todo: generalize this with z_min + # ᶠu₃_with_bcs = @. lazy(ifelse(iszero(ᶠz), sfc_u₃, ᶠY.u₃)) + # _ᶠu₃_with_bcs = similar(Base.materialize(sfc_u₃)) + # @. _ᶠu₃_with_bcs = ifelse(iszero(ᶠz), sfc_u₃, ᶠY.u₃) + # ᶠu₃_with_bcs = @. lazy(ifelse(ᶠz==0, sfc_u₃, ᶠY.u₃)) + # @show Base.materialize(ᶠu₃_with_bcs) + + # return @. lazy(ᶠuₕ³ + CT3(ᶠu₃)) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) + # @show vtt + @show Base.materialize(vtt) + ∑tendencies = lazy.(∑tendencies .+ vtt) # problematic + vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) + # need to improve NullBroadcast support for this. + # neg_vtt_central = @. lazy(vtt_central * (-1)) + neg_vtt_central = vtt_central # TODO: this is a -1 factor bug! + # ∑tendencies = lazy.(∑tendencies .+ neg_vtt_central) + end + + ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) return (;ρe_tot=∑tendencies) end +function compute_ᶠu³(ᶜY, ᶠY) + ᶜρ = ᶜY.ρ + ᶜuₕ = ᶜY.uₕ + ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) + ᶠu₃ = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) + return @. lazy(ᶠuₕ³ + CT3(ᶠu₃)) +end +function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) + :ρq_tot in propertynames(ᶜY) || return () + ∑tendencies = zero(eltype(ᶜY.ρq_tot)) + ᶜJ = Fields.local_geometry_field(ᶜY).J + ᶠJ = Fields.local_geometry_field(ᶠY).J + ᶜρ = ᶜY.ρ + if !(p.atmos.moisture_model isa DryModel) + ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) + cmc = CAP.microphysics_cloud_params(p.params) + cmp = CAP.microphysics_1m_params(p.params) + thp = CAP.thermodynamics_params(p.params) + ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) + end + (; tracer_upwinding) = p.atmos.numerics + if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none) + (; dt) = p + (; ᶜspecific) = p.precomputed + ᶜq_tot = @. lazy(ᶜY.ρq_tot / ᶜY.ρ) + ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), tracer_upwinding) + vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, float(dt), Val(:none)) + ∑tendencies = lazy.(∑tendencies .+ vtt .- vtt_central) + end + + return (;ρq_tot=∑tendencies) +end function ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) :sgsʲs in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.sgsʲs)) @@ -62,6 +232,12 @@ end function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t) :u₃ in propertynames(ᶠY) || return () ∑tendencies = zero(eltype(ᶠY.u₃)) + (; viscous_sponge) = p.atmos + ᶜρ = ᶜY.ρ + ᶜuₕ = ᶜY.uₕ + ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) + ᶠu₃ = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) + ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)) return (;u₃=∑tendencies) end function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) @@ -72,7 +248,12 @@ end NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) device = ClimaComms.device(axes(Y.c)) - p_rt = (;) + (; moisture_model, viscous_sponge, precip_model) = p.atmos + p_rt = (; + atmos = p.atmos, + params = p.params, + dt = p.dt, + ) Operators.columnwise!( device, ᶜremaining_tendency, @@ -92,7 +273,6 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) horizontal_advection_tendency!(Yₜ, Y, p, t) hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) explicit_vertical_advection_tendency!(Yₜ, Y, p, t) - vertical_advection_of_water_tendency!(Yₜ, Y, p, t) additional_tendency!(Yₜ, Y, p, t) return Yₜ end @@ -120,10 +300,6 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) thermo_params = CAP.thermodynamics_params(params) (; ᶜp, sfc_conditions, ᶜts) = p.precomputed - vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) - vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge) - vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge) - rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge) hs_args = (ᶜuₕ, ᶜp, params, sfc_conditions.ts, moisture_model, forcing_type) hs_tendency_uₕ = held_suarez_forcing_tendency_uₕ(hs_args...) hs_tendency_ρe_tot = held_suarez_forcing_tendency_ρe_tot(ᶜρ, hs_args...) @@ -131,13 +307,6 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) lsa_args = (ᶜρ, thermo_params, ᶜts, t, ls_adv) bc_lsa_tend_ρe_tot = large_scale_advection_tendency_ρe_tot(lsa_args...) - # TODO: fuse, once we fix - # https://github.com/CliMA/ClimaCore.jl/issues/2165 - @. Yₜ.c.uₕ += vst_uₕ - @. Yₜ.c.uₕ += rst_uₕ - @. Yₜ.f.u₃.components.data.:1 += vst_u₃ - @. Yₜ.c.ρe_tot += vst_ρe_tot - # TODO: can we write this out explicitly? for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific) χ_name == :e_tot && continue diff --git a/src/prognostic_equations/water_advection.jl b/src/prognostic_equations/water_advection.jl index 1ce7f6504b..735dd1f6ce 100644 --- a/src/prognostic_equations/water_advection.jl +++ b/src/prognostic_equations/water_advection.jl @@ -1,18 +1,2 @@ import ClimaCore: Fields -function vertical_advection_of_water_tendency!(Yₜ, Y, p, t) - - ᶜJ = Fields.local_geometry_field(Y.c).J - ᶠJ = Fields.local_geometry_field(Y.f).J - (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed - - if !(p.atmos.moisture_model isa DryModel) - @. Yₜ.c.ρ -= - ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₜqₜ))) - @. Yₜ.c.ρe_tot -= - ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₕhₜ))) - @. Yₜ.c.ρq_tot -= - ᶜprecipdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜwₜqₜ))) - end - return nothing -end diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index dba3faf059..82b5fd848d 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -48,7 +48,7 @@ sort_files_by_time(files) = permute!(files, sortperm(time_from_filename.(files))) """ - κ .= compute_kinetic(uₕ::Field, uᵥ::Field) + κ .= compute_kinetic(uₕ, uᵥ) Compute the specific kinetic energy at cell centers, resulting in `κ` from individual velocity components: @@ -58,11 +58,13 @@ individual velocity components: cell centers, and - `uᵥ` should be a `Covariant3Vector`-valued field at cell faces. """ -function compute_kinetic(uₕ::Fields.Field, uᵥ::Fields.Field) - @assert eltype(uₕ) <: Union{C1, C2, C12} - @assert eltype(uᵥ) <: C3 +function compute_kinetic(uₕ, uᵥ) + # @assert eltype(uₕ) <: Union{C1, C2, C12} + # @assert eltype(uᵥ) <: C3 + FT = Spaces.undertype(axes(uₕ)) + onehalf = FT(1 / 2) return @. lazy( - 1 / 2 * ( + onehalf * ( dot(C123(uₕ), CT123(uₕ)) + ᶜinterp(dot(C123(uᵥ), CT123(uᵥ))) + 2 * dot(CT123(uₕ), ᶜinterp(C123(uᵥ))) @@ -117,12 +119,12 @@ function compute_strain_rate_face(u::Fields.Field) end """ - g³³_field(field) + g³³_field(space) -Extracts the value of `g³³` from `Fields.local_geometry_field(field)`. +Extracts the value of `g³³` from `Fields.local_geometry_field(space)`. """ -function g³³_field(field) - g_field = Fields.local_geometry_field(field).gⁱʲ.components.data +function g³³_field(space) + g_field = Fields.local_geometry_field(space).gⁱʲ.components.data end_index = fieldcount(eltype(g_field)) # This will be 4 in 2D and 9 in 3D. return g_field.:($end_index) # For both 2D and 3D spaces, g³³ = g[end]. end From 6b1380ca79d25e8124ddd9b588e53c8a1ce79228 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 23 May 2025 16:24:21 -0400 Subject: [PATCH 5/6] wip wip wip --- examples/hybrid/driver.jl | 3 + .../remaining_tendency.jl | 60 ++++++++++++------- 2 files changed, 40 insertions(+), 23 deletions(-) diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index 73fae7553f..f8214dc0e0 100644 --- a/examples/hybrid/driver.jl +++ b/examples/hybrid/driver.jl @@ -10,6 +10,9 @@ import ClimaAtmos as CA import Random Random.seed!(1234) +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: diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 925715cd20..0fbefd716d 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,14 +10,17 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end +print_dlt(x) = type_depth_limit(stdout, string(typeof(x)); maxtypedepth=3) +print_dlt(::Type{T}) where {T} = type_depth_limit(stdout, string(T); maxtypedepth=3) function type_depth_limit(io::IO, s::String; maxtypedepth::Union{Nothing, Int}) sz = get(io, :displaysize, displaysize(io))::Tuple{Int, Int} return Base.type_depth_limit(s, max(sz[2], 120); maxdepth = maxtypedepth) end function assert_concretized(var) - if var isa Base.AbstractBroadcasted - T_lim = type_depth_limit(stdout, string(typeof(var)); maxtypedepth = 3) + val = Operators.getidx(axes(var), var, 1, (1,1,1)) + if val isa Base.AbstractBroadcasted + T_lim = type_depth_limit(stdout, string(typeof(val)); maxtypedepth = 3) error("This should be concretized:\n\n\t$T_lim\n\n") end return nothing @@ -33,13 +36,11 @@ function ᶜremaining_tendency(ᶜY, ᶠY, p, t) ) debug_tendencies = true if debug_tendencies - val = Operators.getidx(axes(ᶜY), nt, 1, (1,1,1)) - assert_concretized(val) - :ρ in propertynames(ᶜY) && assert_concretized(val.ρ) - :uₕ in propertynames(ᶜY) && assert_concretized(val.uₕ) - :ρe_tot in propertynames(ᶜY) && assert_concretized(val.ρe_tot) - :ρq_tot in propertynames(ᶜY) && assert_concretized(val.ρq_tot) - :sgsʲs in propertynames(ᶜY) && assert_concretized(val.sgsʲs) + :ρ in propertynames(ᶜY) && assert_concretized(nt.ρ) + :uₕ in propertynames(ᶜY) && assert_concretized(nt.uₕ) + :ρe_tot in propertynames(ᶜY) && assert_concretized(nt.ρe_tot) + :ρq_tot in propertynames(ᶜY) && assert_concretized(nt.ρq_tot) + :sgsʲs in propertynames(ᶜY) && assert_concretized(nt.sgsʲs) end return nt end @@ -50,10 +51,8 @@ function ᶠremaining_tendency(ᶜY, ᶠY, p, t) ) debug_tendencies = true if debug_tendencies - val = Operators.getidx(axes(ᶠY), nt, 1, (1,1,1)) - assert_concretized(val) - :u₃ in propertynames(ᶠY) && assert_concretized(val.u₃) - :sgsʲs in propertynames(ᶠY) && assert_concretized(val.sgsʲs) + :u₃ in propertynames(ᶠY) && assert_concretized(nt.u₃) + :sgsʲs in propertynames(ᶠY) && assert_concretized(nt.sgsʲs) end return nt end @@ -69,9 +68,6 @@ function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) if !(p.atmos.moisture_model isa DryModel) ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) - cmc = CAP.microphysics_cloud_params(p.params) - cmp = CAP.microphysics_1m_params(p.params) - thp = CAP.thermodynamics_params(p.params) ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) end return (;ρ=∑tendencies) @@ -109,7 +105,7 @@ function compute_ᶜts(ᶜY, ᶠY, p, t) FT = Spaces.undertype(axes(ᶜY)) grav = FT(CAP.grav(p.params)) ᶜΦ = @. lazy(grav * ᶜz) - + ᶜρ = ᶜY.ρ ᶜuₕ = ᶜY.uₕ ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) @@ -155,14 +151,12 @@ function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) if !(p.atmos.moisture_model isa DryModel) ᶜwₕhₜ = compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) - cmc = CAP.microphysics_cloud_params(p.params) - cmp = CAP.microphysics_1m_params(p.params) - thp = CAP.thermodynamics_params(p.params) ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₕhₜ)) # @show "test 1" end if energy_upwinding != Val(:none) # @show energy_upwinding + # error("Done") # @show "test 2" (; dt) = p ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) @@ -175,11 +169,31 @@ function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) # @. _ᶠu₃_with_bcs = ifelse(iszero(ᶠz), sfc_u₃, ᶠY.u₃) # ᶠu₃_with_bcs = @. lazy(ifelse(ᶠz==0, sfc_u₃, ᶠY.u₃)) # @show Base.materialize(ᶠu₃_with_bcs) + # _ᶠu³ = Base.materialize(ᶠu³) + _ᶠu³ = ᶠu³ # works! + # _ᶜh_tot = Base.materialize(ᶜh_tot) + _ᶜh_tot = ᶜh_tot # results in `Killed: 9` (seg fault?) # return @. lazy(ᶠuₕ³ + CT3(ᶠu₃)) - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) - # @show vtt - @show Base.materialize(vtt) + vtt = vertical_transport(ᶜρ, _ᶠu³, _ᶜh_tot, float(dt), energy_upwinding) + # # @show vtt + # ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) + # ᶠu₃_with_bcs = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) + # # @show print_dlt(Base.materialize(ᶠuₕ³)) + # # @show print_dlt(Base.materialize(ᶠu₃_with_bcs)) + # # @show print_dlt(Base.materialize(ᶠu³)) + # @show print_dlt(eltype(_ᶠu³)) + # @show print_dlt(eltype(ᶜh_tot)) + # @show print_dlt(eltype(_ᶜh_tot)) + # @show print_dlt(eltype(vtt)) + # val = Operators.getidx(axes(ᶜY), vtt, 1, (1,1,1)) + # if val isa Base.AbstractBroadcasted + # error("getidx of the input is a broadcasted.") + # else + # @show print_dlt(val) + # end + # error("Done") + ∑tendencies = lazy.(∑tendencies .+ vtt) # problematic vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) # need to improve NullBroadcast support for this. From 741d4b20e2bf90f34d21982b292e91a131509a60 Mon Sep 17 00:00:00 2001 From: Charlie Kawczynski Date: Tue, 27 May 2025 10:08:26 -0700 Subject: [PATCH 6/6] wip --- config/default_configs/default_config.yml | 2 +- examples/hybrid/driver.jl | 21 ++-- .../remaining_tendency.jl | 98 ++++--------------- 3 files changed, 36 insertions(+), 85 deletions(-) diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index 3094e66790..3ba6a8afc8 100644 --- a/config/default_configs/default_config.yml +++ b/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: ~ diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index f8214dc0e0..26c6b8a045 100644 --- a/examples/hybrid/driver.jl +++ b/examples/hybrid/driver.jl @@ -22,14 +22,23 @@ Base.broadcasted(::typeof(+), x, ::NullBroadcasts.NullBroadcasted) = x # Base.broadcasted(::typeof(-), ::NullBroadcasts.NullBroadcasted, x) = x Base.broadcasted(::typeof(-), x, ::NullBroadcasts.NullBroadcasted) = x -if !(@isdefined config) +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 diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 0fbefd716d..07e76dd289 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,51 +10,27 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end -print_dlt(x) = type_depth_limit(stdout, string(typeof(x)); maxtypedepth=3) -print_dlt(::Type{T}) where {T} = type_depth_limit(stdout, string(T); maxtypedepth=3) -function type_depth_limit(io::IO, s::String; maxtypedepth::Union{Nothing, Int}) - sz = get(io, :displaysize, displaysize(io))::Tuple{Int, Int} - return Base.type_depth_limit(s, max(sz[2], 120); maxdepth = maxtypedepth) -end - -function assert_concretized(var) - val = Operators.getidx(axes(var), var, 1, (1,1,1)) - if val isa Base.AbstractBroadcasted - T_lim = type_depth_limit(stdout, string(typeof(val)); maxtypedepth = 3) - error("This should be concretized:\n\n\t$T_lim\n\n") - end - return nothing -end +prognostic_nt(::Val{names}, tends...) where {names} = + NamedTuple{names}(tends) function ᶜremaining_tendency(ᶜY, ᶠY, p, t) - nt = (; + 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_sgsʲs(ᶜY, ᶠY, p, t)..., ) - debug_tendencies = true - if debug_tendencies - :ρ in propertynames(ᶜY) && assert_concretized(nt.ρ) - :uₕ in propertynames(ᶜY) && assert_concretized(nt.uₕ) - :ρe_tot in propertynames(ᶜY) && assert_concretized(nt.ρe_tot) - :ρq_tot in propertynames(ᶜY) && assert_concretized(nt.ρq_tot) - :sgsʲs in propertynames(ᶜY) && assert_concretized(nt.sgsʲs) - end - return nt + return lazy.(prognostic_nt.(Val(names), tends...)) end function ᶠremaining_tendency(ᶜY, ᶠY, p, t) - nt = (; + names = propertynames(ᶠY) + tends = ( ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)..., ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., ) - debug_tendencies = true - if debug_tendencies - :u₃ in propertynames(ᶠY) && assert_concretized(nt.u₃) - :sgsʲs in propertynames(ᶠY) && assert_concretized(nt.sgsʲs) - end - return nt + return lazy.(prognostic_nt.(Val(names), tends...)) end water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = @@ -152,54 +128,15 @@ function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) if !(p.atmos.moisture_model isa DryModel) ᶜwₕhₜ = compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₕhₜ)) - # @show "test 1" end if energy_upwinding != Val(:none) - # @show energy_upwinding - # error("Done") - # @show "test 2" (; dt) = p ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) - # ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) - # ᶠz = Fields.coordinate_field(axes(ᶠY.u₃)).z - # sfc_u₃ = surface_velocity_full(ᶠY.u₃, ᶠuₕ³) - # todo: generalize this with z_min - # ᶠu₃_with_bcs = @. lazy(ifelse(iszero(ᶠz), sfc_u₃, ᶠY.u₃)) - # _ᶠu₃_with_bcs = similar(Base.materialize(sfc_u₃)) - # @. _ᶠu₃_with_bcs = ifelse(iszero(ᶠz), sfc_u₃, ᶠY.u₃) - # ᶠu₃_with_bcs = @. lazy(ifelse(ᶠz==0, sfc_u₃, ᶠY.u₃)) - # @show Base.materialize(ᶠu₃_with_bcs) - # _ᶠu³ = Base.materialize(ᶠu³) - _ᶠu³ = ᶠu³ # works! - # _ᶜh_tot = Base.materialize(ᶜh_tot) - _ᶜh_tot = ᶜh_tot # results in `Killed: 9` (seg fault?) - - # return @. lazy(ᶠuₕ³ + CT3(ᶠu₃)) - vtt = vertical_transport(ᶜρ, _ᶠu³, _ᶜh_tot, float(dt), energy_upwinding) - # # @show vtt - # ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) - # ᶠu₃_with_bcs = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) - # # @show print_dlt(Base.materialize(ᶠuₕ³)) - # # @show print_dlt(Base.materialize(ᶠu₃_with_bcs)) - # # @show print_dlt(Base.materialize(ᶠu³)) - # @show print_dlt(eltype(_ᶠu³)) - # @show print_dlt(eltype(ᶜh_tot)) - # @show print_dlt(eltype(_ᶜh_tot)) - # @show print_dlt(eltype(vtt)) - # val = Operators.getidx(axes(ᶜY), vtt, 1, (1,1,1)) - # if val isa Base.AbstractBroadcasted - # error("getidx of the input is a broadcasted.") - # else - # @show print_dlt(val) - # end - # error("Done") - - ∑tendencies = lazy.(∑tendencies .+ vtt) # problematic + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) + ∑tendencies = lazy.(∑tendencies .+ vtt) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) # need to improve NullBroadcast support for this. - # neg_vtt_central = @. lazy(vtt_central * (-1)) - neg_vtt_central = vtt_central # TODO: this is a -1 factor bug! - # ∑tendencies = lazy.(∑tendencies .+ neg_vtt_central) + ∑tendencies = lazy.(∑tendencies .+ (-1) .* vtt_central) end ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) @@ -263,11 +200,16 @@ end NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) device = ClimaComms.device(axes(Y.c)) (; moisture_model, viscous_sponge, precip_model) = p.atmos - p_rt = (; + p_kernel = (; atmos = p.atmos, params = p.params, dt = p.dt, ) + (localmem_lg, localmem_state) = if device isa ClimaComms.CUDADevice + Val(false), Val(true) + else + Val(false), Val(false) + end Operators.columnwise!( device, ᶜremaining_tendency, @@ -276,10 +218,10 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) Yₜ.f, Y.c, Y.f, - p_rt, + p_kernel, t, - Val(false), - Val(false), + localmem_lg, + localmem_state ) Yₜ_lim .= zero(eltype(Yₜ_lim)) horizontal_tracer_advection_tendency!(Yₜ_lim, Y, p, t)