anemoi-datasets 0.5.20__py3-none-any.whl → 0.5.21__py3-none-any.whl

anemoi/datasets/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.20'
-__version_tuple__ = version_tuple = (0, 5, 20)
+__version__ = version = '0.5.21'
+__version_tuple__ = version_tuple = (0, 5, 21)

anemoi/datasets/commands/grib-index.py

@@ -29,8 +29,6 @@ class GribIndexCmd(Command):
             The command parser to which arguments are added.
         """
 
-        from anemoi.datasets.create.sources.grib_index import KEYS
-
         command_parser.add_argument(
             "--index",
             help="Path to the index file to create or update",
@@ -52,7 +50,6 @@ class GribIndexCmd(Command):
         command_parser.add_argument(
             "--keys",
             help="GRIB keys to add to the index, separated by commas. If the list starts with a +, the keys are added to default list.",
-            default=",".join(KEYS),
         )
 
         command_parser.add_argument(

anemoi/datasets/commands/inspect.py

@@ -401,7 +401,7 @@ class Version:
             return
 
         if self.build_flags is None:
-            print("🪫 Dataset not initialized")
+            print("🪫 Dataset not initialised")
             return
 
         build_flags = self.build_flags
@@ -426,7 +426,7 @@ class Version:
             )
             start = self.initialised
             if self.initialised:
-                print(f"🕰️  Dataset initialized {when(start)}.")
+                print(f"🕰️  Dataset initialised {when(start)}.")
                 if built and latest:
                     speed = (latest - start) / built
                     eta = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + speed * (total - built)

anemoi/datasets/create/__init__.py

@@ -1215,7 +1215,7 @@ class _InitAdditions(Actor, HasRegistryMixin, AdditionsMixin):
         self.tmp_storage = build_storage(directory=self.tmp_storage_path, create=True)
         self.tmp_storage.delete()
         self.tmp_storage.create()
-        LOG.info(f"Dataset {self.tmp_storage_path} additions initialized.")
+        LOG.info(f"Dataset {self.tmp_storage_path} additions initialised.")
 
     def cleanup(self) -> None:
         """Clean up the temporary storage."""

anemoi/datasets/create/sources/forcings.py

@@ -25,7 +25,7 @@ def forcings(context: Any, dates: List[str], template: str, param: str) -> Any:
         The context in which the function is executed.
     dates : list
         List of dates for which data is to be loaded.
-    template : str
+    template : FieldList
         Template for the data source.
     param : str
         Parameter for the data source.

anemoi/datasets/create/sources/grib.py

@@ -17,9 +17,11 @@ from typing import Optional
 from typing import Union
 
 import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_grid
+from anemoi.transform.fields import new_fieldlist_from_list
 from anemoi.transform.flavour import RuleBasedFlavour
+from anemoi.transform.grids import grid_registry
 from earthkit.data import from_source
-from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.data.utils.patterns import Pattern
 
 from .legacy import legacy_source
@@ -27,169 +29,6 @@ from .legacy import legacy_source
 LOG = logging.getLogger(__name__)
 
 
-def _load(context: Any, name: str, record: Dict[str, Any]) -> tuple:
-    """Load data from a given source.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    name : str
-        The name of the data source.
-    record : dict of str to Any
-        The record containing source information.
-
-    Returns
-    -------
-    tuple
-        A tuple containing the data as a numpy array and the UUID of the HGrid.
-    """
-    ds = None
-
-    param = record["param"]
-
-    if "path" in record:
-        context.info(f"Using {name} from {record['path']} (param={param})")
-        ds = from_source("file", record["path"])
-
-    if "url" in record:
-        context.info(f"Using {name} from {record['url']} (param={param})")
-        ds = from_source("url", record["url"])
-
-    ds = ds.sel(param=param)
-
-    assert len(ds) == 1, f"{name} {param}, expected one field, got {len(ds)}"
-    ds = ds[0]
-
-    return ds.to_numpy(flatten=True), ds.metadata("uuidOfHGrid")
-
-
-class Geography:
-    """This class retrieves the latitudes and longitudes of unstructured grids,
-    and checks if the fields are compatible with the grid.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    latitudes : dict of str to Any
-        Latitude information.
-    longitudes : dict of str to Any
-        Longitude information.
-    """
-
-    def __init__(self, context: Any, latitudes: Dict[str, Any], longitudes: Dict[str, Any]) -> None:
-        """Initialize the Geography class.
-
-        Parameters
-        ----------
-        context : Any
-            The context in which the function is executed.
-        latitudes : dict of str to Any
-            Latitude information.
-        longitudes : dict of str to Any
-            Longitude information.
-        """
-        latitudes, uuidOfHGrid_lat = _load(context, "latitudes", latitudes)
-        longitudes, uuidOfHGrid_lon = _load(context, "longitudes", longitudes)
-
-        assert (
-            uuidOfHGrid_lat == uuidOfHGrid_lon
-        ), f"uuidOfHGrid mismatch: lat={uuidOfHGrid_lat} != lon={uuidOfHGrid_lon}"
-
-        context.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}")
-        assert len(latitudes) == len(longitudes)
-
-        self.uuidOfHGrid = uuidOfHGrid_lat
-        self.latitudes = latitudes
-        self.longitudes = longitudes
-        self.first = True
-
-    def check(self, field: Any) -> None:
-        """Check if the field is compatible with the grid.
-
-        Parameters
-        ----------
-        field : Any
-            The field to check.
-        """
-        if self.first:
-            # We only check the first field, for performance reasons
-            assert (
-                field.metadata("uuidOfHGrid") == self.uuidOfHGrid
-            ), f"uuidOfHGrid mismatch: {field.metadata('uuidOfHGrid')} != {self.uuidOfHGrid}"
-            self.first = False
-
-
-class AddGrid:
-    """An earth-kit.data.Field wrapper that adds grid information.
-
-    Parameters
-    ----------
-    field : Any
-        The field to wrap.
-    geography : Geography
-        The geography information.
-    """
-
-    def __init__(self, field: Any, geography: Geography) -> None:
-        """Initialize the AddGrid class.
-
-        Parameters
-        ----------
-        field : Any
-            The field to wrap.
-        geography : Geography
-            The geography information.
-        """
-        self._field = field
-
-        geography.check(field)
-
-        self._latitudes = geography.latitudes
-        self._longitudes = geography.longitudes
-
-    def __getattr__(self, name: str) -> Any:
-        """Get an attribute from the wrapped field.
-
-        Parameters
-        ----------
-        name : str
-            The name of the attribute.
-
-        Returns
-        -------
-        Any
-            The attribute value.
-        """
-        return getattr(self._field, name)
-
-    def __repr__(self) -> str:
-        """Get the string representation of the wrapped field.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        return repr(self._field)
-
-    def grid_points(self) -> tuple:
-        """Get the grid points (latitudes and longitudes).
-
-        Returns
-        -------
-        tuple
-            The latitudes and longitudes.
-        """
-        return self._latitudes, self._longitudes
-
-    @property
-    def resolution(self) -> str:
-        """Get the resolution of the grid."""
-        return "unknown"
-
-
 def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
     """Check if the dataset matches the expected number of fields.
 
@@ -243,9 +82,8 @@ def execute(
     context: Any,
     dates: List[Any],
     path: Union[str, List[str]],
-    latitudes: Optional[Dict[str, Any]] = None,
-    longitudes: Optional[Dict[str, Any]] = None,
     flavour: Optional[Union[str, Dict[str, Any]]] = None,
+    grid_definition: Optional[Dict[str, Any]] = None,
     *args: Any,
     **kwargs: Any,
 ) -> ekd.FieldList:
@@ -255,9 +93,8 @@ def execute(
         context (Any): The context in which the function is executed.
         dates (List[Any]): List of dates.
         path (Union[str, List[str]]): Path or list of paths to the GRIB files.
-        latitudes (Optional[Dict[str, Any]], optional): Latitude information. Defaults to None.
-        longitudes (Optional[Dict[str, Any]], optional): Longitude information. Defaults to None.
         flavour (Optional[Union[str, Dict[str, Any]]], optional): Flavour information. Defaults to None.
+        grid_definition (Optional[Dict[str, Any]]): Grid definition config to create a Grid object.
         *args (Any): Additional arguments.
         **kwargs (Any): Additional keyword arguments.
 
@@ -270,9 +107,10 @@ def execute(
     if flavour is not None:
         flavour = RuleBasedFlavour(flavour)
 
-    geography = None
-    if latitudes is not None and longitudes is not None:
-        geography = Geography(context, latitudes, longitudes)
+    if grid_definition is not None:
+        grid = grid_registry.from_config(grid_definition)
+    else:
+        grid = None
 
     ds = from_source("empty")
     dates = [d.isoformat() for d in dates]
@@ -295,8 +133,8 @@ def execute(
     if kwargs and not context.partial_ok:
         check(ds, given_paths, valid_datetime=dates, **kwargs)
 
-    if geography is not None:
-        ds = FieldArray([AddGrid(_, geography) for _ in ds])
+    if grid is not None:
+        ds = new_fieldlist_from_list([new_field_from_grid(f, grid) for f in ds])
 
     if len(ds) == 0:
         LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")

anemoi/datasets/create/writer.py

@@ -20,7 +20,7 @@ LOG = logging.getLogger(__name__)
 class ViewCacheArray:
     """A class that provides a caching mechanism for writing to a NumPy-like array.
 
