add doc for FourierOperator

Author: foldfelis

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

@@ -41,12 +41,14 @@ end
 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},
@@ -104,17 +106,70 @@ 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)