Merge pull request #908 from ecmwf/feature/xr-engine-new-field

Feature/xr engine new field

Author: sandorkertesz

src/earthkit/data/core/field.py

@@ -33,7 +33,7 @@
     "time.base_datetime",
     "time.step",
     "vertical.level",
-    "vertical.type",
+    "vertical.level_type",
     "ensemble.member",
     "geography.grid_type",
 ]

src/earthkit/data/indexing/tensor.py

@@ -502,8 +502,9 @@ def _subset(self, indexes):
 
     def make_valid_datetime(self, dims_map, dtype="datetime64[ns]"):
         # TODO: make it more general
+        # PW: TODO: make it more general - it could allow to use it when allow_holes=True
 
-        for k in ["valid_datetime", "valid_time"]:
+        for k in ["valid_time", "time.valid_datetime", "metadata.valid_time", "metadata.valid_datetime"]:
             if k in self.user_coords:
                 import datetime
 
@@ -512,10 +513,24 @@ def make_valid_datetime(self, dims_map, dtype="datetime64[ns]"):
         # in the tensor the dims.coords are GRIB keys
         # dims_map is a mapping from dim names to GRIB keys
         DIM_ROLES = {
-            "forecast_reference_time": ("forecast_reference_time", "base_datetime"),
-            "step": ("step_timedelta", "step", "ensStep", "stepRange"),
-            "date": ("date", "dataDate"),
-            "time": ("time", "dataTime"),
+            "forecast_reference_time": (
+                "forecast_reference_time",
+                "time.forecast_reference_time",
+                "time.base_datetime",
+                "metadata.base_datetime",
+                "metadata.indexing_datetime",
+                "metadata.indexing_time",
+            ),
+            "step": (
+                "step",
+                "time.step",
+                "metadata.step_timedelta",
+                "metadata.step",
+                "metadata.endStep",
+                "metadata.stepRange",
+            ),
+            "date": ("date", "metadata.dataDate"),
+            "time": ("time", "metadata.dataTime"),
         }
 
         # map dim roles to keys available in the tensor
@@ -561,7 +576,7 @@ def make_valid_datetime(self, dims_map, dtype="datetime64[ns]"):
                     }
 
                     vals = np.array(
-                        [x for x in self.source.sel(**other_coords).get("valid_datetime")],
+                        [x for x in self.source.sel(**other_coords).get("time.valid_datetime")],
                         dtype=dtype,
                     )
 
@@ -573,7 +588,7 @@ def make_valid_datetime(self, dims_map, dtype="datetime64[ns]"):
                     import numpy as np
 
                     vals = np.array(
-                        [x for x in self.source.get("valid_datetime")],
+                        [x for x in self.source.get("time.valid_datetime")],
                         dtype=dtype,
                     )
 

src/earthkit/data/xr_engine/attrs.py

@@ -107,13 +107,23 @@ class NamespaceAttr(Attr):
 
     def __init__(self, name, ns=None):
         super().__init__(name, remove_prefix=ns is None)
-        self.ns = ns if ns is not None else name
+        if ns is None:
+            ns = name
+
+        if ns.startswith("metadata."):
+            self.ns = "metadata"
+            self.filter = ns.removeprefix("metadata.")
+        else:
+            self.ns = ns
+            self.filter = None
 
     def get(self, field):
-        # PW: TODO
-        # print(f"{self=}, {field=}", flush=True)
-        raise NotImplementedError
-        return {self.name: field.as_namespace(self.ns)}
+        namespace_attr = field._dump_component(self.ns, filter=self.filter)
+        if self.filter is not None:
+            # namespace_attr is of the form `{self.filter: actual_namespace_attr_dict}`
+            # so need to extract the `actual_namespace_attr_dict`
+            namespace_attr = namespace_attr[self.filter]
+        return {self.name: namespace_attr}
 
     def __repr__(self) -> str:
         return f"NamespaceAttr({self.name}, ns={self.ns})"

src/earthkit/data/xr_engine/builder.py

@@ -645,7 +645,7 @@ def parse(self, ds, profile=None, full=False):
         return ds_xr, profile
 
     def grid(self, ds):
-        grids = ds.index("md5GridSection")
+        grids = ds.index("metadata.md5GridSection")
 
         if not grids:
             grid = "_custom_" + str(id(ds))

src/earthkit/data/xr_engine/check.py

@@ -76,8 +76,8 @@ def neighbour_field(self, field_num, index):
         return f_other, index_other
 
     def namespace_diff(self, f, f_other):
-        meta = f.namespace()
-        meta_other = f_other.namespace()
+        meta = f._dump_component()
+        meta_other = f_other._dump_component()
         return DictDiff.diff(meta, meta_other)
 
     def meta_diff(self, f, f_other, coords_keys):
@@ -147,7 +147,7 @@ def check(self, details=False):
                         text_ns += f"{k}:\n {v}\n"
 
                 md = {"coordinates": self.meta(f, coord_keys)}
-                md.update(f.namespace())
+                md.update(f._dump_component())
 
                 text_meta = f"\nField[{index}] metadata:\n"
                 for k, v in md.items():

src/earthkit/data/xr_engine/coord.py

@@ -177,21 +177,30 @@ def attrs(self, profile):
 
 class LevelCoord(Coord):
     def __init__(self, name, vals, dims=None, ds=None, **kwargs):
