anemoi-datasets 0.5.7__py3-none-any.whl → 0.5.11__py3-none-any.whl

anemoi/datasets/create/functions/sources/mars.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 datetime
+import re
 
 from anemoi.utils.humanize import did_you_mean
 from earthkit.data import from_source
@@ -31,6 +33,25 @@ def _date_to_datetime(d):
     return datetime.datetime.fromisoformat(d)
 
 
+def expand_to_by(x):
+
+    if isinstance(x, (str, int)):
+        return expand_to_by(str(x).split("/"))
+
+    if len(x) == 3 and x[1] == "to":
+        start = int(x[0])
+        end = int(x[2])
+        return list(range(start, end + 1))
+
+    if len(x) == 5 and x[1] == "to" and x[3] == "by":
+        start = int(x[0])
+        end = int(x[2])
+        by = int(x[4])
+        return list(range(start, end + 1, by))
+
+    return x
+
+
 def normalise_time_delta(t):
     if isinstance(t, datetime.timedelta):
         assert t == datetime.timedelta(hours=t.hours), t
@@ -42,25 +63,48 @@ def normalise_time_delta(t):
     return t
 
 
+def _normalise_time(t):
+    t = int(t)
+    if t < 100:
+        t * 100
+    return "{:04d}".format(t)
+
+
 def _expand_mars_request(request, date, request_already_using_valid_datetime=False, date_key="date"):
     requests = []
-    step = to_list(request.get("step", [0]))
-    for s in step:
+
+    user_step = to_list(expand_to_by(request.get("step", [0])))
+    user_time = None
+    user_date = None
+
+    if not request_already_using_valid_datetime:
+        user_time = request.get("time")
+        if user_time is not None:
+            user_time = to_list(user_time)
+            user_time = [_normalise_time(t) for t in user_time]
+
+        user_date = request.get(date_key)
+        if user_date is not None:
+            assert isinstance(user_date, str), user_date
+            user_date = re.compile("^{}$".format(user_date.replace("-", "").replace("?", ".")))
+
+    for step in user_step:
         r = request.copy()
 
         if not request_already_using_valid_datetime:
 
-            if isinstance(s, str) and "-" in s:
-                assert s.count("-") == 1, s
+            if isinstance(step, str) and "-" in step:
+                assert step.count("-") == 1, step
+
             # this takes care of the cases where the step is a period such as 0-24 or 12-24
-            hours = int(str(s).split("-")[-1])
+            hours = int(str(step).split("-")[-1])
 
             base = date - datetime.timedelta(hours=hours)
             r.update(
                 {
                     date_key: base.strftime("%Y%m%d"),
                     "time": base.strftime("%H%M"),
-                    "step": s,
+                    "step": step,
                 }
             )
 
