anemoi-datasets 0.5.27__py3-none-any.whl → 0.5.29__py3-none-any.whl

anemoi/datasets/create/sources/grib_index.py

@@ -7,9 +7,12 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import hashlib
+import json
 import logging
 import os
 import sqlite3
+from collections import defaultdict
 from collections.abc import Iterator
 from typing import Any
 
@@ -19,7 +22,8 @@ from anemoi.transform.flavour import RuleBasedFlavour
 from cachetools import LRUCache
 from earthkit.data.indexing.fieldlist import FieldArray
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
@@ -45,8 +49,8 @@ class GribIndex:
         ----------
         database : str
             Path to the SQLite database file.
-        keys : Optional[List[str] | str], optional
-            List of keys or a string of keys to use for indexing, by default None.
+        keys : Optional[list[str] | str], optional
+            list of keys or a string of keys to use for indexing, by default None.
         flavour : Optional[str], optional
             Flavour configuration for mapping fields, by default None.
         update : bool, optional
@@ -160,7 +164,7 @@ class GribIndex:
 
         Returns
         -------
-        List[str]
+        list[str]
             A list of metadata keys stored in the database.
         """
         self.cursor.execute("SELECT key FROM metadata_keys")
@@ -228,7 +232,7 @@ class GribIndex:
 
         Returns
         -------
-        List[str]
+        list[str]
             A list of column names.
         """
         if self._columns is not None:
@@ -244,8 +248,8 @@ class GribIndex:
 
         Parameters
         ----------
-        columns : List[str]
-            List of column names to ensure in the table.
+        columns : list[str]
+            list of column names to ensure in the table.
         """
         assert self.update
 
@@ -363,7 +367,7 @@ class GribIndex:
 
         Returns
         -------
-        List[dict]
+        list[dict]
             A list of GRIB2 parameter information.
         """
         if ("grib2", paramId) in self.cache:
@@ -523,8 +527,8 @@ class GribIndex:
 
         Parameters
         ----------
-        dates : List[Any]
-            List of dates to retrieve data for.
+        dates : list[Any]
+            list of dates to retrieve data for.
         **kwargs : Any
             Additional filtering criteria.
 
@@ -544,6 +548,9 @@ class GribIndex:
         params = dates
 
         for k, v in kwargs.items():
+            if k not in self._columns:
+                LOG.warning(f"Warning : {k} not in database columns, key discarded")
+                continue
             if isinstance(v, list):
                 query += f" AND {k} IN ({', '.join('?' for _ in v)})"
                 params.extend([str(_) for _ in v])
@@ -551,11 +558,14 @@ class GribIndex:
                 query += f" AND {k} = ?"
                 params.append(str(v))
 
-        print("SELECT", query)
-        print("SELECT", params)
+        print("SELECT (query)", query)
+        print("SELECT (params)", params)
 
         self.cursor.execute(query, params)
-        for path_id, offset, length in self.cursor.fetchall():
+
+        fetch = self.cursor.fetchall()
+
+        for path_id, offset, length in fetch:
             if path_id in self.cache:
                 file = self.cache[path_id]
             else:
@@ -569,44 +579,82 @@ class GribIndex:
             yield data
 
 
