16 v2 global embedding by meiertgrootes · Pull Request #49 · ESMValGroup/ClimaNet

meiertgrootes · 2026-05-28T08:58:52Z

This pull request replaces #48 which was corrupted by mal-configured local github setup

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.

…predict. Added use_device to train

SarahAlidoost

@meiertgrootes thanks, the implementation looks good 👍 I left some comments. Most of them are related to code style. If something isnot clear, let me know.

Also, please consider running ruff as it can fixes things automatically and helps saving time in reviewing:

pip install ruff
ruff check --fix your_script.py   # this fixes/shows errors
ruff format --check your_script.py --diff  # this shows formatting issues

SarahAlidoost · 2026-06-01T12:08:27Z

+
+
+


Suggested change

SarahAlidoost · 2026-06-01T12:09:27Z

+        #geo_pos_tensor = self.sh_geo_pos[i: i + ph, j: j + pw] # (H,W, sh_emb_dim) -> (pH, pW, sh_embed_dim)
+
+


Suggested change

#geo_pos_tensor = self.sh_geo_pos[i: i + ph, j: j + pw] # (H,W, sh_emb_dim) -> (pH, pW, sh_embed_dim)

SarahAlidoost · 2026-06-01T12:10:17Z

+        lon_patch = self.lon_coords[j : j + pw] # (W,) -> (pW,)
+
+        #get patch geo pos embedding
+        #geo_pos_embedding_tensor = compute_patch_geo_pos_embedding(geo_pos_tensor, lat_patch)


Suggested change

#geo_pos_embedding_tensor = compute_patch_geo_pos_embedding(geo_pos_tensor, lat_patch)

SarahAlidoost · 2026-06-01T12:10:29Z

+        geo_pos_embedding_tensor = self.patch_geo_embeddings[idx]
+
+        #get scale feature for patch
+        #scale_feature_tensor = compute_patch_scale_features(lat_patch, lon_patch) # -> (10,)


Suggested change

#scale_feature_tensor = compute_patch_scale_features(lat_patch, lon_patch) # -> (10,)

SarahAlidoost · 2026-06-01T13:09:05Z

+        #self.spatial_pe = SpatialPositionalEncoding2D(
+        #    embed_dim=embed_dim, max_H=max_H, max_W=max_W
+        #)


Suggested change

#self.spatial_pe = SpatialPositionalEncoding2D(

# embed_dim=embed_dim, max_H=max_H, max_W=max_W

#)

can you please remove "SpatialPositionalEncoding2D" from the script since we are not using it anymore.

SarahAlidoost · 2026-06-01T13:54:52Z

+
+    # east-west extent
+    dx = earth_radius * cos_lat_c * dlon
+    dx_pix = dx/max(lon_ext -1, 1)


Suggested change

dx_pix = dx/max(lon_ext -1, 1)

dx_pix = dx / max(lon_ext -1, eps)

SarahAlidoost · 2026-06-01T14:01:48Z

+
+
+
+    def _set_geo_pos_table(self, sh_pos_table: str):


this function doesnot return anything, but in __init__ method, it is called as self.geo_pos_table = self._set_geo_pos_table(sh_pos_table). As a result, self.geo_pos_table will be None.

SarahAlidoost · 2026-06-01T14:22:07Z

+        self.sh_order_L = sh_order_L
+
+
+
        # Check that the input data has the expected dimensions


Suggested change

self.sh_order_L = sh_order_L

# Check that the input data has the expected dimensions

self.sh_order_L = sh_order_L

# Check that the input data has the expected dimensions

SarahAlidoost · 2026-06-01T14:23:09Z

+            "land_mask_patch": land_tensor,  # (pH,pW) True=Land
            "daily_timef_patch": daily_timef_tensor, #(M, T=31, 2)
            "padded_days_mask": self.padded_days_tensor,  # (M, T=31) True=padded
+            #"sh_geo_pos_patch": geo_pos_tensor, # (pH, pW, sh_embed_dim)


Suggested change

#"sh_geo_pos_patch": geo_pos_tensor, # (pH, pW, sh_embed_dim)

