anemoi-datasets 0.5.6__py3-none-any.whl → 0.5.10__py3-none-any.whl

anemoi/datasets/data/misc.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import calendar
 import datetime
 import logging
@@ -235,16 +238,21 @@ def _auto_adjust(datasets, kwargs):
             if set(d.variables) != variables:
                 subset_kwargs[i]["select"] = sorted(variables)
 
+    if "start" or "end" in adjust_set:
+        common = datasets[0].dates
+        for d in datasets[0:]:
+            common = np.intersect1d(common, d.dates)
+
     if "start" in adjust_set:
         assert "start" not in kwargs, "Cannot use 'start' in adjust and kwargs"
-        start = max(d.dates[0] for d in datasets).astype(object)
+        start = min(common).astype(object)
         for i, d in enumerate(datasets):
             if start != d.dates[0]:
                 subset_kwargs[i]["start"] = start
 
     if "end" in adjust_set:
         assert "end" not in kwargs, "Cannot use 'end' in adjust and kwargs"
-        end = min(d.dates[-1] for d in datasets).astype(object)
+        end = max(common).astype(object)
         for i, d in enumerate(datasets):
             if end != d.dates[-1]:
                 subset_kwargs[i]["end"] = end
@@ -262,6 +270,7 @@ def _auto_adjust(datasets, kwargs):
 
 
 def _open_dataset(*args, **kwargs):
+
     sets = []
     for a in args:
         sets.append(_open(a))
@@ -302,6 +311,12 @@ def _open_dataset(*args, **kwargs):
         assert not sets, sets
         return concat_factory(args, kwargs).mutate()
 
+    if "merge" in kwargs:
+        from .merge import merge_factory
+
+        assert not sets, sets
+        return merge_factory(args, kwargs).mutate()
+
     if "ensemble" in kwargs:
         from .ensemble import ensemble_factory
 

anemoi/datasets/data/missing.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 

anemoi/datasets/data/rescale.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 

anemoi/datasets/data/select.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 
@@ -76,6 +79,10 @@ class Select(Forwards):
     def variables(self):
         return [self.dataset.variables[i] for i in self.indices]
 
+    @cached_property
+    def variables_metadata(self):
+        return {k: v for k, v in self.dataset.variables_metadata.items() if k in self.variables}
+
     @cached_property
     def name_to_index(self):
         return {k: i for i, k in enumerate(self.variables)}
@@ -108,13 +115,20 @@ class Rename(Forwards):
         super().__init__(dataset)
         for n in rename:
             assert n in dataset.variables, n
+
         self._variables = [rename.get(v, v) for v in dataset.variables]
+        self._variables_metadata = {rename.get(k, k): v for k, v in dataset.variables_metadata.items()}
+
         self.rename = rename
 
     @property
     def variables(self):
         return self._variables
 
+    @property
+    def variables_metadata(self):
+        return self._variables_metadata
+
     @cached_property
     def name_to_index(self):
         return {k: i for i, k in enumerate(self.variables)}

anemoi/datasets/data/statistics.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 

