update manual

Author: yuehhua

docs/src/index.md

@@ -9,6 +9,36 @@ Documentation for [NeuralOperators](https://github.com/foldfelis/NeuralOperators
 ```@index
 ```
 
-```@autodocs
-Modules = [NeuralOperators]
+## Layers
+
+### Spectral convolutional layer
+
+```math
+F(s) = \mathcal{F} \{ v(x) \} \\
+F'(s) = g(F(s)) \\
+v'(x) = \mathcal{F}^{-1} \{ F'(s) \}
+```
+
+where ``v(x)`` and ``v'(x)`` denotes input and output function, ``\mathcal{F} \{ \cdot \}``, ``\mathcal{F}^{-1} \{ \cdot \}`` are Fourier transform, inverse Fourier transform, respectively. Function ``g`` is a linear transform for lowering Fouier modes.
+
+```@docs
+SpectralConv
+```
+
+Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+
+---
+
+### Fourier operator layer
+
+```math
+v_{t+1}(x) = \sigma(W v_t(x) + \mathcal{K} \{ v_t(x) \} )
 ```
+
+where ``v_t(x)`` is the input function for ``t``-th layer and ``\mathcal{K} \{ \cdot \}`` denotes spectral convolutional layer. Activation function ``\sigma`` can be arbitrary non-linear function.
+
+```@docs
+FourierOperator
+```
+
+Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)

src/fourier.jl

@@ -27,13 +27,11 @@ end
         init=c_glorot_uniform, permuted=false, T=ComplexF32
     )
 
-## SpectralConv
+## Arguments
 
-* ``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}``
+* `ch`: Input and output channel size, e.g. `64=>64`.
+* `modes`: The Fourier modes to be preserved.
+* `σ`: Activation function.
 
 ## Example
 
@@ -109,18 +107,11 @@ end
 """
     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
+## Arguments
 
-```
-        ┌ F -> L -> F¯¹ ┐
-v(x) -> ┤               ├ -> + -> σ
-        └      D        ┘
-```
+* `ch`: Input and output channel size for spectral convolution, e.g. `64=>64`.
+* `modes`: The Fourier modes to be preserved for spectral convolution.
+* `σ`: Activation function.
 
 ## Example