feat: introduce spectral losses module

[frazane]

Following #599 and after discussing with @theissenhelen I am opening this PR to start the work on a generic spectral losses module. The focus of this PR is to establish the design of the module and aims to define the right abstractions and tests. Only very basic spectral losses will be actually implemented here for testing purposes. The work will also build on the existing spectral losses implemented in spatial.py. Feedback is of course more than welcome!

Design proposal

A main concern when designing spectral losses is the separation of the logic of the transformation to spectral domain from the logic for computing the error metric itself. Therefore, we introduce two respective abstract classes, SpectralTransform and SpectralLoss for this goal. The former is then assigned as the transform attribute of the latter.

• SpectralTransform: This class requires all subclasses to implement the __call__() method. As an example, I already implemented FFT2D by inheriting from the SpectralTransform class. SHT is added as a placeholder. Note that this is temporarily implemented inside the spectral losses module, but might be moved at some point into anemoi-models. An important assumption of this class is that the input and output tensors will always be assumed to be of shape [batch, time, [ensemble], points, variable], so any reshaping required for the transform happens inside and is undone before returning the value. It's still unclear to me whether this generalizes well (see second point in open questions).

• SpectralLoss: All subclasses that directly inherit from this class will have to implement the forward() method, which will contain the entire logic for the loss computation, including scaling and reduction. For now, only one loss has been implemented (ported from the spatial.py module).

• FunctionalSpectralLoss: Following an existing pattern in the losses package, we also introduce a FunctionalSpectralLoss class that inherits from both FunctionalLoss and SpectralLoss and allows for substantial code deduplication. An example implementation is implemented as SpectralL2Loss.

Open questions

• What will be the general signature of the SHT? Shapes of input and output tensors? Arguments?
• What would be the correct way to exclude some scalers, e.g. scalers applied on the nodes which become meaningless in the spectral domain?
• What do we do where the logic of the transform and the error metric cannot be strictly decoupled? For instance, I imagine some losses might only be computed on 2D arrays in the spectral domain so with FFT2D.
• Shall we implement correctness tests?

[OpheliaMiralles]

I wouldn't say it starts the work on spectral loss since there are some already implemented in the anemoi-core losses...

[frazane]

I wouldn't say it starts the work on spectral loss since there are some already implemented in the anemoi-core losses...

Starts the work for a generic module for spectral losses, which has a broader goal of defining abstractions. But yes absolutely, it also draws on what's already implemented in the spatial.py, you're right to point that out. I adapted the PR message!

[OpheliaMiralles]

I wouldn't say it starts the work on spectral loss since there are some already implemented in the anemoi-core losses...

Starts the work for a generic module for spectral losses, which has a broader goal of defining abstractions. But yes absolutely, it also draws on what's already implemented in the spatial.py, you're right to point that out. I adapted the PR message!

I assigned myself as a reviewer so I'm willing to contribute and give feedback :)

[OpheliaMiralles]

I would also add the log version of the FFT2 norm