-        self._level_type = None
+        self._level_type = {}
         self._cf = None
         if ds is not None:
             self._cf = ds[0].vertical.cf()
-            self._level_type = ds[0].vertical.level_type()
+            for k in ["vertical.level_type", "metadata.typeOfLevel", "metadata.levtype"]:
+                v = ds[0]._get_fast(k)
+                if v is not None:
+                    self._level_type[k] = v
 
         super().__init__(name, vals, dims, **kwargs)
 
     def attrs(self, profile):
         attrs = profile.attrs
         conf = attrs.coord_attrs.get(self.name, {})
-        res = {}
+        _attrs = {}
         if conf:
-            res = conf.get(self._level_type, {})
-        if not res:
-            res = self._cf
-        return res
+            keys = conf["keys"]
+            for key in keys:
+                level_type = self._level_type.get(key)
+                if level_type is None:
+                    continue
+                _attrs = conf.get(level_type)
+                if _attrs is not None:
+                    return _attrs
+        _attrs = self._cf
+        return _attrs
         # raise ValueError(f"Cannot determine level type for coordinate {name}")

src/earthkit/data/xr_engine/defaults.yaml

@@ -83,7 +83,11 @@ coord_attrs:
     # "calendar": "proleptic_gregorian",
     standard_name: time
     long_name: valid_time
-  level:
+  metadata.level: &_level
+    keys:
+    - level_type
+    - vertical.level_type
+    - metadata.typeOfLevel
     depthBelowLand:
       units: m
       positive: down
@@ -151,6 +155,7 @@ coord_attrs:
       positive: up
       long_name: air_potential_temperature
       standard_name: air_potential_temperature
+  _metadata.levelist: &_levelist
     ml:
       units: 1
       positive: down
@@ -180,3 +185,19 @@ coord_attrs:
       units: 1
       long_name: snow layer
       standard_name: unknown
+  metadata.levelist:
+    <<: *_levelist
+    keys:
+    - level_type
+    - vertical.level_type
+    - metadata.levtype
+  levelist:
+    <<: *_levelist
+    keys:
+    - level_type
+    - vertical.level_type
+    - metadata.levtype
+  vertical.level:
+    <<: [*_level, *_levelist]
+  level:
+    <<: [*_level, *_levelist]

src/earthkit/data/xr_engine/diff.py

@@ -27,7 +27,7 @@ def _flatten_dict(d, parent_key="", sep="."):
         flat = {}
         for key, value in d.items():
             new_key = f"{parent_key}{sep}{key}" if parent_key else key
-            if isinstance(value, dict) and len(value) > 0:
+            if isinstance(value, dict):
                 flat.update(DictDiff._flatten_dict(value, parent_key=new_key, sep=sep))
             else:
                 flat[new_key] = value
@@ -89,6 +89,8 @@ def _compare(v1, v2):
             return v1 == v2, ListDiff.VALUE_DIFF
         elif isinstance(v1, datetime.timedelta) and isinstance(v2, datetime.timedelta):
             return v1 == v2, ListDiff.VALUE_DIFF
+        elif v1 is None and v2 is None:
+            return True, ListDiff.VALUE_DIFF
         elif type(v1) is not type(v2):
             return False, ListDiff.TYPE_DIFF
         else:

src/earthkit/data/xr_engine/dim.py

@@ -75,26 +75,24 @@ def _get_metadata_keys(keys):
 _GRIB_MONTH_KEYS = ["forecastMonth", "fcmonth"]
 MONTH_KEYS = _get_metadata_keys(_GRIB_MONTH_KEYS)
 
-_GRIB_VALID_DATETIME_KEYS = ["valid_time", "valid_datetime"]
-_VALID_DATETIME_KEYS = ["time.valid_datetime"]
+_GRIB_VALID_DATETIME_KEYS = ["valid_datetime"]
+_VALID_DATETIME_KEYS = ["valid_datetime"]
 VALID_DATETIME_KEYS = (
     ["valid_time"]
     + _get_component_keys("time", _VALID_DATETIME_KEYS)
     + _get_metadata_keys(_GRIB_VALID_DATETIME_KEYS)
 )
 
 _GRIB_BASE_DATETIME_KEYS = [
-    "forecast_reference_time",
-    "base_time",
     "base_datetime",
-    "reference_time",
-    "reference_datetime",
-    "indexing_time",
+    "indexingDate",
+    "indexingTime",
     "indexing_datetime",
+    "indexing_time",
 ]
 _BASE_DATETIME_KEYS = ["forecast_reference_time", "base_datetime"]
 BASE_DATETIME_KEYS = (
-    ["forecast_reference_time"]
+    ["forecast_reference_time", "indexing_time"]
     + _get_component_keys("time", _BASE_DATETIME_KEYS)
     + _get_metadata_keys(_GRIB_BASE_DATETIME_KEYS)
 )

src/earthkit/data/xr_engine/fieldlist.py

@@ -58,7 +58,7 @@ def index(self, key, maker=None):
         if key not in self._index:
             # # LOG.debug(f"Key={key} not found in IndexDB")
             if maker is not None:
-                self._index[key] = maker(key)[0][key]
+                self._index[key] = maker(key)[key]
             else:
                 raise KeyError(f"Could not find index for {key=}")
         return self._index[key]