From 69256c0600c86b229ef535eb1ca04f94a90760aa Mon Sep 17 00:00:00 2001 From: Florian Pinault Date: Fri, 3 May 2024 13:01:36 +0000 Subject: [PATCH] adding read for statistics tendencies --- src/anemoi/datasets/create/loaders.py | 6 ++++-- src/anemoi/datasets/data/forewards.py | 5 +++++ src/anemoi/datasets/data/join.py | 8 ++++++++ src/anemoi/datasets/data/select.py | 5 +++++ src/anemoi/datasets/data/statistics.py | 5 +++++ src/anemoi/datasets/data/stores.py | 15 +++++++++++++++ src/anemoi/datasets/data/unchecked.py | 4 ++++ 7 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/anemoi/datasets/create/loaders.py b/src/anemoi/datasets/create/loaders.py index 063405f6..780a078e 100644 --- a/src/anemoi/datasets/create/loaders.py +++ b/src/anemoi/datasets/create/loaders.py @@ -550,7 +550,7 @@ def __init__(self, name="", **kwargs): super().__init__(**kwargs) self.name = name - storage_path = os.path.join(self.path + ".tmp_data", name) + storage_path = f"{self.path}.tmp_storage_{name}" self.tmp_storage = build_storage(directory=storage_path, create=True) def initialise(self): @@ -725,6 +725,8 @@ class TendenciesStatisticsDeltaNotMultipleOfFrequency(ValueError): class TendenciesStatisticsAddition(GenericAdditions): + DATASET_NAME_PATTERN = "statistics_tendencies_{delta}" + def __init__(self, path, delta=None, **kwargs): full_ds = open_dataset(path) self.variables = full_ds.variables @@ -739,7 +741,7 @@ def __init__(self, path, delta=None, **kwargs): ) idelta = delta // frequency - super().__init__(path=path, name=f"tendencies_statistics_{delta}h", **kwargs) + super().__init__(path=path, name=self.DATASET_NAME_PATTERN.format(delta=f"{delta}h"), **kwargs) z = zarr.open(self.path, mode="r") start = z.attrs["statistics_start_date"] diff --git a/src/anemoi/datasets/data/forewards.py b/src/anemoi/datasets/data/forewards.py index 078c68ae..197e2fcd 100644 --- a/src/anemoi/datasets/data/forewards.py +++ b/src/anemoi/datasets/data/forewards.py @@ -67,6 +67,11 @@ def variables(self): def statistics(self): return self.forward.statistics + def statistics_tendencies(self, delta=None): + if delta is None: + delta = self.frequency + return self.forward.statistics_tendencies(delta) + @property def shape(self): return self.forward.shape diff --git a/src/anemoi/datasets/data/join.py b/src/anemoi/datasets/data/join.py index fd382ccd..b47f7762 100644 --- a/src/anemoi/datasets/data/join.py +++ b/src/anemoi/datasets/data/join.py @@ -121,6 +121,14 @@ def statistics(self): k: np.concatenate([d.statistics[k] for d in self.datasets], axis=0) for k in self.datasets[0].statistics } + def statistics_tendencies(self, delta=None): + if delta is None: + delta = self.frequency + return { + k: np.concatenate([d.statistics_tendencies(delta)[k] for d in self.datasets], axis=0) + for k in self.datasets[0].statistics_tendencies(delta) + } + def source(self, index): i = index for dataset in self.datasets: diff --git a/src/anemoi/datasets/data/select.py b/src/anemoi/datasets/data/select.py index aaf2f224..feae504b 100644 --- a/src/anemoi/datasets/data/select.py +++ b/src/anemoi/datasets/data/select.py @@ -74,6 +74,11 @@ def name_to_index(self): def statistics(self): return {k: v[self.indices] for k, v in self.dataset.statistics.items()} + def statistics_tendencies(self, delta=None): + if delta is None: + delta = self.frequency + return {k: v[self.indices] for k, v in self.dataset.statistics_tendencies(delta).items()} + def metadata_specific(self, **kwargs): return super().metadata_specific(indices=self.indices, **kwargs) diff --git a/src/anemoi/datasets/data/statistics.py b/src/anemoi/datasets/data/statistics.py index 78f0e1cb..90b3115c 100644 --- a/src/anemoi/datasets/data/statistics.py +++ b/src/anemoi/datasets/data/statistics.py @@ -29,6 +29,11 @@ def __init__(self, dataset, statistic): def statistics(self): return self._statistic.statistics + def statistics_tendencies(self, delta=None): + if delta is None: + delta = self.frequency + return self._statistic.statistics_tendencies(delta) + def metadata_specific(self, **kwargs): return super().metadata_specific( statistics=self._statistic.metadata_specific(), diff --git a/src/anemoi/datasets/data/stores.py b/src/anemoi/datasets/data/stores.py index dc5d3d63..62173e08 100644 --- a/src/anemoi/datasets/data/stores.py +++ b/src/anemoi/datasets/data/stores.py @@ -218,6 +218,21 @@ def statistics(self): minimum=self.z.minimum[:], ) + def statistics_tendencies(self, delta=None): + if delta is None: + delta = self.frequency + if isinstance(delta, int): + delta = f"{delta}h" + from anemoi.datasets.create.loaders import TendenciesStatisticsAddition + + prefix = TendenciesStatisticsAddition.DATASET_NAME_PATTERN.format(delta=delta) + "_" + return dict( + mean=self.z[f"{prefix}mean"][:], + stdev=self.z[f"{prefix}stdev"][:], + maximum=self.z[f"{prefix}maximum"][:], + minimum=self.z[f"{prefix}minimum"][:], + ) + @property def resolution(self): return self.z.attrs["resolution"] diff --git a/src/anemoi/datasets/data/unchecked.py b/src/anemoi/datasets/data/unchecked.py index 7d9b02b1..776a0f92 100644 --- a/src/anemoi/datasets/data/unchecked.py +++ b/src/anemoi/datasets/data/unchecked.py @@ -101,6 +101,10 @@ def variables(self): def statistics(self): raise NotImplementedError() + @check("check_same_variables") + def statistics_tendencies(self, delta=None): + raise NotImplementedError() + @property def shape(self): raise NotImplementedError()