From cba96d106b9dd83554708b91edd3418c967e48d2 Mon Sep 17 00:00:00 2001
From: Samata Bag <2024itb047.samata@students.iiests.ac.in>
Date: Tue, 31 Mar 2026 12:09:28 +0530
Subject: [PATCH] Fix feature coordinate handling in WeatherDataset

---
 neural_lam/weather_dataset.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index 5547fdd4..dfc00c78 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -607,10 +607,16 @@ def _is_listlike(obj):
 
         da_datastore_state = getattr(self, f"da_{category}")
         da_grid_index = da_datastore_state.grid_index
-        da_state_feature = da_datastore_state.state_feature
+        feature_key = f"{category}_feature"
+        if feature_key not in da_datastore_state.coords:
+            raise KeyError(
+                    f"Expected coordinate '{feature_key}' not found for category '{category}'. "
+                    f"Available coords: {list(da_datastore_state.coords.keys())}"
+            )
 
+        da_feature = da_datastore_state.coords[feature_key]
         coords = {
-            f"{category}_feature": da_state_feature,
+            feature_key: da_feature,
             "grid_index": da_grid_index,
         }
         if add_time_as_dim: