Enforce that the input to huracanpy.save is not modified, but sort by track_id for the saved file to ensure the contiguous ragged data makes sense

Author: leosaffin

huracanpy/_data/_load.py

@@ -14,6 +14,8 @@
     lat_track="lat",
     # Names for CHAZ netCDF
     stormID="track_id",
+    # Names for TRACK netCDF
+    TRACK_ID="track_id",
 )
 
 

huracanpy/_data/_netcdf.py

@@ -47,11 +47,10 @@ def save(dataset, filename):
             f"{trajectory_id.name} spans multiple dimensions, should be 1d"
         )
 
-    # np.unique returns a sorted array, so return the index so that the trajectory_ids
-    # can be put back in the same order as they are in the original dataset otherwise
-    # the ordering of data can be messed up if the trajectories ids aren't monotonic
-    trajectory_ids, idx = np.unique(trajectory_id, return_index=True)
-    trajectory_ids = trajectory_id[sorted(idx)].values
+    # Sort by trajectory_id so each track can be described by the first index and
+    # number of elements of the unique trajectory id
+    dataset = dataset.sortby(trajectory_id.name)
+    trajectory_ids = np.unique(trajectory_id)
     rowsize = [np.count_nonzero(trajectory_id == x) for x in trajectory_ids]
 
     dataset[trajectory_id.name] = ("trajectory", trajectory_ids)
@@ -75,9 +74,10 @@ def stretch_trid(dataset):
 
     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
+    dataset[trajectory_id.name] = (sample_dimension, trajectory_id_stretched)
+    # Keep attributes (add cf_role if not already there)
+    dataset[trajectory_id.name].attrs = trajectory_id.attrs
+    dataset[trajectory_id.name].attrs["cf_role"] = "trajectory_id"
 
     return dataset
 
@@ -118,7 +118,6 @@ def _find_trajectory_id(dataset):
         return trajectory_id[0]
     else:
         if "track_id" in dataset:
-            dataset["track_id"].attrs["cf_role"] = "trajectory_id"
             return dataset["track_id"]
         else:
             raise ValueError(

tests/test_huracanpy.py

@@ -65,29 +65,40 @@ def test_save(filename, source, extension, muddle, tmp_path):
     # Caused an issue because they got sorted before
     if muddle:
         data = data.sortby("track_id", ascending=False)
+
     # Copy the data because save modifies the dataset at the moment
-    huracanpy.save(data.copy(), str(tmp_path / f"tmp_file.{extension}"))
+    data_orig = data.copy()
+    huracanpy.save(data, str(tmp_path / f"tmp_file.{extension}"))
+
+    # Check that the original data is not modified by the save function
+    _assert_dataset_identical(data_orig, data)
 
     # Reload the data and check it is still the same
-    data_ = huracanpy.load(str(tmp_path / f"tmp_file.{extension}"))
+    # Saving as netcdf does force sorting by track_id so apply this
+    if extension == "nc":
+        data = data.sortby("track_id")
+    data_reload = huracanpy.load(str(tmp_path / f"tmp_file.{extension}"))
+    _assert_dataset_identical(data, data_reload)
+
 
-    assert len(data.variables) == len(data_.variables)
-    assert len(data.coords) == len(data_.coords)
-    for var in list(data.variables) + list(data.coords):
+def _assert_dataset_identical(ds1, ds2):
+    assert len(ds1.variables) == len(ds2.variables)
+    assert len(ds1.coords) == len(ds2.coords)
+    for var in list(ds1.variables) + list(ds1.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, object, or string
-        if np.issubdtype(data[var].dtype, np.number):
-            if data[var].dtype != data_[var].dtype:
+        if np.issubdtype(ds1[var].dtype, np.number):
+            if ds1[var].dtype != ds2[var].dtype:
                 rtol = 1e-6
             else:
                 rtol = 0
             np.testing.assert_allclose(
-                data[var].data.astype(data_[var].dtype), data_[var].data, rtol=rtol
+                ds1[var].data.astype(ds2[var].dtype), ds2[var].data, rtol=rtol
             )
         else:
-            assert (data[var].data == data_[var].data).all()
+            assert (ds1[var].data == ds2[var].data).all()
 
-    assert len(data.attrs) == len(data_.attrs)
-    for attr in data.attrs:
-        assert data.attrs[attr] == data_.attrs[attr]
+    assert len(ds1.attrs) == len(ds2.attrs)
+    for attr in ds1.attrs:
+        assert ds1.attrs[attr] == ds2.attrs[attr]