Fourier Neural Operators (FNO) in JAX

[@enrikosmutje6201]

Do I understand correctly, that you add the mesh coordinates as a feature to the input data (somewhere in the beginning you concatenate the mesh to a)? Is that really necessary? I image that this will just add a Fourier transform of equally spaced point coordinates to the latent space, which will be the same for all test data.

[@MachineLearningSimulation]

Yes, that's right :). I also think that it is not necessary but I wanted to follow original FNO implementation closely, see here: github.com/neuraloperator/neuraloperator/blob/af93f781d5e013f8ba5c52baa547f2ada304ffb0/fourier_1d.py#L99

[@beincheekym8]

I was asking myself the same question, you are basically adding a constant input to your network which will likely be ignored. for the sake of being close to original implementation I understand, but you can likely just drop this channel entirely and it would work just as well. EDIT: now that I think of it, maybe it helps with the zero-shot upsampling? but at the same time you always train with the same grid so the second channel was probably ignored (as it's constant).

[@beincheekym8]

thank you for the great video, very clean implementation as well, satisfying to watch. eagerly waiting for your video on the use of FNO for next-state prediction / recurrent NN.

[@MachineLearningSimulation]

Thanks a lot 😊 Yes, definitely, the autoregressive prediction of transient problems will be one of the future videos 👍

[@kanishkbhatia95]

Super cool as always. Some feedback to enhance clarity - when writing modules (SpectralConv1d, FNOBlock1d, FNO1d), overlaying the flowchart on the right hand side to show the block to which the code corresponds would be really helpful. I felt a bit lost in these parts.

[@MachineLearningSimulation]

That's a great idea 👍 I will try to include this in a future video.

[@lightfreak999]

Very cool video! The walkthrough write-up of this alternate program of 1D FNO is super useful for newcomers like myself :)

[@MachineLearningSimulation]

Great to hear! 😊 Thanks for the kind feedback ❤️

[@GUYPOLO1]

Thank you for the detailed video! I have a question regarding the zero-shot super-resolution. If we train for 1 second as shown in your video, is it possible to run the test for, let's say, 5 seconds of propagation? Additionally, is it possible to plot the propagation from t=0 to t=1, or can it only provide the result? Since the training data can include the propagation as well, perhaps this information can be utilized in the training, not just the starting point. Thanks for your help!

[@MachineLearningSimulation]

Hi, thanks a lot for the kind feedback :). Here, in this video, the FNO is just trained to predict 1 time unit into the future. The dataset also only consisted of these input-output pairs. I assume you are interested in the state at times between 0 and 1 and beyond 1. There are two ways to obtain spatiotemporal information: either the FNO is set up in a way to return not only the state at t=1, but the state at multiple time levels (t=0.1, 0.2, ... 2.9, 3.0), or we use an FNO with a fixed dt (either the 1.0 of this video or a lower value) to obtain a trajectory autoregressive (kzbin.info/www/bejne/hJ20YoFpjJiKo9U ). For this, a corresponding dataset (or a reference solver to actively create the reference we need) is required.

[@Machine_Learner]

Awesome stuff. I am wondering if you can do a similar video for the new neural spectral methods paper?

[@MachineLearningSimulation]

Thanks 🤗 That's a cool paper. I just skimmed over it. It's probably a good idea to start covering more recent papers. I'll put it on my todo list, still have some other content in the loop I want to do first, but will come back to it later. Thanks for the suggestion 👍

[@amirmahdijafary2734]

Where in the code is the burgers equation used? Is it possible to mix PINN models with FNO models?

[@MachineLearningSimulation]

Hi, thanks for the questions 😊 1) the burgers equation is "within the dataset". It contains the (discretized) state for the Burgers equation at two time levels. So let's say you had a simulator for the burger equations, and you would use it to starting at the first state integrate the dynamics for 1 time unit, you would arrive the second state. We want to learn this mapping. 2) Yes, there are definitely approaches to incorporating PDE constraints via autodiff into the neural operator learning process (see for instance the paper by Wang et al. on physics-informed DeepONets or the Li et al paper on physics-informed neural Operators). This can be helpful but does not have to be. In their original conception FNOs are trained purely data driven.

[@lineherz]

I was looking at the reference code that you mentioned in the jupyter notebook and found that they coded something weird that I can't understand for 2D. out_ft[:, :, : self.mode1, : self.mode2] = self.compl_mul2d(x_ft[:, :, : self.mode1, : self.mode2], self.weights1) out_ft[:, :, -self.mode1 :, : self.mode2] = self.compl_mul2d(x_ft[:, :, -self.mode1 :, : self.mode2], self.weights2) I don't understand why there are two weights (weights1, weights2) and why they take upper mode1 frequencies. Can you explain this? Thanks for your video.

[@gitfted_by_AI]

I have a question about R , only low frequencies are saved but in some case keeping the high frequencies can be relevant.

[@MachineLearningSimulation]

Hi, yes that's indeed a valid concern. In the paper by Li et al. it is argued that the nonlinearities recover the higher modes. However, it is up to experimental evidence how well they do this. For benchmarks against UNets check for instance the pde arena paper ("towards multiscale....") by Gupta and Brandstetter.

[@starshipx1282]

great effort

[@MachineLearningSimulation]

Thanks a lot 😊

[@SahbaZehisaadat]

great. thanks

[@MachineLearningSimulation]

You are welcome! 😊

[@soudaminipanda]

I am having a very hard time understanding your data preparation. It should have been taught more slowly.

[@MachineLearningSimulation]

Hi, thanks for letting me know. I'm sorry to hear that. 😐 What was too fast for you? The speed of me talking or did I spend too little on explaining the involved commands? Can you also pinpoint specific operations that were hard for you to comprehend? Then, I can do them in greater detail in a future video. 😊