-    The is initialized with a NumPy-like array, a shape and a list to reindex the first
+    The is initialised with a NumPy-like array, a shape and a list to reindex the first
     dimension. The array is used to store the final data, while the cache is used to
     temporarily store the data before flushing it to the array.
 

anemoi/datasets/data/misc.py

@@ -677,7 +677,7 @@ def _save_dataset(recipe: Dict[str, Any], zarr_path: str, n_workers: int = 1) ->
     # Workers return (date, subset) tuples.
     root = zarr.open(zarr_path, mode="a")
     initialize_zarr_store(root, full_ds, recipe)
-    print("Zarr store initialized.", flush=True)
+    print("Zarr store initialised.", flush=True)
 
     existing_dates = np.array(sorted(root["dates"]), dtype="datetime64[s]")
     all_dates = full_ds.dates

{anemoi_datasets-0.5.20.dist-info → anemoi_datasets-0.5.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-datasets
-Version: 0.5.20
+Version: 0.5.21
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -225,9 +225,10 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 License-File: LICENSE
-Requires-Dist: anemoi-transform>=0.1.6
-Requires-Dist: anemoi-utils[provenance]>=0.4.18
+Requires-Dist: anemoi-transform>=0.1.9
+Requires-Dist: anemoi-utils[provenance]>=0.4.19
 Requires-Dist: cfunits
+Requires-Dist: numcodecs<0.16
 Requires-Dist: numpy
 Requires-Dist: pyyaml
 Requires-Dist: semantic-version

{anemoi_datasets-0.5.20.dist-info → anemoi_datasets-0.5.21.dist-info}/RECORD

@@ -1,6 +1,6 @@
 anemoi/datasets/__init__.py,sha256=i_wsAT3ezEYF7o5dpqGrpoG4wmLS-QIBug18uJbSYMs,1065
 anemoi/datasets/__main__.py,sha256=ErwAqE3rBc7OaNO2JRsEOhWpB8ldjAt7BFSuRhbnlqQ,936
-anemoi/datasets/_version.py,sha256=0Ib1b9ijC26N78gsH6oTzFbmTu3c8De1ac8TlFGaEZU,513
+anemoi/datasets/_version.py,sha256=xzRAV3vnhRWoa5IGzmE1SpAqijHmFvxk2UzeNXSigRM,513
 anemoi/datasets/grids.py,sha256=Hhj1aOXHvDjmI46M_UlLSjCs1qYqxH-uqd_kapDSdbU,18134
 anemoi/datasets/testing.py,sha256=fy_JzavUwLlK_2rtXAT-UGUyo5gjyQW2y826zf334Wg,2645
 anemoi/datasets/commands/__init__.py,sha256=O5W3yHZywRoAqmRUioAr3zMCh0hGVV18wZYGvc00ioM,698
@@ -11,10 +11,10 @@ anemoi/datasets/commands/copy.py,sha256=UlvW9YIlP7jwKAY7TikVAfhkrjaQ9Kkxqfx8jEut
 anemoi/datasets/commands/create.py,sha256=5BXdPZMO-ZULBnEVgyeSS-IMy4p84HSyFVG855gqj3k,6598
 anemoi/datasets/commands/finalise-additions.py,sha256=2LqU7ke3i-yRQbjkgldX6e2QlyE-tKqp0b6QOhJF19g,1985
 anemoi/datasets/commands/finalise.py,sha256=-YtN9wFFDrM_i_V9YHoXZsajF3eAax-73Zsi4uHAFCI,1709
-anemoi/datasets/commands/grib-index.py,sha256=1aO4Lm_7ZpRF_lRRF-OEMZ37ByYPZfUAIFW7jOsIf7Q,3161
+anemoi/datasets/commands/grib-index.py,sha256=H9snsk1w2tL6ObBhlHVxFXxCm5sB-eUGrX772G8Nb2s,3057
 anemoi/datasets/commands/init-additions.py,sha256=2vMom5L38UvLLopzP2z-R_Fq31fU2uMvKXoDq5d8oI4,1931
 anemoi/datasets/commands/init.py,sha256=5IKyJ_hJA4lLIbpT88XtcGzXccHLSGwSoqVSvVJGxPg,2852
-anemoi/datasets/commands/inspect.py,sha256=dGx_3IQTNaV3iGfwhHp_G1fIecBWeRbOLU9ckrZ7epg,26578
+anemoi/datasets/commands/inspect.py,sha256=kaDHXP8Cv8PsGqEXUF5Yruf5OQHwOIkjCS0SNxMs6eg,26578
 anemoi/datasets/commands/load-additions.py,sha256=zqmRherIHXb5WIB4cnAuCBEsxFJmUpGjafvm6RtQ7Co,2004
 anemoi/datasets/commands/load.py,sha256=U1Y8qByjreu7H9WeX4G8-tyKsz48va10fkW1L4U4wWg,2034
 anemoi/datasets/commands/patch.py,sha256=Q9FDabWxlvK1QaeH4D9zhNpoSGB4h7EliWgcV76iFBs,1599
@@ -22,7 +22,7 @@ anemoi/datasets/commands/publish.py,sha256=7YusLCWYdVLuexZzvyh8ztYoBOBzVmve3uJs-
 anemoi/datasets/commands/scan.py,sha256=e5t_oxSi-II38TVQiMlWMJ8AZhDEBk5PcPD22DDbHfU,4008
 anemoi/datasets/compute/__init__.py,sha256=hCW0QcLHJmE-C1r38P27_ZOvCLNewex5iQEtZqx2ckI,393
 anemoi/datasets/compute/recentre.py,sha256=kwxDB8qpgOCFZSQJvjAmVcpH5zWsfk5FSoIureqNHd4,5915
-anemoi/datasets/create/__init__.py,sha256=D0a5Q-xv5_mtBTzEO_6IaWHCQUbZyH0NjHwEtixMCXs,50699
+anemoi/datasets/create/__init__.py,sha256=_aYkQ-5so9hBRW8rIj7Y-lQpn1ACY7PUpqadffsT0jU,50699
 anemoi/datasets/create/check.py,sha256=FrgyZP3Xyx4qXHl8_ZfM31fgNhcxMqxlE5oLweMDGU0,10003
 anemoi/datasets/create/chunks.py,sha256=kZV3dWoCuv3Bttc0wysJB7OPbXsD99exKyrrj4HGFwQ,4025
 anemoi/datasets/create/config.py,sha256=xrSlaY2p5zssfLIt8A1CP9WwJReSXVWBMQM7bT1aFbU,13448
@@ -34,7 +34,7 @@ anemoi/datasets/create/source.py,sha256=xoV8uH_y6aBSE4_PWuy5w7Q7cX-tGm8e-2xC9flS
 anemoi/datasets/create/testing.py,sha256=FzTSsbv_JBGViGrD1jT6z_T2yaA0KCrbJ3SCgp-rFPQ,2406
 anemoi/datasets/create/typing.py,sha256=Hs2uDG8ZVtQ-Q-5I9-W0Pik2p1hZH5-JPVjpJXRXP7M,484
 anemoi/datasets/create/utils.py,sha256=94JKPYcNSurA62yFAytW5dUFVz-r-fdwiPOkmu121pM,5572
-anemoi/datasets/create/writer.py,sha256=yjgic7FfVWuVKysi3HbaW6i2j4e831oCUdyD48Ua-MA,2203
+anemoi/datasets/create/writer.py,sha256=nZBJvYZ63g_c9FfL65bAeG10Y6bX2R7CgtZvY0kW3fI,2203
 anemoi/datasets/create/zarr.py,sha256=N9PGGD-dYvcc97BZjLVWm5XQeYTiK9gwvhtreRiQzBI,9437
 anemoi/datasets/create/filters/__init__.py,sha256=rhetUFPZKe-vwDIfMY33SbYrJMq909dUJRgBzuX6Q6E,901
 anemoi/datasets/create/filters/empty.py,sha256=Dw1kUnAlFt6b5ds0kmrw9Gak09XjSqF8m1_MpHZNV9I,1013
@@ -78,8 +78,8 @@ anemoi/datasets/create/sources/anemoi_dataset.py,sha256=2xJJTmKlv87F_2ECMKeehaeW
 anemoi/datasets/create/sources/constants.py,sha256=5O6d9tEuAmVjl5vNkNfmkaAjKXFlw1UjeueTsF1GZCI,1528
 anemoi/datasets/create/sources/eccc_fstd.py,sha256=8HK38f444HcWMvBhooP0XqTfMXYoCbN_8G9RI_Ne5rc,659
 anemoi/datasets/create/sources/empty.py,sha256=5mVIVRUwnBfE3zp-bvNA_imXCSpyds-4mewcI8HXAiY,1020
-anemoi/datasets/create/sources/forcings.py,sha256=877OZoXUoJncQ2_AAGSijwWqM-4kJJdxdIa6SFvZBUw,1216
-anemoi/datasets/create/sources/grib.py,sha256=nGjrKTQZB8GQCvfX7SgfdmcVJVSXF9kpb2GmSX7H7bM,8731
+anemoi/datasets/create/sources/forcings.py,sha256=PmcAd-nFZzILRxfvrptb45VLD9Nd2vjy0xeToeoKl1Y,1222
+anemoi/datasets/create/sources/grib.py,sha256=QGAxX103bOQRef8WkamK8g4DGMoo9ATLUfawzYO-vTk,4210
 anemoi/datasets/create/sources/grib_index.py,sha256=Pnm0RLga9lpD4MqVaZr7IqXMBlw1DtTIWZRfz7fq30Q,19026
 anemoi/datasets/create/sources/hindcasts.py,sha256=_4880rgd4AsRxlDXVi6dkh8mlKXrz2i27btVlmlMFjY,2611
 anemoi/datasets/create/sources/legacy.py,sha256=RJce-9TwmUUCFbgC8A3Dp61nSBfB8_lWti8WNoOMIcU,2652
@@ -122,7 +122,7 @@ anemoi/datasets/data/interpolate.py,sha256=-kSYwdjKH7zJtfITdbqdH6KyOFGVZDyHg4TaF
 anemoi/datasets/data/join.py,sha256=ZEHOsCecKBkKKH-vki404Sm7r7cV368ECO7PXPpay3s,9212
 anemoi/datasets/data/masked.py,sha256=giOvHLcGbLf6mZPqZjAxQd1kvydmkepDFh2EqchXLTQ,10213
 anemoi/datasets/data/merge.py,sha256=SvQhJHf-C-Kn7hEjFqomienk-epPPjMtoccRNCJpMtw,8733
-anemoi/datasets/data/misc.py,sha256=l8ZJwEjUf9CoPn-RhnilVDVTGypkVIzJ5xV77h8zzkc,20951
+anemoi/datasets/data/misc.py,sha256=6JOag2VJBJhz-eEEWL9F5nriEJyKz9aZL2P3FYcOSJo,20951
 anemoi/datasets/data/missing.py,sha256=ogfVDponbs0bGHMxps32Fj_fq4gT26R70yEMco5gdK8,12593
 anemoi/datasets/data/rescale.py,sha256=nGfJ5tWCncMJ7NMXkLbmt6z0ELrD6FxpbjJreQ3W91g,7004
 anemoi/datasets/data/select.py,sha256=Xs6uOzJL0CoOGeWA_E5_ukr8Jav2kXbZ41vhk7Vr8PE,8277
@@ -134,9 +134,9 @@ anemoi/datasets/data/xy.py,sha256=-jWzYismrK3eI3YCKIBpU1BCmraRncmVn0_2IUY--lk,75
 anemoi/datasets/dates/__init__.py,sha256=pEArHDQ7w5E0WC8Vvf9ypyKSdm6gnhoN9TmooITB7C4,13617
 anemoi/datasets/dates/groups.py,sha256=IOveL6IyTXZwEdXZEnRAnpu9pINY95VN7LzcpLfJ09E,10105
 anemoi/datasets/utils/__init__.py,sha256=hCW0QcLHJmE-C1r38P27_ZOvCLNewex5iQEtZqx2ckI,393
-anemoi_datasets-0.5.20.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
-anemoi_datasets-0.5.20.dist-info/METADATA,sha256=iGcUiyvv_JjdtsXo_oq_WgwweD1nPE0xAM8n5d0qeQU,16039
-anemoi_datasets-0.5.20.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-anemoi_datasets-0.5.20.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
-anemoi_datasets-0.5.20.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_datasets-0.5.20.dist-info/RECORD,,
+anemoi_datasets-0.5.21.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
+anemoi_datasets-0.5.21.dist-info/METADATA,sha256=wc23I3DuajyKdNV67zvhmAL8X5Uan-2ogszGdc9kVVs,16069
+anemoi_datasets-0.5.21.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+anemoi_datasets-0.5.21.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
+anemoi_datasets-0.5.21.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_datasets-0.5.21.dist-info/RECORD,,