fix wrong dim on permutation

Author: foldfelis

src/fourier.jl

@@ -34,7 +34,7 @@ Base.ndims(::SpectralConv{N}) where {N} = N
 spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m]
 
 function (m::SpectralConv)(𝐱::AbstractArray)
-    𝐱ᵀ = permutedims(Zygote.hook(real, 𝐱), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch]
+    𝐱ᵀ = permutedims(Zygote.hook(real, 𝐱), (2:ndims(m)+1..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch]
     𝐱_fft = fft(𝐱ᵀ, 1:ndims(m)) # [x, in_chs, batch]
 
     # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes]
@@ -48,7 +48,7 @@ function (m::SpectralConv)(𝐱::AbstractArray)
     𝐱_padded = cat(𝐱_shaped, pad, dims=1:ndims(m))
 
     𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch]
-    𝐱_outᵀ = permutedims(real(𝐱_out), (2:ndims(m)+1..., 1, ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch]
+    𝐱_outᵀ = permutedims(real(𝐱_out), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch]
 
     return m.σ.(𝐱_outᵀ)
 end

test/fourier.jl

@@ -1,4 +1,4 @@
-@testset "SpectralConv" begin
+@testset "SpectralConv1d" begin
     modes = (16, )
     ch = 64 => 64
 
@@ -17,7 +17,7 @@
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperator" begin
+@testset "FourierOperator1d" begin
     modes = (16, )
     ch = 64 => 64
 
@@ -33,3 +33,39 @@ end
     data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
+
+@testset "SpectralConv2d" begin
+    modes = (16, 16)
+    ch = 64 => 64
+
+    m = Chain(
+        Dense(1, 64),
+        SpectralConv(ch, modes)
+    )
+    @test ndims(SpectralConv(ch, modes)) == 2
+
+    𝐱, _ , _, _ = get_darcy_flow_data()
+    @test size(m(𝐱)) == (64, 85, 85, 1024)
+
+    T = Float32
+    loss(x, y) = Flux.mse(m(x), y)
+    data = [(T.(𝐱[:, :, :, 1:5]), rand(T, 64, 85, 85, 5))]
+    Flux.train!(loss, params(m), data, Flux.ADAM())
+end
+
+@testset "FourierOperator2d" begin
+    modes = (16, 16)
+    ch = 64 => 64
+
+    m = Chain(
+        Dense(1, 64),
+        FourierOperator(ch, modes)
+    )
+
+    𝐱, _ , _, _ = get_darcy_flow_data()
+    @test size(m(𝐱)) == (64, 85, 85, 1024)
+
+    loss(x, y) = Flux.mse(m(x), y)
+    data = [(Float32.(𝐱[:, :, :, 1:5]), rand(Float32, 64, 85, 85, 5))]
+    Flux.train!(loss, params(m), data, Flux.ADAM())
+end

test/runtests.jl

@@ -3,7 +3,7 @@ using Test
 using Flux
 
 @testset "NeuralOperators.jl" begin
-    include("data.jl")
+    # include("data.jl")
     include("fourier.jl")
-    include("model.jl")
+    # include("model.jl")
 end