Refactor project

Author: foldfelis

src/DeepONet.jl renamed to src/DeepONet/DeepONet.jl

@@ -1,3 +1,7 @@
+export DeepONet
+
+include("subnets.jl")
+
 """
 `DeepONet(architecture_branch::Tuple, architecture_trunk::Tuple,
                         act_branch = identity, act_trunk = identity;

src/subnets.jl renamed to src/DeepONet/subnets.jl

@@ -1,3 +1,5 @@
+export NOMAD
+
 struct NOMAD{T1, T2}
     approximator_net::T1
     decoder_net::T2

src/model.jl renamed to src/FNO/FNO.jl

@@ -1,4 +1,5 @@
 module NeuralOperators
+
 using Flux
 using FFTW
 using Tullio
@@ -10,15 +11,15 @@ using ChainRulesCore
 using GeometricFlux
 using Statistics
 
-export DeepONet
-export NOMAD
-
+# kernels
 include("Transform/Transform.jl")
 include("operator_kernel.jl")
+include("graph_kernel.jl")
 include("loss.jl")
-include("model.jl")
-include("DeepONet.jl")
-include("subnets.jl")
-include("NOMAD.jl")
 
-end
+# models
+include("FNO/FNO.jl")
+include("DeepONet/DeepONet.jl")
+include("NOMAD/NOMAD.jl")
+
+end # module

src/NOMAD.jl renamed to src/NOMAD/NOMAD.jl

@@ -0,0 +1,50 @@
+export GraphKernel
+
+"""
+    GraphKernel(κ, ch, σ=identity)
+
+Graph kernel layer.
+
+## Arguments
+
+* `κ`: A neural network layer for approximation, e.g. a `Dense` layer or a MLP.
+* `ch`: Channel size for linear transform, e.g. `32`.
+* `σ`: Activation function.
+
+## Keyword Arguments
+
+* `init`: Initial function to initialize parameters.
+"""
+struct GraphKernel{A, B, F} <: MessagePassing
+    linear::A
+    κ::B
+    σ::F
+end
+
+function GraphKernel(κ, ch::Int, σ = identity; init = Flux.glorot_uniform)
+    W = init(ch, ch)
+    return GraphKernel(W, κ, σ)
+end
+
+Flux.@functor GraphKernel
+
+function GeometricFlux.message(l::GraphKernel, x_i::AbstractArray, x_j::AbstractArray, e_ij)
+    return l.κ(vcat(x_i, x_j))
+end
+
+function GeometricFlux.update(l::GraphKernel, m::AbstractArray, x::AbstractArray)
+    return l.σ.(GeometricFlux._matmul(l.linear, x) + m)
+end
+
+function (l::GraphKernel)(el::NamedTuple, X::AbstractArray)
+    GraphSignals.check_num_nodes(el.N, X)
+    _, V, _ = GeometricFlux.propagate(l, el, nothing, X, nothing, mean, nothing, nothing)
+    return V
+end
+
+function Base.show(io::IO, l::GraphKernel)
+    channel, _ = size(l.linear)
+    print(io, "GraphKernel(", l.κ, ", channel=", channel)
+    l.σ == identity || print(io, ", ", l.σ)
+    print(io, ")")
+end

src/NeuralOperators.jl

@@ -1,8 +1,7 @@
 export
        OperatorConv,
        SpectralConv,
-       OperatorKernel,
-       GraphKernel
+       OperatorKernel
 
 struct OperatorConv{P, T, S, TT}
     weight::T
@@ -66,8 +65,8 @@ function SpectralConv(ch::Pair{S, S},
                       init = c_glorot_uniform,
                       permuted = false,
                       T::DataType = ComplexF32) where {S <: Integer, N}
-    return OperatorConv(ch, modes, FourierTransform, init = init, permuted = permuted,
-                        T = T)
+    return OperatorConv(ch, modes, FourierTransform,
+                        init = init, permuted = permuted, T = T)
 end
 
 Flux.@functor OperatorConv{true}
@@ -181,55 +180,6 @@ function (m::OperatorKernel)(𝐱)
     return m.σ.(m.linear(𝐱) + m.conv(𝐱))
 end
 
-"""
-    GraphKernel(κ, ch, σ=identity)
-
-Graph kernel layer.
-
-## Arguments
-
-* `κ`: A neural network layer for approximation, e.g. a `Dense` layer or a MLP.
-* `ch`: Channel size for linear transform, e.g. `32`.
-* `σ`: Activation function.
-
-## Keyword Arguments
-
-* `init`: Initial function to initialize parameters.
-"""
-struct GraphKernel{A, B, F} <: MessagePassing
-    linear::A
-    κ::B
-    σ::F
-end
-
-function GraphKernel(κ, ch::Int, σ = identity; init = Flux.glorot_uniform)
-    W = init(ch, ch)
-    return GraphKernel(W, κ, σ)
-end
-
-Flux.@functor GraphKernel
-
-function GeometricFlux.message(l::GraphKernel, x_i::AbstractArray, x_j::AbstractArray, e_ij)
-    return l.κ(vcat(x_i, x_j))
-end
-
-function GeometricFlux.update(l::GraphKernel, m::AbstractArray, x::AbstractArray)
-    return l.σ.(GeometricFlux._matmul(l.linear, x) + m)
-end
-
-function (l::GraphKernel)(el::NamedTuple, X::AbstractArray)
-    GraphSignals.check_num_nodes(el.N, X)
-    _, V, _ = GeometricFlux.propagate(l, el, nothing, X, nothing, mean, nothing, nothing)
-    return V
-end
-
-function Base.show(io::IO, l::GraphKernel)
-    channel, _ = size(l.linear)
-    print(io, "GraphKernel(", l.κ, ", channel=", channel)
-    l.σ == identity || print(io, ", ", l.σ)
-    print(io, ")")
-end
-
 #########
 # utils #
 #########

src/graph_kernel.jl

@@ -28,6 +28,6 @@
           "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
+    # @info mgrad.grads
     @test length(mgrad.grads) == 5
 end