anemoi-datasets 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.0'
-__version_tuple__ = version_tuple = (0, 4, 0)
+__version__ = version = '0.4.2'
+__version_tuple__ = version_tuple = (0, 4, 2)

anemoi/datasets/commands/compare.py

@@ -8,6 +8,10 @@
 # nor does it submit to any jurisdiction.
 #
 
+import numpy as np
+import tqdm
+import zarr
+
 from anemoi.datasets import open_dataset
 
 from . import Command
@@ -19,6 +23,8 @@ class Compare(Command):
     def add_arguments(self, command_parser):
         command_parser.add_argument("dataset1")
         command_parser.add_argument("dataset2")
+        command_parser.add_argument("--data", action="store_true", help="Compare the data.")
+        command_parser.add_argument("--statistics", action="store_true", help="Compare the statistics.")
 
     def run(self, args):
         ds1 = open_dataset(args.dataset1)
@@ -42,5 +48,58 @@ class Compare(Command):
                 f"{ds2.statistics['mean'][ds2.name_to_index[v]]:14g}",
             )
 
+        if args.data:
+            print()
+            print("Data:")
+            print("-----")
+            print()
+
+            diff = 0
+            for a, b in tqdm.tqdm(zip(ds1, ds2)):
+                if not np.array_equal(a, b, equal_nan=True):
+                    diff += 1
+
+            print(f"Number of different rows: {diff}/{len(ds1)}")
+
+        if args.data:
+            print()
+            print("Data 2:")
+            print("-----")
+            print()
+
+            ds1 = zarr.open(args.dataset1, mode="r")
+            ds2 = zarr.open(args.dataset2, mode="r")
+
+            for name in (
+                "data",
+                "count",
+                "sums",
+                "squares",
+                "mean",
+                "stdev",
+                "minimum",
+                "maximum",
+                "latitudes",
+                "longitudes",
+            ):
+                a1 = ds1[name]
+                a2 = ds2[name]
+
+                if len(a1) != len(a2):
+                    print(f"{name}: lengths mismatch {len(a1)} != {len(a2)}")
+                    continue
+
+                diff = 0
+                for a, b in tqdm.tqdm(zip(a1, a2), leave=False):
+                    if not np.array_equal(a, b, equal_nan=True):
+                        if diff == 0:
+                            print(f"\n{name}: first different row:")
+                            print(a[a != b])
+                            print(b[a != b])
+
+                        diff += 1
+
+                print(f"{name}: {diff} different rows out of {len(a1)}")
+
 
 command = Compare

anemoi/datasets/commands/create.py

@@ -1,7 +1,39 @@
-from anemoi.datasets.create import Creator
+import datetime
+import logging
+import time
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import as_completed
+
+import tqdm
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.create.trace import enable_trace
 
 from . import Command
 
+LOG = logging.getLogger(__name__)
+
+
+def task(what, options, *args, **kwargs):
+    """
+    Make sure `import Creator` is done in the sub-processes, and not in the main one.
+    """
+
+    now = datetime.datetime.now()
+    LOG.debug(f"Task {what}({args},{kwargs}) starting")
+
+    from anemoi.datasets.create import Creator
+
+    if "trace" in options:
+        enable_trace(options["trace"])
+
+    c = Creator(**options)
+    result = getattr(c, what)(*args, **kwargs)
+
+    LOG.debug(f"Task {what}({args},{kwargs}) completed ({datetime.datetime.now()-now})")
+    return result
+
 
 class Create(Command):
     """Create a dataset."""
@@ -22,12 +54,61 @@ class Create(Command):
         )
         command_parser.add_argument("config", help="Configuration yaml file defining the recipe to create the dataset.")
         command_parser.add_argument("path", help="Path to store the created data.")
+        group = command_parser.add_mutually_exclusive_group()
+        group.add_argument("--threads", help="Use `n` parallel thread workers.", type=int, default=0)
+        group.add_argument("--processes", help="Use `n` parallel process workers.", type=int, default=0)
+        command_parser.add_argument("--trace", action="store_true")
 
     def run(self, args):
