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

anemoi/datasets/create/functions/filters/rotate_winds.py

@@ -9,61 +9,10 @@
 
 from collections import defaultdict
 
-from climetlab.indexing.fieldset import FieldArray
-
-
-def rotate_winds(lats, lons, x_wind, y_wind, source_projection, target_projection):
-    """Code provided by MetNO"""
-    import numpy as np
-    import pyproj
-
-    if source_projection == target_projection:
-        return x_wind, x_wind
-
-    source_projection = pyproj.Proj(source_projection)
-    target_projection = pyproj.Proj(target_projection)
-
-    transformer = pyproj.transformer.Transformer.from_proj(source_projection, target_projection)
-
-    # To compute the new vector components:
-    # 1) perturb each position in the direction of the winds
-    # 2) convert the perturbed positions into the new coordinate system
-    # 3) measure the new x/y components.
-    #
-    # A complication occurs when using the longlat "projections", since this is not a cartesian grid
-    # (i.e. distances in each direction is not consistent), we need to deal with the fact that the
-    # width of a longitude varies with latitude
-    orig_speed = np.sqrt(x_wind**2 + y_wind**2)
-
-    x0, y0 = source_projection(lons, lats)
-
-    if source_projection.name != "longlat":
-        x1 = x0 + x_wind
-        y1 = y0 + y_wind
-    else:
-        # Reduce the perturbation, since x_wind and y_wind are in meters, which would create
-        # large perturbations in lat, lon. Also, deal with the fact that the width of longitude
-        # varies with latitude.
-        factor = 3600000.0
-        x1 = x0 + x_wind / factor / np.cos(np.deg2rad(lats))
-        y1 = y0 + y_wind / factor
-
-    X0, Y0 = transformer.transform(x0, y0)
-    X1, Y1 = transformer.transform(x1, y1)
-
-    new_x_wind = X1 - X0
-    new_y_wind = Y1 - Y0
-    if target_projection.name == "longlat":
-        new_x_wind *= np.cos(np.deg2rad(lats))
-
-    if target_projection.name == "longlat" or source_projection.name == "longlat":
-        # Ensure the wind speed is not changed (which might not the case since the units in longlat
-        # is degrees, not meters)
-        curr_speed = np.sqrt(new_x_wind**2 + new_y_wind**2)
-        new_x_wind *= orig_speed / curr_speed
-        new_y_wind *= orig_speed / curr_speed
-
-    return new_x_wind, new_y_wind
+import tqdm
+from anemoi.utils.humanize import plural
+from earthkit.data.indexing.fieldlist import FieldArray
+from earthkit.geo.rotate import rotate_vector
 
 
 class NewDataField:
@@ -77,6 +26,9 @@ class NewDataField:
     def __getattr__(self, name):
         return getattr(self.field, name)
 
