Merge branch 'main' into mk/radiation

SpeedyWeather · Dec 2, 2024 · 659691f · 659691f
2 parents 22c4ce2 + f943780
commit 659691f
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 40 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,8 +3,10 @@
 ## Unreleased
 
 - One-band longwave radiation [#624](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/624)
+- compat entry for FiniteDifferences.jl [#620](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/620)
+- update to GPUArrays v11, JLArrays v0.2 and remove CUDA from tests [#590](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/590)
 - Power spectrum for n-dim LowerTriangularArrays [#618](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/618)
-- Added custom EnzymeRules for the SpectralTransform and an extension for compatibility with FiniteDifferences.jl [#589](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/589)
+- Added custom EnzymeRules for the SpectralTransform and an extension for compatibility with FiniteDifferences.jl [#589](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/589) [#625](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/625)
 - Optical depth introduced and array-agnostic ColumnVariables [#606](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/606)
 - Include large-scale condensation tests [#615](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/615)
 - bugfix: large-scale condensation also at <100% [#609](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/609)

diff --git a/Project.toml b/Project.toml
@@ -54,10 +54,11 @@ Dates = "1.10"
 DocStringExtensions = "0.9"
 FFTW = "1"
 FastGaussQuadrature = "0.4, 0.5, 1"
-GPUArrays = "10"
+FiniteDifferences = "0.12"
+GPUArrays = "11"
 GenericFFT = "0.1"
-JLArrays = "0.1"
 GeoMakie = "0.7.6"
+JLArrays = "0.1.4, 0.2"
 JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1.10"

diff --git a/docs/Project.toml b/docs/Project.toml
@@ -6,8 +6,8 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [compat]
-CairoMakie = "0.11, 0.12"
+CairoMakie = "0.12"
 Documenter = "0.26, 0.27, 1"
-GeoMakie = "0.7.5"
+GeoMakie = "0.7.6"
 NCDatasets = "0.12, 0.13, 0.14"
 UnicodePlots = "3"
diff --git a/src/LowerTriangularMatrices/LowerTriangularMatrices.jl b/src/LowerTriangularMatrices/LowerTriangularMatrices.jl
@@ -6,6 +6,7 @@ using DocStringExtensions
 # GPU
 import Adapt
 import GPUArrays
+import KernelAbstractions
 
 # NUMERICS
 import LinearAlgebra: tril!

diff --git 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 GPUArrays.backend(
-    ::Type{LowerTriangularArray{T, N, ArrayType}}
+function KernelAbstractions.get_backend(
+    a::LowerTriangularArray{T, N, ArrayType}
 ) where {T, N, ArrayType <: GPUArrays.AbstractGPUArray}
-    return GPUArrays.backend(ArrayType)
+    return KernelAbstractions.get_backend(a.data)
 end
 
 Adapt.adapt_structure(to, L::LowerTriangularArray) =

diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl
@@ -12,6 +12,7 @@ import LinearAlgebra
 # GPU
 import Adapt
 import GPUArrays
+import KernelAbstractions
 
 # ABSTRACT GRIDS (2D) AND GRIDARRAYS (3D+)
 export  AbstractGridArray,

diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl
@@ -482,10 +482,10 @@ AbstractGPUGridArrayStyle{2, ArrayType, Grid}(::Val{1}) where {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 GPUArrays.backend(
-    ::Type{Grid}
+function KernelAbstractions.get_backend(
+    g::Grid
 ) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray}
-    return GPUArrays.backend(ArrayType)
+    return KernelAbstractions.get_backend(g.data)
 end
 
 function Base.similar(

diff --git a/test/Project.toml b/test/Project.toml
@@ -1,7 +1,6 @@
 [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"

diff --git a/test/spectral_transform_ad_rules.jl b/test/spectral_transform_ad_rules.jl
@@ -7,6 +7,7 @@ import AbstractFFTs
 
 grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms 
 grid_dealiasing = [2, 3]
+fd_tests = [true, false] # to save CI time, only do FiniteDifferences test for one of the grids
 
 # currenlty there's an issue with EnzymeTestUtils not being able to work with structs with undefined fields like FFT plans
 # https://github.com/EnzymeAD/Enzyme.jl/issues/1992
@@ -37,7 +38,7 @@ end
                 # these tests don't pass for reduced grids 
                 # this is likely due to FiniteDifferences and not our EnzymeRules 
                 # see comments in https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/589
-                if !(grid_type <: AbstractReducedGridArray)
+                if !(grid_type <: AbstractReducedGridArray) & fd_tests[i_grid]
                     spectral_grid = SpectralGrid(Grid=grid_type, nlayers=1, trunc=5, dealiasing=grid_dealiasing[i_grid])
                     S = SpectralTransform(spectral_grid)
                     grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
@@ -61,44 +62,47 @@ end
         for (i_grid, grid_type) in enumerate(grid_types)
 
             spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
-
-            # forwards 
             S = SpectralTransform(spectral_grid, one_more_degree=true)
             dS = deepcopy(S)
-            grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-            dgrid = zero(grid)
-            specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)
-
-            # seed
-            dspecs = zero(specs)
-            fill!(dspecs, 1+1im)
 
-            autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))
+            if fd_tests[i_grid]
+
+                # forwards 
+                grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+                dgrid = zero(grid)
+                specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)
+
+                # seed
+                dspecs = zero(specs)
+                fill!(dspecs, 1+1im)
+
+                autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))
 
-            # new seed
-            dspecs2 = zero(specs)
-            fill!(dspecs2, 1+1im)
+                # new seed
+                dspecs2 = zero(specs)
+                fill!(dspecs2, 1+1im)
 
-            # finite difference comparision, seeded with a one adjoint to get the direct gradient
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
-            @test isapprox(dgrid, fd_jvp[1])
+                # finite difference comparision, seeded with a one adjoint to get the direct gradient
+                fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
+                @test isapprox(dgrid, fd_jvp[1])
 
-            ## now backwards, as the input for spec we use the output of the forward transform
+                ## now backwards, as the input for spec we use the output of the forward transform
 
-            fill!(dspecs,0)
-            grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
-            dgrid = similar(grid)
-            fill!(dgrid, 1)
+                fill!(dspecs,0)
+                grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
+                dgrid = similar(grid)
+                fill!(dgrid, 1)
 
-            autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))
+                autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))
 
-            # new seed 
-            dgrid2 = similar(grid)
-            fill!(dgrid2, 1)
+                # new seed 
+                dgrid2 = similar(grid)
+                fill!(dgrid2, 1)
 
-            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
+                fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
 
-            @test isapprox(dspecs, fd_jvp[1])
+                @test isapprox(dspecs, fd_jvp[1])
+            end 
 
             # test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains) 
             # this only holds for exact transforms, like Gaussian grids