diff --git a/examples/example_drifters.py b/examples/example_drifters.py index 1a1cedd..5216df9 100644 --- a/examples/example_drifters.py +++ b/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, diff --git a/tests/test_geographical_selection.py b/tests/test_geographical_selection.py new file mode 100644 index 0000000..2a238a6 --- /dev/null +++ b/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 diff --git a/trajan/traj.py b/trajan/traj.py index 6115a51..e0408d4 100644 --- a/trajan/traj.py +++ b/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.tlonlatmin) & (self.tlatlonmin) & + (self.tlon.max(dim=self.obs_dim)latmin) & + (self.tlat.max(dim=self.obs_dim)