anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/create/functions/__init__.py

@@ -1,66 +0,0 @@
-# (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 importlib
-
-import entrypoints
-
-
-def assert_is_fieldlist(obj):
-    from earthkit.data.indexing.fieldlist import FieldList
-
-    assert isinstance(obj, FieldList), type(obj)
-
-
-def import_function(name, kind):
-
-    from anemoi.transform.filters import filter_registry
-    from anemoi.transform.sources import source_registry
-
-    name = name.replace("-", "_")
-
-    plugins = {}
-    for e in entrypoints.get_group_all(f"anemoi.datasets.{kind}"):
-        plugins[e.name.replace("_", "-")] = e
-
-    if name in plugins:
-        return plugins[name].load()
-
-    try:
-        module = importlib.import_module(
-            f".{kind}.{name}",
-            package=__name__,
-        )
-        return module.execute
-    except ModuleNotFoundError:
-        pass
-
-    if kind == "filters":
-        if filter_registry.lookup(name, return_none=True):
-
-            def proc(context, data, *args, **kwargs):
-                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
-
-    raise ValueError(f"Unknown {kind} '{name}'")

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

@@ -1,9 +0,0 @@
-# (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/create/functions/filters/empty.py

@@ -1,17 +0,0 @@
-# (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 earthkit.data as ekd
-
-
-def execute(context, input, **kwargs):
-    # Useful to create a pipeline that returns an empty result
-    # So we can reference an earlier step in a function like 'constants'
-    return ekd.from_source("empty")

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

@@ -1,58 +0,0 @@
-# (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/create/functions/filters/rename.py

@@ -1,79 +0,0 @@
-# (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 re
-
-from earthkit.data.indexing.fieldlist import FieldArray
-
-
-class RenamedFieldMapping:
-    """Rename a field based on the value of another field.
-
-    Args:
-        field (Field): The field to be renamed.
-        what (str): The name of the field that will be used to rename the field.
-        renaming (dict): A dictionary mapping the values of 'what' to the new names.
-    """
-
-    def __init__(self, field, what, renaming):
-        self.field = field
-        self.what = what
-        self.renaming = {}
-        for k, v in renaming.items():
-            self.renaming[k] = {str(a): str(b) for a, b in v.items()}
-
-    def metadata(self, key=None, **kwargs):
-        if key is None:
-            return self.field.metadata(**kwargs)
-
-        value = self.field.metadata(key, **kwargs)
-        if key == self.what:
-            return self.renaming.get(self.what, {}).get(value, value)
-
-        return value
-
-    def __getattr__(self, name):
-        return getattr(self.field, name)
-
-    def __repr__(self) -> str:
-        return repr(self.field)
-
-
-class RenamedFieldFormat:
-    """Rename a field based on a format string.
-
-    Args:
-        format (str): A string that defines the new name of the field.
-    """
-
-    def __init__(self, field, what, format):
-        self.field = field
-        self.what = what
-        self.format = format
-        self.bits = re.findall(r"{(\w+)}", format)
-
-    def metadata(self, *args, **kwargs):
-        value = self.field.metadata(*args, **kwargs)
-        if args:
-            assert len(args) == 1
-            if args[0] == self.what:
-                bits = {b: self.field.metadata(b, **kwargs) for b in self.bits}
-                return self.format.format(**bits)
-        return value
-
-    def __getattr__(self, name):
-        return getattr(self.field, name)
-
-
-def execute(context, input, what="param", **kwargs):
-    if what in kwargs and isinstance(kwargs[what], str):
-        return FieldArray([RenamedFieldFormat(fs, what, kwargs[what]) for fs in input])
-
-    return FieldArray([RenamedFieldMapping(fs, what, kwargs) for fs in input])

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

