streamline da_grid_reference variable naming

mllam · Dec 20, 2024 · 857f748 · 857f748
1 parent 8cc608d
commit 857f748
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py
@@ -386,11 +386,10 @@ def grid_shape_state(self):
                 "no state data found in datastore"
                 "returning grid shape from forcing data"
             )
-            ds_forcing = self.unstack_grid_coords(self._ds["forcing"])
-            da_x, da_y = ds_forcing.x, ds_forcing.y
+            da_grid_reference = self.unstack_grid_coords(self._ds["forcing"])
         else:
-            ds_state = self.unstack_grid_coords(self._ds["state"])
-            da_x, da_y = ds_state.x, ds_state.y
+            da_grid_reference = self.unstack_grid_coords(self._ds["state"])
+        da_x, da_y = da_grid_reference.x, da_grid_reference.y
         assert da_x.ndim == da_y.ndim == 1
         return CartesianGridShape(x=da_x.size, y=da_y.size)