Merge pull request #187 from knutfrode/dev

Fundamental change in sel and seltime methods - now not masking NaN-v…

Author: knutfrode

examples/example_concat.py

@@ -32,7 +32,7 @@
 #%%
 # Concatenate two 2D datasets (with observation dimension).
 
-dc = xr.concat((d1, d2), dim='trajectory')
+dc = xr.concat((d1, d2), dim='trajectory', join='outer')
 dc = dc.traj.condense_obs()
 print(dc)
 
@@ -51,16 +51,17 @@
 
 #%%
 # Concatenating two 1D datasets will cause a lot of NaNs to be inserted.
-dc = xr.concat((d1, d2), dim='trajectory')
+d1 = d1.drop_duplicates('time')
+d2 = d2.drop_duplicates('time')
+dc = xr.concat((d1, d2), dim='trajectory', join='outer')
 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()
+               dim='trajectory', join='outer').traj.condense_obs()
 print(dc)
 assert dc.sizes['obs'] == ds.sizes['obs']

trajan/traj.py

@@ -752,6 +752,7 @@ def distance_to_next(self):
         distance = xr.DataArray(distance, coords=lonfrom.coords, dims=lon.dims)
         distance = xr.concat((distance, distance.isel({self.obs_dim: -1})),
                              dim=self.obs_dim)  # repeating last time step to
+        distance = distance.assign_coords({self.obs_dim: np.arange(self.ds.sizes[self.obs_dim])})
         return distance
 
     def azimuth_to_next(self):
@@ -787,6 +788,8 @@ def azimuth_to_next(self):
         azimuth_forward = xr.concat(
             (azimuth_forward, azimuth_forward.isel({self.obs_dim: -1})),
             dim=self.obs_dim)  # repeating last time step to
+        azimuth_forward = azimuth_forward.assign_coords(
+            {self.obs_dim: np.arange(self.ds.sizes[self.obs_dim])})
         return azimuth_forward
 
     def velocity_components(self):

trajan/traj1d.py

@@ -270,7 +270,10 @@ def sel(self, *args, **kwargs):
         return self.ds.sel(*args, **kwargs)
 
     def seltime(self, t0=None, t1=None):
-        return self.ds.sel({self.time_varname: slice(t0, t1)})
+        # Preserving NaN in trajectories, as these provide information about gaps / segments
+        subset_indices = np.where((self.ds[self.time_varname] >= pd.to_datetime(t0)) &
+                                  (self.ds[self.time_varname] <= pd.to_datetime(t1)))[0]
+        return self.ds.isel({self.obs_dim: slice(subset_indices.min(), subset_indices.max()+1)})
 
     def iseltime(self, i):
         return self.ds.isel({self.time_varname: i})

trajan/traj2d.py

@@ -135,8 +135,7 @@ def insert_nan_where(self, condition):
 
             nd[varname] = da.astype(var.dtype)
 
-        nd = nd.drop_vars(
-            (self.obs_dim, self.trajectory_dim))  # Remove coordinates
+        nd = nd.assign_coords({self.obs_dim: np.arange(max_obs)})  # New obs index 1..N
 
         return nd
 
@@ -164,11 +163,14 @@ def drop_where(self, condition):
 
         # Ensure all trajectories have equal length by padding with NaN at the end
         trajs = [
-            t.pad(pad_width={self.obs_dim: (0, newlen - t.sizes[self.obs_dim])})
+            t.pad(pad_width={self.obs_dim: (0, newlen - t.sizes[self.obs_dim])}).
+                    assign_coords({self.obs_dim: np.arange(newlen)})  # New obs index 1..N
             for t in trajs
         ]
 
-        return xr.concat(trajs, dim=self.trajectory_dim)
+        ds = xr.concat(trajs, dim=self.trajectory_dim, join='exact')
+
+        return ds
 
     @__require_obs_dim__
     def condense_obs(self) -> xr.Dataset:
@@ -220,8 +222,7 @@ def condense_obs(self) -> xr.Dataset:
         ds = ds.isel({self.obs_dim: slice(0, maxN)})
 
         # Write new observation coordinate.
-        obs = np.arange(0, maxN)
-        ds = ds.assign_coords({self.obs_dim: obs})
+        ds = ds.assign_coords({self.obs_dim: np.arange(0, maxN)})
 
         return ds
 
@@ -246,11 +247,12 @@ def append(self, da, obs_dims=None):
 
     def sel(self, *args, **kwargs):
         return self.ds.groupby(self.trajectory_dim).map(
-            lambda d: ensure_time_dim(d.traj.to_1d().sel(*args, **kwargs), self
-                                      .time_varname).traj.to_2d(self.obs_dim))
+            lambda d: ensure_time_dim(d.traj.to_1d().traj.sel(*args, **kwargs), self.time_varname).traj.to_2d(self.obs_dim))
 
     def seltime(self, t0=None, t1=None):
-        return self.sel({self.time_varname: slice(t0, t1)})
+        # Using TrajAn sel method that allows NaN
+        return self.ds.groupby(self.trajectory_dim).map(
+            lambda d: ensure_time_dim(d.traj.to_1d().traj.seltime(t0, t1), self.time_varname).traj.to_2d(self.obs_dim))
 
     @__require_obs_dim__
     def iseltime(self, i):
@@ -275,18 +277,24 @@ def to_1d(self):
             )
         else:
             ds = self.ds.copy()
-            ds = ds.dropna(self.obs_dim, how='all')
+
+            # Do not remove NaN's since these now have meaning
+            #ds = ds.dropna(self.obs_dim, how='all')
+
+            # For 1D objects, we rename obs-dimension to name of time variable
+            # so that time becomes a coordinate variable,
+            # i.e. typically:   time(traj, obs) -> time(time)
             ds = ds.assign_coords({self.obs_dim: ds[self.time_varname]})
-            ds = ds.drop_vars(self.time_varname).rename(
-                {self.obs_dim: self.time_varname})
-
-            ds[self.time_varname] = ds[self.time_varname].squeeze(
-                self.trajectory_dim)
-            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]})
+            ds = ds.drop_vars(self.time_varname).rename({self.obs_dim: self.time_varname})
+            ds[self.time_varname] = ds[self.time_varname].squeeze(self.trajectory_dim)
+
+            # Do not remove NaN's since these now have meaning
+            #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})
+
+            # Keep trajectory dimension, although always length 1 for 1D objects
+            ds = ds.assign_coords({self.trajectory_dim: ds[self.trajectory_dim]})
 
             return ds