add option to subset and return complete_traj (#332)

Author: philippemiron

.codecov.yml

@@ -8,10 +8,10 @@ coverage:
   status:
     project:
       default:
-        informational: true
+        threshold: 10%
     patch:
       default:
-        informational: true
+        threshold: 10%
 
 comment: off
 

.github/workflows/ci.yml

@@ -26,14 +26,16 @@ jobs:
         with:
           environment-file: environment.yml
           environment-name: clouddrift
-      - name: Testing
+      - name: Run unit tests
         shell: bash -l {0}
         run: |
-              pip install coverage 
-              pip install matplotlib-base cartopy
+              micromamba install coverage matplotlib-base cartopy
               coverage run -m unittest discover -s tests -p "*.py"
+      - name: Create coverage report
+        shell: bash -l {0}
+        run: |
               coverage xml
-      - name: Upload coverage reports to Codecov
+      - name: Upload coverage report to Codecov
         if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.9'
         uses: codecov/codecov-action@v3
         with:

clouddrift/plotting.py

@@ -212,4 +212,8 @@ def plot_ragged(
             start = end
             lines.append(line)
 
+    # set axis limits
+    ax.set_xlim([np.min(longitude), np.max(longitude)])
+    ax.set_ylim([np.min(latitude), np.max(latitude)])
+
     return lines

clouddrift/ragged.py

@@ -543,6 +543,7 @@ def subset(
     rowsize_var_name: str = "rowsize",
     traj_dim_name: str = "traj",
     obs_dim_name: str = "obs",
+    full_trajectories=False,
 ) -> xr.Dataset:
     """Subset the dataset as a function of one or many criteria. The criteria are
     passed as a dictionary, where a variable to subset is assigned to either a
@@ -567,6 +568,10 @@ def subset(
         Name of the trajectory dimension (default is "traj")
     obs_dim_name : str, optional
         Name of the observation dimension (default is "obs")
+    full_trajectories : bool, optional
+        If True, it returns the complete trajectories where at least one observation
+        matches the criteria, rather than just the segments where the criteria are satisfied.
+        Default is False.
 
     Returns
     -------
@@ -582,6 +587,10 @@ def subset(
 
     >>> subset(ds, {"lat": (21, 31), "lon": (-98, -78)})
 
+    The parameter `full_trajectories` can be used to retrieve trajectories passing through a region, for example all trajectories passing through the Gulf of Mexico:
+
+    >>> subset(ds, {"lat": (21, 31), "lon": (-98, -78)}, full_trajectories=True)
+
     Retrieve drogued trajectory segments:
 
     >>> subset(ds, {"drogue_status": True})
@@ -641,6 +650,14 @@ def subset(
         mask_traj, np.in1d(ds[id_var_name], np.unique(ids_with_mask_obs))
     )
 
+    # reset mask_obs to True to keep complete trajectories
+    if full_trajectories:
+        for i in np.where(mask_traj)[0]:
+            mask_obs[slice(traj_idx[i], traj_idx[i + 1])] = True
+        ids_with_mask_obs = np.repeat(
+            ds[id_var_name].values, ds[rowsize_var_name].values
+        )[mask_obs]
+
     if not any(mask_traj):
         warnings.warn("No data matches the criteria; returning an empty dataset.")
         return xr.Dataset()

tests/ragged_tests.py

@@ -667,6 +667,32 @@ def test_arraylike_criterion(self):
         ds_sub = subset(self.ds, {"ID": self.ds["ID"][:2].values})
         self.assertTrue(ds_sub["ID"].size == 2)
 
+    def test_full_trajectories(self):
+        ds_id_rowsize = {
+            i: j for i, j in zip(self.ds.ID.values, self.ds.rowsize.values)
+        }
+
+        ds_sub = subset(self.ds, {"lon": (-125, -111)}, full_trajectories=True)
+        self.assertTrue(all(ds_sub.lon == [-121, -111, 51, 61, 71]))
+
+        ds_sub_id_rowsize = {
+            i: j for i, j in zip(ds_sub.ID.values, ds_sub.rowsize.values)
+        }
+        for k, v in ds_sub_id_rowsize.items():
+            self.assertTrue(ds_id_rowsize[k] == v)
+
+        ds_sub = subset(self.ds, {"lat": (30, 40)}, full_trajectories=True)
+        self.assertTrue(all(ds_sub.lat == [10, 20, 30, 40]))
+
+        ds_sub_id_rowsize = {
+            i: j for i, j in zip(ds_sub.ID.values, ds_sub.rowsize.values)
+        }
+        for k, v in ds_sub_id_rowsize.items():
+            self.assertTrue(ds_id_rowsize[k] == v)
+
+        ds_sub = subset(self.ds, {"time": (4, 5)}, full_trajectories=True)
+        xr.testing.assert_equal(self.ds, ds_sub)
+
 
 class unpack_tests(unittest.TestCase):
     def test_unpack(self):