Merge pull request #13 from yuehhua/doc

Improve documentation and README

Author: foldfelis

README.md

@@ -1,7 +1,5 @@
 # NeuralOperators
 
-## Source code status
-
 | **Documentation** | **Build Status** |
 |:-----------------:|:----------------:|
 | [![doc dev badge]][doc dev link] | [![ci badge]][ci link] [![codecov badge]][codecov link] |
@@ -14,11 +12,82 @@
 [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.
+Neural operator is a novel deep learning architecture. It learns a operator, which is a mapping
+between infinite-dimensional function spaces. It can be used to resolve [partial differential equations (PDE)](https://en.wikipedia.org/wiki/Partial_differential_equation).
+Instead of solving by finite element method, a PDE problem can be resolved by learning a neural network to learn an operator
+mapping from infinite-dimensional space (u, t) to infinite-dimensional space f(u, t). Neural operator learns a continuous function
+between two continuous function spaces. The kernel can be trained on different geometry, which is learned from a graph.
+
+Fourier neural operator learns a neural operator with Dirichlet kernel to form a Fourier transformation. It performs Fourier transformation across infinite-dimensional function spaces and learns better than neural operator.
+
+Currently, `FourierOperator` is provided in this work.
+
+## Usage
+
+```
+function FourierNeuralOperator()
+    modes = (16, )
+    ch = 64 => 64
+    σ = gelu
+
+    return Chain(
+        # project finite-dimensional data to infinite-dimensional space
+        Dense(2, 64),
+        # operator projects data between infinite-dimensional spaces
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes, σ),
+        FourierOperator(ch, modes),
+        # project infinite-dimensional function to finite-dimensional space
+        Dense(64, 128, σ),
+        Dense(128, 1),
+        flatten
+    )
+end
+```
+
+Or you can just call:
+
+```
+fno = FourierNeuralOperator()
+```
+
+And then train as a Flux model.
+
+```
+loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- fno(𝐱)) / size(𝐱)[end]
+opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+Flux.@epochs 50 Flux.train!(loss, params(m), data, opt)
+```
+
+## Examples
+
+PDE training examples are provided in `example` folder.
+
+### One-dimensional Burgers' equation
+
+[Burgers' equation](https://en.wikipedia.org/wiki/Burgers%27_equation) example can be found in `example/burgers.jl`.
+
+### Two-dimensional Darcy flow equation
+
+WIP
+
+### Two-dimensional Navier-Stokes equation
+
+WIP
+
+## Roadmap
+
+- [x] `FourierOperator` layer
+- [x] One-dimensional Burgers' equation example
+- [ ] Two-dimensional Darcy flow equation example
+- [ ] Two-dimensional Navier-Stokes equation example
+- [ ] `NeuralOperator` layer
+- [ ] Poisson equation example
+
+## References
 
-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)
+- [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+  - [zongyi-li/fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator)
+- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://arxiv.org/abs/2003.03485)
+  - [zongyi-li/graph-pde](https://github.com/zongyi-li/graph-pde)

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,13 @@ 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.
+* `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
+    data in the order of `(..., ch, batch)`, otherwise the order is `(ch, ..., batch)`.
 
 ## Example
 
@@ -109,18 +109,13 @@ 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.
+* `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
+    data in the order of `(..., ch, batch)`, otherwise the order is `(ch, ..., batch)`.
 
 ## Example