anemoi-datasets 0.5.11__tar.gz → 0.5.13__tar.gz

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
     - --force-single-line-imports
     - --profile black
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.7.2
+  rev: v0.8.1
   hooks:
   - id: ruff
     args:

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/CHANGELOG.md

@@ -10,12 +10,20 @@ Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-datasets/compare/0.5.8...HEAD)
 
+## Changed
+
+- Fix metadata serialization handling of numpy.integer (#140)
+- Fix negative variance for constant variables (#148)
+- Fix cutout slicing of grid dimension (#145)
+- Use cKDTree instead of KDTree
+- Implement 'complement' feature
+- Add ability to patch xarrays (#160)
 
 ### Added
 
 - Call filters from anemoi-transform
-- make test optional when adls is not installed Pull request #110
-
+- Make test optional when adls is not installed Pull request #110
+- Add wz_to_w, orog_to_z, and sum filters (#149)
 
 ## [0.5.8](https://github.com/ecmwf/anemoi-datasets/compare/0.5.7...0.5.8) - 2024-10-26
 

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: anemoi-datasets
-Version: 0.5.11
+Version: 0.5.13
 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

anemoi_datasets-0.5.13/docs/building/filters/orog_to_z.rst

@@ -0,0 +1,17 @@
+###########
+ orog_to_z
+###########
+
+The ``orog_to_z`` filter converts orography (in meters) to surface
+geopotential height (m^2/s^2) using the equation:
+
+.. math::
+
+   z &= g \cdot \textrm{orog}\\
+   g &= 9.80665\ m \cdot s^{-1}
+
+This filter needs to follow a source that provides orography, which is
+replaced by surface geopotential height.
+
+.. literalinclude:: yaml/orog_to_z.yaml
+   :language: yaml

anemoi_datasets-0.5.13/docs/building/filters/sum.rst

@@ -0,0 +1,13 @@
+#####
+ sum
+#####
+
+The ``sum`` filter computes the sum over multiple variables. This can be
+useful for computing total precipitation from its components (snow,
+rain) or summing the components of total column integrated water. This
+filter needs to follow a source that provides the list of variables to
+be summed up. These variables are removed by the filter and replaced by
+a single summed variable.
+
+.. literalinclude:: yaml/sum.yaml
+   :language: yaml

anemoi_datasets-0.5.13/docs/building/filters/wz_to_w.rst

@@ -0,0 +1,12 @@
+#########
+ wz_to_w
+#########
+
+The ``wz_to_w`` filter converts geometric vertical velocity (provided in
+m/s) to vertical velocity in pressure coordinates (Pa/s). This filter
+needs to follow a source that provides geometric vertical velocity.
+Geometric vertical velocity is removed by the filter and pressure
+vertical velocity is added.
+
+.. literalinclude:: yaml/wz_to_w.yaml
+   :language: yaml

anemoi_datasets-0.5.13/docs/building/filters/yaml/orog_to_z.yaml

@@ -0,0 +1,10 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load an orography variable
+
+    - orog_to_z:
+        orog: orog  # Name of orography (input) variable
+        z: z        # Name of z (output) variable

anemoi_datasets-0.5.13/docs/building/filters/yaml/sum.yaml

@@ -0,0 +1,13 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load the variables to be summed
+
+    - sum:
+        params:  # List of input variables
+          variable1
+          variable2
+          variable3
+        output: variable_total    # Name of output variable

anemoi_datasets-0.5.13/docs/building/filters/yaml/wz_to_w.yaml

@@ -0,0 +1,10 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load geometric vertical velocity
+
+    - wz_to_w:
+        wz: wz  # Name of geometric vertical velocity (input) variable
+        x: z    # Name of pressure vertical velocity (output) variable

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/building/filters.rst

@@ -15,8 +15,11 @@ Filters are used to modify the data or metadata in a dataset.
    :maxdepth: 1
 
    filters/select
+   filters/orog_to_z
    filters/rename
    filters/rotate_winds
+   filters/sum
    filters/unrotate_winds
+   filters/wz_to_w
    filters/noop
    filters/empty

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/building/introduction.rst

@@ -10,7 +10,7 @@ file, which is a YAML file that describes sources of meteorological
 fields as well as the operations to perform on them, before they are
 written to a zarr file. The input of the process is a range of dates and
 some options to control the layout of the output. Statistics will be
-computed as the dataset is build, and stored in the metadata, with other
+computed as the dataset is built, and stored in the metadata, with other
 information such as the the locations of the grid points, the list of
 variables, etc.
 
@@ -24,35 +24,35 @@ variables, etc.
 
 date
    Throughout this document, the term `date` refers to a date and time,
-   not just a date. A training dataset is covers a continuous range of
+   not just a date. A training dataset covers a continuous range of
    dates with a given frequency. Missing dates are still part of the
-   dataset, but the data are missing and marked as such using NaNs.
-   Dates are always in UTC, and refer to date at which the data is
-   valid. For accumulations and fluxes, that would be the end of the
-   accumulation period.
+   dataset, but missing data are marked as such using NaNs. Dates are
+   always in UTC, and refer to date at which the data is valid. For
+   accumulations and fluxes, that would be the end of the accumulation
+   period.
 
 variable
-   A `variable` is meteorological parameter, such as temperature, wind,
-   etc. Multilevel parameters are treated as separate variables, one for
-   each level. For example, temperature at 850 hPa and temperature at
-   500 hPa will be treated as two separate variables (`t_850` and
-   `t_500`).
+   A `variable` is a meteorological parameter, such as temperature,
+   wind, etc. Multilevel parameters are treated as separate variables,
+   one for each level. For example, temperature at 850 hPa and
+   temperature at 500 hPa will be treated as two separate variables
+   (`t_850` and `t_500`).
 
 field
-   A `field` is a variable at a given date. It is represented by a array
-   of values at each grid point.
+   A `field` is a variable at a given date. It is represented by an
+   array of values at each grid point.
 
 source
-   The `source` is a software component that given a list of dates and
-   variables will return the corresponding fields. A example of source
+   The `source` is a software component that, given a list of dates and
+   variables will return the corresponding fields. An example of source
    is ECMWF's MARS archive, a collection of GRIB or NetCDF files, a
    database, etc. See :ref:`sources` for more information.
 
 filter
    A `filter` is a software component that takes as input the output of
-   a source or the output of another filter can modify the fields and/or
-   their metadata. For example, typical filters are interpolations,
-   renaming of variables, etc. See :ref:`filters` for more information.
+   a source or another filter and can modify the fields and/or their
+   metadata. For example, typical filters are interpolations, renaming
+   of variables, etc. See :ref:`filters` for more information.
 
 ************
  Operations
@@ -62,19 +62,20 @@ In order to build a training dataset, sources and filters are combined
 using the following operations:
 
 join
-   The join is the process of combining several sources data. Each
-   source is expected to provide different variables at the same dates.
+   The join is the process of combining several sources of data. Each
+   source is expected to provide different variables for the same of
+   dates.
 
 pipe
    The pipe is the process of transforming fields using filters. The
-   first step of a pipe is typically a source, a join or another pipe.
-   The following steps are filters.
+   first step of a pipe is typically a source, a join, or another pipe.
+   This can subsequently followed by more filters.
 
 concat
    The concatenation is the process of combining different sets of
-   operation that handle different dates. This is typically used to
-   build a dataset that spans several years, when the several sources
-   are involved, each providing a different period.
+   operations that handle different dates. This is typically used to
+   build a dataset that spans several years, when several sources are
+   involved, each providing data for different period.
 
 Each operation is considered as a :ref:`source <sources>`, therefore
 operations can be combined to build complex datasets.
@@ -87,7 +88,7 @@ First recipe
 ============
 
 The simplest `recipe` file must contain a ``dates`` section and an
-``input`` section. The latter must contain a `source` In that case, the
+``input`` section. The latter must contain a `source`. In that case, the
 source is ``mars``
 
 .. literalinclude:: yaml/building1.yaml
@@ -132,15 +133,15 @@ This will build the following dataset:
 Adding some forcing variables
 =============================
 
-When training a data-driven models, some forcing variables may be
+When training a data-driven model, some forcing variables may be
 required such as the solar radiation, the time of day, the day in the
 year, etc.
 
-These are provided by the ``forcings`` source. In that example, we add a
-few of them. The `template` option is used to point to another source,
-in that case the first instance of ``mars``. This source is used to get
-information about the grid points, as some of the forcing variables are
-grid dependent.
+These are provided by the ``forcings`` source. Let us add a few of them
+to the above example. The `template` option is used to point to another
+source, in that case the first instance of ``mars``. This source is used
+to get information about the grid points, as some of the forcing
+variables are grid dependent.
 
 .. literalinclude:: yaml/building3.yaml
    :language: yaml

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/building/sources/yaml/accumulations1.yaml

@@ -1,6 +1,6 @@
 input:
   accumulations:
-    accumulations_period: 6
+    accumulation_period: 6
     class: ea
     param: [tp, cp, sf]
     levtype: sfc

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/building/sources/yaml/accumulations2.yaml

@@ -1,6 +1,6 @@
 input:
   accumulations:
-    accumulations_period: [6, 12]
+    accumulation_period: [6, 12]
     class: od
     param: [tp, cp, sf]
     levtype: sfc

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/index.rst

@@ -27,6 +27,7 @@ datasets <building-introduction>`.
 
 -  :doc:`overview`
 -  :doc:`installing`
+-  :doc:`naming_conventions`
 
 .. toctree::
    :maxdepth: 1
@@ -34,6 +35,7 @@ datasets <building-introduction>`.
 
    overview
    installing
+   naming_conventions
 
 **Using training datasets**
 
@@ -43,6 +45,7 @@ datasets <building-introduction>`.
 -  :doc:`using/subsetting`
 -  :doc:`using/combining`
 -  :doc:`using/selecting`
+-  :doc:`using/ensembles`
 -  :doc:`using/grids`
 -  :doc:`using/zip`
 -  :doc:`using/statistics`
@@ -63,6 +66,7 @@ datasets <building-introduction>`.
    using/subsetting
    using/combining
    using/selecting
+   using/ensembles
    using/grids
    using/zip
    using/statistics

anemoi_datasets-0.5.13/docs/naming_conventions.rst

@@ -0,0 +1,8 @@
+####################
+ Naming conventions
+####################
+
+This is the current naming conventions for datasets :
+
+`Anemoi Naming Conventions
+<https://anemoi-registry.readthedocs.io/en/latest/naming-conventions.html>`_

anemoi_datasets-0.5.13/docs/using/code/complement1_.py

@@ -0,0 +1,6 @@
+open_dataset(
+    complement=dataset1,
+    source=dataset2,
+    what="variables",
+    interpolate="nearest",
+)

anemoi_datasets-0.5.13/docs/using/code/complement2_.py

@@ -0,0 +1,12 @@
+open_dataset(
+    cutout=[
+        {
+            "complement": lam_dataset,
+            "source": global_dataset,
+            "interpolate": "nearest",
+        },
+        {
+            "dataset": global_dataset,
+        },
+    ]
+)

anemoi_datasets-0.5.13/docs/using/code/complement3_.py

@@ -0,0 +1,4 @@
+open_dataset(
+    complement=dataset1,
+    source=dataset2,
+)

anemoi_datasets-0.5.13/docs/using/code/number1_.py

@@ -0,0 +1,4 @@
+ds = open_dataset(
+    dataset,
+    number=1,
+)

anemoi_datasets-0.5.13/docs/using/code/number2_.py

@@ -0,0 +1,4 @@
+ds = open_dataset(
+    dataset,
+    number=[1, 3, 5],
+)

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/using/combining.rst

@@ -182,3 +182,32 @@ The difference can be seen at the boundary between the two grids:
 To debug the combination, you can pass `plot=True` to the `cutout`
 function (when running from a Notebook), of use `plot="prefix"` to save
 the plots to series of PNG files in the current directory.
+
+.. _complement:
+
+************
+ complement
+************
+
+That feature will interpolate the variables of `dataset2` that are not
+in `dataset1` to the grid of `dataset1` , add them to the list of
+variable of `dataset1` and return the result.
+
+.. literalinclude:: code/complement1_.py
+
+Currently ``what`` can only be ``variables`` and can be omitted.
+
+The value for ``interpolate`` can be one of ``none`` (default) or
+``nearest``. In the case of ``none``, the grids of the two datasets must
+match.
+
+This feature was originally designed to be used in conjunction with
+``cutout``, where `dataset1` is the lam, and `dataset2` is the global
+dataset.
+
+.. literalinclude:: code/complement2_.py
+
+Another use case is to simply bring all non-overlapping variables of a
+dataset into an other:
+
+.. literalinclude:: code/complement3_.py

anemoi_datasets-0.5.13/docs/using/ensembles.rst

@@ -0,0 +1,27 @@
+.. _selecting-members:
+
+###################
+ Selecting members
+###################
+
+This section describes how to subset data that are part of an ensemble.
+To combine ensembles, see :ref:`ensembles` in the
+:ref:`combining-datasets` section.
+
+.. _number:
+
+If a dataset is an ensemble, you can select one or more specific members
+using the `number` option. You can also use ``numbers`` (which is an
+alias for ``number``), and ``member`` (or ``members``). The difference
+between the two is that ``number`` is **1-based**, while ``member`` is
+**0-based**.
+
+Select a single element:
+
+.. literalinclude:: code/number1_.py
+   :language: python
+
+... or a list:
+
+.. literalinclude:: code/number2_.py
+   :language: python

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/docs/using/selecting.rst

@@ -67,6 +67,28 @@ You can also rename variables:
 This will be useful when you join datasets and do not want variables
 from one dataset to override the ones from the other.
 
+********
+ number
+********
+
+If a dataset is an ensemble, you can select one or more specific members
+using the `number` option. You can also use ``numbers`` (which is an
+alias for ``number``), and ``member`` (or ``members``). The difference
+between the two is that ``number`` is **1-based**, while ``member`` is
+**0-based**.
+
+Select a single element:
+
+.. literalinclude:: code/number1_.py
+   :language: python
+
+... or a list:
+
+.. literalinclude:: code/number2_.py
+   :language: python
+
+.. _rescale:
+
 *********
  rescale
 *********
@@ -87,7 +109,9 @@ rescale the data.
 .. warning::
 
    When providing units, the library assumes that the mapping between
-   them is a linear transformation. No check is does to ensure this is
+   them is a linear transformation. No check is done to ensure this is
    the case.
 
 .. _cfunits: https://github.com/NCAS-CMS/cfunits
+
+.. _number:

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/src/anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.11'
-__version_tuple__ = version_tuple = (0, 5, 11)
+__version__ = version = '0.5.13'
+__version_tuple__ = version_tuple = (0, 5, 13)

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/src/anemoi/datasets/create/__init__.py

@@ -622,10 +622,14 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
 
         check_shape(cube, dates, dates_in_data)
 
-        def check_dates_in_data(lst, lst2):
-            lst2 = [np.datetime64(_) for _ in lst2]
-            lst = [np.datetime64(_) for _ in lst]
-            assert lst == lst2, ("Dates in data are not the requested ones:", lst, lst2)
+        def check_dates_in_data(dates_in_data, requested_dates):
+            requested_dates = [np.datetime64(_) for _ in requested_dates]
+            dates_in_data = [np.datetime64(_) for _ in dates_in_data]
+            assert dates_in_data == requested_dates, (
+                "Dates in data are not the requested ones:",
+                dates_in_data,
+                requested_dates,
+            )
 
         check_dates_in_data(dates_in_data, dates)
 

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/src/anemoi/datasets/create/check.py

@@ -58,7 +58,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+)-?([a-zA-Z0-9-]+)?$"
+        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h|\d+m)-v(\d+)-?([a-zA-Z0-9-]+)?$"
         match = re.match(pattern, name)
 
         if not match:

{anemoi_datasets-0.5.11 → anemoi_datasets-0.5.13}/src/anemoi/datasets/create/functions/__init__.py

@@ -22,6 +22,7 @@ def assert_is_fieldlist(obj):
 def import_function(name, kind):
 
     from anemoi.transform.filters import filter_registry
+    from anemoi.transform.sources import source_registry
 
     name = name.replace("-", "_")
 
@@ -45,7 +46,20 @@ def import_function(name, kind):
         if filter_registry.lookup(name, return_none=True):
 
             def proc(context, data, *args, **kwargs):
-                return filter_registry.create(name, *args, **kwargs)(data)
+                filter = filter_registry.create(name, *args, **kwargs)
+                filter.context = context
+                # filter = filter_registry.create(context, name, *args, **kwargs)
+                return filter.forward(data)
+
+            return proc
+
+    if kind == "sources":
+        if source_registry.lookup(name, return_none=True):
+
+            def proc(context, data, *args, **kwargs):
+                source = source_registry.create(name, *args, **kwargs)
+                # source = source_registry.create(context, name, *args, **kwargs)
+                return source.forward(data)
 
             return proc
 

anemoi_datasets-0.5.13/src/anemoi/datasets/create/functions/filters/orog_to_z.py

@@ -0,0 +1,58 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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.
+
+
+from collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+class NewDataField:
+    def __init__(self, field, data, new_name):
+        self.field = field
+        self.data = data
+        self.new_name = new_name
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == "param":
+            return self.new_name
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, orog, z="z"):
+    """Convert orography [m] to z (geopotential height)"""
+    result = FieldArray()
+
+    processed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param == orog:
+            key = tuple(key.items())
+
+            if param in processed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            output = f.to_numpy(flatten=True) * 9.80665
+            result.append(NewDataField(f, output, z))
+        else:
+            result.append(f)
+
+    return result

anemoi_datasets-0.5.13/src/anemoi/datasets/create/functions/filters/sum.py

@@ -0,0 +1,71 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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.
+
+
+from collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+class NewDataField:
+    def __init__(self, field, data, new_name):
+        self.field = field
+        self.data = data
+        self.new_name = new_name
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == "param":
+            return self.new_name
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, params, output):
+    """Computes the sum over a set of variables"""
+    result = FieldArray()
+
+    needed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param in params:
+            key = tuple(key.items())
+
+            if param in needed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            needed_fields[key][param] = f
+        else:
+            result.append(f)
+
+    for keys, values in needed_fields.items():
+
+        if len(values) != len(params):
+            raise ValueError("Missing fields")
+
+        s = None
+        for k, v in values.items():
+            c = v.to_numpy(flatten=True)
+            if s is None:
+                s = c
+            else:
+                s += c
+        result.append(NewDataField(values[list(values.keys())[0]], s, output))
+
+    return result