16 global spatial embedding

[meiertgrootes]

This pull request adds fully sphere aware geo position and scale encoding for patches.
Geo position is encoded using real-valued spherical harmonics as basis.
To create embeddings real-valued spherical harmonics at the lat/lon positions of pixels of the input data (i.e. native resolution) are calculated up to a user-defined order L. This results in an (L+1)^2 dimensional embedding vector. L ~ 10 should be fine.
Subsequently a sphere aware area-weighted PCA is performed on the native resolution SH embdding grid, with the requirment that the target diemension for the PCA (the sh_embed_dim) be smaller than (L+1)^2. The ranked PCA components up to sh_embed_dim are retained to be used as basis functions. The SH embedding vectors are then reprojected to this basis and scaled to zero mean and unit variance with tanh based soft-clipping at ~3 sigma to suppress pathological outliers in high order harmonics.
For each patch, a patch position embedding is then constructed as the area weighted mean of the token/pixel embeddings in the patch.

In addition for each patch a scale embedding is constructed consisting of: the patch physical extent in lat and lon directions, the patch area, the anisotropy of the patch extent, the pixel scale in phyical units [m] in lat/lon, anisotropy and isotropized linear scale, as well as finally effective harmonic order cutoff in lat/lon.

Both embeddings are precalculated once for all patches

The 10-dimensional scale embedding is concatenated with the geo position embedding for each patch. During training an trainable MLP is used to project the concatented embeddings into the desired embedding diemension for additive incorporation.

Note, it makes sense to choose the hidden diemension for the projection larger than bith the dimension of the concatenated embeddings, as well as the target embedding dimension.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ meiertgrootes
❌ SarahAlidoost
_{You have signed the CLA already but the status is still pending? Let us recheck it.}