Fix EinopsError in wall-to-wall example due to incorrect tensor shapes

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]

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]
         )

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]
         )

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]

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
+% Currently empty but required by the Jupyter Book configuration

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",