SpeedyTransforms: Differentiable transforms via custom Enzyme rule

.github/workflows/CI.yml

@@ -14,7 +14,6 @@ jobs:
       matrix:
         version:
           - '1.10'
-          - '1.11'
         os:
           - ubuntu-latest
         arch:

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Added custom EnzymeRules for the SpectralTransform and an extension for compatibility with FiniteDifferences.jl [#589](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/589)
 - Also allow SpectralGrid as positional argument to model constructors [#593](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/593)
 - De-interweave SpectralTransform [#587](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/587)
 - Rossby-Haurwitz wave initial conditions [#591](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/591)

Project.toml

@@ -11,8 +11,11 @@ BitInformation = "de688a37-743e-4ac2-a6f0-bd62414d1aa7"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -29,11 +32,15 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
 SpeedyWeatherCUDAExt = "CUDA"
+SpeedyWeatherEnzymeExt = "Enzyme"
+SpeedyWeatherFiniteDifferencesExt = "FiniteDifferences"
 SpeedyWeatherJLArraysExt = "JLArrays"
 SpeedyWeatherMakieExt = "Makie"
 
@@ -50,7 +57,7 @@ FFTW = "1"
 FastGaussQuadrature = "0.4, 0.5, 1"
 GPUArrays = "10"
 GenericFFT = "0.1"
-JLArrays = "0.1.4"
+JLArrays = "0.1"
 JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1.10"
@@ -64,11 +71,3 @@ Statistics = "1.10"
 TOML = "1"
 UnicodePlots = "3.3"
 julia = "1.10"
-
-[extras]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test", "JLArrays", "CUDA"]

ext/SpeedyWeatherEnzymeExt.jl

@@ -0,0 +1,115 @@
+module SpeedyWeatherEnzymeExt
+
+using SpeedyWeather
+using Enzyme
+import .EnzymeRules: reverse, augmented_primal
+using .EnzymeRules
+
+# import all functions for which we define rules
+import SpeedyWeather.SpeedyTransforms: _fourier!
+
+# Rules for SpeedyTransforms
+
+# _fourier! 
+
+# Computes the scale for the adjoint/pullback of all discrete Fourier transforms. 
+function adjoint_scale(S::SpectralTransform)
+    (; nlat_half, nlons, rfft_plans) = S
+    nfreqs = [rfft_plan.osz[1] for rfft_plan in rfft_plans] # TODO: This works with FFTW, but does it with cuFFT as well?
+
+    scale = zeros(Int, maximum(nfreqs), nlat_half) 
+
+    for i=1:nlat_half
+        scale[1:nfreqs[i],i] = rfft_adjoint_scale(nfreqs[i], nlons[i])
+    end 
+
+    # TODO: transfer array to GPU in case we are on GPU
+    return reshape(scale, maximum(nfreqs), 1, nlat_half) # the scratch memory is (Freq x lvl x lat), so we insert 
+                                                         # an additional dimension here for easier matrix multiply
+end 
+
+# Computes the scale for the adjoint/pullback of a real discrete fourier transform.
+function rfft_adjoint_scale(n_freq::Int, n_real::Int)
+    if iseven(n_real)
+        return [1; [2 for i=2:(n_freq-1)]; 1]
+    else 
+        return [1; [2 for i=2:n_freq]]
+    end 
+end 
+
+### Custom rule for _fourier!(f_north, f_north, grid, S)
+function augmented_primal(config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(_fourier!)}, ::Type{<:Const}, 
+    f_north::Duplicated, f_south::Duplicated, grids::Duplicated{<:AbstractGridArray}, S::Union{Const, MixedDuplicated}) 
+
+    func.val(f_north.val, f_south.val, grids.val, S.val) # forward pass
+
+    # save grids in tape if grids will be overwritten
+    if overwritten(config)[4] # TODO: Not sure this is really necessary because grids won't ever get overwritten by this _fourier!
+        tape = copy(grids.val)
+    else
+        tape = nothing
+    end
+
+    return AugmentedReturn(nothing, nothing, tape) # because the function actually returns nothing
+
+end 
+
+function reverse(config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(_fourier!)}, ::Type{<:Const}, tape,
+    f_north::Duplicated, f_south::Duplicated, grids::Duplicated{<:AbstractGridArray}, S::Union{Const, MixedDuplicated})
+
+    # adjoint/jvp of FFT has a different scaling, compute it, apply it later to f_north, f_south
+    scale = adjoint_scale(S.val)
+
+    # retrieve grids value, either from original grids or from tape if grids may have been overwritten.
+    gridsval = overwritten(config)[4] ? tape : grids.val
+
+    # compute the adjoint
+    dgridval = zero(gridsval)
+    _fourier!(dgridval, f_north.dval ./ scale, f_south.dval ./ scale, S.val) # inverse FFT (w/o normalization)
+    grids.dval .+= dgridval 
+
+    # no derivative wrt the f_north and f_south that were input because they are overwritten
+    make_zero!(f_north.dval) 
+    make_zero!(f_south.dval)
+
+    # the function has no return values, so we also return nothing here
+    return (nothing, nothing, nothing, nothing)
+end
+
+### Custom rule for _fourier!(grid, f_north, f_south, S)
+function augmented_primal(config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(_fourier!)}, ::Type{<:Const}, 
+    grids::Duplicated{<:AbstractGridArray}, f_north::Duplicated, f_south::Duplicated, S::Union{Const, MixedDuplicated}) 
+
+    func.val(grids.val, f_north.val, f_south.val, S.val) # forward pass
+
+    # TODO: make an overwritten check here? 
+
+    return AugmentedReturn(nothing, nothing, nothing) # because the function actually returns nothing
+
+end 
+
+function reverse(config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(_fourier!)}, ::Type{<:Const}, tape,
+    grids::Duplicated{<:AbstractGridArray}, f_north::Duplicated, f_south::Duplicated, S::Union{Const, MixedDuplicated})
+
+    # adjoint/vjp of FFT has a different scaling, compute it, apply it later to f_north, f_south
+    scale = adjoint_scale(S.val)
+
+    # TODO: retrieve from tape here if overwritten? 
+
+    # compute the adjoint # TODO: could we reuse the f_north.val for that a well? and not allocate here
+    dfnorthval = zero(f_north.val)
+    dfsouthval = zero(f_south.val)
+
+    _fourier!(dfnorthval, dfsouthval, grids.dval, S.val) # inverse FFT (w/o normalization)
+
+    f_north.dval .+= scale .* dfnorthval
+    f_south.dval .+= scale .* dfsouthval 
+
+    # no derivative wrt the grids that were input because they are overwritten
+    make_zero!(grids.dval) 
+
+    # the function has no return values, so we also return nothing here
+    return (nothing, nothing, nothing, nothing)
+end
+
+end

ext/SpeedyWeatherFiniteDifferencesExt.jl

@@ -0,0 +1,30 @@
+module SpeedyWeatherFiniteDifferencesExt 
+
+using SpeedyWeather 
+import FiniteDifferences
+import FiniteDifferences: to_vec 
+
+# FiniteDifferences needs to be able to convert data structures to Vectors and back 
+# This doesn't work out of the box with our data types, so we'll define those 
+# conversions here.
+function FiniteDifferences.to_vec(x::Grid) where Grid <: AbstractGridArray
+    x_vec, from_vec = FiniteDifferences.to_vec(Array(x))
+
+    function GridArray_from_vec(x_vec)
+        return Grid(reshape(from_vec(x_vec), size(x)), x.nlat_half)
+    end 
+
+    return x_vec, GridArray_from_vec
+end 
+
+function FiniteDifferences.to_vec(x::LTA) where LTA <: LowerTriangularArray
+    x_vec, from_vec = FiniteDifferences.to_vec(x.data)
+
+    function LowerTriangularArray_from_vec(x_vec)
+        return LowerTriangularArray(reshape(from_vec(x_vec), size(x)), x.m, x.n)
+    end 
+
+    return x_vec, LowerTriangularArray_from_vec
+end
+
+end 

test/Project.toml

@@ -0,0 +1,13 @@
+[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"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -2,6 +2,7 @@ using SpeedyWeather
 using Test
 
 # GENERAL
+include("test_transforms_ad_rules.jl") # TODO: put the test somewhere else (and next to the transforms)
 include("utility_functions.jl")
 include("dates.jl")
 include("lower_triangular_matrix.jl")