Closed MLuchmann closed 5 months ago
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.
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?
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.
Feature: Fourier Neural operator
Description
Implementation of the Fourier Neural Operator.
References
https://arxiv.org/abs/2010.08895
Which issue does this PR tackle?
How does it solve the problem?
torch.fft.rfftn
) and its inverse (torch.fft.irfftn
)torch.fft.rfftn
was chosen instead oftorch.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.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?
Checklist for Contributors
feature/title-slug
convention.Bugfix: Title
convention.Checklist for Reviewers: