feat: custom discrete cosine transform

[dietervdb-meteo]

Description

EXPERIMENTAL

The current DCT implementation, which relies on torch-dct, has an issue which might be quite specific to RMI:

1. we work with CERRA, which has a 1069x1069 grid,
2. we work on LUMI which has AMD hardware.

These two combine in a destructive way, since 1069 is a large prime number which does not go well with the HipFFT backend on Rocm.

RuntimeError: cuFFT error: HIPFFT_EXEC_FAILED

For this reason we implement a custom made 2d DCT into anemoi. It currently avoids FFT completely and works via explicit matrix multiplication, so that it will also work for grids with relatively large prime number sizes. Note that the default torch_dct.dct_2d does not seem a very optimized implementation, so it could be the custom implementation here improves on it in general.

This is currentlty a draft PR, performance and sanity still needs to be checked.

Since we want to test this in our particular setup (CERRA based LAM) this is PR towards #1143 rather than main.

Example use in config:

training_loss:
  datasets:
    data:
      _target_: anemoi.training.losses.combined.CombinedLoss
      scalers: ['*']
      loss_weights: [1.0, 0.1]
      losses:
        - _target_: anemoi.training.losses.CRPS
          scalers: ['pressure_level', 'general_variable', 'nan_mask_weights', 'node_weights']
          alpha: 0.95
        - _target_: anemoi.training.losses.spectral.SpectralCRPSLoss
          scalers: ['pressure_level', 'general_variable', 'nan_mask_weights']
          alpha: 0.95
          transform: 'dct2d'
          apply_filter: False
          x_dim: 1069
          y_dim: 1069
          dct_backend: custom_dct   #new
          dct_backend_kwargs:       #new
            norm: unnormalized      #new

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.