Pressure Level Embeddings

[mc4117]

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

The current code does not take into account embeddings in the z-dimension (either through height or pressure levels). We saw with Sara's work that pressure level weightings can have a large impact on results, especially in the upper atmosphere.

This has been done in Aurora (https://arxiv.org/abs/2405.13063) which shows significant improvements to graphcast in the upper atmosphere, and also in https://arxiv.org/abs/2405.06004v1 and https://arxiv.org/pdf/2405.14527.

Describe the solution you'd like

Include pressure level embeddings in our encoder implementation and evaluate the impact of this.

Describe alternatives you've considered

No response

Additional context

No response

Organisation

ECMWF