Add support for anemoi-datasets

src/earthkit/data/core/__init__.py

@@ -98,7 +98,7 @@ def order_by(self, *args, **kwargs):
         """Reorder the elements of the object."""
         self._not_implemented()
 
-    def unique_values(self, *coords, remapping=None, progress_bar=True):
+    def unique_values(self, *coords, remapping=None, patches=None, progress_bar=True):
         """
         Given a list of metadata attributes, such as date, param, levels,
         returns the list of unique values for each attributes
@@ -109,7 +109,7 @@ def unique_values(self, *coords, remapping=None, progress_bar=True):
         assert len(coords)
         assert all(isinstance(k, str) for k in coords), coords
 
-        remapping = build_remapping(remapping)
+        remapping = build_remapping(remapping, patches)
         iterable = self
 
         if progress_bar:
@@ -118,16 +118,16 @@ def unique_values(self, *coords, remapping=None, progress_bar=True):
                 desc=f"Finding coords in dataset for {coords}",
             )
 
-        dic = defaultdict(dict)
+        vals = defaultdict(dict)
         for f in iterable:
             metadata = remapping(f.metadata)
             for k in coords:
-                v = metadata(k)
-                dic[k][v] = True
+                v = metadata(k, default=None)
+                vals[k][v] = True
 
-        dic = {k: tuple(values.keys()) for k, values in dic.items()}
+        vals = {k: tuple(values.keys()) for k, values in vals.items()}
 
-        return dic
+        return vals
 
     # @abstractmethod
     # def to_points(self, *args, **kwargs):

src/earthkit/data/core/fieldlist.py

@@ -12,7 +12,7 @@
 from collections import defaultdict
 
 from earthkit.data.core import Base
-from earthkit.data.core.index import Index
+from earthkit.data.core.index import Index, MaskIndex, MultiIndex
 from earthkit.data.decorators import cached_method, detect_out_filename
 from earthkit.data.utils.array import ensure_backend, numpy_backend
 from earthkit.data.utils.metadata import metadata_argument
@@ -341,6 +341,31 @@ def to_latlon(self, flatten=False, dtype=None):
         lon, lat = self.data(("lon", "lat"), flatten=flatten, dtype=dtype)
         return dict(lat=lat, lon=lon)
 
+    def grid_points(self):
+        r = self.to_latlon(flatten=True)
+        return r["lat"], r["lon"]
+
+    def grid_points_unrotated(self):
+        lat = self._metadata.geography.latitudes_unrotated()
+        lon = self._metadata.geography.longitudes_unrotated()
+        return lat, lon
+
+    @property
+    def rotation(self):
+        return self._metadata.geography.rotation
+
+    @property
+    def resolution(self):
+        return self._metadata.geography.resolution()
+
+    @property
+    def mars_grid(self):
+        return self._metadata.geography.mars_grid()
+
+    @property
+    def mars_area(self):
+        return self._metadata.geography.mars_area()
+
     @property
     def shape(self):
         r"""tuple: Get the shape of the field.
@@ -1420,3 +1445,27 @@ def to_fieldlist(self, array_backend=None, **kwargs):
     def _to_array_fieldlist(self, **kwargs):
         md = [f.metadata() for f in self]
         return self.from_array(self.to_array(**kwargs), md)
+
+    def cube(self, *args, **kwargs):
+        from earthkit.data.indexing.cube import FieldCube
+
+        return FieldCube(self, *args, **kwargs)
+
+    @classmethod
+    def new_mask_index(self, *args, **kwargs):
+        return MaskFieldList(*args, **kwargs)
+
+    @classmethod
+    def merge(cls, sources):
+        assert all(isinstance(_, FieldList) for _ in sources)
+        return MultiFieldList(sources)
+
+
+class MaskFieldList(FieldList, MaskIndex):
+    def __init__(self, *args, **kwargs):
+        MaskIndex.__init__(self, *args, **kwargs)
+
+
+class MultiFieldList(FieldList, MultiIndex):
+    def __init__(self, *args, **kwargs):
+        MultiIndex.__init__(self, *args, **kwargs)

src/earthkit/data/core/geography.py

