# NeuralOperators

Ground TruthInferenced

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.

## 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:

pkg> add NeuralOperators

## Usage

### Fourier Neural Operator

model = Chain(
# lift (d + 1)-dimensional vector field to n-dimensional vector field
# here, d == 1 and n == 64
Dense(2, 64),
# map each hidden representation to the next by integral kernel operator
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),
)

Or one can just call:

model = FourierNeuralOperator(
ch=(2, 64, 64, 64, 64, 64, 128, 1),
modes=(16, ),
σ=gelu
)

And then train as a Flux model.

loss(𝐱, 𝐲) = l₂loss(model(𝐱), 𝐲)
Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)

### 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

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.

loss(xtrain, ytrain, sensor) = Flux.Losses.mse(model(xtrain, sensor), ytrain)
evalcb() = @show(loss(xval, yval, grid))

learning_rate = 0.001
Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb=evalcb)