Add unstructured grid

.gitignore

@@ -187,3 +187,4 @@ _build/
 *.sync
 _version.py
 *.code-workspace
+*.zip

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"]

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))])

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

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)
 
 

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()