Apply JuliaFormatter to codebase

benchmarks/benchFourierLayer.jl

@@ -1,10 +1,10 @@
-# Stolen from Flux as well
-
-for n in [2, 20, 200, 2000]
-    x = randn(Float32, n, 2000, 100)
-    model = FourierLayer(n, n, 2000, 100, 16)
-    println("CPU n=$n")
-    run_benchmark(model, x, cuda=false)
-    println("CUDA n=$n")
-    run_benchmark(model, x, cuda=true)    
-end
+# Stolen from Flux as well
+
+for n in [2, 20, 200, 2000]
+    x = randn(Float32, n, 2000, 100)
+    model = FourierLayer(n, n, 2000, 100, 16)
+    println("CPU n=$n")
+    run_benchmark(model, x, cuda = false)
+    println("CUDA n=$n")
+    run_benchmark(model, x, cuda = true)
+end

benchmarks/bench_utils.jl

@@ -1,51 +1,55 @@
-using BenchmarkTools
-using Flux
-using OperatorLearning
-using CUDA
-using Zygote: pullback
-
-# Stolen from Flux
-
-fw(m, x) = m(x)
-bw(back) = back(1f0)
-fwbw(m, ps, x) = gradient(() -> sum(m(x)), ps)
-
-function run_benchmark(model, x; cuda=true)
-
-    if cuda 
-        model = model |> gpu
-        x = x |> gpu
-    end
-
-    ps = Flux.params(model)
-    y, back = pullback(() -> sum(model(x)), ps)
-
-
-    if cuda
-        CUDA.allowscalar(false)
-        # CUDA.device!(3)
-        println("  forward")
-        fw(model, x); GC.gc(); CUDA.reclaim(); #warmup
-        @btime CUDA.@sync(fw($model, $x)) teardown=(GC.gc(); CUDA.reclaim())
-
-        println("  backward")
-        bw(back); GC.gc(); CUDA.reclaim(); #warmup
-        @btime CUDA.@sync(bw($back)) teardown=(GC.gc(); CUDA.reclaim())
-
-        println("  forw and back")
-        fwbw(model, ps, x); GC.gc(); CUDA.reclaim(); #warmup
-        @btime CUDA.@sync(fwbw($model, $ps, $x)) teardown=(GC.gc(); CUDA.reclaim())
-    else
-        println("  forward")
-        fw(model, x)  #warmup
-        @btime fw($model, $x)
-
-        println("  backward")
-        bw(back)  #warmup
-        @btime bw($back)
-
-        println("  forw and back")
-        fwbw(model, ps, x) # warmup
-        @btime fwbw($model, $ps, $x)
-    end
-end
+using BenchmarkTools
+using Flux
+using OperatorLearning
+using CUDA
+using Zygote: pullback
+
+# Stolen from Flux
+
+fw(m, x) = m(x)
+bw(back) = back(1.0f0)
+fwbw(m, ps, x) = gradient(() -> sum(m(x)), ps)
+
+function run_benchmark(model, x; cuda = true)
+    if cuda
+        model = model |> gpu
+        x = x |> gpu
+    end
+
+    ps = Flux.params(model)
+    y, back = pullback(() -> sum(model(x)), ps)
+
+    if cuda
+        CUDA.allowscalar(false)
+        # CUDA.device!(3)
+        println("  forward")
+        fw(model, x);
+        GC.gc();
+        CUDA.reclaim(); #warmup
+        @btime CUDA.@sync(fw($model, $x)) teardown=(GC.gc(); CUDA.reclaim())
+
+        println("  backward")
+        bw(back);
+        GC.gc();
+        CUDA.reclaim(); #warmup
+        @btime CUDA.@sync(bw($back)) teardown=(GC.gc(); CUDA.reclaim())
+
+        println("  forw and back")
+        fwbw(model, ps, x);
+        GC.gc();
+        CUDA.reclaim(); #warmup
+        @btime CUDA.@sync(fwbw($model, $ps, $x)) teardown=(GC.gc(); CUDA.reclaim())
+    else
+        println("  forward")
+        fw(model, x)  #warmup
+        @btime fw($model, $x)
+
+        println("  backward")
+        bw(back)  #warmup
+        @btime bw($back)
+
+        println("  forw and back")
+        fwbw(model, ps, x) # warmup
+        @btime fwbw($model, $ps, $x)
+    end
+end

