Updated tests and readme

Author: ba2tripleO

README.md

@@ -36,11 +36,15 @@ 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](src/model.jl).
 
 ## Usage
 
+Fourier Neural Operator
+
 ```julia
 model = Chain(
     # lift (d + 1)-dimensional vector field to n-dimensional vector field
@@ -76,6 +80,14 @@ 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)
+```
+You can again specify loss, optimization and training parameters just as you would for a simple neural network with Flux.
+
 ## Examples
 
 PDE training examples are provided in `example` folder.
@@ -113,3 +125,4 @@ PDE training examples are provided in `example` folder.
 - [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)
+- [DeepONet: Learning nonlinear operators for identifying  differential equations based on the universal approximation theorem of operators](https://arxiv.org/abs/1910.03193)

test/deeponet.jl

@@ -33,4 +33,9 @@ mgrad = Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)
 @test !iszero(Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)[1])
 @test !iszero(Flux.Zygote.gradient((x,p)->sum(model(x,p)),a,sensors)[2])
 
-#training
+#Output size of branch and trunk subnets should be same
+branch = Chain(Dense(16, 22), Dense(22, 30))
+trunk = Chain(Dense(1, 16), Dense(16, 24), Dense(24, 32))
+m = DeepONet(branch, trunk)
+@test_throws AssertionError DeepONet((32,64,70), (24,48,72), σ, tanh)
+@test_throws DimensionMismatch m(a, sensors)