-        kwargs = vars(args)
+        now = time.time()
+        if args.threads + args.processes:
+            self.parallel_create(args)
+        else:
+            self.serial_create(args)
+        LOG.info(f"Create completed in {seconds_to_human(time.time()-now)}")
 
-        c = Creator(**kwargs)
+    def serial_create(self, args):
+        from anemoi.datasets.create import Creator
+
+        options = vars(args)
+        c = Creator(**options)
         c.create()
 
+    def parallel_create(self, args):
+        """Some modules, like fsspec do not work well with fork()
+        Other modules may not be thread safe. So we implement
+        parallel loadining using multiprocessing before any
+        of the modules are imported.
+        """
+
+        options = vars(args)
+        parallel = args.threads + args.processes
+        args.use_threads = args.threads > 0
+
+        if args.use_threads:
+            ExecutorClass = ThreadPoolExecutor
+        else:
+            ExecutorClass = ProcessPoolExecutor
+
+        with ExecutorClass(max_workers=1) as executor:
+            total = executor.submit(task, "init", options).result()
+
+        futures = []
+
+        with ExecutorClass(max_workers=parallel) as executor:
+            for n in range(total):
+                futures.append(executor.submit(task, "load", options, parts=f"{n+1}/{total}"))
+
+            for future in tqdm.tqdm(
+                as_completed(futures), desc="Loading", total=len(futures), colour="green", position=parallel + 1
+            ):
+                future.result()
+
+        with ExecutorClass(max_workers=1) as executor:
+            executor.submit(task, "statistics", options).result()
+            executor.submit(task, "additions", options).result()
+            executor.submit(task, "cleanup", options).result()
+            executor.submit(task, "verify", options).result()
+
 
 command = Create

anemoi/datasets/commands/inspect.py

@@ -16,7 +16,7 @@ import numpy as np
 import semantic_version
 import tqdm
 from anemoi.utils.humanize import bytes
-from anemoi.utils.humanize import number
+from anemoi.utils.humanize import bytes_to_human
 from anemoi.utils.humanize import when
 from anemoi.utils.text import dotted_line
 from anemoi.utils.text import progress
@@ -215,9 +215,9 @@ class Version:
             total_size, n = compute_directory_size(self.path)
 
         if total_size is not None:
-            print(f"💽 Size       : {bytes(total_size)} ({number(total_size)})")
+            print(f"💽 Size       : {bytes(total_size)} ({bytes_to_human(total_size)})")
         if n is not None:
-            print(f"📁 Files      : {number(n)}")
+            print(f"📁 Files      : {n:,}")
 
     @property
     def statistics(self):

anemoi/datasets/create/__init__.py

@@ -7,8 +7,15 @@
 # nor does it submit to any jurisdiction.
 #
 
+import logging
 import os
 
+LOG = logging.getLogger(__name__)
+
+
+def _ignore(*args, **kwargs):
+    pass
+
 
 class Creator:
     def __init__(
@@ -16,19 +23,21 @@ class Creator:
         path,
         config=None,
         cache=None,
-        print=print,
+        use_threads=False,
         statistics_tmp=None,
         overwrite=False,
         test=None,
+        progress=None,
         **kwargs,
     ):
         self.path = path  # Output path
         self.config = config
         self.cache = cache
-        self.print = print
+        self.use_threads = use_threads
         self.statistics_tmp = statistics_tmp
         self.overwrite = overwrite
         self.test = test
+        self.progress = progress if progress is not None else _ignore
 
     def init(self, check_name=False):
         # check path
@@ -44,10 +53,11 @@ class Creator:
                 path=self.path,
                 config=self.config,
                 statistics_tmp=self.statistics_tmp,
-                print=self.print,
+                use_threads=self.use_threads,
+                progress=self.progress,
                 test=self.test,
             )
-            obj.initialise(check_name=check_name)
+            return obj.initialise(check_name=check_name)
 
     def load(self, parts=None):
         from .loaders import ContentLoader
@@ -56,7 +66,8 @@ class Creator:
             loader = ContentLoader.from_dataset_config(
                 path=self.path,
                 statistics_tmp=self.statistics_tmp,
-                print=self.print,
+                use_threads=self.use_threads,
+                progress=self.progress,
                 parts=parts,
             )
             loader.load()
