Doc

README.md

@@ -12,13 +12,20 @@
 [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 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.
+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 training 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.
+**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.
+
+**Markov neural operator** learns a neural operator with Fourier operators.
+With only one time step information of learning, it can predict the following few steps with low loss
+by linking the operators into a Markov chain.
 
 Currently, the `FourierOperator` layer is provided in this work.
 As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided. Please take a glance at them [here](src/model.jl).
@@ -41,7 +48,7 @@ model = Chain(
 )
 ```
 
-Or you can just call:
+Or one can just call:
 
 ```julia
 model = FourierNeuralOperator(
@@ -106,3 +113,4 @@ julia> using FlowOverCircle; FlowOverCircle.train()
   - [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)
+- [Markov Neural Operators for Learning Chaotic Systems](https://arxiv.org/abs/2106.06898)

docs/make.jl

@@ -15,6 +15,7 @@ makedocs(;
     ),
     pages=[
         "Home" => "index.md",
+        "APIs" => "apis.md"
     ],
 )
 

docs/src/apis.md

@@ -0,0 +1,61 @@
+## Index
+
+```@index
+```
+
+## 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)
+
+## Models
+
+### Fourier neural operator
+
+```@docs
+FourierNeuralOperator
+```
+
+Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+
+---
+
+### Markov neural operator
+
+```@docs
+MarkovNeuralOperator
+```
+
+Reference: [Markov Neural Operators for Learning Chaotic Systems](https://arxiv.org/abs/2106.06898)

docs/src/index.md

@@ -6,39 +6,75 @@ CurrentModule = NeuralOperators
 
 Documentation for [NeuralOperators](https://github.com/foldfelis/NeuralOperators.jl).
 
-```@index
-```
+| **Ground Truth** | **Inferenced** |
+|:----------------:|:--------------:|
+| ![](https://github.com/foldfelis/NeuralOperators.jl/blob/master/example/FlowOverCircle/gallery/ans.gif?raw=true) | ![](https://github.com/foldfelis/NeuralOperators.jl/blob/master/example/FlowOverCircle/gallery/ans.gif?raw=true) |
 
-## Layers
+The demonstration showing above is Navier-Stokes equation learned by the `MarkovNeuralOperator` with only one time step information.
+Example can be found in [`example/FlowOverCircle`](https://github.com/foldfelis/NeuralOperators.jl/tree/master/example/FlowOverCircle).
+The result is also provided [here](assets/notebook/mno.jl.html)
 
-### Spectral convolutional layer
+## Abstract
 
-```math
-F(s) = \mathcal{F} \{ v(x) \} \\
-F'(s) = g(F(s)) \\
-v'(x) = \mathcal{F}^{-1} \{ F'(s) \}
-```
+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 training 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.
 
-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.
+**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.
 
-```@docs
-SpectralConv
-```
+**Markov neural operator** learns a neural operator with Fourier operators.
+With only one time step information of learning, it can predict the following few steps with low loss
+by linking the operators into a Markov chain.
+
+Currently, the `FourierOperator` layer is provided in this work.
+As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided.
+Please take a glance at them [here](apis.html#Models).
 
-Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+## Quick start
 
----
+The package can be installed with the Julia package manager. From the Julia REPL, type `]` to enter the Pkg REPL mode and run:
 
-### Fourier operator layer
+```julia-repl
+pkg> add NeuralOperators
+```
+
+## Usage
 
-```math
-v_{t+1}(x) = \sigma(W v_t(x) + \mathcal{K} \{ v_t(x) \} )
+```julia
+model = Chain(
+    # project finite-dimensional data to infinite-dimensional space
+    Dense(2, 64),
+    # operator projects data between infinite-dimensional spaces
+    FourierOperator(64=>64, (16, ), gelu),
+    FourierOperator(64=>64, (16, ), gelu),
+    FourierOperator(64=>64, (16, ), gelu),
+    FourierOperator(64=>64, (16, )),
+    # project infinite-dimensional function to finite-dimensional space
+    Dense(64, 128, gelu),
+    Dense(128, 1),
+    flatten
+)
 ```
 
-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.
+Or one can just call:
 
-```@docs
-FourierOperator
+```julia
+model = FourierNeuralOperator(
+    ch=(2, 64, 64, 64, 64, 64, 128, 1),
+    modes=(16, ),
+    σ=gelu
+)
 ```
 
-Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
+And then train as a Flux model.
+
+```julia
+loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- model(𝐱)) / size(𝐱)[end]
+opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)
+```

src/model.jl

@@ -2,6 +2,16 @@ export
     FourierNeuralOperator,
     MarkovNeuralOperator
 
+"""
+    FourierNeuralOperator(;
+        ch=(2, 64, 64, 64, 64, 64, 128, 1),
+        modes=(16, ),
+        σ=gelu
+    )
+
+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.
+"""
 function FourierNeuralOperator(;
     ch=(2, 64, 64, 64, 64, 64, 128, 1),
     modes=(16, ),
@@ -19,6 +29,17 @@ function FourierNeuralOperator(;
     )
 end
 
+"""
+    MarkovNeuralOperator(;
+        ch=(1, 64, 64, 64, 64, 64, 1),
+        modes=(24, 24),
+        σ=gelu
+    )
+
+Markov neural operator learns a neural operator with Fourier operators.
+With only one time step information of learning, it can predict the following few steps with low loss
+by linking the operators into a Markov chain.
+"""
 function MarkovNeuralOperator(;
     ch=(1, 64, 64, 64, 64, 64, 1),
     modes=(24, 24),