better metadata handling

Author: b8raoult

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

@@ -8,6 +8,8 @@
 #
 
 
+import glob
+
 from climetlab import load_source
 from climetlab.utils.patterns import Pattern
 
@@ -22,6 +24,12 @@ def check(ds, paths, **kwargs):
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})")
 
 
+def _expand(paths):
+    for path in paths:
+        for p in glob.glob(path):
+            yield p
+
+
 def execute(context, dates, path, *args, **kwargs):
     given_paths = path if isinstance(path, list) else [path]
 
@@ -35,7 +43,7 @@ def execute(context, dates, path, *args, **kwargs):
             if name in kwargs:
                 raise ValueError(f"MARS interpolation parameter '{name}' not supported")
 
-        for path in paths:
+        for path in _expand(paths):
             context.trace("📁", "PATH", path)
             s = load_source("file", path)
             s = s.sel(valid_datetime=dates, **kwargs)

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

@@ -7,10 +7,29 @@
 # nor does it submit to any jurisdiction.
 #
 
+import glob
+
 from climetlab import load_source
 from climetlab.utils.patterns import Pattern
 
 
+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
+
+        for p in glob.glob(path):
+            yield p
+
+
 def check(what, ds, paths, **kwargs):
     count = 1
     for k, v in kwargs.items():
@@ -32,7 +51,7 @@ def load_netcdfs(emoji, what, context, dates, path, *args, **kwargs):
 
         levels = kwargs.get("level", kwargs.get("levelist"))
 
