fix wrong dim on permutation and shrink test

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
 
@@ -8,16 +8,15 @@
     )
     @test ndims(SpectralConv(ch, modes)) == 1
 
-    𝐱, _ = get_burgers_data(n=1000)
-    @test size(m(𝐱)) == (64, 1024, 1000)
+    𝐱, _ = get_burgers_data(n=5)
+    @test size(m(𝐱)) == (64, 1024, 5)
 
-    T = Float32
     loss(x, y) = Flux.mse(m(x), y)
-    data = [(T.(𝐱[:, :, 1:5]), rand(T, 64, 1024, 5))]
+    data = [(𝐱, rand(Float32, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperator" begin
+@testset "FourierOperator1d" begin
     modes = (16, )
     ch = 64 => 64
 
@@ -26,10 +25,45 @@ end
         FourierOperator(ch, modes)
     )
 
-    𝐱, _ = get_burgers_data(n=1000)
-    @test size(m(𝐱)) == (64, 1024, 1000)
+    𝐱, _ = get_burgers_data(n=5)
+    @test size(m(𝐱)) == (64, 1024, 5)
 
     loss(x, y) = Flux.mse(m(x), y)
-    data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))]
+    data = [(𝐱, 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(n=5, Δsamples=20)
+    @test size(m(𝐱)) == (64, 22, 22, 5)
+
+    loss(x, y) = Flux.mse(m(x), y)
+    data = [(𝐱, rand(Float32, 64, 22, 22, 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(n=5, Δsamples=20)
+    @test size(m(𝐱)) == (64, 22, 22, 5)
+
+    loss(x, y) = Flux.mse(m(x), y)
+    data = [(𝐱, rand(Float32, 64, 22, 22, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end