Implement FNO

Author: foldfelis

src/fourier.jl

@@ -3,26 +3,29 @@ using FFTW
 using Tullio
 
 export
-    SpectralConv1d
+    SpectralConv1d,
+    FNO
 
 struct SpectralConv1d{T,S}
     weight::T
     in_channel::S
     out_channel::S
     modes::S
+    σ::Function
 end
 
 function SpectralConv1d(
     ch::Pair{<:Integer,<:Integer},
-    modes::Integer;
+    modes::Integer,
+    σ::Function=identity;
     init=Flux.glorot_uniform,
     T::DataType=Float32
 )
     in_chs, out_chs = ch
     scale = one(T) / (in_chs * out_chs)
     weights = scale * init(out_chs, in_chs, modes)
 
-    return SpectralConv1d(weights, in_chs, out_chs, modes)
+    return SpectralConv1d(weights, in_chs, out_chs, modes, σ)
 end
 
 Flux.@functor SpectralConv1d
@@ -39,19 +42,37 @@ function (m::SpectralConv1d)(𝐱::AbstractArray)
 
     𝐱_out = irfft(𝐱_padded , size(𝐱, 1), 1)
 
-    return 𝐱_out
+    return m.σ.(𝐱_out)
 end
 
-# function FNO(modes::Integer, width::Integer)
-#     return Chain(
-#         PermutedDimsArray(Dense(2, width),(2,1,3)),
-#         relu(SpectralConv1d(width, width, modes) + Conv(width, width, 1)),
-#         relu(SpectralConv1d(width, width, modes) + Conv(width, width, 1)),
-#         relu(SpectralConv1d(width, width, modes) + Conv(width, width, 1)),
-#         PermutedDimsArray(relu(SpectralConv1d(width, width, modes) + Conv(width, width, 1)), (0, 2, 1)),
-#         Dense(width, 128, relu),
-#         Dense(128, 1)
-#     )
-# end
+function FourierBlock(
+    ch::Pair{<:Integer,<:Integer},
+    modes::Integer,
+    σ::Function=identity
+)
+    return Chain(
+        Parallel(+,
+            Conv((1, ), ch),
+            SpectralConv1d(ch, modes)
+        ),
+        x -> σ.(x)
+    )
+end
+
+function FNO()
+    modes = 16
+    ch = 64 => 64
+
+    return Chain(
+        Conv((1, ), 2=>64),
+        FourierBlock(ch, modes, relu),
+        FourierBlock(ch, modes, relu),
+        FourierBlock(ch, modes, relu),
+        FourierBlock(ch, modes),
+        Conv((1, ), 64=>128, relu),
+        Conv((1, ), 128=>1),
+        flatten
+    )
+end
 
 # loss(m::SpectralConv1d, x, x̂) = sum(abs2, x̂ .- m(x)) / len

test/fourier.jl

@@ -2,13 +2,18 @@ using Flux
 
 @testset "fourier" begin
     modes = 16
-    width = 64
-    ch = width => width
+    ch = 64 => 64
+
     m = Chain(
-        Conv((1, ), 2=>width),
+        Conv((1, ), 2=>64),
         SpectralConv1d(ch, modes)
     )
 
     𝐱, _ = get_data()
-    @show size(m(𝐱))
+    @test size(m(𝐱)) == (1024, 64, 1000)
+end
+
+@testset "FNO" begin
+    𝐱, 𝐲 = get_data()
+    @test size(FNO()(𝐱)) == size(𝐲)
 end