-        for path in paths:
+        for path in _expand(paths):
             context.trace(emoji, what.upper(), path)
             s = load_source("opendap", path)
             s = s.sel(

src/anemoi/datasets/data/concat.py

@@ -12,7 +12,7 @@
 
 from .debug import Node
 from .debug import debug_indexing
-from .forewards import Combined
+from .forwards import Combined
 from .indexing import apply_index_to_slices_changes
 from .indexing import expand_list_indexing
 from .indexing import index_to_slices

src/anemoi/datasets/data/ensemble.py

@@ -8,7 +8,7 @@
 import logging
 
 from .debug import Node
-from .forewards import GivenAxis
+from .forwards import GivenAxis
 from .misc import _auto_adjust
 from .misc import _open
 

src/anemoi/datasets/data/forewards.py renamed to src/anemoi/datasets/data/forwards.py

@@ -91,12 +91,18 @@ def grids(self):
     def metadata_specific(self, **kwargs):
         return super().metadata_specific(
             forward=self.forward.metadata_specific(),
+            **self.subclass_metadata_specific(),
             **kwargs,
         )
 
     def source(self, index):
         return self.forward.source(index)
 
+    def subclass_metadata_specific(self):
+        raise NotImplementedError(
+            f"subclass_metadata_specific() must be implemented in derived class {self.__class__.__name__}"
+        )
+
 
 class Combined(Forwards):
     def __init__(self, datasets):

src/anemoi/datasets/data/grids.py

@@ -12,8 +12,8 @@
 
 from .debug import Node
 from .debug import debug_indexing
-from .forewards import Combined
-from .forewards import GivenAxis
+from .forwards import Combined
+from .forwards import GivenAxis
 from .indexing import apply_index_to_slices_changes
 from .indexing import expand_list_indexing
 from .indexing import index_to_slices
@@ -91,7 +91,7 @@ def tree(self):
         return Node(self, [d.tree() for d in self.datasets])
 
 
-class Grids(GivenAxis):
+class GridsBase(GivenAxis):
     def __init__(self, datasets, axis):
         super().__init__(datasets, axis)
         # Shape: (dates, variables, ensemble, 1d-values)
@@ -106,7 +106,7 @@ def check_same_resolution(self, d1, d2):
         pass
 
 
-class ConcatGrids(Grids):
+class Grids(GridsBase):
     # TODO: select the statistics of the most global grid?
     @property
     def latitudes(self):
@@ -127,7 +127,7 @@ def tree(self):
         return Node(self, [d.tree() for d in self.datasets], mode="concat")
 
 
-class Cutout(Grids):
+class Cutout(GridsBase):
     def __init__(self, datasets, axis):
         from anemoi.datasets.grids import cutout_mask
 
@@ -220,7 +220,7 @@ def grids_factory(args, kwargs):
     datasets = [_open(e) for e in grids]
     datasets, kwargs = _auto_adjust(datasets, kwargs)
 
-    return ConcatGrids(datasets, axis=axis)._subset(**kwargs)
+    return Grids(datasets, axis=axis)._subset(**kwargs)
 
 
 def cutout_factory(args, kwargs):

src/anemoi/datasets/data/join.py

@@ -13,7 +13,7 @@
 from .debug import Node
 from .debug import Source
 from .debug import debug_indexing
-from .forewards import Combined
+from .forwards import Combined
 from .indexing import apply_index_to_slices_changes
 from .indexing import expand_list_indexing
 from .indexing import index_to_slices

src/anemoi/datasets/data/masked.py

@@ -14,7 +14,7 @@
 from .dataset import Dataset
 from .debug import Node
 from .debug import debug_indexing
-from .forewards import Forwards
+from .forwards import Forwards
 from .indexing import apply_index_to_slices_changes
 from .indexing import expand_list_indexing
 from .indexing import index_to_slices
@@ -83,6 +83,9 @@ def __init__(self, forward, thinning, method):
     def tree(self):
         return Node(self, [self.forward.tree()], thinning=self.thinning, method=self.method)
 
+    def subclass_metadata_specific(self):
+        return dict(thinning=self.thinning, method=self.method)
+
 
 class Cropping(Masked):
     def __init__(self, forward, area):
@@ -104,3 +107,6 @@ def __init__(self, forward, area):
 
     def tree(self):
         return Node(self, [self.forward.tree()], area=self.area)
+
+    def metadata_specific(self, **kwargs):
+        return super().metadata_specific(area=self.area, **kwargs)

src/anemoi/datasets/data/select.py

@@ -11,7 +11,7 @@
 from .debug import Node
 from .debug import Source
 from .debug import debug_indexing
-from .forewards import Forwards
+from .forwards import Forwards
 from .indexing import apply_index_to_slices_changes
 from .indexing import expand_list_indexing
 from .indexing import index_to_slices
@@ -88,6 +88,10 @@ def source(self, index):
     def tree(self):
         return Node(self, [self.dataset.tree()], **self.title)
 
+    def subclass_metadata_specific(self):
+        # return dict(indices=self.indices)
+        return {}
+
 
 class Rename(Forwards):
     def __init__(self, dataset, rename):
@@ -105,8 +109,8 @@ def variables(self):
     def name_to_index(self):
         return {k: i for i, k in enumerate(self.variables)}
 
-    def metadata_specific(self, **kwargs):
-        return super().metadata_specific(rename=self.rename, **kwargs)
-
     def tree(self):
         return Node(self, [self.forward.tree()], rename=self.rename)
+
+    def subclass_metadata_specific(self):
+        return dict(rename=self.rename)

src/anemoi/datasets/data/statistics.py

@@ -10,7 +10,7 @@
 
 from . import open_dataset
 from .debug import Node
-from .forewards import Forwards
+from .forwards import Forwards
 
 LOG = logging.getLogger(__name__)
 
@@ -34,11 +34,8 @@ def statistics_tendencies(self, delta=None):
             delta = self.frequency
         return self._statistic.statistics_tendencies(delta)
 
-    def metadata_specific(self, **kwargs):
-        return super().metadata_specific(
-            statistics=self._statistic.metadata_specific(),
-            **kwargs,
-        )
+    def subclass_metadata_specific(self):
+        return dict(statistics=self._statistic.metadata_specific())
 
     def tree(self):
         return Node(self, [self.forward.tree()])