Better error message when thinning non lat/lon fields

ecmwf · Jun 18, 2024 · 1150845 · 1150845
1 parent 2299d36
commit 1150845
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/src/anemoi/datasets/data/masked.py b/src/anemoi/datasets/data/masked.py
@@ -70,8 +70,12 @@ def __init__(self, forward, thinning, method):
         self.thinning = thinning
         self.method = method
 
-        latitudes = forward.latitudes.reshape(forward.field_shape)
-        longitudes = forward.longitudes.reshape(forward.field_shape)
+        shape = forward.field_shape
+        if len(shape) != 2:
+            raise ValueError("Thinning only works latitude/longitude fields")
+
+        latitudes = forward.latitudes.reshape(shape)
+        longitudes = forward.longitudes.reshape(shape)
         latitudes = latitudes[::thinning, ::thinning].flatten()
         longitudes = longitudes[::thinning, ::thinning].flatten()
 

diff --git a/src/anemoi/datasets/data/stores.py b/src/anemoi/datasets/data/stores.py
@@ -235,7 +235,11 @@ def resolution(self):
 
     @property
     def field_shape(self):
-        return tuple(self.z.attrs["field_shape"])
+        try:
+            return tuple(self.z.attrs["field_shape"])
+        except KeyError:
+            LOG.warning("No 'field_shape' in %r, assuming 1D fields", self)
+            return (self.shape[-1],)
 
     @property
     def frequency(self):