Merge pull request #12 from foldfelis/doc

Doc

Author: foldfelis

Project.toml

@@ -18,7 +18,7 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-CUDA = "3.3"
+CUDA = "3.4"
 CUDAKernels = "0.3"
 ChainRulesCore = "1.3"
 DataDeps = "0.7"

README.md

@@ -13,3 +13,12 @@
 [ci link]: https://github.com/foldfelis/NeuralOperators.jl/actions/workflows/CI.yml
 [codecov badge]: https://codecov.io/gh/foldfelis/NeuralOperators.jl/branch/master/graph/badge.svg?token=JQH3MP1Y9R
 [codecov link]: https://codecov.io/gh/foldfelis/NeuralOperators.jl
+
+[Neural Operator](https://github.com/zongyi-li/graph-pde) is a novel deep learning method to learned the mapping
+between infinite-dimensional spaces of functions introduced by [Zongyi Li](https://github.com/zongyi-li)et al.
+
+In this project I temporarily provide the SpectralConv layer and the
+[Fourier Neural Operator](https://github.com/zongyi-li/fourier_neural_operator).
+For more information, please take a look at the
+[Fourier Neural Operator model](src/model.jl) and the [example](example/burgers.jl) of solving
+[Burgers' equation](https://www.wikiwand.com/en/Burgers%27_equation)

docs/Project.toml

@@ -1,3 +1,4 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"

src/fourier.jl

@@ -1,23 +1,55 @@
 export
     SpectralConv,
+    SpectralConvPerm,
     FourierOperator
 
-struct SpectralConv{N, T, S, F}
+abstract type AbstractSpectralConv{N, T, S, F} end
+
+struct SpectralConv{N, T, S, F} <: AbstractSpectralConv{N, T, S, F}
     weight::T
     in_channel::S
     out_channel::S
     modes::NTuple{N, S}
     σ::F
 end
 
-struct SpectralConvPerm{N, T, S, F}
+struct SpectralConvPerm{N, T, S, F} <: AbstractSpectralConv{N, T, S, F}
     weight::T
     in_channel::S
     out_channel::S
     modes::NTuple{N, S}
     σ::F
 end
 
+"""
+    SpectralConv(
+        ch, modes, σ=identity;
+        init=c_glorot_uniform, permuted=false, T=ComplexF32
+    )
+
+## SpectralConv
+
+* ``v(x)``: input
+* ``F``, ``F^{-1}``: Fourier transform, inverse fourier transform
+* ``L``: linear transform on the lower Fouier modes.
+
+``v(x)`` -> ``F`` -> ``L`` -> ``F^{-1}``
+
+## Example
+
+```jldoctest
+julia> SpectralConv(2=>5, (16, ))
+SpectralConv(2 => 5, (16,), σ=identity)
+
+julia> using Flux
+
+julia> SpectralConv(2=>5, (16, ), relu)
+SpectralConv(2 => 5, (16,), σ=relu)
+
+julia> SpectralConv(2=>5, (16, ), relu, permuted=true)
+SpectralConvPerm(2 => 5, (16,), σ=relu)
+```
+"""
 function SpectralConv(
     ch::Pair{S, S},
     modes::NTuple{N, S},
@@ -38,10 +70,14 @@ end
 Flux.@functor SpectralConv
 Flux.@functor SpectralConvPerm
 
-Base.ndims(::SpectralConv{N}) where {N} = N
-Base.ndims(::SpectralConvPerm{N}) where {N} = N
+Base.ndims(::AbstractSpectralConv{N}) where {N} = N
+
+function Base.show(io::IO, l::AbstractSpectralConv)
+    T = (l isa SpectralConv) ? SpectralConv : SpectralConvPerm
+    print(io, "$(string(T))($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)))")
+end
 
-function spectral_conv(m, 𝐱)
+function spectral_conv(m::AbstractSpectralConv, 𝐱::AbstractArray)
     n_dims = ndims(𝐱)
 
     𝐱_fft = fft(Zygote.hook(real, 𝐱), 1:ndims(m)) # [x, in_chs, batch]
@@ -54,33 +90,86 @@ function spectral_conv(m, 𝐱)
     return m.σ.(𝐱_ifft)
 end
 
-function (m::SpectralConv)(𝐱::AbstractArray)
+function (m::SpectralConv)(𝐱)
     𝐱ᵀ = permutedims(𝐱, (ntuple(i->i+1, ndims(m))..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch]
     𝐱_out = spectral_conv(m, 𝐱ᵀ) # [x, out_chs, batch]
     𝐱_outᵀ = permutedims(𝐱_out, (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch]
 
     return 𝐱_outᵀ
 end
 
-function (m::SpectralConvPerm)(𝐱::AbstractArray)
+function (m::SpectralConvPerm)(𝐱)
     return spectral_conv(m, 𝐱) # [x, out_chs, batch]
 end
 
 ############
 # operator #
 ############
 
+"""
+    FourierOperator(ch, modes, σ=identity; permuted=false)
+
+## FourierOperator
+
+* ``v(x)``: input
+* ``F``, ``F^{-1}``: Fourier transform, inverse fourier transform
+* ``L``: linear transform on the lower Fouier modes
+* ``D``: local linear transform
+
+```
+        ┌ F -> L -> F¯¹ ┐
+v(x) -> ┤               ├ -> + -> σ
+        └      D        ┘
+```
+
+## Example
+
+```jldoctest
+julia> FourierOperator(2=>5, (16, ))
+Chain(
+  Parallel(
+    +,
+    Dense(2, 5),                        # 15 parameters
+    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
+  ),
+  NeuralOperators.var"#activation_func#14"{typeof(identity)}(identity),
+)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+
+julia> using Flux
+
+julia> FourierOperator(2=>5, (16, ), relu)
+Chain(
+  Parallel(
+    +,
+    Dense(2, 5),                        # 15 parameters
+    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
+  ),
+  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
+)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+
+julia> FourierOperator(2=>5, (16, ), relu, permuted=true)
+Chain(
+  Parallel(
+    +,
+    Conv((1,), 2 => 5),                 # 15 parameters
+    SpectralConvPerm(2 => 5, (16,), σ=identity),  # 160 parameters
+  ),
+  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
+)                   # Total: 3 arrays, 175 parameters, 1.871 KiB.
+```
+"""
 function FourierOperator(
     ch::Pair{S, S},
     modes::NTuple{N, S},
     σ=identity;
     permuted=false
 ) where {S<:Integer, N}
     short_cut = permuted ? Conv(Tuple(ones(Int, length(modes))), ch) : Dense(ch.first, ch.second)
+    activation_func(x) = σ.(x)
 
     return Chain(
         Parallel(+, short_cut, SpectralConv(ch, modes, permuted=permuted)),
-        x -> σ.(x)
+        activation_func
     )
 end
 

test/fourier.jl

@@ -7,6 +7,7 @@
         SpectralConv(ch, modes)
     )
     @test ndims(SpectralConv(ch, modes)) == 1
+    @test repr(SpectralConv(ch, modes)) == "SpectralConv(64 => 128, (16,), σ=identity)"
 
     𝐱, _ = get_burgers_data(n=5)
     @test size(m(𝐱)) == (128, 1024, 5)