@@ -127,3 +127,15 @@ def gridspec(self):
         :class:`~data.core.gridspec.GridSpec>`
         """
         pass
+
+    @abstractmethod
+    def resolution(self):
+        pass
+
+    @abstractmethod
+    def mars_grid(self):
+        pass
+
+    @abstractmethod
+    def mars_area(self):
+        pass

src/earthkit/data/core/index.py

@@ -104,6 +104,7 @@ def build_actions(self, kwargs):
         raise NotImplementedError()
 
     def compare_elements(self, a, b):
+        assert callable(self.remapping), (type(self.remapping), self.remapping)
         a_metadata = self.remapping(a.metadata)
         b_metadata = self.remapping(b.metadata)
         for k, v in self.actions.items():
@@ -120,18 +121,18 @@ def build_actions(self, kwargs):
         def ascending(a, b):
             if a == b:
                 return 0
-            if a > b:
+            if b is None or a > b:
                 return 1
-            if a < b:
+            if a is None or a < b:
                 return -1
             raise ValueError(f"{a},{b}")
 
         def descending(a, b):
             if a == b:
                 return 0
-            if a > b:
+            if b is None or a > b:
                 return -1
-            if a < b:
+            if a is None or a < b:
                 return 1
             raise ValueError(f"{a},{b}")
 
@@ -169,6 +170,9 @@ def get(self, x):
                     order[int(key)] = i
                 except ValueError:
                     pass
+                except TypeError:
+                    print('Cannot convert "%s" to int (%s)' % (key, type(key)))
+                    raise
                 try:
                     order[float(key)] = i
                 except ValueError:
@@ -273,13 +277,15 @@ def sel(self, *args, remapping=None, **kwargs):
         Using ``remapping`` to specify the selection by a key created from two other keys
         (we created key "param_level" from "param" and "levelist"):
 
-        >>> for f in ds.order_by(
+        >>> subset = ds.sel(
         ...     param_level=["t850", "u1000"],
         ...     remapping={"param_level": "{param}{levelist}"},
-        ... ):
+        ... )
+        >>> for f in subset:
         ...     print(f)
-        GribField(t,850,20180801,1200,0,0)
+        ...
         GribField(u,1000,20180801,1200,0,0)
+        GribField(t,850,20180801,1200,0,0)
         """
         kwargs = normalize_selection(*args, **kwargs)
         kwargs = self._normalize_kwargs_names(**kwargs)
@@ -391,7 +397,7 @@ def isel(self, *args, **kwargs):
 
         return self.sel(**kwargs)
 
-    def order_by(self, *args, remapping=None, **kwargs):
+    def order_by(self, *args, remapping=None, patches=None, **kwargs):
         """Changes the order of the elements in a fieldlist-like object.
 
         Parameters
@@ -488,7 +494,7 @@ def order_by(self, *args, remapping=None, **kwargs):
         kwargs = normalize_order_by(*args, **kwargs)
         kwargs = self._normalize_kwargs_names(**kwargs)
 
-        remapping = build_remapping(remapping)
+        remapping = build_remapping(remapping, patches)
 
         if not kwargs:
             return self
@@ -507,38 +513,39 @@ def cmp(i, j):
 
     def __getitem__(self, n):
         if isinstance(n, slice):
-            return self.from_slice(n)
+            return self._from_slice(n)
         if isinstance(n, (tuple, list)):
-            return self.from_multi(n)
+            return self._from_sequence(n)
         if isinstance(n, dict):
-            return self.from_dict(n)
+            return self._from_dict(n)
         else:
             import numpy as np
 
             if isinstance(n, np.ndarray):
-                return self.from_multi(n)
+                return self._from_ndarray(n)
 
         return self._getitem(n)
 
-    def from_slice(self, s):
+    def _from_slice(self, s):
         indices = range(len(self))[s]
         return self.new_mask_index(self, indices)
 
-    def from_mask(self, lst):
+    def _from_mask(self, lst):
         indices = [i for i, x in enumerate(lst) if x]
         return self.new_mask_index(self, indices)
 
-    def from_multi(self, a):
-        import numpy as np
+    def _from_sequence(self, s):
+        return self.new_mask_index(self, s)
 
-        if not isinstance(a, list):
-            a = list(a)
+    def _from_ndarray(self, a):
+        return self._from_sequence(a.tolist())
+        # import numpy as np
 
-        # will raise IndexError if an index is out of bounds
-        n = len(self)
-        indices = np.arange(0, n if n > 0 else 0)
-        indices = indices[a].tolist()
-        return self.new_mask_index(self, indices)
+        # # will raise IndexError if an index is out of bounds
+        # n = len(self)
+        # indices = np.arange(0, n if n > 0 else 0)
+        # indices = indices[a].tolist()
+        # return self.new_mask_index(self, indices)
 
     def from_dict(self, dic):
         return self.sel(dic)

src/earthkit/data/core/metadata.py

@@ -207,7 +207,7 @@ def _get_custom_key(self, key, default=None, raise_on_missing=True, **kwargs):
         if self._is_custom_key(key):
             try:
                 if key == DATETIME:
-                    return self._valid_datetime()
+                    return self._valid_datetime().isoformat()
                 elif key == GRIDSPEC:
                     return self.grid_spec
             except Exception as e:

src/earthkit/data/core/order.py

@@ -13,16 +13,20 @@
 LOG = logging.getLogger(__name__)
 
 
-class Remapping:
+class Remapping(dict):
+    # inherit from dict to make it serialisable
+
     def __init__(self, remapping):
-        self.remapping = {}
+        super().__init__(remapping)
 
+        self.lists = {}
         for k, v in remapping.items():
-            m = re.split(r"\{([^}]*)\}", v)
-            self.remapping[k] = m
+            if isinstance(v, str):
+                v = re.split(r"\{([^}]*)\}", v)
+            self.lists[k] = v
 
     def __call__(self, func):
-        if self.remapping is None or not self.remapping:
+        if not self:
             return func
 
         class CustomJoiner:
@@ -43,9 +47,12 @@ def wrapped(name, **kwargs):
         return wrapped
 
     def substitute(self, name, joiner, **kwargs):
-        if name in self.remapping:
+        if name in self.lists:
+            if callable(self.lists[name]):
+                return self.lists[name]()
+
             lst = []
-            for i, bit in enumerate(self.remapping[name]):
+            for i, bit in enumerate(self.lists[name]):
                 if i % 2:
                     p = joiner.format_name(bit, **kwargs)
                     if p is not None:
@@ -58,19 +65,61 @@ def substitute(self, name, joiner, **kwargs):
         return joiner.format_name(name, **kwargs)
 
     def as_dict(self):
-        return self.remapping
+        return dict(self)
 
 
-def build_remapping(mapping):
+def _build_remapping(mapping):
     if mapping is None:
         return Remapping({})
 
-    if isinstance(mapping, dict):
+    if not isinstance(mapping, Remapping) and isinstance(mapping, dict):
         return Remapping(mapping)
 
     return mapping
 
 
+class Patch(dict):
+    # inherit from dict to make it serialisable
+
+    def __init__(self, proc, name, patch):
+        self.proc = proc
+        self.name = name
+
+        if isinstance(patch, dict):
+            self.patch = lambda x: patch.get(x, x)
+        elif isinstance(patch, (int, bool, float, str)) or patch is None:
+            self.patch = lambda x: patch
+        else:
+            assert callable(patch)
+            self.patch = patch
+
+        # For JSON, we simply forward to the remapping
+        super().__init__(proc.as_dict())
+
+    def __call__(self, func):
+        next = self.proc(func)
+
+        def wrapped(name, **kwargs):
+            result = next(name, **kwargs)
+            if name == self.name:
+                result = self.patch(result)
+            return result
+
+        return wrapped
+        # assert False, (name, self.proc, self.name, self.patch)
+
+    def as_dict(self):
+        return dict(self)
+
+
+def build_remapping(mapping, patches=None):
+    result = _build_remapping(mapping)
+    if patches:
+        for k, v in patches.items():
+            result = Patch(result, k, v)
+    return result
+
+
 def normalize_order_by(*args, **kwargs):
     _kwargs = {}
     for a in args: