From c7c74f9f8f736285d97ec9ff68521b96d47481df Mon Sep 17 00:00:00 2001 From: Maximilian Gelbrecht Date: Sat, 14 Dec 2024 00:28:56 +0100 Subject: [PATCH 1/4] Roll back GPUArrays upgrade (#636) * roll back GPUArrays upgrade * changelog, GPUArrays downgrade * backend instead of get_backend * KA compat 0.9.29 * KA compat again --- CHANGELOG.md | 1 + Project.toml | 6 +++--- src/LowerTriangularMatrices/lower_triangular_array.jl | 6 +++--- src/RingGrids/general.jl | 8 ++++---- test/Project.toml | 1 + 5 files changed, 12 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 325775f1e..9f113bbda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636) - Change default timestep to 40min at T31 [#623](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/623) ## v0.13 diff --git a/Project.toml b/Project.toml index 0c7b8b0c5..a9745a00f 100644 --- a/Project.toml +++ b/Project.toml @@ -56,12 +56,12 @@ Enzyme = "0.13" FFTW = "1" FastGaussQuadrature = "0.4, 0.5, 1" FiniteDifferences = "0.12" -GPUArrays = "11" +GPUArrays = "10" GenericFFT = "0.1" GeoMakie = "0.7.6" -JLArrays = "0.1.4, 0.2" +JLArrays = "0.1" JLD2 = "0.4, 0.5" -KernelAbstractions = "0.9" +KernelAbstractions = "0.9.0 - 0.9.29" LinearAlgebra = "1.10" Makie = "0.20, 0.21" NCDatasets = "0.12, 0.13, 0.14" diff --git a/src/LowerTriangularMatrices/lower_triangular_array.jl b/src/LowerTriangularMatrices/lower_triangular_array.jl index b38659f66..b0d0cb288 100644 --- a/src/LowerTriangularMatrices/lower_triangular_array.jl +++ b/src/LowerTriangularMatrices/lower_triangular_array.jl @@ -630,10 +630,10 @@ function Base.similar( return LowerTriangularArray{T, N, ArrayType{T,N}}(undef, size(L; as=Matrix)) end -function KernelAbstractions.get_backend( - a::LowerTriangularArray{T, N, ArrayType} +function GPUArrays.backend( + ::Type{LowerTriangularArray{T, N, ArrayType}} ) where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} - return KernelAbstractions.get_backend(a.data) + return GPUArrays.backend(ArrayType) end Adapt.adapt_structure(to, L::LowerTriangularArray) = diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 42d6ca185..cbfc52f4f 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -481,11 +481,11 @@ AbstractGPUGridArrayStyle{2, ArrayType, Grid}(::Val{3}) where {ArrayType, Grid} AbstractGPUGridArrayStyle{2, ArrayType, Grid}(::Val{1}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{2, ArrayType, Grid}() AbstractGPUGridArrayStyle{3, ArrayType, Grid}(::Val{4}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{4, ArrayType, Grid}() AbstractGPUGridArrayStyle{3, ArrayType, Grid}(::Val{2}) where {ArrayType, Grid} = AbstractGPUGridArrayStyle{3, ArrayType, Grid}() - -function KernelAbstractions.get_backend( - g::Grid + +function GPUArrays.backend( + ::Type{Grid} ) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} - return KernelAbstractions.get_backend(g.data) + return GPUArrays.backend(ArrayType) end function Base.similar( diff --git a/test/Project.toml b/test/Project.toml index ede81c90c..838ec629a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" From d24e4d118d68ed560815e95f36a06afe05e0e750 Mon Sep 17 00:00:00 2001 From: milankl Date: Fri, 13 Dec 2024 23:59:14 +0000 Subject: [PATCH 2/4] ConvectiveHeating implemented --- src/physics/convection.jl | 59 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/src/physics/convection.jl b/src/physics/convection.jl index 44c57421b..a0a0bc6dc 100644 --- a/src/physics/convection.jl +++ b/src/physics/convection.jl @@ -430,4 +430,63 @@ function dry_adiabat!( # level of zero buoyancy is reached when the loop stops, but in case it's at the top it's still buoyant level_zero_buoyancy = k + (1-buoyant) return level_zero_buoyancy +end + +export ConvectiveHeating +@kwdef struct ConvectiveHeating{NF} <: AbstractConvection + # DIMENSION + nlat::Int + + "[OPTION] Q_max heating strength as 1K/time_scale" + time_scale::Second = Hour(12) + + "[OPTION] Pressure of maximum heating [hPa]" + p₀::NF = 525 + + "[OPTION] Vertical extent of heating [hPa]" + σₚ::NF = 200 + + "[OPTION] Latitude of heating [˚N]" + θ₀::NF = 0 + + "[OPTION] Latitudinal width of heating [˚]" + σθ::NF = 20 + + # precomputed latitude mask + lat_mask::Vector{NF} = zeros(NF, nlat) +end + +# generator +ConvectiveHeating(SG::SpectralGrid; kwargs...) = ConvectiveHeating{SG.NF}(nlat=SG.nlat; kwargs...) + +function initialize!(C::ConvectiveHeating, model::PrimitiveEquation) + + (; latd) = model.geometry + (; θ₀, σθ) = C + + for (j, θ) in enumerate(latd) + # Lee and Kim, 2003, eq. 2 + C.lat_mask[j] = cosd((θ-θ₀)/σθ)^2 + end +end + +function convection!( + column::ColumnVariables, + scheme::ConvectiveHeating, + model::PrimitiveEquation, +) + # escape immediately if not in the tropics + abs(column.latd) > scheme.σθ && return nothing + + p₀ = scheme.p₀*100 # hPa -> Pa + σₚ = scheme.σₚ*100 # hPa -> Pa + cos²θ_term = scheme.lat_mask[column.jring] + Qmax = 1/Second(scheme.time_scale).value + + for k in eachindex(column) + p = column.pres[k] # Pressure in Pa + + # Lee and Kim, 2003, eq. 2, + column.temp_tend[k] += Qmax*exp(-((p-p₀)/σₚ)^2 / 2)*cos²θ_term + end end \ No newline at end of file From 464a9a5936c028135179c6f0a58c84aefa2f6767 Mon Sep 17 00:00:00 2001 From: milankl Date: Sat, 14 Dec 2024 00:01:23 +0000 Subject: [PATCH 3/4] remove ConvectiveHeating --- src/physics/convection.jl | 59 --------------------------------------- 1 file changed, 59 deletions(-) diff --git a/src/physics/convection.jl b/src/physics/convection.jl index a0a0bc6dc..44c57421b 100644 --- a/src/physics/convection.jl +++ b/src/physics/convection.jl @@ -430,63 +430,4 @@ function dry_adiabat!( # level of zero buoyancy is reached when the loop stops, but in case it's at the top it's still buoyant level_zero_buoyancy = k + (1-buoyant) return level_zero_buoyancy -end - -export ConvectiveHeating -@kwdef struct ConvectiveHeating{NF} <: AbstractConvection - # DIMENSION - nlat::Int - - "[OPTION] Q_max heating strength as 1K/time_scale" - time_scale::Second = Hour(12) - - "[OPTION] Pressure of maximum heating [hPa]" - p₀::NF = 525 - - "[OPTION] Vertical extent of heating [hPa]" - σₚ::NF = 200 - - "[OPTION] Latitude of heating [˚N]" - θ₀::NF = 0 - - "[OPTION] Latitudinal width of heating [˚]" - σθ::NF = 20 - - # precomputed latitude mask - lat_mask::Vector{NF} = zeros(NF, nlat) -end - -# generator -ConvectiveHeating(SG::SpectralGrid; kwargs...) = ConvectiveHeating{SG.NF}(nlat=SG.nlat; kwargs...) - -function initialize!(C::ConvectiveHeating, model::PrimitiveEquation) - - (; latd) = model.geometry - (; θ₀, σθ) = C - - for (j, θ) in enumerate(latd) - # Lee and Kim, 2003, eq. 2 - C.lat_mask[j] = cosd((θ-θ₀)/σθ)^2 - end -end - -function convection!( - column::ColumnVariables, - scheme::ConvectiveHeating, - model::PrimitiveEquation, -) - # escape immediately if not in the tropics - abs(column.latd) > scheme.σθ && return nothing - - p₀ = scheme.p₀*100 # hPa -> Pa - σₚ = scheme.σₚ*100 # hPa -> Pa - cos²θ_term = scheme.lat_mask[column.jring] - Qmax = 1/Second(scheme.time_scale).value - - for k in eachindex(column) - p = column.pres[k] # Pressure in Pa - - # Lee and Kim, 2003, eq. 2, - column.temp_tend[k] += Qmax*exp(-((p-p₀)/σₚ)^2 / 2)*cos²θ_term - end end \ No newline at end of file From b122e2f059c57f7cae75ece746a7cbf70e419325 Mon Sep 17 00:00:00 2001 From: Maximilian Gelbrecht Date: Sat, 14 Dec 2024 01:17:35 +0100 Subject: [PATCH 4/4] RingGrinds more general indexing with Colon (less escapes) (#637) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * RingGrinds more general indexing with Colon (less escapes) * changelog RingGrids colon indexing --------- Co-authored-by: Milan Klöwer --- CHANGELOG.md | 1 + src/RingGrids/general.jl | 2 +- test/grids.jl | 3 +++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f113bbda..dad761940 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- RingGrids indexing with leading Colon should now always return another RingGrid instance [#637](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/637) - Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636) - Change default timestep to 40min at T31 [#623](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/623) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index cbfc52f4f..fedfb1b8b 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -76,7 +76,7 @@ Base.@propagate_inbounds Base.getindex(G::AbstractGridArray, ijk...) = getindex( @inline function Base.getindex( G::GridArray, col::Colon, - k::Integer..., + k..., ) where {GridArray<:AbstractGridArray} GridArray_ = nonparametric_type(GridArray) # obtain parameters from G.data return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) diff --git a/test/grids.jl b/test/grids.jl index a11683b39..bbb1bd936 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -380,6 +380,9 @@ end @test grid[1:2, 1:2, 1:2] == grid.data[1:2, 1:2, 1:2] @test grid[1, 1, :] == grid.data[1, 1, :] + @test SpeedyWeather.RingGrids.nonparametric_type(typeof(grid[:,1:2,1:2])) <: RingGrids.nonparametric_type(G) + @test grid[:, 1:2, 1:2].data == grid.data[:, 1:2, 1:2] + idx = CartesianIndex((1, 2, 3)) @test grid[idx] == grid.data[idx]