implemented statistics with tendencies

Author: floriankrb

src/anemoi/datasets/create/__init__.py

@@ -32,13 +32,13 @@ def init(self, check_name=False):
         # check path
         _, ext = os.path.splitext(self.path)
         assert ext != "zarr", f"Unsupported extension={ext}"
-        from .loaders import InitialiseLoader
+        from .loaders import InitialiserLoader
 
         if self._path_readable() and not self.overwrite:
             raise Exception(f"{self.path} already exists. Use overwrite=True to overwrite.")
 
         with self._cache_context():
-            obj = InitialiseLoader.from_config(
+            obj = InitialiserLoader.from_config(
                 path=self.path,
                 config=self.config,
                 statistics_tmp=self.statistics_tmp,
@@ -59,12 +59,11 @@ def load(self, parts=None):
             loader.load()
 
     def statistics(self, force=False, output=None, start=None, end=None):
-        from .loaders import StatisticsLoader
+        from .loaders import StatisticsAdder
 
-        loader = StatisticsLoader.from_dataset(
+        loader = StatisticsAdder.from_dataset(
             path=self.path,
             print=self.print,
-            force=force,
             statistics_tmp=self.statistics_tmp,
             statistics_output=output,
             recompute=False,
@@ -74,26 +73,72 @@ def statistics(self, force=False, output=None, start=None, end=None):
         loader.run()
 
     def size(self):
-        from .loaders import SizeLoader
+        from .loaders import DatasetHandler
+        from .size import compute_directory_sizes
 
-        loader = SizeLoader.from_dataset(path=self.path, print=self.print)
-        loader.add_total_size()
+        metadata = compute_directory_sizes(self.path)
+        handle = DatasetHandler.from_dataset(path=self.path, print=self.print)
+        handle.update_metadata(**metadata)
 
     def cleanup(self):
-        from .loaders import CleanupLoader
+        from .loaders import DatasetHandlerWithStatistics
 
-        loader = CleanupLoader.from_dataset(
-            path=self.path,
-            print=self.print,
-            statistics_tmp=self.statistics_tmp,
+        cleaner = DatasetHandlerWithStatistics.from_dataset(
+            path=self.path, print=self.print, statistics_tmp=self.statistics_tmp
         )
-        loader.run()
+        cleaner.tmp_statistics.delete()
+        cleaner.registry.clean()
 
     def patch(self, **kwargs):
         from .patch import apply_patch
 
         apply_patch(self.path, **kwargs)
 
+    def init_additions(self, delta=[1, 3, 6, 12]):
+        from .loaders import StatisticsAddition
+        from .loaders import TendenciesStatisticsAddition
+        from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
+
+        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+        a.initialise()
+
+        for d in delta:
+            try:
+                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a.initialise()
+            except TendenciesStatisticsDeltaNotMultipleOfFrequency:
+                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+
+    def run_additions(self, parts=None, delta=[1, 3, 6, 12]):
+        from .loaders import StatisticsAddition
+        from .loaders import TendenciesStatisticsAddition
+        from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
+
+        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+        a.run(parts)
+
+        for d in delta:
+            try:
+                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a.run(parts)
+            except TendenciesStatisticsDeltaNotMultipleOfFrequency:
+                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+
+    def finalise_additions(self, delta=[1, 3, 6, 12]):
+        from .loaders import StatisticsAddition
+        from .loaders import TendenciesStatisticsAddition
+        from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
+
+        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+        a.finalise()
+
+        for d in delta:
+            try:
+                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a.finalise()
+            except TendenciesStatisticsDeltaNotMultipleOfFrequency:
+                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+
     def finalise(self, **kwargs):
         self.statistics(**kwargs)
         self.size()
@@ -102,8 +147,14 @@ def create(self):
         self.init()
         self.load()
         self.finalise()
+        self.additions()
         self.cleanup()
 
+    def additions(self):
+        self.init_additions()
+        self.run_additions()
+        self.finalise_additions()
+
     def _cache_context(self):
         from .utils import cache_context
 

src/anemoi/datasets/create/check.py

@@ -8,29 +8,14 @@
 #
 
 import logging
-import os
 import re
 import warnings
 
 import numpy as np
-import tqdm
 
 LOG = logging.getLogger(__name__)
 
 
-def compute_directory_size(path):
-    if not os.path.isdir(path):
-        return None
-    size = 0
-    n = 0
-    for dirpath, _, filenames in tqdm.tqdm(os.walk(path), desc="Computing size", leave=False):
-        for filename in filenames:
-            file_path = os.path.join(dirpath, filename)
-            size += os.path.getsize(file_path)
-            n += 1
-    return size, n
-
-
 class DatasetName:
     def __init__(
         self,

src/anemoi/datasets/create/chunks.py

@@ -0,0 +1,78 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+import logging
+import warnings
+
+LOG = logging.getLogger(__name__)
+
+ALL = object()
+
+
+class ChunkFilter:
+    def __init__(self, *, parts, total):
+        self.total = total
+
+        if isinstance(parts, list):
+            if len(parts) == 1:
+                parts = parts[0]
+            elif len(parts) == 0:
+                parts = None
+            else:
+                raise ValueError(f"Invalid parts format: {parts}. Must be in the form 'i/n'.")
+
+        if not parts:
+            parts = "all"
+
+        assert isinstance(parts, str), f"Argument parts must be a string, got {parts}."
+
+        if parts.lower() == "all" or parts == "*":
+            self.allowed = ALL
+            return
+
+        assert "/" in parts, f"Invalid parts format: {parts}. Must be in the form 'i/n'."
+
+        i, n = parts.split("/")
+        i, n = int(i), int(n)
+
+        assert i > 0, f"Chunk number {i} must be positive."
+        assert i <= n, f"Chunk number {i} must be less than total chunks {n}."
+        if n > total:
+            warnings.warn(
+                f"Number of chunks {n} is larger than the total number of chunks: {total}. "
+                "Some chunks will be empty."
+            )
+
+        chunk_size = total / n
+        parts = [x for x in range(total) if x >= (i - 1) * chunk_size and x < i * chunk_size]
+
+        for i in parts:
+            if i < 0 or i >= total:
+                raise AssertionError(f"Invalid chunk number {i}. Must be between 0 and {total - 1}.")
+        if not parts:
+            warnings.warn(f"Nothing to do for chunk {i}/{n}.")
+
+        LOG.info(f"Running parts: {parts}")
+
+        self.allowed = parts
+
+    def __call__(self, i):
+        if i < 0 or i >= self.total:
+            raise AssertionError(f"Invalid chunk number {i}. Must be between 0 and {self.total - 1}.")
+
+        if self.allowed == ALL:
+            return True
+        return i in self.allowed
+
+    def __iter__(self):
+        for i in range(self.total):
+            if self(i):
+                yield i
+
+    def __len__(self):
+        return len([_ for _ in self])

src/anemoi/datasets/create/functions/sources/perturbations.py

@@ -58,6 +58,7 @@ def load_if_needed(context, dates, dict_or_dataset):
 def perturbations(context, dates, members, center, remapping={}, patches={}):
     members = load_if_needed(context, dates, members)
     center = load_if_needed(context, dates, center)
+    # return perturbations(member, centers....)
 
     keys = ["param", "level", "valid_datetime", "date", "time", "step", "number"]