benchmarks/runbenchmarks.jl

@@ -1,7 +1,7 @@
-using OperatorLearning, Flux
-
-versioninfo()
-include("bench_utils.jl")
-
-@info "Benchmark FourierLayer"
-include("benchFourierLayer.jl")
+using OperatorLearning, Flux
+
+versioninfo()
+include("bench_utils.jl")
+
+@info "Benchmark FourierLayer"
+include("benchFourierLayer.jl")

docs/make.jl

@@ -1,29 +1,28 @@
 using OperatorLearning
 using Documenter, DocumenterTools
 
-DocMeta.setdocmeta!(OperatorLearning, :DocTestSetup, :(using OperatorLearning); recursive=true)
+DocMeta.setdocmeta!(OperatorLearning, :DocTestSetup, :(using OperatorLearning); recursive = true)
 
 makedocs(;
-    modules=[OperatorLearning],
-    authors="Patrick Zimbrod <patrick.zimbrod@gmail.com> and contributors",
-    repo="https://github.com/pzimbrod/OperatorLearning.jl/blob/{commit}{path}#{line}",
-    sitename="OperatorLearning.jl",
-    format=Documenter.HTML(;
-        prettyurls=get(ENV, "CI", "false") == "true",
-        canonical="https://pzimbrod.github.io/OperatorLearning.jl",
-        assets=String[],
+    modules = [OperatorLearning],
+    authors = "Patrick Zimbrod <patrick.zimbrod@gmail.com> and contributors",
+    repo = "https://github.com/pzimbrod/OperatorLearning.jl/blob/{commit}{path}#{line}",
+    sitename = "OperatorLearning.jl",
+    format = Documenter.HTML(;
+        prettyurls = get(ENV, "CI", "false") == "true",
+        canonical = "https://pzimbrod.github.io/OperatorLearning.jl",
+        assets = String[]
     ),
-    pages=[
+    pages = [
         "Home" => "index.md",
-        "Examples" =>
-            ["Burgers Equation with FNO" => "examples/burgers_FNO.md",
-             "Burgers Equation with DeepONet" => "examples/burgers_DeepONet.md"],
+        "Examples" => ["Burgers Equation with FNO" => "examples/burgers_FNO.md",
+            "Burgers Equation with DeepONet" => "examples/burgers_DeepONet.md"],
         "Frequently Asked Questions" => "faq.md",
-        "Module Reference" => "reference.md",
-    ],
+        "Module Reference" => "reference.md"
+    ]
 )
 
 deploydocs(;
-    repo="github.com/pzimbrod/OperatorLearning.jl",
-    #devbranch="main",
+    repo = "github.com/pzimbrod/OperatorLearning.jl",
+#devbranch="main",
 )

docs/src/reference.md

@@ -1,3 +1,3 @@
 ```@autodocs
 Modules = [OperatorLearning]
-```
+```

examples/burgers_DeepONet.jl

@@ -21,22 +21,22 @@ subsample = 2^3;
 # create the x training array, according to our desired grid size
 xtrain = vars["a"][1:1000, 1:subsample:end]' |> device;
 # create the x test array
-xtest = vars["a"][end-99:end, 1:subsample:end]' |> device;
+xtest = vars["a"][(end - 99):end, 1:subsample:end]' |> device;
 
 # Create the y training array
 ytrain = vars["u"][1:1000, 1:subsample:end] |> device;
 # Create the y test array
-ytest = vars["u"][end-99:end, 1:subsample:end] |> device;
+ytest = vars["u"][(end - 99):end, 1:subsample:end] |> device;
 
 # The data is missing grid data, so we create it
 # `collect` converts data type `range` into an array
-grid = collect(range(0, 1, length=1024))' |> device
+grid = collect(range(0, 1, length = 1024))' |> device
 
 # Create the DeepONet:
 # IC is given on grid of 1024 points, and we solve for a fixed time t in one
 # spatial dimension x, making the branch input of size 1024 and trunk size 1
 # We choose GeLU activation for both subnets
-model = DeepONet((1024,1024,1024),(1,1024,1024),gelu,gelu) |> device
+model = DeepONet((1024, 1024, 1024), (1, 1024, 1024), gelu, gelu) |> device
 
 # We use the ADAM optimizer for training
 learning_rate = 0.001
@@ -48,12 +48,12 @@ parameters = params(model)
 # The loss function
 # We can't use the "vanilla" implementation of the mse here since we have
 # two distinct inputs to our DeepONet, so we wrap them into a tuple
-loss(xtrain,ytrain,sensor) = Flux.Losses.mse(model(xtrain,sensor),ytrain)
+loss(xtrain, ytrain, sensor) = Flux.Losses.mse(model(xtrain, sensor), ytrain)
 
 # Define a callback function that gives some output during training
-evalcb() = @show(loss(xtest,ytest,grid))
+evalcb() = @show(loss(xtest, ytest, grid))
 # Print the callback only every 5 seconds
 throttled_cb = throttle(evalcb, 5)
 
 # Do the training loop
-Flux.@epochs 500 train!(loss, parameters, [(xtrain,ytrain,grid)], opt, cb = evalcb)
+Flux.@epochs 500 train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb = evalcb)

examples/burgers_FNO.jl

@@ -13,55 +13,57 @@ subsample = 2^3;
 # create the x training array, according to our desired grid size
 xtrain = vars["a"][1:1000, 1:subsample:end] |> device;
 # create the x test array
-xtest = vars["a"][end-99:end, 1:subsample:end] |> device;
+xtest = vars["a"][(end - 99):end, 1:subsample:end] |> device;
 
 # Create the y training array
 ytrain = vars["u"][1:1000, 1:subsample:end] |> device;
 # Create the y test array
-ytest = vars["u"][end-99:end, 1:subsample:end] |> device;
+ytest = vars["u"][(end - 99):end, 1:subsample:end] |> device;
 
 # The data is missing grid data, so we create it
 # `collect` converts data type `range` into an array
-grid = collect(range(0, 1, length=length(xtrain[1,:]))) |> device
+grid = collect(range(0, 1, length = length(xtrain[1, :]))) |> device
 
 # Merge the created grid with the data
 # Output has the dims: batch x grid points x 2  (a(x), x)
 # First, reshape the data to a 3D tensor,
 # Then, create a 3D tensor from the synthetic grid data
 # and concatenate them along the newly created 3rd dim
-xtrain = cat(reshape(xtrain,(1000,1024,1)),
-            reshape(repeat(grid,1000),(1000,1024,1));
-            dims=3) |> device
-ytrain = cat(reshape(ytrain,(1000,1024,1)),
-            reshape(repeat(grid,1000),(1000,1024,1));
-            dims=3) |> device
+xtrain = cat(reshape(xtrain, (1000, 1024, 1)),
+    reshape(repeat(grid, 1000), (1000, 1024, 1));
+    dims = 3) |> device
+ytrain = cat(reshape(ytrain, (1000, 1024, 1)),
+    reshape(repeat(grid, 1000), (1000, 1024, 1));
+    dims = 3) |> device
 # Same treatment with the test data
-xtest = cat(reshape(xtest,(100,1024,1)),
-            reshape(repeat(grid,100),(100,1024,1));
-            dims=3) |> device
-ytest = cat(reshape(ytest,(100,1024,1)),
-            reshape(repeat(grid,100),(100,1024,1));
-            dims=3) |> device
+xtest = cat(reshape(xtest, (100, 1024, 1)),
+    reshape(repeat(grid, 100), (100, 1024, 1));
+    dims = 3) |> device
+ytest = cat(reshape(ytest, (100, 1024, 1)),
+    reshape(repeat(grid, 100), (100, 1024, 1));
+    dims = 3) |> device
 
 # Our net wants the input in the form (2,grid,batch), though,
 # So we permute
-xtrain, xtest = permutedims(xtrain,(3,2,1)), permutedims(xtest,(3,2,1)) |> device
-ytrain, ytest = permutedims(ytrain,(3,2,1)), permutedims(ytest,(3,2,1)) |> device
+xtrain, xtest = permutedims(xtrain, (3, 2, 1)), permutedims(xtest, (3, 2, 1)) |> device
+ytrain, ytest = permutedims(ytrain, (3, 2, 1)), permutedims(ytest, (3, 2, 1)) |> device
 
 # Pass the data to the Flux DataLoader and give it a batch of 20
-train_loader = Flux.Data.DataLoader((xtrain, ytrain), batchsize=20, shuffle=true) |> device
-test_loader = Flux.Data.DataLoader((xtest, ytest), batchsize=20, shuffle=false) |> device
+train_loader = Flux.Data.DataLoader((xtrain, ytrain), batchsize = 20, shuffle = true) |>
+               device
+test_loader = Flux.Data.DataLoader((xtest, ytest), batchsize = 20, shuffle = false) |>
+              device
 
 # Set up the Fourier Layer
 # 128 in- and outputs, batch size 20 as given above, grid size 1024
 # 16 modes to keep, σ activation on the gpu
-layer = FourierLayer(128,128,1024,16,gelu,bias_fourier=false) |> device
+layer = FourierLayer(128, 128, 1024, 16, gelu, bias_fourier = false) |> device
 
 # The whole architecture
 # linear transform into the latent space, 4 Fourier Layers,
 # then transform it back
-model = Chain(Dense(2,128;bias=false), layer, layer, layer, layer,
-                Dense(128,2;bias=false)) |> device
+model = Chain(Dense(2, 128; bias = false), layer, layer, layer, layer,
+    Dense(128, 2; bias = false)) |> device
 
 # We use the ADAM optimizer for training
 learning_rate = 0.001
@@ -71,21 +73,21 @@ opt = ADAM(learning_rate)
 parameters = params(model)
 
 # The loss function
-loss(x,y) = Flux.Losses.mse(model(x),y)
+loss(x, y) = Flux.Losses.mse(model(x), y)
 
 # Define a callback function that gives some output during training
-evalcb() = @show(loss(xtest,ytest))
+evalcb() = @show(loss(xtest, ytest))
 # Print the callback only every 5 seconds, 
 throttled_cb = throttle(evalcb, 5)
 
 # Do the training loop
 Flux.@epochs 500 train!(loss, parameters, train_loader, opt, cb = throttled_cb)
 
 # Accuracy metrics
-val_loader = Flux.Data.DataLoader((xtest, ytest), batchsize=1, shuffle=false) |> device
+val_loader = Flux.Data.DataLoader((xtest, ytest), batchsize = 1, shuffle = false) |> device
 loss = 0.0 |> device
 
-for (x,y) in val_loader
+for (x, y) in val_loader
     ŷ = model(x)
-    loss += Flux.Losses.mse(ŷ,y)
+    loss += Flux.Losses.mse(ŷ, y)
 end

src/ComplexWeights.jl

@@ -4,7 +4,9 @@ cglorot_uniform([rng=GLOBAL_RNG], dims...)
 A modification of the `glorot_uniform` function provided by `Flux` to accommodate Complex numbers.
 This is necessary since the parameters of the global convolution operator in the Fourier Layer generally has complex weights.
 """
-cglorot_uniform(rng::AbstractRNG, dims...) = (rand(rng, ComplexF32, dims...) .- 0.5f0) .* sqrt(24.0f0 / sum(nfan(dims...)))
+function cglorot_uniform(rng::AbstractRNG, dims...)
+    (rand(rng, ComplexF32, dims...) .- 0.5f0) .* sqrt(24.0f0 / sum(nfan(dims...)))
+end
 cglorot_uniform(dims...) = cglorot_uniform(Random.GLOBAL_RNG, dims...)
 cglorot_uniform(rng::AbstractRNG) = (dims...) -> cglorot_uniform(rng, dims...)
 
@@ -14,6 +16,8 @@ cglorot_normal([rng=GLOBAL_RNG], dims...)
 A modification of the `glorot_normal` function provided by `Flux` to accommodate Complex numbers.
 This is necessary since the parameters of the global convolution operator in the Fourier Layer generally has complex weights.
 """
-cglorot_normal(rng::AbstractRNG, dims...) = randn(rng, ComplexF32, dims...) .* sqrt(2.0f0 / sum(nfan(dims...)))
+function cglorot_normal(rng::AbstractRNG, dims...)
+    randn(rng, ComplexF32, dims...) .* sqrt(2.0f0 / sum(nfan(dims...)))
+end
 cglorot_normal(dims...) = cglorot_normal(Random.GLOBAL_RNG, dims...)
-cglorot_normal(rng::AbstractRNG) = (dims...) -> cglorot_normal(rng, dims...)
+cglorot_normal(rng::AbstractRNG) = (dims...) -> cglorot_normal(rng, dims...)