work on tests

src/anemoi/datasets/create/functions/sources/xarray/time.py

@@ -8,9 +8,6 @@
 #
 
 
-import datetime
-
-
 class Time:
     @classmethod
     def from_coordinates(cls, coordinates):
@@ -19,16 +16,16 @@ def from_coordinates(cls, coordinates):
         date_coordinate = [c for c in coordinates if c.is_date]
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 1:
-            return ForecasstFromValidTimeAndStep(step_coordinate[0])
+            return ForecasstFromValidTimeAndStep(time_coordinate[0], step_coordinate[0])
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
-            return Analysis()
+            return Analysis(time_coordinate[0])
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 0 and len(step_coordinate) == 0:
             return Constant()
 
         if len(date_coordinate) == 1 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
-            return ForecastFromValidTimeAndBaseTime(date_coordinate[0])
+            return ForecastFromValidTimeAndBaseTime(date_coordinate[0], time_coordinate[0])
 
         if len(date_coordinate) == 1 and len(time_coordinate) == 0 and len(step_coordinate) == 1:
             return ForecastFromBaseTimeAndDate(date_coordinate[0], step_coordinate[0])
@@ -38,62 +35,73 @@ def from_coordinates(cls, coordinates):
 
 class Constant(Time):
 
-    def fill_time_metadata(self, time, metadata):
-        metadata["date"] = time.strftime("%Y%m%d")
-        metadata["time"] = time.strftime("%H%M")
-        metadata["step"] = 0
+    def fill_time_metadata(self, coords_values, metadata):
+        raise NotImplementedError("Constant time not implemented")
+        # print("Constant", coords_values, metadata)
+        # metadata["date"] = time.strftime("%Y%m%d")
+        # metadata["time"] = time.strftime("%H%M")
+        # metadata["step"] = 0
 
 
 class Analysis(Time):
 
-    def fill_time_metadata(self, time, metadata):
-        metadata["date"] = time.strftime("%Y%m%d")
-        metadata["time"] = time.strftime("%H%M")
+    def __init__(self, time_coordinate):
+        self.time_coordinate_name = time_coordinate.variable.name
+
+    def fill_time_metadata(self, coords_values, metadata):
+        valid_datetime = coords_values[self.time_coordinate_name]
+
+        metadata["date"] = valid_datetime.strftime("%Y%m%d")
+        metadata["time"] = valid_datetime.strftime("%H%M")
         metadata["step"] = 0
 
+        return valid_datetime
+
 
 class ForecasstFromValidTimeAndStep(Time):
     def __init__(self, step_coordinate):
         self.step_name = step_coordinate.variable.name
 
-    def fill_time_metadata(self, time, metadata):
-        step = metadata.pop(self.step_name)
-        assert isinstance(step, datetime.timedelta)
-        base = time - step
+    def fill_time_metadata(self, coords_values, metadata):
+        raise NotImplementedError("Constant time not implemented")
+        # step = metadata.pop(self.step_name)
+        # assert isinstance(step, datetime.timedelta)
+        # base = time - step
 
-        hours = step.total_seconds() / 3600
-        assert int(hours) == hours
+        # hours = step.total_seconds() / 3600
+        # assert int(hours) == hours
 
-        metadata["date"] = base.strftime("%Y%m%d")
-        metadata["time"] = base.strftime("%H%M")
-        metadata["step"] = int(hours)
+        # metadata["date"] = base.strftime("%Y%m%d")
+        # metadata["time"] = base.strftime("%H%M")
+        # metadata["step"] = int(hours)
 
 
 class ForecastFromValidTimeAndBaseTime(Time):
     def __init__(self, date_coordinate):
         self.date_coordinate = date_coordinate
 
-    def fill_time_metadata(self, time, metadata):
+    def fill_time_metadata(self, coords_values, metadata):
+        raise NotImplementedError("Constant time not implemented")
 
-        step = time - self.date_coordinate
+        # step = time - self.date_coordinate
 
-        hours = step.total_seconds() / 3600
-        assert int(hours) == hours
+        # hours = step.total_seconds() / 3600
+        # assert int(hours) == hours
 
-        metadata["date"] = self.date_coordinate.single_value.strftime("%Y%m%d")
-        metadata["time"] = self.date_coordinate.single_value.strftime("%H%M")
-        metadata["step"] = int(hours)
+        # metadata["date"] = self.date_coordinate.single_value.strftime("%Y%m%d")
+        # metadata["time"] = self.date_coordinate.single_value.strftime("%H%M")
+        # metadata["step"] = int(hours)
 
 
 class ForecastFromBaseTimeAndDate(Time):
     def __init__(self, date_coordinate, step_coordinate):
-        self.date_coordinate = date_coordinate.name
-        self.step_coordinate = step_coordinate.name
+        self.date_coordinate_name = date_coordinate.name
+        self.step_coordinate_name = step_coordinate.name
 
     def fill_time_metadata(self, coords_values, metadata):
 
-        date = coords_values[self.date_coordinate]
-        step = coords_values[self.step_coordinate]
+        date = coords_values[self.date_coordinate_name]
+        step = coords_values[self.step_coordinate_name]
 
         metadata["date"] = date.strftime("%Y%m%d")
         metadata["time"] = date.strftime("%H%M")

tests/xarray/test_zarr.py

@@ -22,7 +22,7 @@ def test_arco_era5():
     print(len(fs))
 
     print(fs[-1].metadata())
-    print(fs[-1].to_numpy())
+    # print(fs[-1].to_numpy())
 
     assert len(fs) == 128677526
 
@@ -47,8 +47,9 @@ def test_weatherbench():
 
     assert len(fs) == 2430240
 
-    assert fs[0].metadata("valid_datetime") == "2020-01-01T00:00:00", fs[0].metadata("valid_datetime")
+    assert fs[0].metadata("valid_datetime") == "2020-01-01T06:00:00", fs[0].metadata("valid_datetime")
+    assert fs[-1].metadata("valid_datetime") == "2021-01-10T12:00:00", fs[-1].metadata("valid_datetime")
 
 
 if __name__ == "__main__":
-    test_weatherbench()
+    test_arco_era5()