@@ -69,12 +113,28 @@ def _expand_mars_request(request, date, request_already_using_valid_datetime=Fal
                 if isinstance(r[pproc], (list, tuple)):
                     r[pproc] = "/".join(str(x) for x in r[pproc])
 
+        if user_date is not None:
+            if not user_date.match(r[date_key]):
+                continue
+
+        if user_time is not None:
+            # It time is provided by the user, we only keep the requests that match the time
+            if r["time"] not in user_time:
+                continue
+
         requests.append(r)
 
+    # assert requests, requests
+
     return requests
 
 
-def factorise_requests(dates, *requests, request_already_using_valid_datetime=False, date_key="date"):
+def factorise_requests(
+    dates,
+    *requests,
+    request_already_using_valid_datetime=False,
+    date_key="date",
+):
     updates = []
     for req in requests:
         # req = normalise_request(req)
@@ -87,6 +147,9 @@ def factorise_requests(dates, *requests, request_already_using_valid_datetime=Fa
                 date_key=date_key,
             )
 
+    if not updates:
+        return
+
     compressed = Availability(updates)
     for r in compressed.iterate():
         for k, v in r.items():
@@ -177,13 +240,24 @@ MARS_KEYS = [
 ]
 
 
-def mars(context, dates, *requests, request_already_using_valid_datetime=False, date_key="date", **kwargs):
+def mars(
+    context,
+    dates,
+    *requests,
+    request_already_using_valid_datetime=False,
+    date_key="date",
+    **kwargs,
+):
+
     if not requests:
         requests = [kwargs]
 
     for r in requests:
+        param = r.get("param", [])
+        if not isinstance(param, (list, tuple)):
+            param = [param]
         # check for "Norway bug" where yaml transforms 'no' into False, etc.
-        for p in r.get("param", []):
+        for p in param:
             if p is False:
                 raise ValueError(
                     "'param' cannot be 'False'. If you wrote 'param: no' or 'param: off' in yaml, you may want to use quotes?"
@@ -197,12 +271,18 @@ def mars(context, dates, *requests, request_already_using_valid_datetime=False,
                     "'param' cannot be 'True'. If you wrote 'param: on' in yaml, you may want to use quotes?"
                 )
 
-    requests = factorise_requests(
-        dates,
-        *requests,
-        request_already_using_valid_datetime=request_already_using_valid_datetime,
-        date_key=date_key,
-    )
+    if len(dates) == 0:  # When using `repeated_dates`
+        assert len(requests) == 1, requests
+        assert "date" in requests[0], requests[0]
+        if isinstance(requests[0]["date"], datetime.date):
+            requests[0]["date"] = requests[0]["date"].strftime("%Y%m%d")
+    else:
+        requests = factorise_requests(
+            dates,
+            *requests,
+            request_already_using_valid_datetime=request_already_using_valid_datetime,
+            date_key=date_key,
+        )
 
     requests = list(requests)
 

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

@@ -1,11 +1,12 @@
-# (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.
-#
+
 
 from .xarray import load_many
 

anemoi/datasets/create/functions/sources/opendap.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.
-#
 
 
 from .xarray import load_many

anemoi/datasets/create/functions/sources/recentre.py

@@ -1,11 +1,12 @@
-# (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.
-#
+
 from copy import deepcopy
 
 from anemoi.datasets.compute.recentre import recentre as _recentre

anemoi/datasets/create/functions/sources/source.py

@@ -1,11 +1,12 @@
-# (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.
-#
+
 from earthkit.data import from_source
 
 from anemoi.datasets.create.utils import to_datetime_list

anemoi/datasets/create/functions/sources/tendencies.py

@@ -1,11 +1,12 @@
-# (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 datetime
 from collections import defaultdict
 

anemoi/datasets/create/functions/sources/xarray/__init__.py

@@ -1,18 +1,17 @@
-# (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 logging
 
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.data.stores import name_to_zarr_store
-from anemoi.datasets.utils.fields import NewMetadataField as NewMetadataField
 
 from .. import iterate_patterns
 from .fieldlist import XarrayFieldList
@@ -46,6 +45,12 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
 
     if isinstance(dataset, str) and ".zarr" in dataset:
         data = xr.open_zarr(name_to_zarr_store(dataset), **options)
+    elif "planetarycomputer" in dataset:
+        store = name_to_zarr_store(dataset)
+        if "store" in store:
+            data = xr.open_zarr(**store)
+        if "filename_or_obj" in store:
+            data = xr.open_dataset(**store)
     else:
         data = xr.open_dataset(dataset, **options)
 

anemoi/datasets/create/functions/sources/xarray/coordinates.py

@@ -1,11 +1,12 @@
-# (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 datetime
 import logging

anemoi/datasets/create/functions/sources/xarray/field.py

@@ -1,14 +1,16 @@
-# (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 datetime
 import logging
+from functools import cached_property
 
 from earthkit.data.core.fieldlist import Field
 from earthkit.data.core.fieldlist import math
@@ -45,8 +47,6 @@ class XArrayField(Field):
             This is actually a nD object, but the first dimensions are always 1.
             The other two dimensions are latitude and longitude.
         """
-        super().__init__(owner.array_backend)
-
         self.owner = owner
         self.selection = selection
 
@@ -85,7 +85,8 @@ class XArrayField(Field):
 
         return values  # .reshape(self.shape)
 
-    def _make_metadata(self):
+    @cached_property
+    def _metadata(self):
         return XArrayMetadata(self)
 
     def grid_points(self):

anemoi/datasets/create/functions/sources/xarray/fieldlist.py

@@ -1,11 +1,12 @@
-# (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 json
 import logging
@@ -13,10 +14,8 @@ import logging
 import yaml
 from earthkit.data.core.fieldlist import FieldList
 
-from .coordinates import is_scalar as is_scalar
 from .field import EmptyFieldList
 from .flavour import CoordinateGuesser
-from .metadata import XArrayMetadata as XArrayMetadata
 from .time import Time
 from .variable import FilteredVariable
 from .variable import Variable
@@ -60,6 +59,15 @@ class XarrayFieldList(FieldList):
                 else:
                     flavour = json.load(f)
 
+        if isinstance(flavour, dict):
+            flavour_coords = [coords["name"] for coords in flavour["rules"].values()]
+            ds_dims = [dim for dim in ds._dims]
+            for dim in ds_dims:
+                if dim in flavour_coords and dim not in ds._coord_names:
+                    ds = ds.assign_coords({dim: ds[dim]})
+                else:
+                    pass
+
         guess = CoordinateGuesser.from_flavour(ds, flavour)
 
         skip = set()

anemoi/datasets/create/functions/sources/xarray/flavour.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.
-#
 
 
 import logging

anemoi/datasets/create/functions/sources/xarray/grid.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.
-#
 
 
 import logging

anemoi/datasets/create/functions/sources/xarray/metadata.py

@@ -1,11 +1,12 @@
-# (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 logging
 from functools import cached_property

anemoi/datasets/create/functions/sources/xarray/time.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.
-#
 
 
 import datetime

anemoi/datasets/create/functions/sources/xarray/variable.py

@@ -1,18 +1,18 @@
-# (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 logging
 import math
 from functools import cached_property
 
 import numpy as np
-from earthkit.data.utils.array import ensure_backend
 
 from .field import XArrayField
 
@@ -29,7 +29,6 @@ class Variable:
         grid,
         time,
         metadata,
-        array_backend=None,
     ):
         self.ds = ds
         self.variable = variable
@@ -47,7 +46,6 @@ class Variable:
         self.by_name = {c.variable.name: c for c in coordinates}
 
         self.length = math.prod(self.shape)
-        self.array_backend = ensure_backend(array_backend)
 
     @property
     def name(self):
@@ -82,9 +80,8 @@ class Variable:
         )
 
     def __getitem__(self, i):
-        """
-        Get a 2D field from the variable
-        """
+        """Get a 2D field from the variable"""
+
         if i >= self.length:
             raise IndexError(i)
 
@@ -117,7 +114,7 @@ class Variable:
 
         variable = Variable(
             ds=self.ds,
-            var=self.variable.isel({k: i}),
+            variable=self.variable.isel({k: i}),
             coordinates=coordinates,
             grid=self.grid,
             time=self.time,

anemoi/datasets/create/functions/sources/xarray_kerchunk.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.
-#
 
 
 from earthkit.data.core.fieldlist import MultiFieldList

anemoi/datasets/create/functions/sources/xarray_zarr.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.
-#
 
 
 from .xarray import load_many

anemoi/datasets/create/functions/sources/zenodo.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.
-#
 
 
 from earthkit.data.core.fieldlist import MultiFieldList

anemoi/datasets/create/input/__init__.py

@@ -1,28 +1,14 @@
-# (C) Copyright 2023 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 datetime
-import itertools
+
 import logging
-import math
-import time
-from collections import defaultdict
 from copy import deepcopy
-from functools import cached_property
-from functools import wraps
-
-import numpy as np
-from anemoi.utils.dates import as_datetime as as_datetime
-from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
-
-from anemoi.datasets.dates import DatesProvider as DatesProvider
-from anemoi.datasets.fields import FieldArray as FieldArray
-from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
 
 from .trace import trace_select
 

anemoi/datasets/create/input/action.py

@@ -1,22 +1,17 @@
-# (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 logging
 from copy import deepcopy
 
-from anemoi.utils.dates import as_datetime as as_datetime
-from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
 from earthkit.data.core.order import build_remapping
 
-from anemoi.datasets.dates import DatesProvider as DatesProvider
-from anemoi.datasets.fields import FieldArray as FieldArray
-from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
-
 from .context import Context
 from .misc import is_function
 

anemoi/datasets/create/input/concat.py

@@ -1,11 +1,12 @@
-# (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 logging
 from copy import deepcopy
 from functools import cached_property

anemoi/datasets/create/input/context.py

@@ -1,22 +1,17 @@
-# (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 logging
 import textwrap
 
-from anemoi.utils.dates import as_datetime as as_datetime
-from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
 from anemoi.utils.humanize import plural
 
-from anemoi.datasets.dates import DatesProvider as DatesProvider
-from anemoi.datasets.fields import FieldArray as FieldArray
-from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
-
 from .trace import step
 from .trace import trace