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