use rfft instead of fft

Author: tknopp

src/Transform/chebyshev_transform.jl

@@ -14,7 +14,7 @@ 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}
+function inverse(t::ChebyshevTransform{N}, 𝐱̂::AbstractArray, M) where {N}
     normalized_𝐱̂ = 𝐱̂ ./ (prod(2 .* (size(𝐱̂)[1:N] .- 1)))
     return FFTW.r2r(normalized_𝐱̂, FFTW.REDFT01, 1:N) # [size(x)..., in_chs, batch]
 end

src/Transform/fourier_transform.jl

@@ -7,7 +7,7 @@ end
 Base.ndims(::FourierTransform{N}) where {N} = N
 
 function transform(ft::FourierTransform, 𝐱::AbstractArray)
-    return fft(Zygote.hook(real, 𝐱), 1:ndims(ft)) # [size(x)..., in_chs, batch]
+    return rfft(Zygote.hook(real, 𝐱), 1:ndims(ft)) # [size(x)..., in_chs, batch]
 end
 
 function low_pass(ft::FourierTransform, 𝐱_fft::AbstractArray)
@@ -16,6 +16,6 @@ end
 
 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]
+function inverse(ft::FourierTransform, 𝐱_fft::AbstractArray, M)
+    return real(irfft(𝐱_fft, M[1], 1:ndims(ft))) # [size(x_fft)..., out_chs, batch]
 end

src/operator_kernel.jl

@@ -92,7 +92,7 @@ function operator_conv(m::OperatorConv, 𝐱::AbstractArray)
     𝐱_padded = pad_modes(𝐱_applied_pattern,
                          (size(𝐱_transformed)[1:(end - 2)]...,
                           size(𝐱_applied_pattern)[(end - 1):end]...)) # [size(x)..., out_chs, batch] <- [modes..., out_chs, batch]
-    𝐱_inversed = inverse(m.transform, 𝐱_padded)
+    𝐱_inversed = inverse(m.transform, 𝐱_padded, size(𝐱))
 
     return 𝐱_inversed
 end

test/Transform/chebyshev_transform.jl

@@ -8,8 +8,8 @@
     @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)
+    @test size(inverse(t, truncate_modes(t, transform(t, 𝐱)), size(𝐱))) == (3, 4, 5, ch, batch)
 
-    g = gradient(x -> sum(inverse(t, truncate_modes(t, transform(t, x)))), 𝐱)
+    g = gradient(x -> sum(inverse(t, truncate_modes(t, transform(t, x)), size(𝐱))), 𝐱)
     @test size(g[1]) == (30, 40, 50, ch, batch)
 end

test/Transform/fourier_transform.jl

@@ -7,8 +7,8 @@
 
     @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)
+    @test size(inverse(ft, truncate_modes(ft, transform(ft, 𝐱)), size(𝐱))) == (3, 4, 5, ch, batch)
 
-    g = Zygote.gradient(x -> sum(inverse(ft, truncate_modes(ft, transform(ft, x)))), 𝐱)
+    g = Zygote.gradient(x -> sum(inverse(ft, truncate_modes(ft, transform(ft, x)), size(𝐱))), 𝐱)
     @test size(g[1]) == (30, 40, 50, ch, batch)
 end