From 7a4131c0ded60c32d501a65f3ff2bb23c1204807 Mon Sep 17 00:00:00 2001 From: Leo Saffin Date: Fri, 2 Aug 2024 15:08:16 +0100 Subject: [PATCH] Remove trajectory_id variable before adding "track_id" when loading from a netcdf. These will be the same variable if saved by huracanpy resulting in it being dropped. Fix save test to catch this --- huracanpy/_data/_netcdf.py | 4 +++- tests/test_huracanpy.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/huracanpy/_data/_netcdf.py b/huracanpy/_data/_netcdf.py index dc71009..0931aaf 100644 --- a/huracanpy/_data/_netcdf.py +++ b/huracanpy/_data/_netcdf.py @@ -24,11 +24,13 @@ def load(filename, **kwargs): for npoints, tr_id in zip(rowsize.data, trajectory_id.data): trajectory_id_stretched.extend([tr_id] * npoints) + dataset = dataset.drop_vars([trajectory_id.name, rowsize.name]) + dataset["track_id"] = (sample_dimension, trajectory_id_stretched) # Keep attributes (including cf_role) dataset["track_id"].attrs = trajectory_id.attrs - return dataset.drop_vars([trajectory_id.name, rowsize.name]) + return dataset def save(dataset, filename): diff --git a/tests/test_huracanpy.py b/tests/test_huracanpy.py index 82b4b0c..51f51c3 100644 --- a/tests/test_huracanpy.py +++ b/tests/test_huracanpy.py @@ -64,7 +64,7 @@ def test_save(filename, tracker, extension, tmp_path): # Reload the data and check it is still the same data_ = huracanpy.load(str(tmp_path / f"tmp_file.{extension}")) - for var in list(data_.variables) + list(data_.coords): + for var in list(data.variables) + list(data.coords): # Work around for xarray inconsistent loading the data as float or double # depending on fill_value and scale_factor # np.testing.assert_allclose doesn't work for datetime64