feat: Consolidate and expand spectral losses #788
Conversation
…, torch-harmonics) (HEALPix transform by M. Price et al., S2FFT).
for more information, see https://pre-commit.ci
… feat/spectral-losses-groundwork
for more information, see https://pre-commit.ci
…moi-core into feat/harmonize-spectral-losses
… dim (as well as predictions with ensemble dim) Signed-off-by: evenmn <[email protected]>
…nt that is already passed from KernelCRPS and AlmostFairKernalCRPS Signed-off-by: evenmn <[email protected]>
…configuration Signed-off-by: evenmn <[email protected]>
…moi-core into feat/harmonize-spectral-losses
| def __init__( | ||
| self, | ||
| nlat: int, | ||
| lmax: int | None = None, |
There was a problem hiding this comment.
We need to document these options better in a docstring here. Also, a section in the documentation for the spectral losses would be necessary.
There was a problem hiding this comment.
Should we harmonize it with the FFT2D class? The apply_filter argument serves the same (or similar) purpose (compared to the lmax/mmax arguments here) - "apply low-pass filter to ignore frequencies beyond the Nyquist limit".
Please notice that the folding argument isn't even implemented, it will raise NotImplementedError (supposedly, it should account for the frequency extension and folding, but I don't fully understand it).
There was a problem hiding this comment.
We have decided to finalise this harmonisation when we have decided which method should go into main.
| class EcTransOctahedralSHT(SpectralTransform): | ||
| def __init__( | ||
| self, | ||
| truncation: int, |
There was a problem hiding this comment.
We need to document these options better in a docstring here. Also, a section in the documentation for the spectral losses would be necessary.
There was a problem hiding this comment.
@sahahner, @PortillaS-Predictia and I had a chat earlier, and we wondered if we should just delete all HEALPix-related stuff from this file. We don't think there's a use case for it yet, and we have our hands full with the other spherical harmonic transform implementations. On top of that, it's currently a "stub" with no wrapper class in spectral_transforms.py, and also only the inverse is implemented so it can't even be used for spectral loss calculations. Shall we get rid of it and potentially resurrect it later when we have a use case?
There was a problem hiding this comment.
Yes I think that sounds sensible!
There was a problem hiding this comment.
Should we also remove the "Cartesian" (Regular) transforms and just add torch-harmonics as a dependency? These classes are mostly the original code from said library. I only ever added type hints and some minor changes (I also solved a minor bug with some dtypes, that's why there's a type change in the forward method of the inverse transform). However, in the meantime, they also added some other features, like properly guessing the truncation frequency for these Grids (@samhatfield).
There was a problem hiding this comment.
I think we could possibly remove the cartesian/regular transform, as I am not aware that anyone is actively using these grids. The code would still be available on the closed PR #729. We could also point to torch-harmonics in the documentation for implementations for other grids and add the dependency when someone starts working with these grids.
There was a problem hiding this comment.
So in fact this PR would only implement octahedral and reduced SHTs?
There was a problem hiding this comment.
And spectral CRPS used in bris models, it’s already a decent expansion
| return legpoly(mmax, lmax, np.cos(t), norm=norm, inverse=inverse, csphase=csphase) | ||
|
|
||
|
|
||
| class CartesianRealSHT(Module): |
There was a problem hiding this comment.
What's the meaning of Real in this context, and below for OctahedralRealSHT? If it's referring to the fact that the inputs must be real valued, I would say this is not necessary to mention as that will always be the case.
There was a problem hiding this comment.
Also, following discussion with @PortillaS-Predictia, I would recommend to rename this and the inverse, along with the wrapper CartesianSHT to RegularSHT to denote the fact that it takes a regular grid (i.e. a latitude-longitude grid with no reduction towards the poles) as input.
There was a problem hiding this comment.
Yes, Real only refers to the reality of the input fields, we can remove it. Renaming the class is also fine by me, but I left another comment related to this (here). Let me know what you think.
|
Two more ideas that would simplify this PR:
|
|
I am very happy to see that we are already at a place where we can train global models with spectral losses. I have played around with the loss and uploaded some results to mlflow. Either way, I feel there are still a few open points:
|
Description
On top of #678, this PR introduces SpectralCRPS, adapt Octahedral losses by @theissenhelen and @PortillaS-Predictia (PRs #709 and #729) to match skeleton introduced in #678 by @frazane. The PR also adds tests for all of the losses. We can decide later which Octahedral loss we want or leave both options.
What problem does this change solve?
New feature/Consolidation
What issue or task does this change relate to?
#678 #709 #729
Additional notes
To be merged after #678 and if/when we decide to harmonise something about #709 and #729.
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.