Updated docs and readme

Author: ba2tripleO

README.md

@@ -86,8 +86,26 @@ DeepONet
 #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)
+```
+
 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))
+
+learning_rate = 0.001
+opt = ADAM(learning_rate)
+parameters = params(model)
+Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain,ytrain,grid)], opt, cb = evalcb)
+```
+
 ## Examples
 
 PDE training examples are provided in `example` folder.
@@ -96,6 +114,10 @@ PDE training examples are provided in `example` folder.
 
 [Burgers' equation](example/Burgers)
 
+### DeepONet implementation for solving Burgers' equation
+
+[Burgers' equation](example/Burgers/Burgers_deeponet)
+
 ### Two-dimensional Fourier Neural Operator
 
 [Double Pendulum](example/DoublePendulum)

docs/src/index.md

@@ -10,7 +10,7 @@ Documentation for [NeuralOperators](https://github.com/foldfelis/NeuralOperators
 |:----------------:|:--------------:|
 | ![](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/inferenced.gif?raw=true) |
 
-The demonstration showing above is Navier-Stokes equation learned by the `MarkovNeuralOperator` with only one time step information.
+The demonstration shown 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).
 
 ## Abstract
@@ -30,6 +30,8 @@ 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.
+
 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).
@@ -44,6 +46,8 @@ pkg> add NeuralOperators
 
 ## Usage
 
+### Fourier Neural Operator
+
 ```julia
 model = Chain(
     # lift (d + 1)-dimensional vector field to n-dimensional vector field
@@ -78,3 +82,31 @@ 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)
 ```
+
+### 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)
+```
+
+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)
+```
+
+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))
+
+learning_rate = 0.001
+opt = ADAM(learning_rate)
+parameters = params(model)
+Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain,ytrain,grid)], opt, cb = evalcb)
+```
+A more complete example using DeepONet architecture to solve Burgers can be found in the [examples]()