From 5158b0def7efa563f280a88537c529b9d825559a Mon Sep 17 00:00:00 2001 From: Aouatef Merghni Date: Sun, 26 Oct 2025 20:37:42 +0000 Subject: [PATCH] Support UKV variable name change #364: Conditionally rename x,y dims for backward compatibility --- ocf_data_sampler/load/nwp/providers/ukv.py | 23 +++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/ocf_data_sampler/load/nwp/providers/ukv.py b/ocf_data_sampler/load/nwp/providers/ukv.py index 18014199..e1ea3420 100755 --- a/ocf_data_sampler/load/nwp/providers/ukv.py +++ b/ocf_data_sampler/load/nwp/providers/ukv.py @@ -21,14 +21,19 @@ def open_ukv(zarr_path: str | list[str]) -> xr.DataArray: """ ds = open_zarr_paths(zarr_path, backend="tensorstore") - ds = ds.rename( - { - "init_time": "init_time_utc", - "variable": "channel", - "x": "x_osgb", - "y": "y_osgb", - }, - ) + # Build rename dictionary conditionally based on existing dimensions + rename_dict = { + "init_time": "init_time_utc", + "variable": "channel", + } + + # Only rename x,y to x_osgb,y_osgb if the legacy names exist + if "x" in ds.dims: + rename_dict["x"] = "x_osgb" + if "y" in ds.dims: + rename_dict["y"] = "y_osgb" + + ds = ds.rename(rename_dict) check_time_unique_increasing(ds.init_time_utc) @@ -37,4 +42,4 @@ def open_ukv(zarr_path: str | list[str]) -> xr.DataArray: ds = ds.transpose("init_time_utc", "step", "channel", "x_osgb", "y_osgb") # TODO: should we control the dtype of the DataArray? - return get_xr_data_array_from_xr_dataset(ds) + return get_xr_data_array_from_xr_dataset(ds) \ No newline at end of file