From b880509737995a419f9b465e5a0ceefa7a3de4b4 Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Thu, 17 Oct 2024 10:13:43 +0000
Subject: [PATCH] Work on grid

---
 src/anemoi/transform/grids/unstructured.py | 25 ++++++++++++++--------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/anemoi/transform/grids/unstructured.py b/src/anemoi/transform/grids/unstructured.py
index 1d394e4..cde6008 100644
--- a/src/anemoi/transform/grids/unstructured.py
+++ b/src/anemoi/transform/grids/unstructured.py
@@ -25,8 +25,8 @@ class Geography:
 
     def __init__(self, latitudes, longitudes, uuidOfHGrid=None):
 
-        assert isinstance(latitudes, np.ndarray)
-        assert isinstance(longitudes, np.ndarray)
+        assert isinstance(latitudes, np.ndarray), type(latitudes)
+        assert isinstance(longitudes, np.ndarray), type(longitudes)
 
         LOG.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}")
         assert len(latitudes) == len(longitudes)
@@ -35,6 +35,9 @@ def __init__(self, latitudes, longitudes, uuidOfHGrid=None):
         self.latitudes = latitudes
         self.longitudes = longitudes
 
+    def shape(self):
+        return self.latitudes.shape
+
 
 def _load(url_or_path, param):
     parsed = urlparse(url_or_path)
@@ -56,16 +59,17 @@ class UnstructuredGridField:
     """An unstructured field."""
 
     def __init__(self, geography):
-        self._latitudes = geography.latitudes
-        self._longitudes = geography.longitudes
+        self.geography = geography
+
+    def metadata(self, *args, default=None, **kwargs):
+
+        if len(args) == 0 and len(kwargs) == 0:
+            return self
 
-    def metadata(self, name, default=None):
-        if name == "uuidOfHGrid":
-            return self.geography.uuidOfHGrid
         return default
 
     def grid_points(self):
-        return self._latitudes, self._longitudes
+        return self.geography.latitudes, self.geography.longitudes
 
     @property
     def resolution(self):
@@ -73,7 +77,10 @@ def resolution(self):
 
     @property
     def shape(self):
-        return (len(self._latitudes),)
+        return self.geography.shape()
+
+    def to_latlon(self, flatten=False):
+        return dict(lat=self.geography.latitudes, lon=self.geography.longitudes)
 
 
 class UnstructuredGridFieldList(FieldArray):