anemoi-datasets 0.3.10__py3-none-any.whl → 0.4.2__py3-none-any.whl

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

@@ -7,21 +7,13 @@
 # nor does it submit to any jurisdiction.
 #
 import datetime
-import warnings
-from copy import deepcopy
-
-import climetlab as cml
-import numpy as np
-from climetlab.core.temporary import temp_file
-from climetlab.readers.grib.output import new_grib_output
-from climetlab.utils.availability import Availability
 
 from anemoi.datasets.create.functions.sources.mars import mars
 
 DEBUG = True
 
 
-def member(field):
+def _member(field):
     # Bug in eccodes has number=0 randomly
     number = field.metadata("number")
     if number is None:
@@ -29,355 +21,12 @@ def member(field):
     return number
 
 
-class Accumulation:
-    def __init__(self, out, /, param, date, time, number, step, frequency, **kwargs):
-        self.out = out
-        self.param = param
-        self.date = date
-        self.time = time
-        self.steps = step
-        self.number = number
-        self.values = None
-        self.seen = set()
-        self.startStep = None
-        self.endStep = None
-        self.done = False
-        self.frequency = frequency
-        self._check = None
-
-    @property
-    def key(self):
-        return (self.param, self.date, self.time, self.steps, self.number)
-
-    def check(self, field):
-        if self._check is None:
-            self._check = field.as_mars()
-
-            assert self.param == field.metadata("param"), (self.param, field.metadata("param"))
-            assert self.date == field.metadata("date"), (self.date, field.metadata("date"))
-            assert self.time == field.metadata("time"), (self.time, field.metadata("time"))
-            assert self.number == member(field), (self.number, member(field))
-
-            return
-
-        mars = field.as_mars()
-        keys1 = sorted(self._check.keys())
-        keys2 = sorted(mars.keys())
-
-        assert keys1 == keys2, (keys1, keys2)
-
-        for k in keys1:
-            if k not in ("step",):
-                assert self._check[k] == mars[k], (k, self._check[k], mars[k])
-
-    def write(self, template):
-
-        assert self.startStep != self.endStep, (self.startStep, self.endStep)
-        assert np.all(self.values >= 0), (np.amin(self.values), np.amax(self.values))
-
-        self.out.write(
-            self.values,
-            template=template,
-            stepType="accum",
-            startStep=self.startStep,
-            endStep=self.endStep,
-        )
-        self.values = None
-        self.done = True
-
-    def add(self, field, values):
-
-        self.check(field)
-
-        step = field.metadata("step")
-        if step not in self.steps:
-            return
-
-        if not np.all(values >= 0):
-            warnings.warn(f"Negative values for {field}: {np.amin(values)} {np.amax(values)}")
-
-        assert not self.done, (self.key, step)
-        assert step not in self.seen, (self.key, step)
-
-        startStep = field.metadata("startStep")
-        endStep = field.metadata("endStep")
-
-        if self.buggy_steps and startStep == endStep:
-            startStep = 0
-
-        assert step == endStep, (startStep, endStep, step)
-
-        self.compute(values, startStep, endStep)
-
-        self.seen.add(step)
-
-        if len(self.seen) == len(self.steps):
-            self.write(template=field)
-
-    @classmethod
-    def mars_date_time_steps(cls, dates, step1, step2, frequency, base_times, adjust_step):
-
-        # assert step1 > 0, (step1, step2, frequency)
-
-        for valid_date in dates:
-            base_date = valid_date - datetime.timedelta(hours=step2)
-            add_step = 0
-            if base_date.hour not in base_times:
-                if not adjust_step:
-                    raise ValueError(
-                        f"Cannot find a base time in {base_times} that validates on {valid_date.isoformat()} for step={step2}"
-                    )
-
-                while base_date.hour not in base_times:
-                    # print(f'{base_date=}, {base_times=}, {add_step=} {frequency=}')
-                    base_date -= datetime.timedelta(hours=1)
-                    add_step += 1
-
-            yield cls._mars_date_time_step(base_date, step1, step2, add_step, frequency)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self.key})"
-
-
-class AccumulationFromStart(Accumulation):
-    buggy_steps = True
-
-    def compute(self, values, startStep, endStep):
-
-        assert startStep == 0, startStep
-
-        if self.values is None:
-
-            self.values = np.copy(values)
-            self.startStep = 0
-            self.endStep = endStep
-
-        else:
-            assert endStep != self.endStep, (self.endStep, endStep)
-
-            if endStep > self.endStep:
-                # assert endStep - self.endStep == self.stepping, (self.endStep, endStep, self.stepping)
-                self.values = values - self.values
-                self.startStep = self.endStep
-                self.endStep = endStep
-            else:
-                # assert self.endStep - endStep == self.stepping, (self.endStep, endStep, self.stepping)
-                self.values = self.values - values
-                self.startStep = endStep
-
-            if not np.all(self.values >= 0):
-                warnings.warn(f"Negative values for {self.param}: {np.amin(self.values)} {np.amax(self.values)}")
-                self.values = np.maximum(self.values, 0)
-
-    @classmethod
-    def _mars_date_time_step(cls, base_date, step1, step2, add_step, frequency):
-        assert not frequency, frequency
-
-        steps = (step1 + add_step, step2 + add_step)
-        if steps[0] == 0:
-            steps = (steps[1],)
-
-        return (
-            base_date.year * 10000 + base_date.month * 100 + base_date.day,
-            base_date.hour * 100 + base_date.minute,
-            steps,
-        )
-
-
-class AccumulationFromLastStep(Accumulation):
-    buggy_steps = False
-
-    def compute(self, values, startStep, endStep):
-
-        assert endStep - startStep == self.frequency, (startStep, endStep, self.frequency)
-
-        if self.startStep is None:
-            self.startStep = startStep
-        else:
-            self.startStep = min(self.startStep, startStep)
-
-        if self.endStep is None:
-            self.endStep = endStep
-        else:
-            self.endStep = max(self.endStep, endStep)
-
-        if self.values is None:
-            self.values = np.zeros_like(values)
-
-        self.values += values
-
-    @classmethod
-    def _mars_date_time_step(cls, base_date, step1, step2, add_step, frequency):
-        assert frequency > 0, frequency
-        # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
-
-        steps = []
-        for step in range(step1 + frequency, step2 + frequency, frequency):
-            steps.append(step + add_step)
-        return (
-            base_date.year * 10000 + base_date.month * 100 + base_date.day,
-            base_date.hour * 100 + base_date.minute,
-            tuple(steps),
-        )
-
-
-def identity(x):
-    return x
-
-
-def compute_accumulations(
-    dates,
-    request,
-    user_accumulation_period=6,
-    data_accumulation_period=None,
-    patch=identity,
-    base_times=None,
-):
-    adjust_step = isinstance(user_accumulation_period, int)
-
-    if not isinstance(user_accumulation_period, (list, tuple)):
-        user_accumulation_period = (0, user_accumulation_period)
-
-    assert len(user_accumulation_period) == 2, user_accumulation_period
-    step1, step2 = user_accumulation_period
-    assert step1 < step2, user_accumulation_period
-
-    if base_times is None:
-        base_times = [0, 6, 12, 18]
-
-    base_times = [t // 100 if t > 100 else t for t in base_times]
-
-    AccumulationClass = AccumulationFromStart if data_accumulation_period in (0, None) else AccumulationFromLastStep
-
-    mars_date_time_steps = AccumulationClass.mars_date_time_steps(
-        dates,
-        step1,
-        step2,
-        data_accumulation_period,
-        base_times,
-        adjust_step,
-    )
-
-    request = deepcopy(request)
-
-    param = request["param"]
-    if not isinstance(param, (list, tuple)):
-        param = [param]
-
-    number = request.get("number", [0])
-    assert isinstance(number, (list, tuple))
-
-    frequency = data_accumulation_period
-
-    type_ = request.get("type", "an")
-    if type_ == "an":
-        type_ = "fc"
-
-    request.update({"type": type_, "levtype": "sfc"})
-
-    tmp = temp_file()
-    path = tmp.path
-    out = new_grib_output(path)
-
-    requests = []
-
-    accumulations = {}
-
-    for date, time, steps in mars_date_time_steps:
-        for p in param:
-            for n in number:
-                requests.append(
-                    patch(
-                        {
-                            "param": p,
-                            "date": date,
-                            "time": time,
-                            "step": sorted(steps),
-                            "number": n,
-                        }
-                    )
-                )
-
-    compressed = Availability(requests)
-    ds = cml.load_source("empty")
-    for r in compressed.iterate():
-        request.update(r)
-        print("🌧️", request)
-        ds = ds + cml.load_source("mars", **request)
-
-    accumulations = {}
-    for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:
-        for s in a.steps:
-            key = (a.param, a.date, a.time, s, a.number)
-            accumulations.setdefault(key, []).append(a)
-
-    for field in ds:
-        key = (
-            field.metadata("param"),
-            field.metadata("date"),
-            field.metadata("time"),
-            field.metadata("step"),
-            member(field),
-        )
-        values = field.values  # optimisation
-        assert accumulations[key], key
-        for a in accumulations[key]:
-            a.add(field, values)
-
-    for acc in accumulations.values():
-        for a in acc:
-            assert a.done, (a.key, a.seen, a.steps)
-
-    out.close()
-
-    ds = cml.load_source("file", path)
-
-    assert len(ds) / len(param) / len(number) == len(dates), (
-        len(ds),
-        len(param),
-        len(dates),
-    )
-    ds._tmp = tmp
-
-    return ds
-
-
-def to_list(x):
+def _to_list(x):
     if isinstance(x, (list, tuple)):
         return x
     return [x]
 
 
-def normalise_time_to_hours(r):
-    r = deepcopy(r)
-    if "time" not in r:
-        return r
-
-    times = []
-    for t in to_list(r["time"]):
-        assert len(t) == 4, r
-        assert t.endswith("00"), r
-        times.append(int(t) // 100)
-    r["time"] = tuple(times)
-    return r
-
-
-def normalise_number(r):
-    if "number" not in r:
-        return r
-    number = r["number"]
-    number = to_list(number)
-
-    if len(number) > 4 and (number[1] == "to" and number[3] == "by"):
-        return list(range(int(number[0]), int(number[2]) + 1, int(number[4])))
-
-    if len(number) > 2 and number[1] == "to":
-        return list(range(int(number[0]), int(number[2]) + 1))
-
-    r["number"] = number
-    return r
-
-
 class HindcastCompute:
     def __init__(self, base_times, available_steps, request):
         self.base_times = base_times
@@ -385,22 +34,34 @@ class HindcastCompute:
         self.request = request
 
     def compute_hindcast(self, date):
-        for step in self.available_steps:
+        result = []
+        for step in sorted(self.available_steps):  # Use the shortest step
             start_date = date - datetime.timedelta(hours=step)
             hours = start_date.hour
             if hours in self.base_times:
-                r = deepcopy(self.request)
+                r = self.request.copy()
                 r["date"] = start_date
                 r["time"] = f"{start_date.hour:02d}00"
                 r["step"] = step
-                return r
-        raise ValueError(
-            f"Cannot find data for {self.request} for {date} (base_times={self.base_times}, available_steps={self.available_steps})"
-        )
+                result.append(r)
+
+        if not result:
+            raise ValueError(
+                f"Cannot find data for {self.request} for {date} (base_times={self.base_times}, "
+                f"available_steps={self.available_steps})"
+            )
+
+        if len(result) > 1:
+            raise ValueError(
+                f"Multiple requests for {self.request} for {date} (base_times={self.base_times}, "
+                f"available_steps={self.available_steps})"
+            )
+
+        return result[0]
 
 
 def use_reference_year(reference_year, request):
-    request = deepcopy(request)
+    request = request.copy()
     hdate = request.pop("date")
     date = datetime.datetime(reference_year, hdate.month, hdate.day)
     request.update(date=date.strftime("%Y-%m-%d"), hdate=hdate.strftime("%Y-%m-%d"))
@@ -408,15 +69,15 @@ def use_reference_year(reference_year, request):
 
 
 def hindcasts(context, dates, **request):
-    request["param"] = to_list(request["param"])
-    request["step"] = to_list(request["step"])
+    request["param"] = _to_list(request["param"])
+    request["step"] = _to_list(request["step"])
     request["step"] = [int(_) for _ in request["step"]]
 
     if request.get("stream") == "enfh" and "base_times" not in request:
         request["base_times"] = [0]
 
     available_steps = request.pop("step")
-    available_steps = to_list(available_steps)
+    available_steps = _to_list(available_steps)
 
     base_times = request.pop("base_times")
 
@@ -431,7 +92,14 @@ def hindcasts(context, dates, **request):
         req = use_reference_year(reference_year, req)
 
         requests.append(req)
-    return mars(context, dates, *requests, date_key="hdate")
+
+    return mars(
+        context,
+        dates,
+        *requests,
+        date_key="hdate",
+        request_already_using_valid_datetime=True,
+    )
 
 
 execute = hindcasts

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

@@ -7,11 +7,10 @@
 # nor does it submit to any jurisdiction.
 #
 import datetime
-from copy import deepcopy
 
 from anemoi.utils.humanize import did_you_mean
-from climetlab import load_source
-from climetlab.utils.availability import Availability
+from earthkit.data import from_source
+from earthkit.data.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
@@ -43,25 +42,27 @@ def normalise_time_delta(t):
     return t
 
 
-def _expand_mars_request(request, date, date_key="date"):
+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:
-        r = deepcopy(request)
-
-        if isinstance(s, str) and "-" in s:
-            assert s.count("-") == 1, s
-        # 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])
-
-        base = date - datetime.timedelta(hours=hours)
-        r.update(
-            {
-                date_key: base.strftime("%Y%m%d"),
-                "time": base.strftime("%H%M"),
-                "step": s,
-            }
-        )
+        r = request.copy()
+
+        if not request_already_using_valid_datetime:
+
+            if isinstance(s, str) and "-" in s:
+                assert s.count("-") == 1, s
+            # 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])
+
+            base = date - datetime.timedelta(hours=hours)
+            r.update(
+                {
+                    date_key: base.strftime("%Y%m%d"),
+                    "time": base.strftime("%H%M"),
+                    "step": s,
+                }
+            )
 
         for pproc in ("grid", "rotation", "frame", "area", "bitmap", "resol"):
             if pproc in r:
@@ -73,13 +74,18 @@ def _expand_mars_request(request, date, date_key="date"):
     return requests
 
 
-def factorise_requests(dates, *requests, 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)
 
         for d in dates:
-            updates += _expand_mars_request(req, date=d, date_key=date_key)
+            updates += _expand_mars_request(
+                req,
+                date=d,
+                request_already_using_valid_datetime=request_already_using_valid_datetime,
+                date_key=date_key,
+            )
 
     compressed = Availability(updates)
     for r in compressed.iterate():
@@ -171,12 +177,33 @@ MARS_KEYS = [
 ]
 
 
-def mars(context, dates, *requests, date_key="date", **kwargs):
+def mars(context, dates, *requests, request_already_using_valid_datetime=False, date_key="date", **kwargs):
     if not requests:
         requests = [kwargs]
 
-    requests = factorise_requests(dates, *requests, date_key=date_key)
-    ds = load_source("empty")
+    for r in requests:
+        # check for "Norway bug" where yaml transforms 'no' into False, etc.
+        for p in r.get("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?"
+                )
+            if p is None:
+                raise ValueError(
+                    "'param' cannot be 'None'. If you wrote 'param: no' in yaml, you may want to use quotes?"
+                )
+            if p is True:
+                raise ValueError(
+                    "'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,
+    )
+    ds = from_source("empty")
     for r in requests:
         r = {k: v for k, v in r.items() if v != ("-",)}
 
@@ -184,14 +211,18 @@ def mars(context, dates, *requests, date_key="date", **kwargs):
             r = use_grib_paramid(r)
 
         if DEBUG:
-            context.trace("✅", f"load_source(mars, {r}")
+            context.trace("✅", f"from_source(mars, {r}")
 
         for k, v in r.items():
             if k not in MARS_KEYS:
                 raise ValueError(
                     f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
                 )
-        ds = ds + load_source("mars", **r)
+        try:
+            ds = ds + from_source("mars", **r)
+        except Exception as e:
+            if "File is empty:" not in str(e):
+                raise
     return ds
 
 

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

@@ -7,66 +7,8 @@
 # nor does it submit to any jurisdiction.
 #
 
-import glob
-
-from climetlab import load_source
-from climetlab.utils.patterns import Pattern
-
-
-def _expand(paths):
-    for path in paths:
-        if path.startswith("file://"):
-            path = path[7:]
-
-        if path.startswith("http://"):
-            yield path
-            continue
-
-        if path.startswith("https://"):
-            yield path
-            continue
-
-        for p in glob.glob(path):
-            yield p
-
-
-def check(what, ds, paths, **kwargs):
-    count = 1
-    for k, v in kwargs.items():
-        if isinstance(v, (tuple, list)):
-            count *= len(v)
-
-    if len(ds) != count:
-        raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, {what}s={paths})")
-
-
-def load_netcdfs(emoji, what, context, dates, path, *args, **kwargs):
-    given_paths = path if isinstance(path, list) else [path]
-
-    dates = [d.isoformat() for d in dates]
-    ds = load_source("empty")
-
-    for path in given_paths:
-        paths = Pattern(path, ignore_missing_keys=True).substitute(*args, date=dates, **kwargs)
-
-        levels = kwargs.get("level", kwargs.get("levelist"))
-
-        for path in _expand(paths):
-            context.trace(emoji, what.upper(), path)
-            s = load_source("opendap", path)
-            s = s.sel(
-                valid_datetime=dates,
-                param=kwargs["param"],
-                step=kwargs.get("step", 0),
-            )
-            if levels:
-                s = s.sel(levelist=levels)
-            ds = ds + s
-
-    check(what, ds, given_paths, valid_datetime=dates, **kwargs)
-
-    return ds
+from .xarray import load_many
 
 
 def execute(context, dates, path, *args, **kwargs):
-    return load_netcdfs("📁", "path", context, dates, path, *args, **kwargs)
+    return load_many("📁", context, dates, path, *args, **kwargs)

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

@@ -7,8 +7,9 @@
 # nor does it submit to any jurisdiction.
 #
 
-from .netcdf import load_netcdfs
+
+from .xarray import load_many
 
 
 def execute(context, dates, url, *args, **kwargs):
-    return load_netcdfs("🌐", "url", context, dates, url, *args, **kwargs)
+    return load_many("🌐", context, dates, url, *args, **kwargs)

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

@@ -6,7 +6,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
-from climetlab import load_source
+from earthkit.data import from_source
 
 from anemoi.datasets.create.utils import to_datetime_list
 
@@ -15,12 +15,12 @@ DEBUG = True
 
 def source(context, dates, **kwargs):
     name = kwargs.pop("name")
-    context.trace("✅", f"load_source({name}, {dates}, {kwargs}")
+    context.trace("✅", f"from_source({name}, {dates}, {kwargs}")
     if kwargs["date"] == "$from_dates":
         kwargs["date"] = list({d.strftime("%Y%m%d") for d in dates})
     if kwargs["time"] == "$from_dates":
         kwargs["time"] = list({d.strftime("%H%M") for d in dates})
-    return load_source(name, **kwargs)
+    return from_source(name, **kwargs)
 
 
 execute = source