Spectral loss

[theissenhelen]

Is your feature request related to a problem? Please describe.

Anemoi-training already provides a spectral loss in the form of a 2d FFT, mostly relevant for limited area models. A spectral loss for global models requires a differentiable implementation of accurate spherical harmonics in torch.

Describe the solution you'd like

We need the following initial steps for this work:

1. An accurate differentiable harmonic transform in torch that

• deals with loading pre-computed and/or computing the Legendre polynomials

• sufficiently efficient to allow computations on standard grid resolutions both with respect to runtime and memory

Our baseline here is the spherical harmonic implementation of ectrans.

2. A loss module in anemoi-training that handles the spectral transform and computes a loss by e.g. using the power spectrum.

Describe alternatives you've considered

Torch-harmonics:

Torch-harmonics does not currently support reduced gaussian grids.

Additional context

No response

Organisation

No response

[frazane]

In the case of ICON regional meshes, we cannot use the standard 2d FFT (unless we perform a regridding step) nor an implementation based on spherical harmonics. We might need to resort to some form of non-uniform Fourier transform.

[PortillaS-Predictia]

Hi @theissenhelen, we are very much interested in this. We've implemented some stuff here and there, I hope you find it useful. Your feedback is welcomed (especially, regarding the "loss module in anemoi-training that handles ...").

Following C. Subich, we've built a basic spectral loss module. The "amse" loss from the article and a "standard" power loss are included.

We also expanded torch-harmonics to support Octahedral and HEALPix grids.

[frazane]

Hi @theissenhelen, I am working on the spectral losses that are currently implemented in anemoi. I am still applying them on a regional domain, so using the 2D FFT. However, I am considering a refactor of the code, and I would like to do it in away that paves the way for the future integration of spherical harmonics. Is this something you are already working on or can I move forward with this? Also, would be nice to have a chat in any case!

Related question: I wonder if it would make sense to implement the spectral computation functionalities in anemoi-models instead of anemoi-training?

[theissenhelen]

Hi @PortillaS-Predictia,

I had a look at your code, very neat! It is very similar to what I am doing except that I am using ECMWF’s ectrans to compute the polynomials and adhere to ectrans conventions for the SHT. Yes, that paper is exactly what we have been looking at as well.

[theissenhelen]

Hi @frazane,

very good point. I think the spectral functionalities should live in anemoi-models rather than anemoi-training except for the actual loss modules.

[theissenhelen]

I opened a PR for the spherical harmonic transform based on ectrans (#709).

[PortillaS-Predictia]

Thanks @theissenhelen, I'm sure the Fortran code wasn't an easy ride. Could you attach a "minimal working example" of the class? I can't figure it out, I tried the following for an O320:

sht = EcTransOctahedralSHT(319, torch.float32)
x = torch.ones(1, 1, 421120, 1)

spectrum = sht(x)
excited = torch.where(spectrum.abs() > 1e-3)

print(excited)

Only the 0-th harmonic should be excited (because x is constant) ... however, there are a few more excited harmonics. Admittedly, there may be something wrong with my installation.