02_transform.md
One important aspect of the 2D RBC problem is that in the vertical direction, the mesh is not uniform nor periodic because of the walls in the top and bottom. Although, this does not prevent to use spectral discretization methods in space : a Chebyshev grid distribution is used in the vertical direction, such that we get a finer representation of the field near the wall, and can still represent it in spectral space through a Discrete Cosinus Transform, that is based on FFT.
Hence we proposed the following Neural Operator design, that replaces the Fourier Transform in the vertical direction by a DCT, as this is the transformation used by the spectral discretization during simulation. Since the DCT produces a real output, we can do it first, and then apply a RFFT transform in the Fourier Layer. The CFNO layer defined that way represent then a proper convolution layer considering our grid that doesn't have a uniform grid point distribution in the vertical direction.
Below is a graph summarizing the new design :
We call this a Chebyshev Fourier Neural Operator (CFNO).
Considering the same model structure, we train a CFNO model and another classical FNO update-predicting model. Comparing the loss during training for the CFNO (left) and the FNO (right) :
We observe that the loss goes down to a similar value, although training for the CFNO seems to require way more epochs to converge, compared to the FNO. Furthermore, the combination of the DCT and RFFT being more expensive compared to the optimized 2D FFT of pytorch, training time per epoch is also bigger for the CFNO.
Then, looking at prediction errors :
| CFNO | FNO | |
|---|---|---|
| v_x | 5.2e-05 | 4.0e-05 |
| v_z | 6.8e-05 | 6.4e-05 |
| b | 4.6e-05 | 4.1e-05 |
| p | 1.5e-05 | 1.5e-05 |
| avg | 4.5e-05 | 4.0e-05 |
we can see that both model produce similar accuracy, the FNO being slightly more accurate since the loss was able to converge faster and the CFNO could have benefit for more training (not done here as the 60000 epochs required already a lot of computation time).
It is interesting also to look at the absolute error for buoyancy for the CFNO (left) and the FNO (right) :
Note that the errors are very similar and quite small. This suggests that even if we are not doing a proper convolution when using the classical FNO (due to the Fourier transform applied on the non uniform z direction), it still manages to "correct" the grid non uniformity and produces an acceptable solution after training.
This motivates at the end to keep the classical FNO design, as it produces a more efficient ML model in terms of training and inference computation time.