+    def __repr__(self) -> str:
+        return repr(self.field)
+
 
 def execute(
     context,
@@ -88,13 +40,15 @@ def execute(
 ):
     from pyproj import CRS
 
+    context.trace("🔄", "Rotating winds (extracting winds from ", plural(len(input), "field"))
+
     result = FieldArray()
 
     wind_params = (x_wind, y_wind)
     wind_pairs = defaultdict(dict)
 
     for f in input:
-        key = f.as_mars()
+        key = f.metadata(namespace="mars")
         param = key.pop("param")
 
         if param not in wind_params:
@@ -108,7 +62,9 @@ def execute(
 
         wind_pairs[key][param] = f
 
-    for _, pairs in wind_pairs.items():
+    context.trace("🔄", "Rotating", plural(len(wind_pairs), "wind"), "(speed will likely include data download)")
+
+    for _, pairs in tqdm.tqdm(list(wind_pairs.items())):
         if len(pairs) != 2:
             raise ValueError("Missing wind component")
 
@@ -118,11 +74,11 @@ def execute(
         assert x.grid_mapping == y.grid_mapping
 
         lats, lons = x.grid_points()
-        x_new, y_new = rotate_winds(
+        x_new, y_new = rotate_vector(
             lats,
             lons,
-            x.to_numpy(reshape=False),
-            y.to_numpy(reshape=False),
+            x.to_numpy(flatten=True),
+            y.to_numpy(flatten=True),
             (source_projection if source_projection is not None else CRS.from_cf(x.grid_mapping)),
             target_projection,
         )

anemoi/datasets/create/functions/filters/unrotate_winds.py

@@ -9,60 +9,9 @@
 
 from collections import defaultdict
 
-import numpy as np
-from climetlab.indexing.fieldset import FieldArray
-
-
-def normalise(x):
-    return max(min(x, 1.0), -1.0)
-
-
-def normalise_longitude(lon, minimum):
-    while lon < minimum:
-        lon += 360
-
-    while lon >= minimum + 360:
-        lon -= 360
-
-    return lon
-
-
-def rotate_winds(
-    lats,
-    lons,
-    raw_lats,
-    raw_lons,
-    x_wind,
-    y_wind,
-    south_pole_latitude,
-    south_pole_longitude,
-    south_pole_rotation_angle=0,
-):
-    # Code from MIR
-    assert south_pole_rotation_angle == 0
-    C = np.deg2rad(90 - south_pole_latitude)
-    cos_C = np.cos(C)
-    sin_C = np.sin(C)
-
-    new_x = np.zeros_like(x_wind)
-    new_y = np.zeros_like(y_wind)
-
-    for i, (vx, vy, lat, lon, raw_lat, raw_lon) in enumerate(zip(x_wind, y_wind, lats, lons, raw_lats, raw_lons)):
-        lonRotated = south_pole_longitude - lon
-        lon_rotated = normalise_longitude(lonRotated, -180)
-        lon_unrotated = raw_lon
-
-        a = np.deg2rad(lon_rotated)
-        b = np.deg2rad(lon_unrotated)
-        q = 1 if (sin_C * lon_rotated < 0.0) else -1.0  # correct quadrant
-
-        cos_c = normalise(np.cos(a) * np.cos(b) + np.sin(a) * np.sin(b) * cos_C)
-        sin_c = q * np.sqrt(1.0 - cos_c * cos_c)
-
-        new_x[i] = cos_c * vx + sin_c * vy
-        new_y[i] = -sin_c * vx + cos_c * vy
-
-    return new_x, new_y
+# import numpy as np
+from earthkit.data.indexing.fieldlist import FieldArray
+from earthkit.geo.rotate import unrotate_vector
 
 
 class NewDataField:
@@ -85,7 +34,7 @@ def execute(context, input, u, v):
     wind_pairs = defaultdict(dict)
 
     for f in input:
-        key = f.as_mars()
+        key = f.metadata(namespace="mars")
         param = key.pop("param")
 
         if param not in wind_params:
@@ -107,18 +56,19 @@ def execute(context, input, u, v):
         y = pairs[v]
 
         lats, lons = x.grid_points()
-        raw_lats, raw_longs = x.grid_points_raw()
+        raw_lats, raw_longs = x.grid_points_unrotated()
 
         assert x.rotation == y.rotation
 
-        u_new, v_new = rotate_winds(
+        u_new, v_new = unrotate_vector(
             lats,
             lons,
-            raw_lats,
-            raw_longs,
-            x.to_numpy(reshape=False),
-            y.to_numpy(reshape=False),
-            *x.rotation,
+            x.to_numpy(flatten=True),
+            y.to_numpy(flatten=True),
+            *x.rotation[:2],
+            south_pole_rotation_angle=x.rotation[2],
+            lat_unrotated=raw_lats,
+            lon_unrotated=raw_longs,
         )
 
         result.append(NewDataField(x, u_new))
@@ -128,9 +78,9 @@ def execute(context, input, u, v):
 
 
 if __name__ == "__main__":
-    from climetlab import load_source
+    from earthkit.data import from_source
 
-    source = load_source(
+    source = from_source(
         "mars",
         date=-1,
         param="10u/10v",

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

@@ -6,3 +6,42 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
+
+import glob
+import logging
+
+from earthkit.data.utils.patterns import Pattern
+
+LOG = logging.getLogger(__name__)
+
+
+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
+
+        cnt = 0
+        for p in glob.glob(path):
+            yield p
+            cnt += 1
+        if cnt == 0:
+            yield path
+
+
+def iterate_patterns(path, dates, **kwargs):
+    given_paths = path if isinstance(path, list) else [path]
+
+    dates = [d.isoformat() for d in dates]
+
+    for path in given_paths:
+        paths = Pattern(path, ignore_missing_keys=True).substitute(date=dates, **kwargs)
+        for path in _expand(paths):
+            yield path, dates

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

@@ -11,11 +11,11 @@ import logging
 import warnings
 from copy import deepcopy
 
-import climetlab as cml
+import earthkit.data as ekd
 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 earthkit.data.core.temporary import temp_file
+from earthkit.data.readers.grib.output import new_grib_output
+from earthkit.data.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
@@ -24,9 +24,9 @@ from .mars import use_grib_paramid
 LOG = logging.getLogger(__name__)
 
 
-def member(field):
+def _member(field):
     # Bug in eccodes has number=0 randomly
-    number = field.metadata("number")
+    number = field.metadata("number", default=0)
     if number is None:
         number = 0
     return number
@@ -54,16 +54,25 @@ class Accumulation:
 
     def check(self, field):
         if self._check is None:
-            self._check = field.as_mars()
+            self._check = field.metadata(namespace="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))
+            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()
+        mars = field.metadata(namespace="mars")
         keys1 = sorted(self._check.keys())
         keys2 = sorted(mars.keys())
 
@@ -196,7 +205,11 @@ class AccumulationFromLastStep(Accumulation):
 
     def compute(self, values, startStep, endStep):
 
-        assert endStep - startStep == self.frequency, (startStep, endStep, self.frequency)
+        assert endStep - startStep == self.frequency, (
+            startStep,
+            endStep,
+            self.frequency,
+        )
 
         if self.startStep is None:
             self.startStep = startStep
@@ -228,17 +241,17 @@ class AccumulationFromLastStep(Accumulation):
         )
 
 
-def identity(x):
+def _identity(x):
     return x
 
 
-def compute_accumulations(
+def _compute_accumulations(
     context,
     dates,
     request,
     user_accumulation_period=6,
     data_accumulation_period=None,
-    patch=identity,
+    patch=_identity,
     base_times=None,
 ):
     adjust_step = isinstance(user_accumulation_period, int)
@@ -307,14 +320,13 @@ def compute_accumulations(
                 )
 
     compressed = Availability(requests)
-    ds = cml.load_source("empty")
+    ds = ekd.from_source("empty")
     for r in compressed.iterate():
         request.update(r)
         if context.use_grib_paramid and "param" in request:
             request = use_grib_paramid(request)
         print("🌧️", request)
-
-        ds = ds + cml.load_source("mars", **request)
+        ds = ds + ekd.from_source("mars", **request)
 
     accumulations = {}
     for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:
@@ -328,7 +340,7 @@ def compute_accumulations(
             field.metadata("date"),
             field.metadata("time"),
             field.metadata("step"),
-            member(field),
+            _member(field),
         )
         values = field.values  # optimisation
         assert accumulations[key], key
@@ -341,7 +353,7 @@ def compute_accumulations(
 
     out.close()
 
-    ds = cml.load_source("file", path)
+    ds = ekd.from_source("file", path)
 
     assert len(ds) / len(param) / len(number) == len(dates), (
         len(ds),
@@ -353,43 +365,13 @@ def compute_accumulations(
     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
-
-
-def scda(request):
+def _scda(request):
     if request["time"] in (6, 18, 600, 1800):
         request["stream"] = "scda"
     else:
@@ -398,14 +380,14 @@ def scda(request):
 
 
 def accumulations(context, dates, **request):
-    to_list(request["param"])
+    _to_list(request["param"])
     class_ = request.get("class", "od")
     stream = request.get("stream", "oper")
 
     user_accumulation_period = request.pop("accumulation_period", 6)
 
     KWARGS = {
-        ("od", "oper"): dict(patch=scda),
+        ("od", "oper"): dict(patch=_scda),
         ("od", "elda"): dict(base_times=(6, 18)),
         ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
         ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
@@ -415,7 +397,7 @@ def accumulations(context, dates, **request):
 
     context.trace("🌧️", f"accumulations {request} {user_accumulation_period} {kwargs}")
 
-    return compute_accumulations(
+    return _compute_accumulations(
         context,
         dates,
         request,

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

@@ -6,15 +6,22 @@
 # 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
 
 
 def constants(context, dates, template, param):
     from warnings import warn
 
-    warn("The source `constants` is deprecated, use `forcings` instead.", DeprecationWarning, stacklevel=2)
-    context.trace("✅", f"load_source(constants, {template}, {param}")
-    return load_source("constants", source_or_dataset=template, date=dates, param=param)
+    warn(
+        "The source `constants` is deprecated, use `forcings` instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    context.trace("✅", f"from_source(constants, {template}, {param}")
+    if len(template) == 0:
+        raise ValueError("Forcings template is empty.")
+
+    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
 
 
 execute = constants

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

@@ -7,8 +7,8 @@
 # nor does it submit to any jurisdiction.
 #
 
-import climetlab as cml
+import earthkit.data as ekd
 
 
 def execute(context, dates, **kwargs):
-    return cml.load_source("empty")
+    return ekd.from_source("empty")

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

@@ -6,12 +6,12 @@
 # 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
 
 
 def forcings(context, dates, template, param):
-    context.trace("✅", f"load_source(forcings, {template}, {param}")
-    return load_source("constants", source_or_dataset=template, date=dates, param=param)
+    context.trace("✅", f"from_source(forcings, {template}, {param}")
+    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
 
 
 execute = forcings

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

@@ -10,8 +10,8 @@
 
 import glob
 
-from climetlab import load_source
-from climetlab.utils.patterns import Pattern
+from earthkit.data import from_source
+from earthkit.data.utils.patterns import Pattern
 
 
 def check(ds, paths, **kwargs):
@@ -26,14 +26,18 @@ def check(ds, paths, **kwargs):
 
 def _expand(paths):
     for path in paths:
+        cnt = 0
         for p in glob.glob(path):
             yield p
+            cnt += 1
+        if cnt == 0:
+            yield path
 
 
 def execute(context, dates, path, *args, **kwargs):
     given_paths = path if isinstance(path, list) else [path]
 
-    ds = load_source("empty")
+    ds = from_source("empty")
     dates = [d.isoformat() for d in dates]
 
     for path in given_paths:
@@ -45,7 +49,7 @@ def execute(context, dates, path, *args, **kwargs):
 
         for path in _expand(paths):
             context.trace("📁", "PATH", path)
-            s = load_source("file", path)
+            s = from_source("file", path)
             s = s.sel(valid_datetime=dates, **kwargs)
             ds = ds + s