From d2605fe02b9090b4dd8d83ee87c8158a90d0d05b Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Mon, 2 Jun 2025 13:09:03 -0400 Subject: [PATCH 01/21] Start using columnwise --- .../remaining_tendency.jl | 117 +++++++++++++++++- 1 file changed, 116 insertions(+), 1 deletion(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index a8929739a3..df440d4763 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,9 +10,124 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end +function prognostic_nt(::Val{names}; tends...) where {names} + nt_ordered = NamedTuple{names}(rzero(values(tends))) + nt_values = NamedTuple{keys(tends)}(values(tends)) + return merge(nt_ordered, nt_values) +end + +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_sgs⁰(ᶜ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 +using ClimaCore.RecursiveApply: rzero +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_sgsʲs(ᶜY, ᶠY, p, t) + :sgsʲs in propertynames(ᶜY) || return () + ∑tendencies = rzero(eltype(ᶜY.sgsʲs)) + return (; sgsʲs = ∑tendencies) +end +function ᶜremaining_tendency_sgs⁰(ᶜY, ᶠY, p, t) + :sgs⁰ in propertynames(ᶜY) || return () + ∑tendencies = rzero(eltype(ᶜY.sgs⁰)) + return (; sgs⁰ = ∑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 = rzero(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 = (;) + if :sfc in propertynames(Y) # columnwise! does not yet handle .sfc + parent(Yₜ.sfc) .= zero(Spaces.undertype(axes(Y.c))) + end + 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) From c15d293c442dd4738a08a6014fd083530bb20d96 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Tue, 3 Jun 2025 10:26:30 -0400 Subject: [PATCH 02/21] Add some tendencies via columnwise --- .../precipitation_precomputed_quantities.jl | 70 ++++++++ src/prognostic_equations/advection.jl | 13 -- .../remaining_tendency.jl | 155 ++++++++++++++++-- src/prognostic_equations/water_advection.jl | 16 -- 4 files changed, 207 insertions(+), 47 deletions(-) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index ca172846d6..a0b89826bd 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/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ᵣ = 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(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 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_velocities!(Y, p, moisture_model, precip_model) diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 304a048869..09d85c15a8 100644 --- a/src/prognostic_equations/advection.jl +++ b/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 diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index df440d4763..9a1e589646 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -44,22 +44,84 @@ using ClimaCore.RecursiveApply: rzero function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) :ρ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρ)) - return (; ρ = ∑tendencies) + ᶜ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) + ∑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ₕ)) - return (; uₕ = ∑tendencies) + (; viscous_sponge, rayleigh_sponge) = p.atmos + ᶜuₕ = ᶜY.uₕ + ∑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 ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) :ρe_tot in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρe_tot)) - return (; ρe_tot = ∑tendencies) + + (; 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) + ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₕhₜ)) + end + if energy_upwinding != Val(:none) + (; dt) = p + ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) + 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. + ∑tendencies = lazy.(∑tendencies .+ (-1) .* vtt_central) + end + + ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) + 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) + ᶜ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_ρq_liq(ᶜY, ᶠY, p, t) :ρq_liq in propertynames(ᶜY) || return () @@ -94,7 +156,13 @@ end function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t) :u₃ in propertynames(ᶠY) || return () ∑tendencies = zero(eltype(ᶠY.u₃)) - return (; u₃ = ∑tendencies) + (; 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) :sgsʲs in propertynames(ᶠY) || return () @@ -103,6 +171,64 @@ function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) end +water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = + @. lazy(ᶜprecipdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜχ)))) + +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 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 + NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) Yₜ_lim .= zero(eltype(Yₜ_lim)) device = ClimaComms.device(axes(Y.c)) @@ -111,7 +237,12 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) else Val(false), Val(false) end - p_kernel = (;) + (; moisture_model, viscous_sponge, precip_model) = p.atmos + p_kernel = (; + atmos = p.atmos, + params = p.params, + dt = p.dt, + ) if :sfc in propertynames(Y) # columnwise! does not yet handle .sfc parent(Yₜ.sfc) .= zero(Spaces.undertype(axes(Y.c))) end @@ -133,7 +264,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 @@ -161,10 +291,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...) @@ -172,13 +298,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 From b0c1d2db42ffc75c56618348e1ad1ca3b5b18a88 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Tue, 3 Jun 2025 15:15:42 -0400 Subject: [PATCH 03/21] Fix NullBroadcasts issues --- .../remaining_tendency.jl | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 9a1e589646..8618e978fc 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -49,17 +49,23 @@ function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) if !(p.atmos.moisture_model isa DryModel) ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) - ∑tendencies = lazy.(∑tendencies .- water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) + ∑tendencies = sub_tend(∑tendencies, water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) end return (;ρ=∑tendencies) end + +add_tend(∑tends, t) = lazy.(∑tends .+ t) +add_tend(∑tends, ::NullBroadcasted) = ∑tends +sub_tend(∑tends, ::NullBroadcasted) = ∑tends +sub_tend(∑tends, t) = lazy.(∑tends .- t) + function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t) :uₕ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.uₕ)) (; viscous_sponge, rayleigh_sponge) = p.atmos ᶜuₕ = ᶜY.uₕ - ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)) - ∑tendencies = lazy.(∑tendencies .+ rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)) + ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)) + ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)) return (;uₕ=∑tendencies) end @@ -82,19 +88,19 @@ 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ₜ)) + ∑tendencies = sub_tend(∑tendencies, water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₕhₜ)) end if energy_upwinding != Val(:none) (; dt) = p ᶠu³ = compute_ᶠu³(ᶜY, ᶠY) vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) - ∑tendencies = lazy.(∑tendencies .+ vtt) + ∑tendencies = add_tend(∑tendencies, vtt) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) # need to improve NullBroadcast support for this. - ∑tendencies = lazy.(∑tendencies .+ (-1) .* vtt_central) + ∑tendencies = sub_tend(∑tendencies, vtt_central) end - ∑tendencies = lazy.(∑tendencies .+ viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) + ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) return (;ρe_tot=∑tendencies) end function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) @@ -108,7 +114,7 @@ function ᶜremaining_tendency_ρq_tot(ᶜ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ₜ)) + ∑tendencies = sub_tend(∑tendencies, water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) end (; tracer_upwinding) = p.atmos.numerics if !(p.atmos.moisture_model isa DryModel) && tracer_upwinding != Val(:none) @@ -118,7 +124,8 @@ function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) ᶠ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) + ∑tendencies = add_tend(∑tendencies, vtt) + ∑tendencies = sub_tend(∑tendencies, vtt_central) end return (;ρq_tot=∑tendencies) @@ -161,7 +168,7 @@ function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t) ᶜ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)) + ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)) return (;u₃=∑tendencies) end function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) From 55ffc8cf624c64713f79c3149e678bafa5328e44 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 08:31:47 -0400 Subject: [PATCH 04/21] Fix broadcasting --- .../precipitation_precomputed_quantities.jl | 28 +++++++++---------- .../remaining_tendency.jl | 24 +++++++++++----- 2 files changed, 31 insertions(+), 21 deletions(-) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index a0b89826bd..d31f0c4744 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -34,31 +34,31 @@ function compute_precip_velocities(ᶜY, ᶠY, p, t) 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( + ᶜ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ₗ = CMNe.terminal_velocity( + ᶜwₗ = @. lazy(CMNe.terminal_velocity( cmc.liquid, cmc.Ch2022.rain, ᶜY.ρ, max(zero(ᶜY.ρ), ᶜY.ρq_liq / ᶜY.ρ), - ) - ᶜwᵢ = CMNe.terminal_velocity( + )) + ᶜ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 @@ -73,12 +73,12 @@ function compute_ᶜwₜqₜ(ᶜ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 Geometry.WVector( + return @. lazy(Geometry.WVector( ᶜwₗ * ᶜY.ρq_liq + ᶜwᵢ * ᶜY.ρq_ice + ᶜwᵣ * ᶜY.ρq_rai + ᶜwₛ * ᶜY.ρq_sno, - ) / ᶜY.ρ + ) / ᶜY.ρ) end compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) = diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 8618e978fc..a9d93c2b34 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -10,12 +10,20 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) end -function prognostic_nt(::Val{names}; tends...) where {names} - nt_ordered = NamedTuple{names}(rzero(values(tends))) - nt_values = NamedTuple{keys(tends)}(values(tends)) - return merge(nt_ordered, nt_values) +@generated function sorted_nt(::Val{snames}, ::Val{unames}, vals...) where {snames,unames} + svals_exprs = [] + for sn in snames + i = findfirst(un -> un == sn, unames)::Int + push!(svals_exprs, :(getfield(vals, $i))) + end + return quote + NamedTuple{snames}(($(svals_exprs...),)) + end end +prognostic_nt(::Val{names}, ::Val{K}, vals...) where {names, K} = + sorted_nt(Val(names), Val(K), vals...) + function ᶜremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶜY) tends = (; @@ -30,7 +38,7 @@ function ᶜremaining_tendency(ᶜY, ᶠY, p, t) ᶜremaining_tendency_sgs⁰(ᶜY, ᶠY, p, t)..., ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., ) - return lazy.(prognostic_nt.(Val(names); tends...)) + return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) end function ᶠremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶠY) @@ -38,7 +46,7 @@ function ᶠremaining_tendency(ᶜY, ᶠY, p, t) ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)..., ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., ) - return lazy.(prognostic_nt.(Val(names); tends...)) + return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) end using ClimaCore.RecursiveApply: rzero function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) @@ -46,6 +54,7 @@ function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) ∑tendencies = zero(eltype(ᶜY.ρ)) ᶜ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) @@ -178,7 +187,8 @@ function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) end -water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = +water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ::Real) = zero(eltype(ᶜρ)) +water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = # only valid when ᶜχ is a field @. lazy(ᶜprecipdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜχ)))) function surface_velocity_full(ᶠu₃, ᶠuₕ³) From e3f10da5899320d621acc2644d26db66e791f089 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 11:14:45 -0400 Subject: [PATCH 05/21] Increase mem request --- .buildkite/pipeline.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 997026e384..70c1825967 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -328,6 +328,7 @@ steps: artifact_paths: "baroclinic_wave_equil/output_active/*" agents: slurm_constraint: icelake|cascadelake|skylake|epyc + slurm_mem: 20GB # columnwise! - label: ":computer: held suarez" key: held_suarez From ac8867cf0071eddda660300378e78c3ccc013ea8 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 11:33:14 -0400 Subject: [PATCH 06/21] Try with avoiding gpu-incompatible remapper logic --- .buildkite/Manifest-v1.11.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index d88a3bd902..596e8e02fe 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.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 = "01c9e1119f235499f01e213da778593de55d02f9" +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"] From 9a9c93ff80f9627db54a49cc1a6a323f73d16997 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 11:33:41 -0400 Subject: [PATCH 07/21] Fix buildkite pipeline --- .buildkite/pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 70c1825967..2ed93984fc 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -328,7 +328,7 @@ steps: artifact_paths: "baroclinic_wave_equil/output_active/*" agents: slurm_constraint: icelake|cascadelake|skylake|epyc - slurm_mem: 20GB # columnwise! + slurm_mem: 20GB # columnwise! - label: ":computer: held suarez" key: held_suarez From 4f552df8d8c7e6bb25d700d0a9734376b688ea93 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 11:55:43 -0400 Subject: [PATCH 08/21] CC fixes --- .buildkite/Manifest-v1.11.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index 596e8e02fe..12de14743e 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -381,7 +381,7 @@ 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 = "01c9e1119f235499f01e213da778593de55d02f9" +git-tree-sha1 = "9f97f9e9b2957a8651c83c37c59dd21ef2964f7f" repo-rev = "ck/fixes" repo-url = "https://github.com/CliMA/ClimaCore.jl.git" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" From bcda95d8384115c5a0169d00f375bcababc32390 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 12:17:16 -0400 Subject: [PATCH 09/21] Fix undefined var --- src/prognostic_equations/remaining_tendency.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index a9d93c2b34..a9a46e659f 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -87,6 +87,8 @@ function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) thermo_params = CAP.thermodynamics_params(p.params) thermo_args = (thermo_params, moisture_model, precip_model) ᶜz = Fields.coordinate_field(ᶜY).z + ᶜJ = Fields.local_geometry_field(ᶜY).J + ᶠJ = Fields.local_geometry_field(ᶠY).J ᶜρ = ᶜY.ρ ᶜuₕ = ᶜY.uₕ ᶜρe_tot = ᶜY.ρe_tot From 380d808136cb51b84141a66a373efcad01e33337 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 12:20:22 -0400 Subject: [PATCH 10/21] Fix scoping issue --- .buildkite/Manifest-v1.11.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index 12de14743e..f118588ffd 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -381,7 +381,7 @@ 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 = "9f97f9e9b2957a8651c83c37c59dd21ef2964f7f" +git-tree-sha1 = "1fcf9c8267bc3e7aa0016f7c52c5f3f94ea5938b" repo-rev = "ck/fixes" repo-url = "https://github.com/CliMA/ClimaCore.jl.git" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" From 494cf02c352676f83363d6afe042a9be377bb871 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 14:14:02 -0400 Subject: [PATCH 11/21] Fix InputOutput writers --- .buildkite/Manifest-v1.11.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index f118588ffd..2998c7df4c 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -381,7 +381,7 @@ 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 = "1fcf9c8267bc3e7aa0016f7c52c5f3f94ea5938b" +git-tree-sha1 = "4085d9d00635b3780533539e12d2e2af5ffc6cb2" repo-rev = "ck/fixes" repo-url = "https://github.com/CliMA/ClimaCore.jl.git" uuid = "d414da3d-4745-48bb-8d80-42e94e092884" From dcacdaf1693e89613077489b3e05244332ac6199 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 14:21:34 -0400 Subject: [PATCH 12/21] Fix undefined precomputed --- src/prognostic_equations/remaining_tendency.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index a9a46e659f..df34f0a3f2 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -130,7 +130,6 @@ function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) (; 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) From 326d6f772e842abf28b6bd1528b850d497fe3abc Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 15:19:48 -0400 Subject: [PATCH 13/21] Fix zero call --- src/cache/precipitation_precomputed_quantities.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index d31f0c4744..09eb38e219 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -84,7 +84,7 @@ 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.ρ) +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) From 062073e2ad7e648c9c73cc590a3b6ea16122cd54 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 15:45:50 -0400 Subject: [PATCH 14/21] Increase mem request --- .buildkite/pipeline.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2ed93984fc..fd84872068 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -328,7 +328,8 @@ steps: artifact_paths: "baroclinic_wave_equil/output_active/*" agents: slurm_constraint: icelake|cascadelake|skylake|epyc - slurm_mem: 20GB # columnwise! + # slurm_mem: 20GB # columnwise! insufficient mem + slurm_mem: 60GB # columnwise! - label: ":computer: held suarez" key: held_suarez From 9cd8d0789ec9164559ed768f602713b45739c2ae Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 15:47:28 -0400 Subject: [PATCH 15/21] Try baro wave on cpu --- .buildkite/pipeline.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index fd84872068..05a82d65a3 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -902,11 +902,11 @@ steps: --config_file $CONFIG_PATH/baroclinic_wave.yml --job_id baroclinic_wave_gpu artifact_paths: "baroclinic_wave_gpu/output_active/*" - env: - CLIMACOMMS_DEVICE: "CUDA" + # env: + # CLIMACOMMS_DEVICE: "CUDA" agents: slurm_gpus: 1 - slurm_mem: 16GB + slurm_mem: 60GB - label: "GPU: compare BW with CPU" command: > From e8ba9ddca655960a4921cf44be4de0ddcf0cacf1 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 4 Jun 2025 22:26:22 -0400 Subject: [PATCH 16/21] Try with cuda + more mem --- .buildkite/pipeline.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 05a82d65a3..7ac2d86542 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -902,8 +902,8 @@ steps: --config_file $CONFIG_PATH/baroclinic_wave.yml --job_id baroclinic_wave_gpu artifact_paths: "baroclinic_wave_gpu/output_active/*" - # env: - # CLIMACOMMS_DEVICE: "CUDA" + env: + CLIMACOMMS_DEVICE: "CUDA" agents: slurm_gpus: 1 slurm_mem: 60GB From 6c232affc2920579edb2c1186df6111165a86a67 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 5 Jun 2025 09:53:12 -0400 Subject: [PATCH 17/21] Disable ClimaCore.Fields.error_mismatched_spaces --- .buildkite/ci_driver.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.buildkite/ci_driver.jl b/.buildkite/ci_driver.jl index 3207bcd3a0..23d4db12b9 100644 --- a/.buildkite/ci_driver.jl +++ b/.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 From e04ca2701874176139d5722e72c99f24976e9361 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 5 Jun 2025 13:08:10 -0400 Subject: [PATCH 18/21] Try to fix gpu inference --- .../remaining_tendency.jl | 101 ++++++++++-------- 1 file changed, 59 insertions(+), 42 deletions(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index df34f0a3f2..6fb0365d3f 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -24,32 +24,49 @@ end prognostic_nt(::Val{names}, ::Val{K}, vals...) where {names, K} = sorted_nt(Val(names), Val(K), vals...) +@generated function construct_tends(::Val{names}, f, ᶜY, ᶠY, p, t) where {names} + calls = [] + for name in names + push!(calls, :(f(Val($name), ᶜY, ᶠY, p, t))) + end + return quote + $(calls...) + end +end + 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_sgs⁰(ᶜY, ᶠY, p, t)..., - ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)..., - ) - return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) + tends = construct_tends(Val(names), ᶜremaining_tendency, ᶜY, ᶠY, p, t) + # tends = (; + # ᶜremaining_tendency(Val(:ρ), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:uₕ), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρe_tot), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρq_tot), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρq_liq), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρq_ice), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρq_rai), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:ρq_sno), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:sgs⁰), ᶜY, ᶠY, p, t)..., + # ᶜremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t)..., + # ) + + # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) + prog_nt = (NamedTuple{names} ∘ tuple) + return lazy.(prog_nt.(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), Val(keys(tends)), values(tends)...)) + tends = construct_tends(Val(names), ᶠremaining_tendency, ᶜY, ᶠY, p, t) + # tends = (; + # ᶠremaining_tendency(Val(:u₃), ᶜY, ᶠY, p, t)..., + # ᶠremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t)..., + # ) + # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) + prog_nt = (NamedTuple{names} ∘ tuple) + return lazy.(prog_nt.(tends)) end using ClimaCore.RecursiveApply: rzero -function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρ}, ᶜY, ᶠY, p, t) :ρ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρ)) ᶜJ = Fields.local_geometry_field(ᶜY).J @@ -60,7 +77,7 @@ function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t) ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t) ∑tendencies = sub_tend(∑tendencies, water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜwₜqₜ)) end - return (;ρ=∑tendencies) + return ∑tendencies end add_tend(∑tends, t) = lazy.(∑tends .+ t) @@ -68,7 +85,7 @@ add_tend(∑tends, ::NullBroadcasted) = ∑tends sub_tend(∑tends, ::NullBroadcasted) = ∑tends sub_tend(∑tends, t) = lazy.(∑tends .- t) -function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:uₕ}, ᶜY, ᶠY, p, t) :uₕ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.uₕ)) (; viscous_sponge, rayleigh_sponge) = p.atmos @@ -76,9 +93,9 @@ function ᶜremaining_tendency_uₕ(ᶜY, ᶠY, p, t) ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)) ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)) - return (;uₕ=∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρe_tot}, ᶜY, ᶠY, p, t) :ρe_tot in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρe_tot)) @@ -112,9 +129,9 @@ function ᶜremaining_tendency_ρe_tot(ᶜY, ᶠY, p, t) end ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge)) - return (;ρe_tot=∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρ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 @@ -138,39 +155,39 @@ function ᶜremaining_tendency_ρq_tot(ᶜY, ᶠY, p, t) ∑tendencies = sub_tend(∑tendencies, vtt_central) end - return (;ρq_tot=∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρq_liq(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρq_liq}, ᶜY, ᶠY, p, t) :ρq_liq in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρq_liq)) - return (; ρq_liq = ∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρq_ice(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρq_ice}, ᶜY, ᶠY, p, t) :ρq_ice in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρq_ice)) - return (; ρq_ice = ∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρq_rai(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρq_rai}, ᶜY, ᶠY, p, t) :ρq_rai in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρq_rai)) - return (; ρq_rai = ∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_ρq_sno(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:ρq_sno}, ᶜY, ᶠY, p, t) :ρq_sno in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.ρq_sno)) - return (; ρq_sno = ∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t) :sgsʲs in propertynames(ᶜY) || return () ∑tendencies = rzero(eltype(ᶜY.sgsʲs)) - return (; sgsʲs = ∑tendencies) + return ∑tendencies end -function ᶜremaining_tendency_sgs⁰(ᶜY, ᶠY, p, t) +function ᶜremaining_tendency(::Val{:sgs⁰}, ᶜY, ᶠY, p, t) :sgs⁰ in propertynames(ᶜY) || return () ∑tendencies = rzero(eltype(ᶜY.sgs⁰)) - return (; sgs⁰ = ∑tendencies) + return ∑tendencies end -function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t) +function ᶠremaining_tendency(::Val{:u₃}, ᶜY, ᶠY, p, t) :u₃ in propertynames(ᶠY) || return () ∑tendencies = zero(eltype(ᶠY.u₃)) (; viscous_sponge) = p.atmos @@ -179,12 +196,12 @@ function ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t) ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ) ᶠu₃ = compute_ᶠu₃_with_bcs(ᶠY.u₃, ᶠuₕ³) ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)) - return (;u₃=∑tendencies) + return ∑tendencies end -function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t) +function ᶠremaining_tendency(::Val{:sgsʲs}, ᶜY, ᶠY, p, t) :sgsʲs in propertynames(ᶠY) || return () ∑tendencies = rzero(eltype(ᶠY.sgsʲs)) - return (; sgsʲs = ∑tendencies) + return ∑tendencies end From b6ce4c4a680144deec1043978c718260e83d7d5b Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 5 Jun 2025 14:18:50 -0400 Subject: [PATCH 19/21] wip --- .../remaining_tendency.jl | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 6fb0365d3f..7759b278cc 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -27,33 +27,38 @@ prognostic_nt(::Val{names}, ::Val{K}, vals...) where {names, K} = @generated function construct_tends(::Val{names}, f, ᶜY, ᶠY, p, t) where {names} calls = [] for name in names - push!(calls, :(f(Val($name), ᶜY, ᶠY, p, t))) + push!(calls, :(f(Val($(QuoteNode(name))), ᶜY, ᶠY, p, t))) end return quote - $(calls...) + ($(calls...),) end end function ᶜremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶜY) tends = construct_tends(Val(names), ᶜremaining_tendency, ᶜY, ᶠY, p, t) - # tends = (; - # ᶜremaining_tendency(Val(:ρ), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:uₕ), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρe_tot), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρq_tot), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρq_liq), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρq_ice), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρq_rai), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:ρq_sno), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:sgs⁰), ᶜY, ᶠY, p, t)..., - # ᶜremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t)..., + @show length(tends) + @show length(names) + # tends = ( + # ᶜremaining_tendency(Val(:ρ), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:uₕ), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρe_tot), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρq_tot), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρq_liq), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρq_ice), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρq_rai), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:ρq_sno), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:sgs⁰), ᶜY, ᶠY, p, t), + # ᶜremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t), # ) # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) - prog_nt = (NamedTuple{names} ∘ tuple) - return lazy.(prog_nt.(tends)) + # prog_nt = (NamedTuple{names} ∘ tuple) + # return lazy.(prog_nt.(tends)) + return lazy.(foo.(Val(names), tends...)) end +foo(::Val{names}, tends...) where {names} = NamedTuple{names}(tends) + function ᶠremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶠY) tends = construct_tends(Val(names), ᶠremaining_tendency, ᶜY, ᶠY, p, t) @@ -62,8 +67,9 @@ function ᶠremaining_tendency(ᶜY, ᶠY, p, t) # ᶠremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t)..., # ) # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) - prog_nt = (NamedTuple{names} ∘ tuple) - return lazy.(prog_nt.(tends)) + # prog_nt = (NamedTuple{names} ∘ tuple) + # return lazy.(prog_nt.(tends)) + return lazy.(foo.(Val(names), tends...)) end using ClimaCore.RecursiveApply: rzero function ᶜremaining_tendency(::Val{:ρ}, ᶜY, ᶠY, p, t) From f67ed0a7ff2e36d11f98dc0ff24b9a9fa5ce0665 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 5 Jun 2025 15:44:36 -0400 Subject: [PATCH 20/21] Remove show calls --- .../remaining_tendency.jl | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 7759b278cc..1499b4dfdc 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -37,24 +37,6 @@ end function ᶜremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶜY) tends = construct_tends(Val(names), ᶜremaining_tendency, ᶜY, ᶠY, p, t) - @show length(tends) - @show length(names) - # tends = ( - # ᶜremaining_tendency(Val(:ρ), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:uₕ), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρe_tot), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρq_tot), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρq_liq), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρq_ice), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρq_rai), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:ρq_sno), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:sgs⁰), ᶜY, ᶠY, p, t), - # ᶜremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t), - # ) - - # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) - # prog_nt = (NamedTuple{names} ∘ tuple) - # return lazy.(prog_nt.(tends)) return lazy.(foo.(Val(names), tends...)) end foo(::Val{names}, tends...) where {names} = NamedTuple{names}(tends) @@ -62,13 +44,6 @@ foo(::Val{names}, tends...) where {names} = NamedTuple{names}(tends) function ᶠremaining_tendency(ᶜY, ᶠY, p, t) names = propertynames(ᶠY) tends = construct_tends(Val(names), ᶠremaining_tendency, ᶜY, ᶠY, p, t) - # tends = (; - # ᶠremaining_tendency(Val(:u₃), ᶜY, ᶠY, p, t)..., - # ᶠremaining_tendency(Val(:sgsʲs), ᶜY, ᶠY, p, t)..., - # ) - # return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...)) - # prog_nt = (NamedTuple{names} ∘ tuple) - # return lazy.(prog_nt.(tends)) return lazy.(foo.(Val(names), tends...)) end using ClimaCore.RecursiveApply: rzero From aafb36a38ff1d29bd4ce9d9f16c9d1d75643cef3 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 5 Jun 2025 16:11:44 -0400 Subject: [PATCH 21/21] Put zmax into cache, for rayleigh sponge --- src/parameterized_tendencies/sponge/rayleigh_sponge.jl | 8 +++++++- src/prognostic_equations/remaining_tendency.jl | 6 +++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/parameterized_tendencies/sponge/rayleigh_sponge.jl b/src/parameterized_tendencies/sponge/rayleigh_sponge.jl index ae69849958..defdd71325 100644 --- a/src/parameterized_tendencies/sponge/rayleigh_sponge.jl +++ b/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 diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 1499b4dfdc..c0fab53a90 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -69,10 +69,13 @@ sub_tend(∑tends, t) = lazy.(∑tends .- t) function ᶜremaining_tendency(::Val{:uₕ}, ᶜY, ᶠY, p, t) :uₕ in propertynames(ᶜY) || return () ∑tendencies = zero(eltype(ᶜY.uₕ)) + (; zmax) = p (; viscous_sponge, rayleigh_sponge) = p.atmos + ᶜz = Fields.coordinate_field(ᶜY).z + ᶠz = Fields.coordinate_field(ᶠY).z ᶜuₕ = ᶜY.uₕ ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)) - ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)) + ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge, ᶜz, ᶠz, zmax)) return ∑tendencies end @@ -255,6 +258,7 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t) end (; moisture_model, viscous_sponge, precip_model) = p.atmos p_kernel = (; + zmax = Spaces.z_max(axes(Y.f)), atmos = p.atmos, params = p.params, dt = p.dt,