Fourier Neural Operators (FNO) in JAX

Neural Operators are mappings between (discretized) function spaces, like from the IC of a PDE to its solution at a later point in time. FNOs do so by employing a spectral convolution that allows for multiscale properties. Let's code a simple example in JAX: https://github.com/Ceyron/machine-lea...

-------

👉 This educational series is supported by the world-leaders in integrating machine learning and artificial intelligence with simulation and scientific computing, Pasteur Labs and Institute for Simulation Intelligence. Check out https://simulation.science/ for more on their pursuit of 'Nobel-Turing' technologies (https://arxiv.org/abs/2112.03235 ), and for partnership or career opportunities.

-------

📝 : Check out the GitHub Repository of the channel, where I upload all the handwritten notes and source-code files (contributions are very welcome): https://github.com/Ceyron/machine-lea...

📢 : Follow me on LinkedIn or Twitter for updates on the channel and other cool Machine Learning & Simulation stuff:   / felix-koehler   and   / felix_m_koehler  

💸 : If you want to support my work on the channel, you can become a Patreon here:   / mlsim  

🪙: Or you can make a one-time donation via PayPal: https://www.paypal.com/paypalme/Felix...

-------

⚙️ My Gear:

(Below are affiliate links to Amazon. If you decide to purchase the product or something else on Amazon through this link, I earn a small commission.)

- 🎙️ Microphone: Blue Yeti: https://amzn.to/3NU7OAs
- ⌨️ Logitech TKL Mechanical Keyboard: https://amzn.to/3JhEtwp
- 🎨 Gaomon Drawing Tablet (similar to a WACOM Tablet, but cheaper, works flawlessly under Linux): https://amzn.to/37katmf
- 🔌 Laptop Charger: https://amzn.to/3ja0imP
- 💻 My Laptop (generally I like the Dell XPS series): https://amzn.to/38xrABL
- 📱 My Phone: Fairphone 4 (I love the sustainability and repairability aspect of it): https://amzn.to/3Jr4ZmV

If I had to purchase these items again, I would probably change the following:

- 🎙️ Rode NT: https://amzn.to/3NUIGtw
- 💻 Framework Laptop (I do not get a commission here, but I love the vision of Framework. It will definitely be my next Ultrabook): https://frame.work

As an Amazon Associate I earn from qualifying purchases.

-------

Timestamps:
00:00 Intro
01:09 What are Neural Operators?
03:11 About FNOs and their multiscale property
05:05 About Spectral Convolutions
09:17 A "Fourier Layer"
10:18 Stacking Layers with Lifting & Projection
11:01 Our Example: Solving the 1d Burgers equation
12:04 Minor technicalities
13:07 Installing and Importing packages
14:02 Obtaining the dataset and reading it in
15:44 Plot and Discussion of the dataset
17:51 Prepare training & test data
22:23 Implementing Spectral Convolution
34:25 Implementing a Fourier Layer/Block
37:48 Implementing the full FNO
43:14 A simple dataloader in JAX
44:18 Loss Function & Training Loop
52:34 Visualize loss history
53:31 Test prediction with trained FNO
57:13 Zero-Shot superresolution
59:59 Compute error as reported in FNO paper
01:03:45 Summary
01:05:46 Outro