Neural operators is a new deep-learning framework which allows learning direct mapping between infinite function spaces on bounded domains. The trained model can predict output of the PDE systems on arbitrary time steps with arbitrary spatial resolutions, and it converges to the continuous solution operator when the resolution of the discretisation decreases.

Finding the numerical result in the traditional way, FEM/FDTD etc, especially with high resolutions, can be very computationally expensive. Moreover, every time the parameters, the initial condition, or the boundary condition of the PDE systems changes, even though the governing PDEs are still the same, one needs to run the simulation again and expect the same compute cost for each simulation.

Researchers have been working on data-driven methods that could potentially solve the PDE systems faster. There are efforts on learning projections between finite dimensional Euclidean spaces or finite sets. It generally discretises the input and learning the vector-to-vector mapping to the solution on the mesh, therefore, the trained neural network can only inference well on the mesh it was trained on. 

Recently, Li et al [https://zongyi-li.github.io/neural-operator/] tested several computational approaches of the neural operators and shows that they can learn the solution operator to a 1D Poisson Equation, a second order elliptic equation, 1D viscous Burgers' equation, and a 2D Navier-Stokes equation for a viscous, incompressible fluid.

At NCI, we have installed one of the implementation using Julia, NeuralOperators.jl and replicate the 1D viscous Burgers' equation example following the code provided by the author of the paper [https://github.com/neuraloperator/graph-pde] and the developer of the Julia package [https://github.com/SciML/NeuralOperators.jl].

Use the Shared Environment

$ module use /g/data/dk92/apps/Modules/modulefiles/
$ module load NCI-ai-ml/23.03
$ julia
 
julia> sdir="/g/data/dk92/notebooks/examples-julia/NeuralOperators/"
julia> pushfirst!(DEPOT_PATH,sdir*"SharedEnv")
julia> pushfirst!(DEPOT_PATH,mkpath(joinpath(ENV["TMPDIR"],"env")))
julia> using NeuralOperators

Run the tested notebook

  1. Go to ARE site: are.nci.org.au
  2. Fill out the JupyterLab request form
    1. Walltime (hours): 2
    2. Queue: normal
    3. Compute Size: small
    4. Project:  <xy01>
    5. Storage: gdata/dk92+scratch/<xy01>
  3. Click Advanced options and fill in the following fields
    1. Module directories: /g/data/dk92/apps/Modules/modulefiles/
    2. Modules: NCI-ai-ml/23.03
    3. Jobfs size: 10GB  
  4. Launch the session to run the tested notebook

 

Note on the tested notebook

  • copy the tested notebook from "/g/data/dk92/notebooks/examples-julia/NeuralOperators/FNO.ipynb" to your own working directory
  • the notebook FNO.ipynb contains an example to train a Fourier neural operator with the simulation data generated .
  • if necessary, modify DEPOT_PATH in the notebook to use another environment. The default one points to "/g/data/dk92/notebooks/examples-julia/NeuralOperators/SharedEnv"


Resolution-invariant

We train a Fourier neural operator to learn the 1D Burge's equation with the resolution of Δsamples=8 and inference with input data of Δsamples=[32,16,4,1]. The following figure shows the inference L2 loss with respect to the ground truth is similar to the training. 


  • No labels