Merge pull request #5 from foldfelis/refactor

Refactor

Author: foldfelis

.github/workflows/CI.yml

@@ -1,4 +1,6 @@
 name: CI
+env:
+  DATADEPS_ALWAYS_ACCEPT: true
 on:
   - push
   - pull_request

src/NeuralOperators.jl

@@ -1,8 +1,18 @@
 module NeuralOperators
+    using DataDeps
+    using Fetch
+    using MAT
+
+    using Flux
+    using FFTW
+    using Tullio
+    using Zygote
+
     function __init__()
         register_datasets()
     end
 
-    include("preprocess.jl")
+    include("data.jl")
     include("fourier.jl")
+    include("model.jl")
 end

src/preprocess.jl renamed to src/data.jl

@@ -1,9 +1,5 @@
-using DataDeps
-using Fetch
-using MAT
-
 export
-    get_data
+    get_burgers_data
 
 function register_datasets()
     register(DataDep(
@@ -19,7 +15,7 @@ function register_datasets()
     ))
 end
 
-function get_data(; n=1000, Δsamples=2^3, grid_size=div(2^13, Δsamples))
+function get_burgers_data(; n=1000, Δsamples=2^3, grid_size=div(2^13, Δsamples))
     file = matopen(joinpath(datadep"BurgersR10", "burgers_data_R10.mat"))
     x_data = collect(read(file, "a")[1:n, 1:Δsamples:end]')
     y_data = collect(read(file, "u")[1:n, 1:Δsamples:end]')

src/fourier.jl

@@ -1,8 +1,3 @@
-using Flux
-using FFTW
-using Tullio
-using Zygote
-
 export
     SpectralConv1d,
     FourierOperator,
@@ -68,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/preprocess.jl renamed to test/data.jl

@@ -1,5 +1,5 @@
-@testset "get data" begin
-    xs, ys = get_data()
+@testset "get burgers data" begin
+    xs, ys = get_burgers_data()
 
     @test size(xs) == (2, 1024, 1000)
     @test size(ys) == (1024, 1000)

test/fourier.jl

@@ -1,5 +1,3 @@
-using Flux
-
 @testset "SpectralConv1d" begin
     modes = 16
     ch = 64 => 64
@@ -9,7 +7,7 @@ using Flux
         SpectralConv1d(ch, modes)
     )
 
-    𝐱, _ = get_data()
+    𝐱, _ = get_burgers_data()
     @test size(m(𝐱)) == (64, 1024, 1000)
 
     T = Float32
@@ -27,21 +25,10 @@ end
         FourierOperator(ch, modes)
     )
 
-    𝐱, _ = get_data()
+    𝐱, _ = get_burgers_data()
     @test size(m(𝐱)) == (64, 1024, 1000)
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
-
-@testset "FNO" begin
-    𝐱, 𝐲 = get_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

@@ -1,9 +1,9 @@
 using NeuralOperators
 using Test
-
-ENV["DATADEPS_ALWAYS_ACCEPT"] = true
+using Flux
 
 @testset "NeuralOperators.jl" begin
-    include("preprocess.jl")
+    include("data.jl")
     include("fourier.jl")
+    include("model.jl")
 end