Feature: Fourier Layer for Fourier Neural Operator

[MLuchmann]

Feature: Fourier Neural operator

Description

Implementation of the Fourier Neural Operator.

References

https://arxiv.org/abs/2010.08895

Which issue does this PR tackle?

• FNOs were missing.

How does it solve the problem?

• Implements FourierLayer class which represents a kernel integral operation in Fourier Space.
• The Layer performs a N-dimensional fast Fourier transformation (torch.fft.rfftn) and its inverse (torch.fft.irfftn)
• torch.fft.rfftn was chosen instead of torch.fft.fftn because it is more space efficient. When transforming a real-valued function we do not need to compute the functional values for the negative frequencies because: $f(x)\in \mathcal{R} \qquad \Rightarrow \qquad \mathcal{F} = f(\omega) = -f^{*}(\omega)$. Therefore, we only need to evaluate and store half of the functional evaluations in Fourier space.
  • The same argument goes for the kernel. Because we assume that the kernel is real-valued, we can parameterize the fourier transformed kernel as fourier_kernel.shape = (num-sensors // 2 + 1, num-sensors // 2 + 1, v-dim, u-dim). +1 because of the 0-mode.

How are the changes tested?

• test_fno.py
• fourierneuraloperator notebook

Checklist for Contributors

• [x] Scope: This PR tackles exactly one problem.
• [ ] Conventions: The branch follows the feature/title-slug convention.
• [x] Conventions: The PR title follows the Bugfix: Title convention.
• [x] Coding style: The code passes all pre-commit hooks.
• [x] Documentation: All changes are well-documented.
• [x] Tests: New features are tested and all tests pass successfully.
• [x] Changelog: Updated CHANGELOG.md for new features or breaking changes.
• [x] Review: A suitable reviewer has been assigned.

Checklist for Reviewers:

• [x] The PR solves the issue it claims to solve and only this one.
• [x] Changes are tested sufficiently and all tests pass.
• [x] Documentation is complete and well-written.
• [x] Changelog has been updated, if necessary.

[MLuchmann]

Hey @samuelburbulla,

this is a draft for the Fourier layer in N-dimensions. I kept the 1d implementation in there for now because its a bit easier to understand.

I added a few tests (test_fourierlayer.py) (for a 1d example and a 2d example). I implemented two tests. One simply checks whether the FFT works as expected ( = convolutional operation in real-space). The other just trains an operator and checks the MSE.

I also added a notebook to play around with the operators. I experimented with different input functions. However, I tried to avoid the sin / cos examples because I was worried that the frequency spectrum could cause issues. The Fourier transformed sin-function should give a delta distribution. Because we use the FFT and not the continuous FT, this only leads to a peak at some freququency. Still, for the notebook I wanted to have a more regular frequency spectrum.

However, I am stuck right now. While the 1d example works very well, the 2D example is quite bad. While it does learn something, I don't manage to improve the performance. I tried to stack several layer on top of each other (see notebook) and also played around with the hyperparameters. It could be an issue with the implementation but I am not sure. What do you think? I implemented the tests for that reason but they only taught me that their are boundary effects when performing a convolution...

Edit: This is solved now. The 1d and 2d case work now as expected and give reasonable performance.

[MLuchmann]

Hi @samuelburbulla , this is the current implementation the FourierLayer. I think it could be merged in this state. Feel free to move around or change things. What do you think?

[samuelburbulla]

Thanks, @MLuchmann, looks good!

I cleaned up a little bit and will merge the Fourier layer first, there are quite some conflicts with the main. Afterwards, I will add a true FNO and convert your notebook to an example.