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

**URL:** <https://discourse.julialang.org/t/ann-neuraloperators-jl-learning-the-infinite-dimensional-operator-for-partial-differential-equations/68779>\
**Category:** Package Announcements\
**Tags:** package, flux, pde\
**Created:** [September 26, 2021, 2:33pm UTC](https://discourse.julialang.org/t/ann-neuraloperators-jl-learning-the-infinite-dimensional-operator-for-partial-differential-equations/68779 "2021-09-26T14:33:00Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![foldfelis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/foldfelis/32/29494_2.png) [@foldfelis](https://discourse.julialang.org/u/foldfelis)\
**Post date:** [September 26, 2021, 2:33pm UTC](https://discourse.julialang.org/t/ann-neuraloperators-jl-learning-the-infinite-dimensional-operator-for-partial-differential-equations/68779/1 "2021-09-26T14:33:00Z")

</div>

I’m pleased to announce that [NeuralOperator.jl](https://foldfelis.github.io/NeuralOperators.jl/dev/) 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](https://arxiv.org/abs/2010.08895) 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

```julia
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.

```julia
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.

```julia
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))

```

---

<div class="post-metadata">

**Author:** ![bad\_at\_math](https://avatars.discourse-cdn.com/v4/letter/b/51bf81/32.png) [@bad\_at\_math](https://discourse.julialang.org/u/bad_at_math)\
**Post date:** [September 28, 2021, 4:25pm UTC](https://discourse.julialang.org/t/ann-neuraloperators-jl-learning-the-infinite-dimensional-operator-for-partial-differential-equations/68779/2 "2021-09-28T16:25:01Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![foldfelis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/foldfelis/32/29494_2.png) [@foldfelis](https://discourse.julialang.org/u/foldfelis)\
**Post date:** [September 28, 2021, 7:06pm UTC](https://discourse.julialang.org/t/ann-neuraloperators-jl-learning-the-infinite-dimensional-operator-for-partial-differential-equations/68779/3 "2021-09-28T19:06:18Z")

</div>

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](https://arxiv.org/abs/2003.03485). 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.