Update docs

Author: foldfelis

docs/src/apis.md

@@ -5,7 +5,7 @@
 
 ## Layers
 
-### Spectral convolutional layer
+### Operator convolutional layer
 
 ```math
 F(s) = \mathcal{F} \{ v(x) \} \\
@@ -25,7 +25,7 @@ Reference: [Fourier Neural Operator for Parametric Partial Differential Equation
 
 ---
 
-### Fourier operator layer
+### Operator kernel layer
 
 ```math
 v_{t+1}(x) = \sigma(W v_t(x) + \mathcal{K} \{ v_t(x) \} )

docs/src/index.md

@@ -30,7 +30,9 @@ It performs Fourier transformation across infinite-dimensional function spaces a
 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.
 
-**DeepONet operator** (Deep Operator Network)learns a neural operator with the help of two sub-neural net structures described as the branch and the trunk network. The branch network is fed the initial conditions data, whereas the trunk is fed with the locations where the target(output) is evaluated from the corresponding initial conditions. It is important that the output size of the branch and trunk subnets is same so that a dot product can be performed between them.
+**DeepONet operator** (Deep Operator Network) learns a neural operator with the help of two sub-neural net structures described as the branch and the trunk network.
+The branch network is fed the initial conditions data, whereas the trunk is fed with the locations where the target(output) is evaluated from the corresponding initial conditions.
+It is important that the output size of the branch and trunk subnets is same so that a dot product can be performed between them.
 
 Currently, the `OperatorKernel` layer is provided in this work.
 As for model, there are `FourierNeuralOperator` and `MarkovNeuralOperator` provided.
@@ -54,10 +56,10 @@ model = Chain(
     # here, d == 1 and n == 64
     Dense(2, 64),
     # map each hidden representation to the next by integral kernel operator
-    OperatorKernel(64=>64, (16, ), gelu),
-    OperatorKernel(64=>64, (16, ), gelu),
-    OperatorKernel(64=>64, (16, ), gelu),
-    OperatorKernel(64=>64, (16, )),
+    OperatorKernel(64=>64, (16, ), FourierTransform, gelu),
+    OperatorKernel(64=>64, (16, ), FourierTransform, gelu),
+    OperatorKernel(64=>64, (16, ), FourierTransform, gelu),
+    OperatorKernel(64=>64, (16, ), FourierTransform),
     # project back to the scalar field of interest space
     Dense(64, 128, gelu),
     Dense(128, 1),
@@ -86,27 +88,28 @@ Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)
 ### DeepONet
 
 ```julia
-#tuple of Ints for branch net architecture and then for trunk net, followed by activations for branch and trunk respectively
-model = DeepONet((32,64,72), (24,64,72), σ, tanh)
+# tuple of Ints for branch net architecture and then for trunk net,
+# followed by activations for branch and trunk respectively
+model = DeepONet((32, 64, 72), (24, 64, 72), σ, tanh)
 ```
 
 Or specify branch and trunk as separate `Chain` from Flux and pass to `DeepONet`
 
 ```julia
-branch = Chain(Dense(32,64,σ), Dense(64,72,σ))
-trunk = Chain(Dense(24,64,tanh), Dense(64,72,tanh))
-model = DeepONet(branch,trunk)
+branch = Chain(Dense(32, 64, σ), Dense(64, 72, σ))
+trunk = Chain(Dense(24, 64, tanh), Dense(64, 72, tanh))
+model = DeepONet(branch, trunk)
 ```
 
 You can again specify loss, optimization and training parameters just as you would for a simple neural network with Flux.
 
 ```julia
-loss(xtrain,ytrain,sensor) = Flux.Losses.mse(model(xtrain,sensor),ytrain)
-evalcb() = @show(loss(xval,yval,grid))
+loss(xtrain, ytrain, sensor) = Flux.Losses.mse(model(xtrain, sensor), ytrain)
+evalcb() = @show(loss(xval, yval, grid))
 
 learning_rate = 0.001
 opt = ADAM(learning_rate)
 parameters = params(model)
-Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain,ytrain,grid)], opt, cb = evalcb)
+Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb=evalcb)
 ```
 A more complete example using DeepONet architecture to solve Burgers' equation can be found in the [examples](../../example/Burgers/src/Burgers_deeponet.jl)