Merge branch 'main' into fix/bug-set!

CHANGELOG.md

@@ -3,6 +3,8 @@
 ## Unreleased
 
 - Number format flexibility with set! [#634](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/634)
+- 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)
 
 ## v0.13

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"

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

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)
@@ -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(

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"

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]