anemoi/datasets/data/stores.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
@@ -69,7 +71,7 @@ class S3Store(ReadOnlyStore):
     """
 
     def __init__(self, url, region=None):
-        from anemoi.utils.s3 import s3_client
+        from anemoi.utils.remote.s3 import s3_client
 
         _, _, self.bucket, self.key = url.split("/", 3)
         self.s3 = s3_client(self.bucket, region=region)
@@ -83,6 +85,41 @@ class S3Store(ReadOnlyStore):
         return response["Body"].read()
 
 
+class PlanetaryComputerStore(ReadOnlyStore):
+    """We write our own Store to access catalogs on Planetary Computer,
+    as it requires some extra arguements to use xr.open_zarr.
+    """
+
+    def __init__(self, data_catalog_id):
+        self.data_catalog_id = data_catalog_id
+
+    def __getitem__(self):
+        import planetary_computer
+        import pystac_client
+
+        catalog = pystac_client.Client.open(
+            "https://planetarycomputer.microsoft.com/api/stac/v1/",
+            modifier=planetary_computer.sign_inplace,
+        )
+        collection = catalog.get_collection(self.data_catalog_id)
+
+        asset = collection.assets["zarr-abfs"]
+
+        if "xarray:storage_options" in asset.extra_fields:
+            store = {
+                "store": asset.href,
+                "storage_options": asset.extra_fields["xarray:storage_options"],
+                **asset.extra_fields["xarray:open_kwargs"],
+            }
+        else:
+            store = {
+                "filename_or_obj": asset.href,
+                **asset.extra_fields["xarray:open_kwargs"],
+            }
+
+        return store
+
+
 class DebugStore(ReadOnlyStore):
     """A store to debug the zarr loading."""
 
@@ -119,6 +156,9 @@ def name_to_zarr_store(path_or_url):
         if len(bits) == 5 and (bits[1], bits[3], bits[4]) == ("s3", "amazonaws", "com"):
             s3_url = f"s3://{bits[0]}{parsed.path}"
             store = S3Store(s3_url, region=bits[2])
+        elif store.startswith("https://planetarycomputer.microsoft.com/"):
+            data_catalog_id = store.rsplit("/", 1)[-1]
+            store = PlanetaryComputerStore(data_catalog_id).__getitem__()
         else:
             store = HTTPStore(store)
 
@@ -302,6 +342,17 @@ class Zarr(Dataset):
             )
         ]
 
+    @cached_property
+    def constant_fields(self):
+        result = self.z.attrs.get("constant_fields")
+        if result is None:
+            LOG.warning("No 'constant_fields' attribute in %r, computing them", self)
+        return self.computed_constant_fields()
+
+    @property
+    def variables_metadata(self):
+        return self.z.attrs.get("variables_metadata", {})
+
     def __repr__(self):
         return self.path
 
@@ -313,6 +364,7 @@ class Zarr(Dataset):
             attrs=dict(self.z.attrs),
             chunks=self.chunks,
             dtype=str(self.dtype),
+            path=self.path,
         )
 
     def source(self, index):
@@ -331,6 +383,12 @@ class Zarr(Dataset):
         name, _ = os.path.splitext(os.path.basename(self.path))
         names.add(name)
 
+    def collect_supporting_arrays(self, collected, *path):
+        pass
+
+    def collect_input_sources(self, collected):
+        pass
+
 
 class ZarrWithMissingDates(Zarr):
     """A zarr dataset with missing dates."""
@@ -339,7 +397,7 @@ class ZarrWithMissingDates(Zarr):
         super().__init__(path)
 
         missing_dates = self.z.attrs.get("missing_dates", [])
-        missing_dates = set([np.datetime64(x) for x in missing_dates])
+        missing_dates = set([np.datetime64(x, "s") for x in missing_dates])
         self.missing_to_dates = {i: d for i, d in enumerate(self.dates) if d in missing_dates}
         self.missing = set(self.missing_to_dates)
 
@@ -392,6 +450,9 @@ class ZarrWithMissingDates(Zarr):
         return "zarr*"
 
 
+QUIET = set()
+
+
 def zarr_lookup(name, fail=True):
 
     if name.endswith(".zarr") or name.endswith(".zip"):
@@ -400,6 +461,9 @@ def zarr_lookup(name, fail=True):
     config = load_config()["datasets"]
 
     if name in config["named"]:
+        if name not in QUIET:
+            LOG.info("Opening `%s` as `%s`", name, config["named"][name])
+            QUIET.add(name)
         return config["named"][name]
 
     tried = []
@@ -413,6 +477,9 @@ def zarr_lookup(name, fail=True):
             if z is not None:
                 # Cache for next time
                 config["named"][name] = full
+                if name not in QUIET:
+                    LOG.info("Opening `%s` as `%s`", name, full)
+                    QUIET.add(name)
                 return full
         except zarr.errors.PathNotFoundError:
             pass

anemoi/datasets/data/subset.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 
@@ -132,6 +135,8 @@ class Subset(Forwards):
     @cached_property
     def frequency(self):
         dates = self.dates
+        if len(dates) < 2:
+            raise ValueError(f"Cannot determine frequency of a subset with less than two dates ({self.dates}).")
         return frequency_to_timedelta(dates[1].astype(object) - dates[0].astype(object))
 
     def source(self, index):

anemoi/datasets/data/unchecked.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 from functools import wraps
@@ -91,6 +94,11 @@ class Unchecked(Combined):
     def variables(self):
         raise NotImplementedError()
 
+    @property
+    @check("check_same_variables")
+    def variables_metadata(self):
+        raise NotImplementedError()
+
     @property
     @check("check_same_variables")
     def statistics(self):

anemoi/datasets/data/xy.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.
 
+
 import logging
 from functools import cached_property
 
@@ -18,15 +21,19 @@ LOG = logging.getLogger(__name__)
 
 class ZipBase(Combined):
 
+    def __init__(self, datasets, check_compatibility=True):
+        self._check_compatibility = check_compatibility
+        super().__init__(datasets)
+
     def swap_with_parent(self, parent):
         new_parents = [parent.clone(ds) for ds in self.datasets]
         return self.clone(new_parents)
 
     def clone(self, datasets):
-        return self.__class__(datasets)
+        return self.__class__(datasets, check_compatibility=self._check_compatibility)
 
     def tree(self):
-        return Node(self, [d.tree() for d in self.datasets])
+        return Node(self, [d.tree() for d in self.datasets], check_compatibility=self._check_compatibility)
 
     def __len__(self):
         return min(len(d) for d in self.datasets)
@@ -86,6 +93,10 @@ class ZipBase(Combined):
     def name_to_index(self):
         return tuple(d.name_to_index for d in self.datasets)
 
+    def check_compatibility(self, d1, d2):
+        if self._check_compatibility:
+            super().check_compatibility(d1, d2)
+
 
 class Zip(ZipBase):
     pass
@@ -110,7 +121,9 @@ def xy_factory(args, kwargs):
 
     assert len(datasets) == 2
 
-    return XY(datasets)._subset(**kwargs)
+    check_compatibility = kwargs.pop("check_compatibility", True)
+
+    return XY(datasets, check_compatibility=check_compatibility)._subset(**kwargs)
 
 
 def zip_factory(args, kwargs):
@@ -122,4 +135,6 @@ def zip_factory(args, kwargs):
     datasets = [_open(e) for e in zip]
     datasets, kwargs = _auto_adjust(datasets, kwargs)
 
-    return Zip(datasets)._subset(**kwargs)
+    check_compatibility = kwargs.pop("check_compatibility", True)
+
+    return Zip(datasets, check_compatibility=check_compatibility)._subset(**kwargs)

anemoi/datasets/dates/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts.
+# (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.
@@ -8,6 +10,8 @@
 
 import datetime
 import warnings
+from functools import reduce
+from math import gcd
 
 # from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import DateTimes
@@ -193,18 +197,16 @@ class HindcastsDates(DatesProvider):
 
         dates = sorted(dates)
 
-        mindelta = None
+        deltas = set()
         for a, b in zip(dates, dates[1:]):
             delta = b - a
             assert isinstance(delta, datetime.timedelta), delta
-            if mindelta is None:
-                mindelta = delta
-            else:
-                mindelta = min(mindelta, delta)
+            deltas.add(delta)
 
+        mindelta_seconds = reduce(gcd, [int(delta.total_seconds()) for delta in deltas])
+        mindelta = datetime.timedelta(seconds=mindelta_seconds)
         self.frequency = mindelta
         assert mindelta.total_seconds() > 0, mindelta
-
         print("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", dates[0], dates[-1], mindelta)
 
         # Use all values between start and end by frequency, and set the ones that are missing

anemoi/datasets/dates/groups.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts.
+# (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.

anemoi/datasets/fields.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

anemoi/datasets/grids.py

@@ -1,12 +1,14 @@
-# (C) Copyright 2024 ECMWF.
+# (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.
-#
 
+
+import base64
 import logging
 
 import numpy as np
@@ -295,6 +297,88 @@ def thinning_mask(
     return np.array([i for i in indices])
 
 
+def outline(lats, lons, neighbours=5):
+    from scipy.spatial import KDTree
+
+    xyx = latlon_to_xyz(lats, lons)
+    grid_points = np.array(xyx).transpose()
+
+    # Use a KDTree to find the nearest points
+    _, indices = KDTree(grid_points).query(grid_points, k=neighbours)
+
+    # Centre of the Earth
+    zero = np.array([0.0, 0.0, 0.0])
+
+    outside = []
+
+    for i, (point, index) in enumerate(zip(grid_points, indices)):
+        inside = False
+        for j in range(1, neighbours):
+            t = Triangle3D(
+                grid_points[index[j]],
+                grid_points[index[(j + 1) % neighbours]],
+                grid_points[index[(j + 2) % neighbours]],
+            )
+            inside = t.intersect(zero, point)
+            if inside:
+                break
+
+        if not inside:
+            outside.append(i)
+
+    return outside
+
+
+def deserialise_mask(encoded):
+    import pickle
+    import zlib
+
+    packed = pickle.loads(zlib.decompress(base64.b64decode(encoded)))
+
+    mask = []
+    value = False
+    for count in packed:
+        mask.extend([value] * count)
+        value = not value
+    return np.array(mask, dtype=bool)
+
+
+def _serialise_mask(mask):
+    import pickle
+    import zlib
+
+    assert len(mask.shape) == 1
+    assert len(mask)
+
+    packed = []
+    last = mask[0]
+    count = 1
+
+    for value in mask[1:]:
+        if value == last:
+            count += 1
+        else:
+            packed.append(count)
+            last = value
+            count = 1
+
+    packed.append(count)
+
+    # We always start with an 'off' value
+    # So if the first value is 'on', we need to add a zero
+    if mask[0]:
+        packed.insert(0, 0)
+
+    return base64.b64encode(zlib.compress(pickle.dumps(packed))).decode("utf-8")
+
+
+def serialise_mask(mask):
+    result = _serialise_mask(mask)
+    # Make sure we can deserialise it
+    assert np.all(mask == deserialise_mask(result))
+    return result
+
+
 if __name__ == "__main__":
     global_lats, global_lons = np.meshgrid(
         np.linspace(90, -90, 90),

anemoi/datasets/testing.py

@@ -0,0 +1,60 @@
+# (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.
+
+
+# A collection of functions to support pytest testing
+
+import logging
+
+LOG = logging.getLogger(__name__)
+
+
+def assert_field_list(fs, size=None, start=None, end=None, constant=False, skip=None):
+    import numpy as np
+
+    if size is None:
+        assert len(fs) > 0, fs
+    else:
+        assert len(fs) == size, (len(fs), size)
+
+    first = fs[0]
+    last = fs[-1]
+
+    if constant:
+        # TODO: add a check for constant fields
+        pass
+    else:
+        assert start is None or first.metadata("valid_datetime") == start, (first.metadata("valid_datetime"), start)
+        assert end is None or last.metadata("valid_datetime") == end, (last.metadata("valid_datetime"), end)
+        print(first.datetime())
+
+    print(last.metadata())
+
+    first = first
+    latitudes, longitudes = first.grid_points()
+
+    assert len(latitudes.shape) == 1, latitudes.shape
+    assert len(longitudes.shape) == 1, longitudes.shape
+
+    assert len(latitudes) == len(longitudes), (len(latitudes), len(longitudes))
+    data = first.to_numpy(flatten=True)
+
+    assert len(data) == len(latitudes), (len(data), len(latitudes))
+
+    north = np.max(latitudes)
+    south = np.min(latitudes)
+    east = np.max(longitudes)
+    west = np.min(longitudes)
+
+    assert north >= south, (north, south)
+    assert east >= west, (east, west)
+    assert north <= 90, north
+    assert south >= -90, south
+    assert east <= 360, east
+    assert west >= -180, west

anemoi/datasets/utils/__init__.py

@@ -0,0 +1,8 @@
+# (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.

anemoi/datasets/utils/fields.py

@@ -1,11 +1,11 @@
-# (C) Copyright 2024 ECMWF.
+# (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.
-#
 
 
 class WrappedField: