From 11508458344dca0e761102e4ac193eec25fdb482 Mon Sep 17 00:00:00 2001 From: Baudouin Raoult Date: Tue, 18 Jun 2024 10:09:34 +0000 Subject: [PATCH] Better error message when thinning non lat/lon fields --- src/anemoi/datasets/data/masked.py | 8 ++++++-- src/anemoi/datasets/data/stores.py | 6 +++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/anemoi/datasets/data/masked.py b/src/anemoi/datasets/data/masked.py index e44344bb..002bf929 100644 --- 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 index 7c59f758..8e8e6ec8 100644 --- 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):