[ANN] NeuralOperators.jl - Learning the Infinite-dimensional operator for partial differential equations

Package Announcements
, ,
1

I’m pleased to announce that NeuralOperator.jl is now available. Neural operator is a novel deep learning architecture introduced by Zongyi-Li et al. It learns a operator, which is a mapping between infinite-dimensional function spaces. Instead of solving by finite element method, a PDE problem can be resolved by training a neural network to learn an operator mapping from infinite-dimensional space (u, t) to infinite-dimensional space f(u, t). Neural operator learns a continuous function between two continuous function spaces.

Fourier Neural Operator learns a operator based on Fourier transformation. It maps a time-domain continuous function to another frequency-domain continuous function and back. Modes are truncated by user-specified parameter. It could be applied to solve PDE problems with initial conditions.

Example

There is a simple example for solving Burgers’ equation. It is available for CUDA. Here how we construct the Fourier neural operator model

if has_cuda()
    @info "CUDA is on"
    device = gpu
    CUDA.allowscalar(false)
else
    device = cpu
end

model = FourierNeuralOperator(
    ch=(2, 64, 64, 64, 64, 64, 128, 1),
    modes=(16, ),
    Οƒ=gelu
) |> device

Then we define loss function and load data from simple interface.

loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- model(𝐱)) / size(𝐱)[end]

loader_train, loader_test = get_dataloader()

function validate()
    validation_losses = [loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test]
    @info "loss: $(sum(validation_losses)/length(loader_test))"
end

Finally, we train it directly with Flux.jl facilities.

data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
Flux.@epochs 500 @time(Flux.train!(loss, params(model), data, opt, cb=call_back))
27 Likes
2

Very exciting! Is this project still WIP, or are the current implementations reasonably stable? Have you benchmarked at all against the python implementations in terms of time and accuracy? The current tests seem focused on verifying validity of code.

3

Hi, although there are still some tasks remain to be done, the layer FourierOperator and the models FourierNeuralOperator MarkovNeuralOperator are ready to learn some new stuff \OwO/.

In this project, we decided to support the general Neural Operator. The framework design is work in slow progress due to my master thesis defence (at the end of January). Also, the benchmarks will be done in the near future.

3 Likes