From fb532dfe524a9b32d7404d3b2d9b259a67bf909b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Oct 2025 23:08:12 +0000 Subject: [PATCH 1/5] Initial plan From 455e788488fe3af6ca9903509e1f6db7d2fb2f55 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Oct 2025 23:08:59 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/references.bib | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/references.bib b/docs/references.bib index 2b5c34cd..e9cee75d 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -1,2 +1,2 @@ % This file contains bibliography references for the Clay Foundation Model documentation -% Currently empty but required by the Jupyter Book configuration \ No newline at end of file +% Currently empty but required by the Jupyter Book configuration From 266891de849c193ebe858967097c6c4eb71ceb49 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Oct 2025 23:13:55 +0000 Subject: [PATCH 3/5] Fix tensor shape issue in wall-to-wall example Co-authored-by: brunosan <434029+brunosan@users.noreply.github.com> --- docs/tutorials/wall-to-wall.ipynb | 4 +- test_wall_to_wall_fix.py | 87 +++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 2 deletions(-) create mode 100644 test_wall_to_wall_fix.py diff --git a/docs/tutorials/wall-to-wall.ipynb b/docs/tutorials/wall-to-wall.ipynb index 0b79f64f..da527826 100644 --- a/docs/tutorials/wall-to-wall.ipynb +++ b/docs/tutorials/wall-to-wall.ipynb @@ -482,12 +482,12 @@ "datacube = {\n", " \"platform\": platform,\n", " \"time\": torch.tensor(\n", - " np.hstack((week_norm, hour_norm)),\n", + " np.column_stack((week_norm, hour_norm)),\n", " dtype=torch.float32,\n", " device=device,\n", " ),\n", " \"latlon\": torch.tensor(\n", - " np.hstack((lat_norm, lon_norm)), dtype=torch.float32, device=device\n", + " np.column_stack((lat_norm, lon_norm)), dtype=torch.float32, device=device\n", " ),\n", " \"pixels\": pixels.to(device),\n", " \"gsd\": torch.tensor(stack.gsd.values, device=device),\n", diff --git a/test_wall_to_wall_fix.py b/test_wall_to_wall_fix.py new file mode 100644 index 00000000..df343e05 --- /dev/null +++ b/test_wall_to_wall_fix.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +""" +Test script to verify the wall-to-wall tensor shape fix. +""" + +import math +import numpy as np +import torch + + +def normalize_timestamp(date): + """Normalize timestamp to week and hour components.""" + week = 1 * 2 * np.pi / 52 # week 1 as example + hour = 12 * 2 * np.pi / 24 # noon as example + return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour)) + + +def normalize_latlon(lat, lon): + """Normalize lat/lon coordinates.""" + lat = lat * np.pi / 180 + lon = lon * np.pi / 180 + return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon)) + + +def test_tensor_shapes(): + """Test that the tensor shapes are correct after the fix.""" + # Simulate multiple dates (like 12 in the notebook) + num_samples = 12 + lat, lon = 37.30939, -8.57207 # From the notebook + + times = [normalize_timestamp(None) for _ in range(num_samples)] + week_norm = [dat[0] for dat in times] + hour_norm = [dat[1] for dat in times] + + latlons = [normalize_latlon(lat, lon)] * len(times) + lat_norm = [dat[0] for dat in latlons] + lon_norm = [dat[1] for dat in latlons] + + # Test the old (broken) way + print("=== Old (broken) way ===") + try: + time_old = torch.tensor(np.hstack((week_norm, hour_norm)), dtype=torch.float32) + latlon_old = torch.tensor(np.hstack((lat_norm, lon_norm)), dtype=torch.float32) + print(f"time_old shape: {time_old.shape}") + print(f"latlon_old shape: {latlon_old.shape}") + + # Try to use it in the model's expected format + time_latlon_old = torch.hstack((time_old, latlon_old)) + print(f"time_latlon_old shape: {time_latlon_old.shape}") + print("This would fail with EinopsError when trying 'B D -> B L D'") + except Exception as e: + print(f"Error with old way: {e}") + + print("\n=== New (fixed) way ===") + try: + time_new = torch.tensor(np.column_stack((week_norm, hour_norm)), dtype=torch.float32) + latlon_new = torch.tensor(np.column_stack((lat_norm, lon_norm)), dtype=torch.float32) + print(f"time_new shape: {time_new.shape}") + print(f"latlon_new shape: {latlon_new.shape}") + + # Try to use it in the model's expected format + time_latlon_new = torch.hstack((time_new, latlon_new)) + print(f"time_latlon_new shape: {time_latlon_new.shape}") + print("This should work correctly with 'B D -> B L D'") + + # Test the einops operation that was failing + from einops import repeat + B, D = time_latlon_new.shape + L = 1024 # From the error message + result = repeat(time_latlon_new, "B D -> B L D", L=L) + print(f"After einops repeat: {result.shape}") + print("āœ… SUCCESS: Einops operation works correctly!") + + except Exception as e: + print(f"Error with new way: {e}") + return False + + return True + + +if __name__ == "__main__": + print("Testing wall-to-wall tensor shape fix...") + success = test_tensor_shapes() + if success: + print("\nšŸŽ‰ All tests passed! The fix should work.") + else: + print("\nāŒ Tests failed. The fix needs more work.") \ No newline at end of file From 704950302252ae2c66bd2d935172f0820ab7f074 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Oct 2025 23:14:07 +0000 Subject: [PATCH 4/5] Remove temporary test file Co-authored-by: brunosan <434029+brunosan@users.noreply.github.com> --- test_wall_to_wall_fix.py | 87 ---------------------------------------- 1 file changed, 87 deletions(-) delete mode 100644 test_wall_to_wall_fix.py diff --git a/test_wall_to_wall_fix.py b/test_wall_to_wall_fix.py deleted file mode 100644 index df343e05..00000000 --- a/test_wall_to_wall_fix.py +++ /dev/null @@ -1,87 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script to verify the wall-to-wall tensor shape fix. -""" - -import math -import numpy as np -import torch - - -def normalize_timestamp(date): - """Normalize timestamp to week and hour components.""" - week = 1 * 2 * np.pi / 52 # week 1 as example - hour = 12 * 2 * np.pi / 24 # noon as example - return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour)) - - -def normalize_latlon(lat, lon): - """Normalize lat/lon coordinates.""" - lat = lat * np.pi / 180 - lon = lon * np.pi / 180 - return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon)) - - -def test_tensor_shapes(): - """Test that the tensor shapes are correct after the fix.""" - # Simulate multiple dates (like 12 in the notebook) - num_samples = 12 - lat, lon = 37.30939, -8.57207 # From the notebook - - times = [normalize_timestamp(None) for _ in range(num_samples)] - week_norm = [dat[0] for dat in times] - hour_norm = [dat[1] for dat in times] - - latlons = [normalize_latlon(lat, lon)] * len(times) - lat_norm = [dat[0] for dat in latlons] - lon_norm = [dat[1] for dat in latlons] - - # Test the old (broken) way - print("=== Old (broken) way ===") - try: - time_old = torch.tensor(np.hstack((week_norm, hour_norm)), dtype=torch.float32) - latlon_old = torch.tensor(np.hstack((lat_norm, lon_norm)), dtype=torch.float32) - print(f"time_old shape: {time_old.shape}") - print(f"latlon_old shape: {latlon_old.shape}") - - # Try to use it in the model's expected format - time_latlon_old = torch.hstack((time_old, latlon_old)) - print(f"time_latlon_old shape: {time_latlon_old.shape}") - print("This would fail with EinopsError when trying 'B D -> B L D'") - except Exception as e: - print(f"Error with old way: {e}") - - print("\n=== New (fixed) way ===") - try: - time_new = torch.tensor(np.column_stack((week_norm, hour_norm)), dtype=torch.float32) - latlon_new = torch.tensor(np.column_stack((lat_norm, lon_norm)), dtype=torch.float32) - print(f"time_new shape: {time_new.shape}") - print(f"latlon_new shape: {latlon_new.shape}") - - # Try to use it in the model's expected format - time_latlon_new = torch.hstack((time_new, latlon_new)) - print(f"time_latlon_new shape: {time_latlon_new.shape}") - print("This should work correctly with 'B D -> B L D'") - - # Test the einops operation that was failing - from einops import repeat - B, D = time_latlon_new.shape - L = 1024 # From the error message - result = repeat(time_latlon_new, "B D -> B L D", L=L) - print(f"After einops repeat: {result.shape}") - print("āœ… SUCCESS: Einops operation works correctly!") - - except Exception as e: - print(f"Error with new way: {e}") - return False - - return True - - -if __name__ == "__main__": - print("Testing wall-to-wall tensor shape fix...") - success = test_tensor_shapes() - if success: - print("\nšŸŽ‰ All tests passed! The fix should work.") - else: - print("\nāŒ Tests failed. The fix needs more work.") \ No newline at end of file From 0325d3dc105d9af48eec96a1ba15966465de0406 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Oct 2025 23:18:46 +0000 Subject: [PATCH 5/5] Fix inconsistent tensor shape comments across codebase Co-authored-by: brunosan <434029+brunosan@users.noreply.github.com> --- claymodel/finetune/embedder/factory.py | 4 ++-- claymodel/finetune/regression/factory.py | 4 ++-- claymodel/finetune/segment/factory.py | 4 ++-- claymodel/model.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/claymodel/finetune/embedder/factory.py b/claymodel/finetune/embedder/factory.py index 529cbfbb..708b2141 100644 --- a/claymodel/finetune/embedder/factory.py +++ b/claymodel/finetune/embedder/factory.py @@ -107,8 +107,8 @@ def add_encodings(self, patches, time, latlon, gsd): def forward(self, datacube): cube, time, latlon, gsd, waves = ( datacube["pixels"], # [B C H W] - datacube["time"], # [B 2] - datacube["latlon"], # [B 2] + datacube["time"], # [B 4] # week_sin, week_cos, hour_sin, hour_cos + datacube["latlon"], # [B 4] # lat_sin, lat_cos, lon_sin, lon_cos datacube["gsd"], # 1 datacube["waves"], # [N] ) # [B C H W] diff --git a/claymodel/finetune/regression/factory.py b/claymodel/finetune/regression/factory.py index 023eac71..05bc928d 100644 --- a/claymodel/finetune/regression/factory.py +++ b/claymodel/finetune/regression/factory.py @@ -100,8 +100,8 @@ def forward(self, datacube): """ cube, time, latlon, gsd, waves = ( datacube["pixels"], # [B C H W] - datacube["time"], # [B 2] - datacube["latlon"], # [B 2] + datacube["time"], # [B 4] # week_sin, week_cos, hour_sin, hour_cos + datacube["latlon"], # [B 4] # lat_sin, lat_cos, lon_sin, lon_cos datacube["gsd"], # 1 datacube["waves"], # [N] ) diff --git a/claymodel/finetune/segment/factory.py b/claymodel/finetune/segment/factory.py index dcc07448..fee33a29 100644 --- a/claymodel/finetune/segment/factory.py +++ b/claymodel/finetune/segment/factory.py @@ -106,8 +106,8 @@ def forward(self, datacube): """ cube, time, latlon, gsd, waves = ( datacube["pixels"], # [B C H W] - datacube["time"], # [B 2] - datacube["latlon"], # [B 2] + datacube["time"], # [B 4] # week_sin, week_cos, hour_sin, hour_cos + datacube["latlon"], # [B 4] # lat_sin, lat_cos, lon_sin, lon_cos datacube["gsd"], # 1 datacube["waves"], # [N] ) diff --git a/claymodel/model.py b/claymodel/model.py index 300eb823..50c22b55 100644 --- a/claymodel/model.py +++ b/claymodel/model.py @@ -161,8 +161,8 @@ def mask_out(self, patches): def forward(self, datacube): cube, time, latlon, gsd, waves = ( datacube["pixels"], # [B C H W] - datacube["time"], # [B 2] - datacube["latlon"], # [B 2] + datacube["time"], # [B 4] # week_sin, week_cos, hour_sin, hour_cos + datacube["latlon"], # [B 4] # lat_sin, lat_cos, lon_sin, lon_cos datacube["gsd"], # 1 datacube["waves"], # [N] ) # [B C H W]