example: concatenating trajectories

Author: gauteh

examples/example_concat.py

@@ -0,0 +1,66 @@
+"""
+Concatenating drifter datasets
+==============================
+"""
+import numpy as np
+import lzma
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import xarray as xr
+import pandas as pd
+import trajan as ta
+
+#%%
+# Importing a dataset with two drifters in the Barents Sea
+with lzma.open('barents.nc.xz') as barents:
+    ds = xr.open_dataset(barents)
+    ds.load()
+
+assert 'obs' in ds.dims
+
+#%%
+# Split into two observational datasets for this example
+
+ds = ds.rename(drifter_names='trajectory').traj.condense_obs()
+print(ds)
+
+d1 = ds.isel(trajectory=0).traj.to_1d().traj.to_2d()
+d2 = ds.isel(trajectory=1).traj.to_1d().traj.to_2d()
+print("d1=", d1)
+print("d2=", d2)
+
+#%%
+# Concatenate two 2D datasets (with observation dimension).
+
+dc = xr.concat((d1, d2), dim='trajectory')
+dc = dc.traj.condense_obs()
+print(dc)
+
+assert np.all(ds.lat.values[~np.isnan(ds.lat.values)] ==
+              dc.lat.values[~np.isnan(dc.lat.values)])
+
+#%%
+# Concatenating two 1D datasets, with observations at different times.
+
+d1 = d1.traj.to_1d(
+)  # trivial conversion since `d1` only contains a single trajectory. No need for gridtime.
+d2 = d2.traj.to_1d()  # Also trivial for d2.
+print(d1)
+
+assert 'obs' not in d1.dims
+
+#%%
+# Concatenating two 1D datasets will cause a lot of NaNs to be inserted.
+dc = xr.concat((d1, d2), dim='trajectory')
+print(dc)
+
+assert np.all(ds.lat.values[~np.isnan(ds.lat.values)] ==
+              dc.lat.values[~np.isnan(dc.lat.values)])
+
+#%%
+# Converting to 2D and condensing the dataset will give a cleaner result.
+
+dc = xr.concat((d1.traj.to_2d(), d2.traj.to_2d()),
+               dim='trajectory').traj.condense_obs()
+print(dc)
+assert dc.sizes['obs'] == ds.sizes['obs']

trajan/traj1d.py

@@ -34,7 +34,9 @@ def to_2d(self, obs_dim='obs'):
         time = ds[self.time_varname].rename({
             self.time_varname: obs_dim
         }).expand_dims(
-            dim={self.trajectory_dim: ds.sizes[self.trajectory_dim]})
+            dim={
+                self.trajectory_dim: ds.sizes[self.trajectory_dim]
+            }).assign_coords({self.trajectory_dim: ds[self.trajectory_dim]})
         # TODO should also add cf_role here
         ds = ds.rename({self.time_varname: obs_dim})
         ds[self.time_varname] = time
@@ -119,19 +121,23 @@ def skill_along_trajectory(self, expected, **kwargs):
 
         def skill_matching(traj, expected):
             traj = traj.where(np.isfinite(traj.lon), drop=True)
-            expected_overlap = expected.sel(time=slice(traj.time[0], traj.time[-1]))
+            expected_overlap = expected.sel(
+                time=slice(traj.time[0], traj.time[-1]))
             mask = False
-            if traj.sizes[self.obs_dim] != expected_overlap.sizes[expected_overlap.traj.obs_dim]:
-                traj = traj.sel(time=slice(expected_overlap.time[0], expected_overlap.time[-1]))
+            if traj.sizes[self.obs_dim] != expected_overlap.sizes[
+                    expected_overlap.traj.obs_dim]:
+                traj = traj.sel(time=slice(expected_overlap.time[0],
+                                           expected_overlap.time[-1]))
                 mask = True
             s = traj.traj.skill(expected_overlap, **kwargs)
             if mask is True:
-                s = s*np.nan  # Mask out the skill score if the trajectories are not of equal length
+                s = s * np.nan  # Mask out the skill score if the trajectories are not of equal length
             s['start_time'] = traj.time[0]
             s['end_time'] = traj.time[-1]
             return s
 
-        s = self.ds.groupby(self.trajectory_dim).apply(skill_matching, expected=expected)
+        s = self.ds.groupby(self.trajectory_dim).apply(skill_matching,
+                                                       expected=expected)
         s['start_lon'] = start_lon
         s['start_lat'] = start_lat
         s['end_lon'] = end_lon
@@ -151,8 +157,7 @@ def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
         if numtraj_self > 1 and numtraj_expected > 1 and numtraj_self != numtraj_expected:
             raise ValueError(
                 'Datasets must have the same number of trajectories, or a single trajectory. '
-                f'This dataset: {numtraj_self}, expected: {numtraj_expected}.'
-            )
+                f'This dataset: {numtraj_self}, expected: {numtraj_expected}.')
 
         numobs_self = self.ds.sizes[self.obs_dim]
         numobs_expected = expected.ds.sizes[expected.obs_dim]
@@ -162,9 +167,10 @@ def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
             )
 
         diff = np.max(
-            np.abs((self.ds[self.obs_dim] -
-                    expected.ds[expected.obs_dim]).astype('timedelta64[s]').astype(
-                        np.float64)))
+            np.abs((
+                self.ds[self.obs_dim] -
+                expected.ds[expected.obs_dim]).astype('timedelta64[s]').astype(
+                    np.float64)))
         if not np.isclose(diff, 0):
             raise ValueError(
                 f"The two datasets must have approximately equal time coordinates, maximum difference: {diff} seconds. Consider using `gridtime` to interpolate one of the datasets."
@@ -193,7 +199,8 @@ def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
         else:
             raise ValueError(f"Unknown skill-score method: {method}.")
 
-        s = skill_method(expected.traj.tlon, expected.traj.tlat, ds.traj.tlon, ds.traj.tlat, **kwargs)
+        s = skill_method(expected.traj.tlon, expected.traj.tlat, ds.traj.tlon,
+                         ds.traj.tlat, **kwargs)
 
         newcoords = dict(ds.lon.sizes)
         newcoords.pop('time')

trajan/traj2d.py

@@ -254,6 +254,8 @@ def to_1d(self):
             ds = ds.loc[{self.time_varname: ~pd.isna(ds[self.time_varname])}]
             _, ui = np.unique(ds[self.time_varname], return_index=True)
             ds = ds.isel({self.time_varname: ui})
+            ds = ds.assign_coords(
+                {self.trajectory_dim: ds[self.trajectory_dim]})
 
             return ds