refactor model

Author: foldfelis

src/NeuralOperators.jl

@@ -14,4 +14,5 @@ module NeuralOperators
 
     include("data.jl")
     include("fourier.jl")
+    include("model.jl")
 end

src/fourier.jl

@@ -63,20 +63,3 @@ function FourierOperator(
         x -> σ.(x)
     )
 end
-
-function FNO()
-    modes = 16
-    ch = 64 => 64
-    σ = relu
-
-    return Chain(
-        Dense(2, 64),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes),
-        Dense(64, 128, σ),
-        Dense(128, 1),
-        flatten
-    )
-end

src/model.jl

@@ -0,0 +1,19 @@
+export
+    FourierNeuralOperator
+
+function FourierNeuralOperator()
+    modes = 16
+    ch = 64 => 64
+    σ = relu
+
+    return Chain(
+        Dense(2, 64),
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes),
+        Dense(64, 128, σ),
+        Dense(128, 1),
+        flatten
+    )
+end

test/fourier.jl

@@ -32,14 +32,3 @@ end
     data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
-
-@testset "FNO" begin
-    𝐱, 𝐲 = get_burgers_data()
-    𝐱, 𝐲 = Float32.(𝐱), Float32.(𝐲)
-    @test size(FNO()(𝐱)) == size(𝐲)
-
-    m = FNO()
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
-    data = [(𝐱[:, :, 1:5], 𝐲[:, 1:5])]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
-end

test/model.jl

@@ -0,0 +1,11 @@
+@testset "FourierNeuralOperator" begin
+    m = FourierNeuralOperator()
+
+    𝐱, 𝐲 = get_burgers_data()
+    𝐱, 𝐲 = Float32.(𝐱), Float32.(𝐲)
+    @test size(m(𝐱)) == size(𝐲)
+
+    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+    data = [(𝐱[:, :, 1:5], 𝐲[:, 1:5])]
+    Flux.train!(loss, params(m), data, Flux.ADAM())
+end

test/runtests.jl

@@ -5,4 +5,5 @@ using Flux
 @testset "NeuralOperators.jl" begin
     include("data.jl")
     include("fourier.jl")
+    include("model.jl")
 end