@@ -66,7 +77,8 @@ class Creator:
 
         loader = StatisticsAdder.from_dataset(
             path=self.path,
-            print=self.print,
+            use_threads=self.use_threads,
+            progress=self.progress,
             statistics_tmp=self.statistics_tmp,
             statistics_output=output,
             recompute=False,
@@ -74,20 +86,22 @@ class Creator:
             statistics_end=end,
         )
         loader.run()
+        assert loader.ready()
 
     def size(self):
         from .loaders import DatasetHandler
         from .size import compute_directory_sizes
 
         metadata = compute_directory_sizes(self.path)
-        handle = DatasetHandler.from_dataset(path=self.path, print=self.print)
+        handle = DatasetHandler.from_dataset(path=self.path, use_threads=self.use_threads)
         handle.update_metadata(**metadata)
+        assert handle.ready()
 
     def cleanup(self):
         from .loaders import DatasetHandlerWithStatistics
 
         cleaner = DatasetHandlerWithStatistics.from_dataset(
-            path=self.path, print=self.print, statistics_tmp=self.statistics_tmp
+            path=self.path, use_threads=self.use_threads, progress=self.progress, statistics_tmp=self.statistics_tmp
         )
         cleaner.tmp_statistics.delete()
         cleaner.registry.clean()
