Apply format

Author: foldfelis

example/Burgers/src/Burgers.jl

@@ -67,7 +67,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 500)
     return learner
 end
 
-function train_nomad(; n=300, cuda=true, learning_rate=0.001, epochs=400)
+function train_nomad(; n = 300, cuda = true, learning_rate = 0.001, epochs = 400)
     if cuda && has_cuda()
         @info "Training on GPU"
         device = gpu
@@ -76,32 +76,31 @@ function train_nomad(; n=300, cuda=true, learning_rate=0.001, epochs=400)
         device = cpu
     end
 
-    x, y = get_data_don(n=n)
+    x, y = get_data_don(n = n)
 
     xtrain = x[1:280, :]'
     ytrain = y[1:280, :]
 
-    xval = x[end-19:end, :]' |> device
-    yval = y[end-19:end, :] |> device
+    xval = x[(end - 19):end, :]' |> device
+    yval = y[(end - 19):end, :] |> device
 
     # grid = collect(range(0, 1, length=1024)') |> device
     grid = rand(collect(0:0.001:1), (280, 1024)) |> device
     gridval = rand(collect(0:0.001:1), (20, 1024)) |> device
 
-
     opt = ADAM(learning_rate)
 
-    m = NOMAD((1024,1024), (2048,1024), gelu, gelu) |> device
+    m = NOMAD((1024, 1024), (2048, 1024), gelu, gelu) |> device
 
     loss(X, y, sensor) = Flux.Losses.mse(m(X, sensor), y)
     evalcb() = @show(loss(xval, yval, gridval))
 
     data = [(xtrain, ytrain, grid)] |> device
-    Flux.@epochs epochs Flux.train!(loss, params(m), data, opt, cb=evalcb)
+    Flux.@epochs epochs Flux.train!(loss, params(m), data, opt, cb = evalcb)
     ỹ = m(xval |> device, gridval |> device)
 
     diffvec = vec(abs.(cpu(yval) .- cpu(ỹ)))
-    mean_diff = sum(diffvec)/length(diffvec)
+    mean_diff = sum(diffvec) / length(diffvec)
     return mean_diff
 end
 

example/Burgers/test/runtests.jl

@@ -15,6 +15,6 @@ using Test
 end
 
 @testset "Burger: NOMAD Training Accuracy" begin
-    ϵ = Burgers.train_nomad(; cuda=true, epochs=100)
+    ϵ = Burgers.train_nomad(; cuda = true, epochs = 100)
     @test ϵ < 0.4 # epoch=100 returns 0.233
-end
+end

src/NOMAD.jl

@@ -49,18 +49,18 @@ Approximator net: (Chain(Dense(2 => 128), Dense(128 => 64)))
 Decoder net: (Chain(Dense(72 => 24), Dense(24 => 12)))
 """
 function NOMAD(architecture_approximator::Tuple, architecture_decoder::Tuple,
-                act_approximator = identity, act_decoder=true;
-                init_approximator = Flux.glorot_uniform,
-                init_decoder = Flux.glorot_uniform,
-                bias_approximator=true, bias_decoder=true)
-
+               act_approximator = identity, act_decoder = true;
+               init_approximator = Flux.glorot_uniform,
+               init_decoder = Flux.glorot_uniform,
+               bias_approximator = true, bias_decoder = true)
     approximator_net = construct_subnet(architecture_approximator, act_approximator;
-                                    init=init_approximator, bias=bias_approximator)
+                                        init = init_approximator, bias = bias_approximator)
 
     decoder_net = construct_subnet(architecture_decoder, act_decoder;
-                                    init=init_decoder, bias=bias_decoder)
+                                   init = init_decoder, bias = bias_decoder)
 
-    return NOMAD{typeof(approximator_net), typeof(decoder_net)}(approximator_net, decoder_net)
+    return NOMAD{typeof(approximator_net), typeof(decoder_net)}(approximator_net,
+                                                                decoder_net)
 end
 
 Flux.@functor NOMAD
@@ -69,12 +69,12 @@ function (a::NOMAD)(x::AbstractArray, y::AbstractVecOrMat)
     # Assign the parameters
     approximator, decoder = a.approximator_net, a.decoder_net
 
-    return decoder(cat(approximator(x), y', dims=1))'
+    return decoder(cat(approximator(x), y', dims = 1))'
 end
 
 # Print nicely
 function Base.show(io::IO, l::NOMAD)
-    print(io, "NOMAD with\nApproximator net: (",l.approximator_net)
+    print(io, "NOMAD with\nApproximator net: (", l.approximator_net)
     print(io, ")\n")
     print(io, "Decoder net: (", l.decoder_net)
     print(io, ")\n")

test/nomad.jl

@@ -1,28 +1,31 @@
 @testset "NOMAD" begin
     @testset "proper construction" begin
-        nomad = NOMAD((32,64,72), (24,48,72), σ, tanh)
+        nomad = NOMAD((32, 64, 72), (24, 48, 72), σ, tanh)
         # approximator net
-        @test size(nomad.approximator_net.layers[end].weight) == (72,64)
+        @test size(nomad.approximator_net.layers[end].weight) == (72, 64)
         @test size(nomad.approximator_net.layers[end].bias) == (72,)
         # decoder net
-        @test size(nomad.decoder_net.layers[end].weight) == (72,48)
+        @test size(nomad.decoder_net.layers[end].weight) == (72, 48)
         @test size(nomad.decoder_net.layers[end].bias) == (72,)
     end
 
     # Accept only Int as architecture parameters
-    @test_throws MethodError NOMAD((32.5,64,72), (24,48,72), σ, tanh)
-    @test_throws MethodError NOMAD((32,64,72), (24.1,48,72))
+    @test_throws MethodError NOMAD((32.5, 64, 72), (24, 48, 72), σ, tanh)
+    @test_throws MethodError NOMAD((32, 64, 72), (24.1, 48, 72))
 
     # Just the first 16 datapoints from the Burgers' equation dataset
     a = [0.83541104, 0.83479851, 0.83404712, 0.83315711, 0.83212979, 0.83096755,
-         0.82967374, 0.82825263, 0.82670928, 0.82504949, 0.82327962, 0.82140651,
-         0.81943734, 0.81737952, 0.8152405, 0.81302771]
-    sensors = collect(range(0, 1, length=16)')
-    model = NOMAD((length(a), 22, length(a)), (length(a) + length(sensors), length(sensors)), σ, tanh; init_approximator=Flux.glorot_normal, bias_decoder=false)
+        0.82967374, 0.82825263, 0.82670928, 0.82504949, 0.82327962, 0.82140651,
+        0.81943734, 0.81737952, 0.8152405, 0.81302771]
+    sensors = collect(range(0, 1, length = 16)')
+    model = NOMAD((length(a), 22, length(a)),
+                  (length(a) + length(sensors), length(sensors)), σ, tanh;
+                  init_approximator = Flux.glorot_normal, bias_decoder = false)
     y = model(a, sensors)
     @test size(y) == (1, 16)
     # Check if model description is printed, when defined
-    @test repr(model) == "NOMAD with\nApproximator net: (Chain(Dense(16 => 22, σ), Dense(22 => 16, σ)))\nDecoder net: (Chain(Dense(32 => 16, tanh; bias=false)))\n"
+    @test repr(model) ==
+          "NOMAD with\nApproximator net: (Chain(Dense(16 => 22, σ), Dense(22 => 16, σ)))\nDecoder net: (Chain(Dense(32 => 16, tanh; bias=false)))\n"
 
     mgrad = Flux.Zygote.gradient(() -> sum(model(a, sensors)), Flux.params(model))
     @info mgrad.grads