SarahAlidoost · 2026-06-01T15:10:01Z

+        geo_emb = geo_emb[:, None, None, :]   # (B,1,1,E)
+
+        x = agg_latent + geo_emb # (B, M, Hp*Wp, E)


Suggested change

geo_emb = geo_emb[:, None, None, :] # (B,1,1,E)

x = agg_latent + geo_emb # (B, M, Hp*Wp, E)

# Broadcasting: same geo embedding for all M months and all Hp*Wp locations

# we use weighted mean patch embedding, see `geo_embedding_utils.py`

geo_emb = geo_emb[:, None, None, :] # (B,1,1,E)

x = agg_latent + geo_emb # (B, M, Hp*Wp, E)

SarahAlidoost

@meiertgrootes thanks, the implementation looks good 👍 I left some comments. Most of them are related to code style. If something isnot clear, let me know.

Also, please consider running ruff as it can fixes things automatically and helps saving time in reviewing:

pip install ruff
ruff check --fix your_script.py   # this fixes/shows errors
ruff format --check your_script.py --diff  # this shows formatting issues

rogerkuou

Hi @meiertgrootes , besides Sarah's comments, I only have two very minor comments from my side. So I will already approve this PR. Please go ahead and merge after adapting Sarah's comment.

rogerkuou · 2026-06-05T08:54:44Z

+                batch["land_mask_patch"].to(device, non_blocking=use_cuda),
+                batch["geo_pos_embedding_patch"].to(device, non_blocking=use_cuda),
+                batch["scale_feature_patch"].to(device, non_blocking=use_cuda),
+                batch["padded_days_mask"].to(device, non_blocking=use_cuda)   ,


Suggested change

batch["padded_days_mask"].to(device, non_blocking=use_cuda) ,

batch["padded_days_mask"].to(device, non_blocking=use_cuda),

rogerkuou · 2026-06-05T08:59:18Z

@@ -35,7 +35,7 @@
   "metadata": {},
   "outputs": [],
   "source": [
-    "data_folder = Path(\"./eso4clima\")\n",
+    "data_folder = Path(\"/Users/mwgrootes/Projects/REPOS/ESO4CLIMA/data/output\") #(\"./eso4clima\")\n",


For the generality of the note book, maybe let's keep the original relative path ./eso4clima?

meiertgrootes added 6 commits May 28, 2026 10:52

added test suite for geo embeddings

cd6f099

updated dependencies to include scipy

31b2d3c

added geoembedding utilities

ceacbd4

implemented geo position and scale mebddings in model

dd05e43

added support for updated model including geo embedding to train and …

431df27

…predict. Added use_device to train

updated example notebook. Included histogram of errors and raw residuals

0f92c94

meiertgrootes requested review from SarahAlidoost and rogerkuou May 28, 2026 08:59

SarahAlidoost requested changes Jun 1, 2026

View reviewed changes

rogerkuou approved these changes Jun 5, 2026

View reviewed changes

		#geo_pos_tensor = self.sh_geo_pos[i: i + ph, j: j + pw] # (H,W, sh_emb_dim) -> (pH, pW, sh_embed_dim)

	#self.spatial_pe = SpatialPositionalEncoding2D(
	# embed_dim=embed_dim, max_H=max_H, max_W=max_W
	#)

	dx_pix = dx/max(lon_ext -1, 1)
	dx_pix = dx / max(lon_ext -1, eps)

		self.sh_order_L = sh_order_L



		# Check that the input data has the expected dimensions

		geo_emb = geo_emb[:, None, None, :] # (B,1,1,E)

		x = agg_latent + geo_emb # (B, M, Hp*Wp, E)

	batch["padded_days_mask"].to(device, non_blocking=use_cuda) ,
	batch["padded_days_mask"].to(device, non_blocking=use_cuda),

Conversation

meiertgrootes commented May 28, 2026

Uh oh!

SarahAlidoost left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SarahAlidoost left a comment

Choose a reason for hiding this comment

Uh oh!

rogerkuou left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants