Merge pull request #63 from yuehhua/chebyshev

Add Chebyshev transform

Author: foldfelis

src/Transform/Transform.jl

@@ -17,3 +17,4 @@ export
 abstract type AbstractTransform end
 
 include("fourier_transform.jl")
+include("chebyshev_transform.jl")

src/Transform/chebyshev_transform.jl

@@ -0,0 +1,50 @@
+export ChebyshevTransform
+
+struct ChebyshevTransform{N, S} <: AbstractTransform
+    modes::NTuple{N, S} # N == ndims(x)
+end
+
+Base.ndims(::ChebyshevTransform{N}) where {N} = N
+
+function transform(t::ChebyshevTransform{N}, 𝐱::AbstractArray) where {N}
+    return FFTW.r2r(𝐱, FFTW.REDFT10, 1:N) # [size(x)..., in_chs, batch]
+end
+
+function truncate_modes(t::ChebyshevTransform, 𝐱̂::AbstractArray)
+    return view(𝐱̂, map(d -> 1:d, t.modes)..., :, :) # [t.modes..., in_chs, batch]
+end
+
+function inverse(t::ChebyshevTransform{N}, 𝐱̂::AbstractArray) where {N}
+    normalized_𝐱̂ = 𝐱̂ ./ (prod(2 .* (size(𝐱̂)[1:N] .- 1)))
+    return FFTW.r2r(normalized_𝐱̂, FFTW.REDFT01, 1:N) # [size(x)..., in_chs, batch]
+end
+
+function ChainRulesCore.rrule(::typeof(FFTW.r2r), x::AbstractArray, kind, dims)
+    y = FFTW.r2r(x, kind, dims)
+    r2r_pullback(Δ) = (NoTangent(), ∇r2r(unthunk(Δ), kind, dims), NoTangent(), NoTangent())
+    return y, r2r_pullback
+end
+
+function ∇r2r(Δ::AbstractArray{T}, kind, dims) where {T}
+    # derivative of r2r turns out to be r2r
+    Δx = FFTW.r2r(Δ, kind, dims)
+
+    # rank 4 correction: needs @bischtob to elaborate the reason using this. 
+    # (M,) = size(Δ)[dims]
+    # a1 = fill!(similar(Δ, M), one(T))
+    # CUDA.@allowscalar a1[1] = a1[end] = zero(T)
+
+    # a2 = fill!(similar(Δ, M), one(T))
+    # a2[1:2:end] .= -one(T)
+    # CUDA.@allowscalar a2[1] = a2[end] = zero(T)
+
+    # e1 = fill!(similar(Δ, M), zero(T))
+    # CUDA.@allowscalar e1[1] = one(T)
+
+    # eN = fill!(similar(Δ, M), zero(T))
+    # CUDA.@allowscalar eN[end] = one(T)
+
+    # Δx .+= @. a1' * sum(e1' .* Δ, dims=2) - a2' * sum(eN' .* Δ, dims=2)
+    # Δx .+= @. eN' * sum(a2' .* Δ, dims=2) - e1' * sum(a1' .* Δ, dims=2)
+    return Δx
+end

src/Transform/fourier_transform.jl

@@ -14,7 +14,7 @@ function low_pass(ft::FourierTransform, 𝐱_fft::AbstractArray)
     return view(𝐱_fft, map(d -> 1:d, ft.modes)..., :, :) # [ft.modes..., in_chs, batch]
 end
 
-truncate_modes(args...) = low_pass(args...)
+truncate_modes(ft::FourierTransform, 𝐱_fft::AbstractArray) = low_pass(ft, 𝐱_fft)
 
 function inverse(ft::FourierTransform, 𝐱_fft::AbstractArray)
     return real(ifft(𝐱_fft, 1:ndims(ft))) # [size(x_fft)..., out_chs, batch]

test/Transform/Transform.jl

@@ -1 +1,4 @@
-@testset "Transform" begin include("fourier_transform.jl") end
+@testset "Transform" begin
+    include("fourier_transform.jl")
+    include("chebyshev_transform.jl")
+end

test/Transform/chebyshev_transform.jl

@@ -0,0 +1,15 @@
+@testset "Chebyshev transform" begin
+    ch = 6
+    batch = 7
+    𝐱 = rand(30, 40, 50, ch, batch)
+
+    t = ChebyshevTransform((3, 4, 5))
+
+    @test ndims(t) == 3
+    @test size(transform(t, 𝐱)) == (30, 40, 50, ch, batch)
+    @test size(truncate_modes(t, transform(t, 𝐱))) == (3, 4, 5, ch, batch)
+    @test size(inverse(t, truncate_modes(t, transform(t, 𝐱)))) == (3, 4, 5, ch, batch)
+
+    g = gradient(x -> sum(inverse(t, truncate_modes(t, transform(t, x)))), 𝐱)
+    @test size(g[1]) == (30, 40, 50, ch, batch)
+end

test/Transform/fourier_transform.jl

@@ -1,9 +1,14 @@
-@testset "fourier transform" begin
-    𝐱 = rand(30, 40, 50, 6, 7) # where ch == 6 and batch == 7
+@testset "Fourier transform" begin
+    ch = 6
+    batch = 7
+    𝐱 = rand(30, 40, 50, ch, batch)
 
     ft = FourierTransform((3, 4, 5))
 
-    @test size(transform(ft, 𝐱)) == (30, 40, 50, 6, 7)
-    @test size(truncate_modes(ft, transform(ft, 𝐱))) == (3, 4, 5, 6, 7)
-    @test size(inverse(ft, truncate_modes(ft, transform(ft, 𝐱)))) == (3, 4, 5, 6, 7)
+    @test size(transform(ft, 𝐱)) == (30, 40, 50, ch, batch)
+    @test size(truncate_modes(ft, transform(ft, 𝐱))) == (3, 4, 5, ch, batch)
+    @test size(inverse(ft, truncate_modes(ft, transform(ft, 𝐱)))) == (3, 4, 5, ch, batch)
+
+    g = Zygote.gradient(x -> sum(inverse(ft, truncate_modes(ft, transform(ft, x)))), 𝐱)
+    @test size(g[1]) == (30, 40, 50, ch, batch)
 end