Subset_fix (#261)

* make test fail

* fix subset

* cleanup sample_ragged_array creation

* lint

---------

Co-authored-by: Philippe Miron <[email protected]>

Author: philippemiron

clouddrift/analysis.py

@@ -1045,9 +1045,10 @@ def subset(
     else:
         # apply the filtering for both dimensions
         ds_sub = ds.isel({traj_dim_name: mask_traj, obs_dim_name: mask_obs})
-        _, ds_sub[rowsize_var_name].values = np.unique(
-            ids_with_mask_obs, return_counts=True
+        _, unique_idx, sorted_rowsize = np.unique(
+            ids_with_mask_obs, return_index=True, return_counts=True
         )
+        ds_sub[rowsize_var_name].values = sorted_rowsize[np.argsort(unique_idx)]
         return ds_sub
 
 

tests/analysis_tests.py

@@ -25,17 +25,17 @@
 
 
 def sample_ragged_array() -> RaggedArray:
-    drifter_id = [1, 2, 3]
-    rowsize = [5, 4, 2]
-    longitude = [[-121, -111, 51, 61, 71], [12, 22, 32, 42], [103, 113]]
-    latitude = [[-90, -45, 45, 90, 0], [10, 20, 30, 40], [10, 20]]
-    t = [[1, 2, 3, 4, 5], [2, 3, 4, 5], [4, 5]]
-    ids = [[1, 1, 1, 1, 1], [2, 2, 2, 2], [3, 3]]
+    drifter_id = [1, 3, 2]
+    longitude = [[-121, -111, 51, 61, 71], [103, 113], [12, 22, 32, 42]]
+    latitude = [[-90, -45, 45, 90, 0], [10, 20], [10, 20, 30, 40]]
+    t = [[1, 2, 3, 4, 5], [4, 5], [2, 3, 4, 5]]
     test = [
         [True, True, True, False, False],
-        [True, False, False, False],
         [False, False],
+        [True, False, False, False],
     ]
+    rowsize = [len(x) for x in longitude]
+    ids = [[d] * rowsize[i] for i, d in enumerate(drifter_id)]
     nb_obs = np.sum(rowsize)
     nb_traj = len(drifter_id)
     attrs_global = {
@@ -796,14 +796,17 @@ def test_select(self):
     def test_range(self):
         # positive
         ds_sub = subset(self.ds, {"lon": (0, 180)})
+        print(ds_sub)
         traj_idx = np.insert(np.cumsum(ds_sub["rowsize"].values), 0, 0)
         self.assertTrue(
             all(ds_sub.lon[slice(traj_idx[0], traj_idx[1])] == [51, 61, 71])
         )
+
+        self.assertTrue(all(ds_sub.lon[slice(traj_idx[1], traj_idx[2])] == [103, 113]))
+
         self.assertTrue(
-            all(ds_sub.lon[slice(traj_idx[1], traj_idx[2])] == [12, 22, 32, 42])
+            all(ds_sub.lon[slice(traj_idx[2], traj_idx[3])] == [12, 22, 32, 42])
         )
-        self.assertTrue(all(ds_sub.lon[slice(traj_idx[2], traj_idx[3])] == [103, 113]))
 
         # negative range
         ds_sub = subset(self.ds, {"lon": (-180, 0)})