Merge pull request #46 from SciML/transforms

Separate out transforms from OperatorKernels

Author: foldfelis

README.md

@@ -36,9 +36,11 @@ 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 `FourierOperator` layer is provided in this work.
+Currently, the `OperatorKernel` 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
@@ -51,10 +53,10 @@ model = Chain(
     # here, d == 1 and n == 64
     Dense(2, 64),
     # map each hidden representation to the next by integral kernel operator
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(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),
@@ -83,27 +85,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)
 ```
 
 ## Examples
@@ -130,16 +133,6 @@ PDE training examples are provided in `example` folder.
 
 [Super resolution on time dependent Navier-Stokes equation](example/SuperResolution)
 
-## Roadmap
-
-- [x] `FourierOperator` layer
-- [x] One-dimensional Burgers' equation example
-- [x] Two-dimensional with time Navier-Stokes equations example
-- [x] `MarkovNeuralOperator` model
-- [x] Flow over a circle prediction example
-- [ ] `NeuralOperator` layer
-- [ ] Poisson equation example
-
 ## References
 
 - [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)

docs/src/apis.md

@@ -5,7 +5,7 @@
 
 ## Layers
 
-### Spectral convolutional layer
+### Operator convolutional layer
 
 ```math
 F(s) = \mathcal{F} \{ v(x) \} \\
@@ -18,14 +18,14 @@ where ``v(x)`` and ``v'(x)`` denotes input and output function,
 Function ``g`` is a linear transform for lowering Fouier modes.
 
 ```@docs
-SpectralConv
+OperatorConv
 ```
 
 Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)
 
 ---
 
-### Fourier operator layer
+### Operator kernel layer
 
 ```math
 v_{t+1}(x) = \sigma(W v_t(x) + \mathcal{K} \{ v_t(x) \} )
@@ -35,7 +35,7 @@ where ``v_t(x)`` is the input function for ``t``-th layer and ``\mathcal{K} \{ \
 Activation function ``\sigma`` can be arbitrary non-linear function.
 
 ```@docs
-FourierOperator
+OperatorKernel
 ```
 
 Reference: [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)

docs/src/index.md

@@ -30,9 +30,11 @@ 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 `FourierOperator` layer is provided in this work.
+Currently, the `OperatorKernel` 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).
 
@@ -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
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(64=>64, (16, ), gelu),
-    FourierOperator(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)

example/Burgers/src/Burgers.jl

@@ -21,12 +21,13 @@ function train()
     modes = (16, )
     ch = 64 => 64
     σ = gelu
+    Transform = FourierTransform
     m = Chain(
         Dense(2, 64),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes, σ),
-        FourierOperator(ch, modes),
+        OperatorKernel(ch, modes, Transform, σ),
+        OperatorKernel(ch, modes, Transform, σ),
+        OperatorKernel(ch, modes, Transform, σ),
+        OperatorKernel(ch, modes, Transform),
         Dense(64, 128, σ),
         Dense(128, 1),
         flatten

example/Burgers/src/Burgers_deeponet.jl

@@ -1,15 +1,15 @@
 function train_don()
-    if has_cuda()
-        @info "CUDA is on"
-        device = gpu
-        CUDA.allowscalar(false)
-    else
+    # if has_cuda()
+    #     @info "CUDA is on"
+    #     device = gpu
+    #     CUDA.allowscalar(false)
+    # else
         device = cpu
-    end
+    # end
 
     x, y = get_data_don(n=300)
     xtrain = x[1:280, :]' |> device
-    xval = x[end-19:end, :]' |device
+    xval = x[end-19:end, :]' |> device
 
     ytrain = y[1:280, :] |> device
     yval = y[end-19:end, :] |> device
@@ -20,7 +20,7 @@ function train_don()
     opt = ADAM(learning_rate)
 
     m = DeepONet((1024,1024,1024),(1,1024,1024),gelu,gelu)
-    loss(xtrain,ytrain,sensor) = Flux.Losses.mse(model(xtrain,sensor),ytrain)
+    loss(xtrain,ytrain,sensor) = Flux.Losses.mse(m(xtrain,sensor),ytrain)
     evalcb() = @show(loss(xval,yval,grid))
 
     Flux.@epochs 400 Flux.train!(loss, params(m), [(xtrain,ytrain,grid)], opt, cb = evalcb)

example/DoublePendulum/src/DoublePendulum.jl

@@ -26,10 +26,10 @@ function train(; Δt=1)
 
     m = Chain(
         Dense(2, 64),
-        FourierOperator(64=>64, (4, 16), gelu),
-        FourierOperator(64=>64, (4, 16), gelu),
-        FourierOperator(64=>64, (4, 16), gelu),
-        FourierOperator(64=>64, (4, 16)),
+        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
+        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
+        OperatorKernel(64=>64, (4, 16), FourierTransform, gelu),
+        OperatorKernel(64=>64, (4, 16), FourierTransform),
         Dense(64, 128, gelu),
         Dense(128, 2),
     ) |> device

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -24,10 +24,10 @@ function train()
 
     m = Chain(
         Dense(1, 64),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
+        OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
+        OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
+        OperatorKernel(64=>64, (24, 24), FourierTransform, gelu),
         Dense(64, 1),
     ) |> device
 

example/SuperResolution/src/SuperResolution.jl

@@ -24,10 +24,10 @@ function train()
 
     m = Chain(
         Dense(1, 64),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
-        FourierOperator(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
         Dense(64, 1),
     ) |> device
 

src/NeuralOperators.jl

@@ -10,7 +10,8 @@ module NeuralOperators
 
     export DeepONet
 
-    include("fourier.jl")
+    include("Transform/Transform.jl")
+    include("operator_kernel.jl")
     include("model.jl")
     include("DeepONet.jl")
     include("subnets.jl")

src/Transform/Transform.jl

@@ -0,0 +1,19 @@
+export
+    AbstractTransform,
+    transform,
+    truncate_modes,
+    inverse
+
+"""
+    AbstractTransform
+
+## Interface
+
+* `Base.ndims(<:AbstractTransform)`: N dims of modes
+* `transform(<:AbstractTransform, 𝐱::AbstractArray)`: Apply the transform to 𝐱
+* `truncate_modes(<:AbstractTransform, 𝐱_transformed::AbstractArray)`: Truncate modes that contribute to the noise
+* `inverse(<:AbstractTransform, 𝐱_transformed::AbstractArray)`: Apply the inverse transform to 𝐱_transformed
+"""
+abstract type AbstractTransform end
+
+include("fourier_transform.jl")