@@ -1,78 +0,0 @@
-# (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
-
-import numpy as np
-from earthkit.data.indexing.fieldlist import FieldArray
-from earthkit.meteo.wind.array import polar_to_xy
-
-
-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, wind_speed, wind_dir, u_component="u", v_component="v", in_radians=False):
-
-    result = FieldArray()
-
-    wind_params = (wind_speed, wind_dir)
-    wind_pairs = defaultdict(dict)
-
-    for f in input:
-        key = f.metadata(namespace="mars")
-        param = key.pop("param")
-
-        if param not in wind_params:
-            result.append(f)
-            continue
-
-        key = tuple(key.items())
-
-        if param in wind_pairs[key]:
-            raise ValueError(f"Duplicate wind component {param} for {key}")
-
-        wind_pairs[key][param] = f
-
-    for _, pairs in wind_pairs.items():
-        if len(pairs) != 2:
-            raise ValueError("Missing wind component")
-
-        magnitude = pairs[wind_speed]
-        direction = pairs[wind_dir]
-
-        # assert speed.grid_mapping == dir.grid_mapping
-        if in_radians:
-            direction = np.rad2deg(direction)
-
-        u, v = polar_to_xy(magnitude.to_numpy(flatten=True), direction.to_numpy(flatten=True))
-
-        result.append(NewDataField(magnitude, u, u_component))
-        result.append(NewDataField(direction, v, v_component))
-
-    return result

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

@@ -1,56 +0,0 @@
-# (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
-
-import numpy as np
-from earthkit.data.indexing.fieldlist import FieldArray
-from earthkit.meteo.wind.array import xy_to_polar
-
-from anemoi.datasets.create.functions.filters.speeddir_to_uv import NewDataField
-
-
-def execute(context, input, u_component, v_component, wind_speed, wind_dir, in_radians=False):
-    result = FieldArray()
-
-    wind_params = (u_component, v_component)
-    wind_pairs = defaultdict(dict)
-
-    for f in input:
-        key = f.metadata(namespace="mars")
-        param = key.pop("param")
-
-        if param not in wind_params:
-            result.append(f)
-            continue
-
-        key = tuple(key.items())
-
-        if param in wind_pairs[key]:
-            raise ValueError(f"Duplicate wind component {param} for {key}")
-
-        wind_pairs[key][param] = f
-
-    for _, pairs in wind_pairs.items():
-        if len(pairs) != 2:
-            raise ValueError("Missing wind component")
-
-        u = pairs[u_component]
-        v = pairs[v_component]
-
-        # assert speed.grid_mapping == dir.grid_mapping
-        magnitude, direction = xy_to_polar(u.to_numpy(flatten=True), v.to_numpy(flatten=True))
-        if in_radians:
-            direction = np.deg2rad(direction)
-
-        result.append(NewDataField(u, magnitude, wind_speed))
-        result.append(NewDataField(v, direction, wind_dir))
-
-    return result

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

@@ -1,15 +0,0 @@
-# (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 earthkit.data as ekd
-
-
-def execute(context, dates, **kwargs):
-    return ekd.from_source("empty")

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