-@legacy_source(__file__)
-def execute(
-    context: Any,
-    dates: list[Any],
-    indexdb: str,
-    flavour: str | None = None,
-    **kwargs: Any,
-) -> FieldArray:
-    """Execute the GRIB data retrieval process.
-
-    Parameters
-    ----------
-    context : Any
-        The execution context.
-    dates : List[Any]
-        List of dates to retrieve data for.
-    indexdb : str
-        Path to the GRIB index database.
-    flavour : Optional[str], optional
-        Flavour configuration for mapping fields, by default None.
-    **kwargs : Any
-        Additional filtering criteria.
-
-    Returns
-    -------
-    FieldArray
-        An array of retrieved GRIB fields.
-    """
-    index = GribIndex(indexdb)
-    result = []
-
-    if flavour is not None:
-        flavour = RuleBasedFlavour(flavour)
-
-    for grib in index.retrieve(dates, **kwargs):
-        field = ekd.from_source("memory", grib)[0]
-        if flavour:
-            field = flavour.apply(field)
-        result.append(field)
-
-    return FieldArray(result)
+@source_registry.register("grib-index")
+class GribIndexSource(LegacySource):
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[Any],
+        indexdb: str,
+        flavour: str | None = None,
+        **kwargs: Any,
+    ) -> FieldArray:
+        """Execute the GRIB data retrieval process.
+
+        Parameters
+        ----------
+        context : Any
+            The execution context.
+        dates : List[Any]
+            List of dates to retrieve data for.
+        indexdb : str
+            Path to the GRIB index database.
+        flavour : Optional[str], optional
+            Flavour configuration for mapping fields, by default None.
+        **kwargs : Any
+            Additional filtering criteria.
+
+        Returns
+        -------
+        FieldArray
+            An array of retrieved GRIB fields.
+        """
+        index = GribIndex(indexdb)
+
+        if flavour is not None:
+            flavour = RuleBasedFlavour(flavour)
+
+        if hasattr(dates, "date_to_intervals"):
+            # When using accumulate source
+            full_requests = []
+            for d, interval in dates.intervals:
+                context.trace("🌧️", "interval:", interval)
+                valid_date, request, _ = dates._adjust_request_to_interval(interval, kwargs)
+                context.trace("🌧️", "  request =", request)
+                full_requests.append(([valid_date], request))
+        else:
+            # Normal case, without accumulate source
+            full_requests = [(dates, kwargs)]
+
+        full_requests = factorise(full_requests)
+        context.trace("🌧️", f"number of (factorised) requests: {len(full_requests)}")
+        for valid_dates, request in full_requests:
+            context.trace("🌧️", f"  dates: {valid_dates}, request: {request}")
+
+        result = []
+        for valid_dates, request in full_requests:
+            for grib in index.retrieve(valid_dates, **request):
+                field = ekd.from_source("memory", grib)[0]
+                if flavour:
+                    field = flavour.apply(field)
+                result.append(field)
+
+        return FieldArray(result)
+
+
+def factorise(lst):
+    """Factorise a list of (dates, request) tuples by merging dates with identical requests."""
+    content = dict()
+
+    d = defaultdict(list)
+    for dates, request in lst:
+        assert isinstance(request, dict), type(request)
+        key = hashlib.md5(json.dumps(request, sort_keys=True).encode()).hexdigest()
+        content[key] = request
+        d[key] += dates
+
+    res = []
+    for key, dates in d.items():
+        dates = list(sorted(set(dates)))
+        res.append((dates, content[key]))
+    return res

anemoi/datasets/create/sources/hindcasts.py

@@ -14,7 +14,8 @@ from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.mars import mars
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOGGER = logging.getLogger(__name__)
 
@@ -37,57 +38,57 @@ def _to_list(x: list | tuple | Any) -> list[Any]:
     return [x]
 
 
