Physics-Informed Neural Networks in JAX (with Equinox & Optax)

PINNs are a modern approach to solving (partial) differential equations (=PDEs) using neural networks based on minimizing a residuum formed with automatic differentiation. This video is a simple example of the 1D Poisson problem. Here is the code: 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 simulation.science/ for more on their pursuit of 'Nobel-Turing' technologies (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): 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: 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: amzn.to/3NU7OAs
- ⌨️ Logitech TKL Mechanical Keyboard: amzn.to/3JhEtwp
- 🎨 Gaomon Drawing Tablet (similar to a WACOM Tablet, but cheaper, works flawlessly under Linux): amzn.to/37katmf
- 🔌 Laptop Charger: amzn.to/3ja0imP
- 💻 My Laptop (generally I like the Dell XPS series): amzn.to/38xrABL
- 📱 My Phone: Fairphone 4 (I love the sustainability and repairability aspect of it): amzn.to/3Jr4ZmV
If I had to purchase these items again, I would probably change the following:
- 🎙️ Rode NT: 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): frame.work
As an Amazon Associate I earn from qualifying purchases.
-------
Timestamps:
00:00 Intro
01:02 What are PINNs?
01:23 1D Poisson Problem with homogeneous Dirichlet BCs
02:11 Training PINNs by residuum losses
04:28 How autodiff comes into play
06:33 Finite Differences as a reference
07:08 Considered forcing function
07:23 Imports
07:57 Constants/Hyperparameters
09:20 Defining and initializing the MLP architecture
11:38 Querying initial PINN state at some points
13:52 Computing reference solution by Finite Differences
18:07 Plot true solution and initial PINN guess
20:41 Defining PDE residuum using automatic differentiation
24:10 Total loss function
28:16 Training loop (including the third autodiff pass)
32:06 Plot Final PINN solution and discussion
34:30 Advantages of having a trained PINN
35:18 Summary
37:23 Potential improvements
37:52 Outro

[@basicmachines]

Not only is this a great introduction to PINNs but it's a nice intro to Jax as well. Thanks!

[@MachineLearningSimulation]

Thank you very much :). Welcome to the world of JAX ;).

[@MathPhysicsFunwithGus]

Amazing expectation thank you so much!

[@MachineLearningSimulation]

You're welcome 🤗
Thanks for the kind words.

[@khawar0o7]

Great video! I wish I could learn from you in person.

[@MachineLearningSimulation]

Thanks for this kind comment and the generous donation 😊

[@paulpetersik8123]

Really great video!

[@MachineLearningSimulation]

Thanks a lot 😊

Nice lesson. Thanks!

[@MachineLearningSimulation]

Glad you liked it! :) thanks for the kind comment.

[@auk6938]

Nice vedio !! i believe pinns can make a great difference on PDE problem

[@MachineLearningSimulation]

Thanks for the kind feedback 😊
PINNs are a nice approach, but my experience is that for more sophisticated problems (than this 1D Poisson PDE), they require additional tricks to train properly. Check for instance this paper: t.co/b0mSyWQQ7d or this one arxiv.org/abs/2109.01050
A very promising (albeit less interesting for classical engineering problems) Problem Domain is high-dimensional PDEs. PINNs can be proven to theoretically overcome the curse of dimensionality: proceedings.neurips.cc/paper_files/paper/2022/hash/46f0114c06524debc60ef2a72769f7a9-Abstract-Conference.html
However, the problem with robustly Training them still remains.
Would be interesting in hearing your opinion 😊 My experience with PINNs is rather limited. My PhD research focuses on "image-to-image"- neural nets in combination with differentiable physics.

[@ivankissiov]

Thanks!

[@MachineLearningSimulation]

Thanks again for the donation, Ivan 😊

[@AndreasEnnemoser]

So we’ll explained!

[@MachineLearningSimulation]

Thanks a lot 😊

[@ahmedshakiraliali397]

Great. Can it solve heterogeneous problems? Like 1D diffusion equation with a heterogeneous domain?

[@MachineLearningSimulation]

Thanks :)
By heterogeneous domain, do you mean a diffusion coefficient that varies over the domain? I think the framework of PINNs is very flexible. Just be aware that for more complicated PDE problems, adaptations of the learning process might be necessary to learn meaningful networks.

[@ahmedshakiraliali397]

Yes, for example, each half of the domain with different values of diffusion coefficient.

[@shadmansakiefhridoy794]

Will you make any tutorial video on inverse problem (estimating any unknown parameter) using PINN . For example estimating thermal conductivity (k) from a boundary value problem of transient heat conduction. I actually didn’t find that much good explanation on this topic.

[@MachineLearningSimulation]

Hi,
Thanks for the suggestion. :) I will definitely have more videos on PINNs in the future. I don't have a fixed recording schedule, it's mostly related to what I am the most interested in at the moment. But stay tuned for the next year!

[@adityak449]

Hello Felix, great video again!
Where can I get your email? I have something to discuss

[@MachineLearningSimulation]

Thanks 🙏
You can send me a connection request with your message on LinkedIn and add the note that you wrote a comment.