@@ -103,15 +117,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.initialise()
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.initialise()
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.info(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
     def run_additions(self, parts=None, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
@@ -119,15 +135,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.run(parts)
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.run(parts)
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.debug(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
     def finalise_additions(self, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
@@ -135,15 +153,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.finalise()
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.finalise()
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.debug(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
     def finalise(self, **kwargs):
         self.statistics(**kwargs)
@@ -174,3 +194,10 @@ class Creator:
             return True
         except zarr.errors.PathNotFoundError:
             return False
+
+    def verify(self):
+        from .loaders import DatasetVerifier
+
+        handle = DatasetVerifier.from_dataset(path=self.path, use_threads=self.use_threads)
+
+        handle.verify()

anemoi/datasets/create/check.py

@@ -56,7 +56,7 @@ class DatasetName:
             raise ValueError(self.error_message)
 
     def _parse(self, name):
-        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?(.*)$"
+        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?([a-zA-Z0-9-]+)$"
         match = re.match(pattern, name)
 
         assert match, (name, pattern)
@@ -136,18 +136,19 @@ class StatisticsValueError(ValueError):
     pass
 
 
-def check_data_values(arr, *, name: str, log=[], allow_nan=False):
-    if allow_nan is False:
-        allow_nan = lambda x: False  # noqa: E731
+def check_data_values(arr, *, name: str, log=[], allow_nans=False):
 
-    if allow_nan(name):
+    if (isinstance(allow_nans, (set, list, tuple, dict)) and name in allow_nans) or allow_nans:
         arr = arr[~np.isnan(arr)]
 
+    assert arr.size > 0, (name, *log)
+
     min, max = arr.min(), arr.max()
     assert not (np.isnan(arr).any()), (name, min, max, *log)
 
     if min == 9999.0:
         warnings.warn(f"Min value 9999 for {name}")
+
     if max == 9999.0:
         warnings.warn(f"Max value 9999 for {name}")
 

anemoi/datasets/create/chunks.py

@@ -57,7 +57,7 @@ class ChunkFilter:
         if not parts:
             warnings.warn(f"Nothing to do for chunk {i}/{n}.")
 
-        LOG.info(f"Running parts: {parts}")
+        LOG.debug(f"Running parts: {parts}")
 
         self.allowed = parts
 

anemoi/datasets/create/config.py

@@ -12,10 +12,10 @@ import os
 from copy import deepcopy
 
 import yaml
+from anemoi.utils.config import DotDict
+from anemoi.utils.config import load_any_dict_format
 from earthkit.data.core.order import normalize_order_by
 
-from .utils import load_json_or_yaml
-
 LOG = logging.getLogger(__name__)
 
 
@@ -43,31 +43,10 @@ def check_dict_value_and_set(dic, key, value):
         if dic[key] == value:
             return
         raise ValueError(f"Cannot use {key}={dic[key]}. Must use {value}.")
-    print(f"Setting {key}={value} in config")
+    LOG.info(f"Setting {key}={value} in config")
     dic[key] = value
 
 
-class DictObj(dict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        for key, value in self.items():
-            if isinstance(value, dict):
-                self[key] = DictObj(value)
-                continue
-            if isinstance(value, list):
-                self[key] = [DictObj(item) if isinstance(item, dict) else item for item in value]
-                continue
-
-    def __getattr__(self, attr):
-        try:
-            return self[attr]
-        except KeyError:
-            raise AttributeError(attr)
-
-    def __setattr__(self, attr, value):
-        self[attr] = value
-
-
 def resolve_includes(config):
     if isinstance(config, list):
         return [resolve_includes(c) for c in config]
@@ -79,11 +58,11 @@ def resolve_includes(config):
     return config
 
 
-class Config(DictObj):
+class Config(DotDict):
     def __init__(self, config=None, **kwargs):
         if isinstance(config, str):
             self.config_path = os.path.realpath(config)
-            config = load_json_or_yaml(config)
+            config = load_any_dict_format(config)
         else:
             config = deepcopy(config if config is not None else {})
         config = resolve_includes(config)

anemoi/datasets/create/functions/filters/rename.py

@@ -26,15 +26,23 @@ class RenamedFieldMapping:
         self.what = what
         self.renaming = renaming
 
-    def metadata(self, key, **kwargs):
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
         value = self.field.metadata(key, **kwargs)
         if key == self.what:
             return self.renaming.get(value, value)
+
         return value
 
     def __getattr__(self, name):
         return getattr(self.field, name)
 
+    def __repr__(self) -> str:
+        return repr(self.field)
+        return f"{self.field} -> {self.what} -> {self.renaming}"
+
 
 class RenamedFieldFormat:
     """Rename a field based on a format string.

anemoi/datasets/create/functions/filters/rotate_winds.py

@@ -9,6 +9,8 @@
 
 from collections import defaultdict
 
+import tqdm
+from anemoi.utils.humanize import plural
 from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.geo.rotate import rotate_vector
 
@@ -24,6 +26,9 @@ class NewDataField:
     def __getattr__(self, name):
         return getattr(self.field, name)
 
+    def __repr__(self) -> str:
+        return repr(self.field)
+
 
 def execute(
     context,
@@ -35,6 +40,8 @@ def execute(
 ):
     from pyproj import CRS
 
+    context.trace("🔄", "Rotating winds (extracting winds from ", plural(len(input), "field"))
+
     result = FieldArray()
 
     wind_params = (x_wind, y_wind)
@@ -55,7 +62,9 @@ def execute(
 
         wind_pairs[key][param] = f
 
-    for _, pairs in wind_pairs.items():
+    context.trace("🔄", "Rotating", plural(len(wind_pairs), "wind"), "(speed will likely include data download)")
+
+    for _, pairs in tqdm.tqdm(list(wind_pairs.items())):
         if len(pairs) != 2:
             raise ValueError("Missing wind component")
 

anemoi/datasets/create/functions/sources/__init__.py

@@ -6,3 +6,42 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
+
+import glob
+import logging
+
+from earthkit.data.utils.patterns import Pattern
+
+LOG = logging.getLogger(__name__)
+
+
+def _expand(paths):
+    for path in paths:
+        if path.startswith("file://"):
+            path = path[7:]
+
+        if path.startswith("http://"):
+            yield path
+            continue
+
+        if path.startswith("https://"):
+            yield path
+            continue
+
+        cnt = 0
+        for p in glob.glob(path):
+            yield p
+            cnt += 1
+        if cnt == 0:
+            yield path
+
+
+def iterate_patterns(path, dates, **kwargs):
+    given_paths = path if isinstance(path, list) else [path]
+
+    dates = [d.isoformat() for d in dates]
+
+    for path in given_paths:
+        paths = Pattern(path, ignore_missing_keys=True).substitute(date=dates, **kwargs)
+        for path in _expand(paths):
+            yield path, dates