@@ -1,150 +0,0 @@
-# (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 glob
-import logging
-
-from earthkit.data import from_source
-from earthkit.data.indexing.fieldlist import FieldArray
-from earthkit.data.utils.patterns import Pattern
-
-LOG = logging.getLogger(__name__)
-
-
-def _load(context, name, record):
-    ds = None
-
-    param = record["param"]
-
-    if "path" in record:
-        context.info(f"Using {name} from {record['path']} (param={param})")
-        ds = from_source("file", record["path"])
-
-    if "url" in record:
-        context.info(f"Using {name} from {record['url']} (param={param})")
-        ds = from_source("url", record["url"])
-
-    ds = ds.sel(param=param)
-
-    assert len(ds) == 1, f"{name} {param}, expected one field, got {len(ds)}"
-    ds = ds[0]
-
-    return ds.to_numpy(flatten=True), ds.metadata("uuidOfHGrid")
-
-
-class Geography:
-    """This class retrieve the latitudes and longitudes of unstructured grids,
-    and checks if the fields are compatible with the grid.
-    """
-
-    def __init__(self, context, latitudes, longitudes):
-
-        latitudes, uuidOfHGrid_lat = _load(context, "latitudes", latitudes)
-        longitudes, uuidOfHGrid_lon = _load(context, "longitudes", longitudes)
-
-        assert (
-            uuidOfHGrid_lat == uuidOfHGrid_lon
-        ), f"uuidOfHGrid mismatch: lat={uuidOfHGrid_lat} != lon={uuidOfHGrid_lon}"
-
-        context.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}")
-        assert len(latitudes) == len(longitudes)
-
-        self.uuidOfHGrid = uuidOfHGrid_lat
-        self.latitudes = latitudes
-        self.longitudes = longitudes
-        self.first = True
-
-    def check(self, field):
-        if self.first:
-            # We only check the first field, for performance reasons
-            assert (
-                field.metadata("uuidOfHGrid") == self.uuidOfHGrid
-            ), f"uuidOfHGrid mismatch: {field.metadata('uuidOfHGrid')} != {self.uuidOfHGrid}"
-            self.first = False
-
-
-class AddGrid:
-    """An earth-kit.data.Field wrapper that adds grid information."""
-
-    def __init__(self, field, geography):
-        self._field = field
-
-        geography.check(field)
-
-        self._latitudes = geography.latitudes
-        self._longitudes = geography.longitudes
-
-    def __getattr__(self, name):
-        return getattr(self._field, name)
-
-    def __repr__(self) -> str:
-        return repr(self._field)
-
-    def grid_points(self):
-        return self._latitudes, self._longitudes
-
-    @property
-    def resolution(self):
-        return "unknown"
-
-
-def check(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}, paths={paths})")
-
-
-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, latitudes=None, longitudes=None, *args, **kwargs):
-    given_paths = path if isinstance(path, list) else [path]
-
-    geography = None
-    if latitudes is not None and longitudes is not None:
-        geography = Geography(context, latitudes, longitudes)
-
-    ds = from_source("empty")
-    dates = [d.isoformat() for d in dates]
-
-    for path in given_paths:
-        paths = Pattern(path, ignore_missing_keys=True).substitute(*args, date=dates, **kwargs)
-
-        for name in ("grid", "area", "rotation", "frame", "resol", "bitmap"):
-            if name in kwargs:
-                raise ValueError(f"MARS interpolation parameter '{name}' not supported")
-
-        for path in _expand(paths):
-            context.trace("📁", "PATH", path)
-            s = from_source("file", path)
-            s = s.sel(valid_datetime=dates, **kwargs)
-            ds = ds + s
-
-    if kwargs and not context.partial_ok:
-        check(ds, given_paths, valid_datetime=dates, **kwargs)
-
-    if geography is not None:
-        ds = FieldArray([AddGrid(_, geography) for _ in ds])
-
-    if len(ds) == 0:
-        LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
-
-    return ds

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

@@ -1,15 +0,0 @@
-# (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
-
-
-def execute(context, dates, path, *args, **kwargs):
-    return load_many("📁", context, dates, path, *args, **kwargs)

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

@@ -1,15 +0,0 @@
-# (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
-
-
-def execute(context, dates, url, *args, **kwargs):
-    return load_many("🌐", context, dates, url, *args, **kwargs)

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

@@ -1,60 +0,0 @@
-# (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
-
-from .mars import mars
-
-
-def to_list(x):
-    if isinstance(x, (list, tuple)):
-        return x
-    if isinstance(x, str):
-        return x.split("/")
-    return [x]
-
-
-def normalise_number(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))
-
-    return number
-
-
-def normalise_request(request):
-    request = deepcopy(request)
-    if "number" in request:
-        request["number"] = normalise_number(request["number"])
-    if "time" in request:
-        request["time"] = to_list(request["time"])
-    request["param"] = to_list(request["param"])
-    return request
-
-
-def load_if_needed(context, dates, dict_or_dataset):
-    if isinstance(dict_or_dataset, dict):
-        dict_or_dataset = normalise_request(dict_or_dataset)
-        dict_or_dataset = mars(context, dates, dict_or_dataset)
-    return dict_or_dataset
-
-
-def recentre(context, dates, members, centre, alpha=1.0, remapping={}, patches={}):
-    members = load_if_needed(context, dates, members)
-    centre = load_if_needed(context, dates, centre)
-    return _recentre(members=members, centre=centre, alpha=alpha)
-
-
-execute = recentre