-@legacy_source(__file__)
-def hindcasts(context: Any, dates: list[Any], **request: dict[str, Any]) -> MultiFieldList:
-    """Generates hindcast requests based on the provided dates and request parameters.
-
-    Parameters
-    ----------
-    context : Any
-        The context containing the dates provider and trace method.
-    dates : List[Any]
-        A list of dates for which to generate hindcast requests.
-    request : Dict[str, Any]
-        Additional request parameters.
-
-    Returns
-    -------
-    MultiFieldList
-        A MultiFieldList containing the hindcast data.
-    """
-    from anemoi.datasets.dates import HindcastsDates
-
-    provider = context.dates_provider
-    assert isinstance(provider, HindcastsDates)
-
-    context.trace("H️", f"hindcasts {len(dates)=}")
-
-    request["param"] = _to_list(request["param"])
-    request["step"] = _to_list(request.get("step", 0))
-    request["step"] = [int(_) for _ in request["step"]]
-
-    context.trace("H️", f"hindcast {request}")
-
-    requests = []
-    for d in dates:
-        r = request.copy()
-        hindcast = provider.mapping[d]
-        r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
-        r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
-        r["time"] = hindcast.refdate.strftime("%H")
-        r["step"] = hindcast.step
-        requests.append(r)
-
-    if len(requests) == 0:
-        return MultiFieldList([])
-
-    return mars(
-        context,
-        dates,
-        *requests,
-        date_key="hdate",
-        request_already_using_valid_datetime=True,
-    )
-
-
-execute = hindcasts
+@source_registry.register("hindcasts")
+class HindcastsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[Any], **request: dict[str, Any]) -> MultiFieldList:
+        """Generates hindcast requests based on the provided dates and request parameters.
+
+        Parameters
+        ----------
+        context : Any
+            The context containing the dates provider and trace method.
+        dates : List[Any]
+            A list of dates for which to generate hindcast requests.
+        request : Dict[str, Any]
+            Additional request parameters.
+
+        Returns
+        -------
+        MultiFieldList
+            A MultiFieldList containing the hindcast data.
+        """
+        from anemoi.datasets.dates import HindcastsDates
+
+        provider = context.dates_provider
+        assert isinstance(provider, HindcastsDates)
+
+        context.trace("H️", f"hindcasts {len(dates)=}")
+
+        request["param"] = _to_list(request["param"])
+        request["step"] = _to_list(request.get("step", 0))
+        request["step"] = [int(_) for _ in request["step"]]
+
+        context.trace("H️", f"hindcast {request}")
+
+        requests = []
+        for d in dates:
+            r = request.copy()
+            hindcast = provider.mapping[d]
+            r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
+            r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
+            r["time"] = hindcast.refdate.strftime("%H")
+            r["step"] = hindcast.step
+            requests.append(r)
+
+        if len(requests) == 0:
+            return MultiFieldList([])
+
+        return mars(
+            context,
+            dates,
+            *requests,
+            date_key="hdate",
+            request_already_using_valid_datetime=True,
+        )

anemoi/datasets/create/sources/legacy.py

@@ -8,16 +8,13 @@
 # nor does it submit to any jurisdiction.
 
 
-import inspect
 import logging
-import os
-from collections.abc import Callable
+from abc import abstractmethod
 from typing import Any
 
-from anemoi.datasets.create.input.template import resolve
+from anemoi.datasets.create.input.context import Context
 
 from ..source import Source
-from . import source_registry
 
 LOG = logging.getLogger(__name__)
 
@@ -27,7 +24,7 @@ class LegacySource(Source):
 
     Parameters
     ----------
-    context : Any
+    context : Context
         The context in which the source is created.
     *args : tuple
         Positional arguments.
@@ -35,64 +32,15 @@ class LegacySource(Source):
         Keyword arguments.
     """
 
-    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, context: Context, *args: Any, **kwargs: Any) -> None:
         super().__init__(context, *args, **kwargs)
         self.args = args
         self.kwargs = kwargs
 
+    @staticmethod
+    @abstractmethod
+    def _execute(context, *args, **kwargs):
+        pass
 
-class legacy_source:
-    """A decorator class for legacy sources.
-
-    Parameters
-    ----------
-    name : str
-        The name of the legacy source.
-    """
-
-    def __init__(self, name: str) -> None:
-        name, _ = os.path.splitext(os.path.basename(name))
-        self.name = name
-
-    def __call__(self, execute: Callable) -> Callable:
-        """Call method to wrap the execute function.
-
-        Parameters
-        ----------
-        execute : function
-            The execute function to be wrapped.
-
-        Returns
-        -------
-        function
-            The wrapped execute function.
-        """
-        this = self
-        name = f"Legacy{self.name.title()}Source"
-        source = ".".join([execute.__module__, execute.__name__])
-
-        def execute_wrapper(self, dates) -> Any:
-            """Wrapper method to call the execute function."""
-
-            args, kwargs = resolve(self.context, (self.args, self.kwargs))
-
-            try:
-                return execute(self.context, dates, *args, **kwargs)
-            except TypeError:
-                LOG.error(f"Error executing source {this.name} from {source}")
-                LOG.error(f"Function signature is: {inspect.signature(execute)}")
-                LOG.error(f"Arguments are: {args=}, {kwargs=}")
-                raise
-
-        klass = type(
-            name,
-            (LegacySource,),
-            {
-                "execute": execute_wrapper,
-                "_source": source,
-            },
-        )
-
-        source_registry.register(self.name)(klass)
-
-        return execute
+    def execute(self, dates: Any) -> Any:
+        return self._execute(self.context, dates, *self.args, **self.kwargs)