fix: fft shift (#67)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/transform.jl

@@ -43,27 +43,30 @@ end
 Base.ndims(T::FourierTransform) = length(T.modes)
 
 function transform(ft::FourierTransform, x::AbstractArray)
-    res = Lux.Utils.eltype(x) <: Complex ? fft(x, 1:ndims(ft)) : rfft(x, 1:ndims(ft))
+    complex_data = Lux.Utils.eltype(x) <: Complex
+    res = complex_data ? fft(x, 1:ndims(ft)) : rfft(x, 1:ndims(ft))
     if ft.shift && ndims(ft) > 1
-        res = fftshift(res, 1:ndims(ft))
+        res = fftshift(res, (1 + !complex_data):ndims(ft))
     end
     return res
 end
 
 function low_pass(ft::FourierTransform, x_fft::AbstractArray)
-    return view(x_fft,(map(d -> 1:d, ft.modes)...),:,:)
+    return view(x_fft, (map(d -> 1:d, ft.modes)...), :, :)
 end
 
 truncate_modes(ft::FourierTransform, x_fft::AbstractArray) = low_pass(ft, x_fft)
 
 function inverse(
     ft::FourierTransform, x_fft::AbstractArray{T,N}, x::AbstractArray{T2,N}
 ) where {T,T2,N}
+    complex_data = Lux.Utils.eltype(x) <: Complex
+
     if ft.shift && ndims(ft) > 1
-        x_fft = fftshift(x_fft, 1:ndims(ft))
+        x_fft = fftshift(x_fft, (1 + !complex_data):ndims(ft))
     end
 
-    if Lux.Utils.eltype(x) <: Complex
+    if complex_data
         return ifft(x_fft, 1:ndims(ft))
     else
         return real(irfft(x_fft, size(x, 1), 1:ndims(ft)))