diff --git a/src/anemoi/datasets/create/functions/sources/xarray/time.py b/src/anemoi/datasets/create/functions/sources/xarray/time.py index 98cb57d7..a5f2e9be 100644 --- a/src/anemoi/datasets/create/functions/sources/xarray/time.py +++ b/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") diff --git a/tests/xarray/test_zarr.py b/tests/xarray/test_zarr.py index d138d371..8016d108 100644 --- a/tests/xarray/test_zarr.py +++ b/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()