Made method to crop dataset to a given region. Added tests and demons…

examples/example_drifters.py

@@ -25,6 +25,13 @@
 ds.traj.plot()
 plt.show()
 
+#%%
+# Cropping to a smaller geographical area
+d_cropped = ds.traj.crop(lonmin=24, lonmax=28, latmin=76.5, latmax=77)
+d_cropped.traj.plot()
+plt.title('Cropped dataset')
+plt.show()
+
 #%%
 # The figure can be customized, combining functionality from Trajan, Xarray and Matplotlib
 ds.isel(trajectory=1).traj.plot(color='b', label=ds.trajectory[1].values,

tests/test_geographical_selection.py

@@ -0,0 +1,42 @@
+import trajan as ta
+import xarray as xr
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def test_clip(barents):
+
+    # Clipping dataset to:
+    lonmin = 24
+    lonmax = 28
+    latmin = 76.5
+    latmax = 77
+
+    # Check that original coverage is larger than bounding box
+    assert barents.traj.tlon.min(skipna=True) < lonmin
+    assert barents.traj.tlon.max(skipna=True) > lonmax
+    assert barents.traj.tlat.min(skipna=True) < latmin
+    assert barents.traj.tlat.max(skipna=True) > latmax
+
+    bc = barents.traj.crop(lonmin=lonmin, lonmax=lonmax, latmin=latmin, latmax=latmax)
+
+    # Check that cropped coverage is within bounding box
+    assert bc.traj.tlon.min(skipna=True) > lonmin
+    assert bc.traj.tlon.max(skipna=True) < lonmax
+    assert bc.traj.tlat.min(skipna=True) > latmin
+    assert bc.traj.tlat.max(skipna=True) < latmax
+
+def test_contained_in(barents):
+
+    # Check that both trajectories are within a large area
+    bc = barents.traj.contained_in(lonmin=0, lonmax=100, latmin=50, latmax=80)
+    assert bc.sizes[bc.traj.trajectory_dim] == 2
+
+    # Check that only one of the trajectories are fully within given area
+    bc = barents.traj.contained_in(lonmin=15, lonmax=28, latmin=72, latmax=77.2)
+    assert bc.sizes[bc.traj.trajectory_dim] == 1
+
+    # Check that none of the trajectories are fully within a small area (although both intersects)
+    bc = barents.traj.contained_in(lonmin=22, lonmax=27, latmin=76.5, latmax=77.2)
+    assert bc.sizes[bc.traj.trajectory_dim] == 0

trajan/traj.py

@@ -881,3 +881,65 @@ def to_2d(self, obs_dim='obs') -> xr.Dataset:
         """
         Convert dataset into a 2D dataset from.
         """
+
+    def crop(self, lonmin=-360, lonmax=360, latmin=-90, latmax=90, shape='not_yet_implemented'):
+        """
+        Remove parts of trajectories outside of given geographical bounds.
+
+        Parameters
+        ----------
+
+        lonmin : float
+            Minimum longitude
+
+        lonmax : float
+            Maximum longitude
+
+        latmin : float
+            Minimum latitude
+
+        latmax : float
+            Maximum latitude
+
+        Returns
+        -------
+
+        Dataset
+            A new Xarray Dataset containing only given area
+        """
+
+        return self.ds.where((self.tlon>lonmin) & (self.tlon<lonmax) &
+                             (self.tlat>latmin) & (self.tlat<latmax))
+
+    def contained_in(self, lonmin=-360, lonmax=360, latmin=-90, latmax=90, shape='not_yet_implemented'):
+        """
+        Return only trajectories fully within given geographical bounds.
+
+        Parameters
+        ----------
+
+        lonmin : float
+            Minimum longitude
+
+        lonmax : float
+            Maximum longitude
+
+        latmin : float
+            Minimum latitude
+
+        latmax : float
+            Maximum latitude
+
+        Returns
+        -------
+
+        Dataset
+            A new Xarray Dataset containing only the trajectories fully within given area.
+        """
+
+        condition = ((self.tlon.min(dim=self.obs_dim)>lonmin) &
+                     (self.tlon.max(dim=self.obs_dim)<lonmax) &
+                     (self.tlat.min(dim=self.obs_dim)>latmin) &
+                     (self.tlat.max(dim=self.obs_dim)<latmax))
+
+        return self.ds.isel({self.trajectory_dim: condition})