From d093e1bfd35df9550a75f8650a95b3a894194ba9 Mon Sep 17 00:00:00 2001 From: Baudouin Raoult Date: Thu, 17 Oct 2024 06:54:49 +0000 Subject: [PATCH] Add unstructured grid --- .gitignore | 1 + src/anemoi/transform/grids/__init__.py | 11 +++ src/anemoi/transform/grids/unstructured.py | 98 +++++++++++++++++++++ src/anemoi/transform/variables/__init__.py | 7 +- src/anemoi/transform/variables/variables.py | 19 ++-- tests/test_grids.py | 42 +++++++++ 6 files changed, 171 insertions(+), 7 deletions(-) create mode 100644 src/anemoi/transform/grids/__init__.py create mode 100644 src/anemoi/transform/grids/unstructured.py create mode 100644 tests/test_grids.py diff --git a/.gitignore b/.gitignore index d4a8551..bfcf3b2 100644 --- a/.gitignore +++ b/.gitignore @@ -187,3 +187,4 @@ _build/ *.sync _version.py *.code-workspace +*.zip diff --git a/src/anemoi/transform/grids/__init__.py b/src/anemoi/transform/grids/__init__.py new file mode 100644 index 0000000..1a4a17b --- /dev/null +++ b/src/anemoi/transform/grids/__init__.py @@ -0,0 +1,11 @@ +# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts. +# 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 .unstructured import UnstructuredGridFieldList + +__all__ = ["UnstructuredGridFieldList"] diff --git a/src/anemoi/transform/grids/unstructured.py b/src/anemoi/transform/grids/unstructured.py new file mode 100644 index 0000000..1d394e4 --- /dev/null +++ b/src/anemoi/transform/grids/unstructured.py @@ -0,0 +1,98 @@ +# (C) Copyright 2024 ECMWF. +# +# 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 urllib.parse import urlparse + +import numpy as np +from earthkit.data import from_source +from earthkit.data.indexing.fieldlist import FieldArray + +LOG = logging.getLogger(__name__) + + +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, latitudes, longitudes, uuidOfHGrid=None): + + assert isinstance(latitudes, np.ndarray) + assert isinstance(longitudes, np.ndarray) + + LOG.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}") + assert len(latitudes) == len(longitudes) + + self.uuidOfHGrid = uuidOfHGrid + self.latitudes = latitudes + self.longitudes = longitudes + + +def _load(url_or_path, param): + parsed = urlparse(url_or_path) + if parsed.scheme: + source = "url" + else: + source = "file" + + ds = from_source(source, url_or_path) + ds = ds.sel(param=param) + + assert len(ds) == 1, f"{url_or_path} {param}, expected one field, got {len(ds)}" + ds = ds[0] + + return ds.to_numpy(flatten=True), ds.metadata("uuidOfHGrid") + + +class UnstructuredGridField: + """An unstructured field.""" + + def __init__(self, geography): + self._latitudes = geography.latitudes + self._longitudes = geography.longitudes + + def metadata(self, name, default=None): + if name == "uuidOfHGrid": + return self.geography.uuidOfHGrid + return default + + def grid_points(self): + return self._latitudes, self._longitudes + + @property + def resolution(self): + return "unknown" + + @property + def shape(self): + return (len(self._latitudes),) + + +class UnstructuredGridFieldList(FieldArray): + @classmethod + def from_grib(cls, latitudes_url_or_path, longitudes_url_or_path, latitudes_param="tlat", longitudes_params="tlon"): + latitudes, latitudes_uuid = _load(latitudes_url_or_path, latitudes_param) + longitudes, longitudes_uuid = _load(longitudes_url_or_path, longitudes_params) + + if latitudes_uuid != longitudes_uuid: + raise ValueError(f"uuidOfHGrid mismatch: lat={latitudes_uuid} != lon={longitudes_uuid}") + + return cls([UnstructuredGridField(Geography(latitudes, longitudes))]) + + @classmethod + def from_values(cls, latitudes, longitudes): + if isinstance(latitudes, (list, tuple)): + latitudes = np.array(latitudes) + + if isinstance(longitudes, (list, tuple)): + longitudes = np.array(longitudes) + + return cls([UnstructuredGridField(Geography(latitudes, longitudes))]) diff --git a/src/anemoi/transform/variables/__init__.py b/src/anemoi/transform/variables/__init__.py index bb49775..b9d6703 100644 --- a/src/anemoi/transform/variables/__init__.py +++ b/src/anemoi/transform/variables/__init__.py @@ -30,7 +30,7 @@ def from_earthkit(cls, field): return VariableFromEarthkit(field) def __repr__(self) -> str: - return f"{self.__class__.__name__}({self.name})" + return self.name def __hash__(self) -> int: return hash(self.name) @@ -49,3 +49,8 @@ def is_pressure_level(self): @abstractmethod def level(self): pass + + @property + @abstractmethod + def is_constant_in_time(self): + pass diff --git a/src/anemoi/transform/variables/variables.py b/src/anemoi/transform/variables/variables.py index 53c4303..2adfc7e 100644 --- a/src/anemoi/transform/variables/variables.py +++ b/src/anemoi/transform/variables/variables.py @@ -5,6 +5,8 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. +import json + from . import Variable @@ -14,24 +16,29 @@ class VariableFromMarsVocabulary(Variable): def __init__(self, name, data: dict) -> None: super().__init__(name) self.data = data + print(json.dumps(data, indent=4)) + if "mars" in self.data: + self.mars = self.data["mars"] + else: + self.mars = self.data @property def is_pressure_level(self): - return self.data.get("levtype", None) == "pl" + return self.mars.get("levtype", None) == "pl" @property def level(self): - return self.data.get("levelist", None) + return self.mars.get("levelist", None) + + @property + def is_constant_in_time(self): + return self.data.get("is_constant_in_time", False) class VariableFromDict(VariableFromMarsVocabulary): """A variable that is defined by a user provided dictionary.""" def __init__(self, name, data: dict) -> None: - - if "mars" in data: - data = data["mars"] - super().__init__(name, data) diff --git a/tests/test_grids.py b/tests/test_grids.py new file mode 100644 index 0000000..36558d9 --- /dev/null +++ b/tests/test_grids.py @@ -0,0 +1,42 @@ +# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts. +# 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 + +from anemoi.transform.grids import UnstructuredGridFieldList + +latitude_url = "http://icon-downloads.mpimet.mpg.de/grids/public/edzw/icon_extpar_0026_R03B07_G_20150805.g2" +tlat = "tlat" + +longitudes_url = "http://icon-downloads.mpimet.mpg.de/grids/public/edzw/icon_extpar_0026_R03B07_G_20150805.g2" +tlon = "tlon" + + +def test_unstructured_from_url(): + ds = UnstructuredGridFieldList.from_grib(latitude_url, longitudes_url, tlat, tlon) + + assert len(ds) == 1 + + lats, lons = ds[0].grid_points() + + assert len(lats) == len(lons) + + forcings = ekd.from_source( + "forcings", + ds, + date="2015-08-05", + param=["cos_latitude", "sin_latitude"], + ) + + assert len(forcings) == 2 + + +if __name__ == "__main__": + for name, obj in list(globals().items()): + if name.startswith("test_") and callable(obj): + print(f"Running {name}...") + obj()