Merge pull request #143 from knutfrode/dev

Trajectory dimension is not anymore hardcoded as 'trajectory', but is…

Author: knutfrode

examples/example_parcels.py

@@ -8,8 +8,18 @@
 import trajan as ta
 
 ds = xr.open_dataset('../tests/test_data/parcels.zarr', engine='zarr')
+#%%
+# Print Xarray dataset
 print(ds)
-ds.traj.plot(land='mask', margin=2)
+
+#%%
+# Print trajectory specific information about dataset
+print(ds.traj)
+
+#%%
+# Basic plot
+ds.traj.plot(land='mask', margin=1)
+# TODO: we must allow no time dimension for the below to work
 #ds.mean('trajectory', skipna=True).traj.plot(color='r', label='Mean trajectory')
 
 plt.show()

tests/test_convert_datalayout.py

@@ -6,7 +6,7 @@
 
 def test_to2d(barents):
     # print(barents)
-    gr = barents.traj.gridtime('1H')
+    gr = barents.traj.gridtime('1h')
     # print(gr)
 
     assert gr.traj.is_1d()

tests/test_repr.py

@@ -5,7 +5,7 @@ def test_repr_1d(opendrift_sim):
     repr = str(opendrift_sim.traj)
     assert '2015-11-16T00:00' in repr
     assert 'Timestep:       1:00:00' in repr
-    assert "67 timesteps    time['time'] (1D)" in repr
+    assert "67 timesteps      [obs_dim: time]" in repr
 
 def test_repr_2d(test_data):
     ds = xr.open_dataset(test_data / 'bug32.nc')

tests/test_skill_score.py

@@ -38,10 +38,10 @@ def test_barents_align(barents):
     b0 = barents.isel(trajectory=0)
 
     # assert b0.sizes['trajectory'] == 1
-    assert barents.sizes['trajectory'] == 2
+    assert barents.sizes[barents.traj.trajectory_dim] == 2
 
     (b01, _) = xr.broadcast(b0, barents)
-    b01 = b01.transpose('trajectory', ...)
+    b01 = b01.transpose(b01.traj.trajectory_dim, ...)
 
     np.testing.assert_allclose(b01.isel(trajectory=0).lon, barents.isel(trajectory=0).lon)
     np.testing.assert_allclose(b01.isel(trajectory=1).lon, barents.isel(trajectory=0).lon)

trajan/accessor.py

@@ -22,18 +22,54 @@ def detect_time_variable(ds, obs_dim):
 
     raise ValueError("No time variable detected")
 
+def detect_trajectory_dim(ds):
+    logger.debug('Detecting trajectory dimension')
+    if 'trajectory_id' in ds.cf.cf_roles:  # This is the proper CF way
+        trajectory_var = ds.cf['trajectory_id']
+        if trajectory_var.name in trajectory_var.sizes:  # Check that this is a dimension
+            return trajectory_var.name
+        else:
+            if len(trajectory_var.dims) > 1:
+                logger.warning(f'trajectory_id is {trajectory_var.name}, but dimensions '
+                               'have other names: {str(list(trajectory_var.sizes))}')
+            elif len(trajectory_var.dims) == 1:  # Using the single dimension name
+                return list(trajectory_var.sizes)[0]
+            else:
+                logger.warning('Single trajectory, a trajectory dimension will be added')
+                return None
+
+    logger.warning('No trajectory_id attribute/variable found, trying to identify by name.')
+    tx = detect_tx_variable(ds)
+    for tdn in ['trajectory', 'traj']:  # Common names of trajectory dimension
+        if tdn in tx.dims:
+            return tdn
+
+    return None  # Did not succeed in detecting trajectory dimension
+
 
 @xr.register_dataset_accessor("traj")
 class TrajA(Traj):
     def __new__(cls, ds):
-        if 'traj' in ds.dims:
-            logger.info(
-                'Normalizing dimension name from "traj" to "trajectory".')
-            ds = ds.rename({'traj': 'trajectory'})
 
-        if 'trajectory' not in ds.dims:  # Add empty trajectory dimension, if single trajectory
-            ds = ds.expand_dims({'trajectory': 1})
-            ds['trajectory'].attrs['cf_role'] = 'trajectory_id'
+        trajectory_dim = detect_trajectory_dim(ds)
+
+        if trajectory_dim is None:
+            if 'trajectory_id' in ds.cf.cf_roles:
+                trajectory_id = ds.cf.cf_roles['trajectory_id']
+                if len(trajectory_id) > 1:
+                        raise ValueError(f'Dataset has several trajectory_id variables: {trajectory_id}')
+                else:
+                    trajectory_id = trajectory_id[0]
+                    logger.warning(f'Using trajectory_id variable name ({trajectory_id}) '
+                                   'as trajectory dimension name')
+                    trajectory_dim = trajectory_id
+                    ds = ds.set_coords(trajectory_dim)
+                    ds = ds.expand_dims(trajectory_dim, create_index_for_new_dim=False)
+            else:
+                logger.debug('Creating new trajectory dimension "trajectory"')
+                trajectory_dim = 'trajectory'
+                ds = ds.expand_dims({trajectory_dim: 1})
+                ds[trajectory_dim].attrs['cf_role'] = 'trajectory_id'
 
         obs_dim = None
         time_varname = None
@@ -65,8 +101,7 @@ def __new__(cls, ds):
                 else:
                     raise ValueError(f"cannot deduce the timecoord; we have the following candidates: {with_standard_name_time_and_dim_index = }")
 
-                # discover the trajectorycoord variable name #################
-                trajectorycoord = ds.cf["trajectory_id"].name
+                # KFD TODO: the below detection should be generalized to dynamic dimension names
 
                 # discover the "rowsize" variable name #######################
                 # NOTE: this is probably not standard; something to point to the CF conventions? should we need a standard_name for this, instead of the following heuristics?
@@ -82,10 +117,10 @@ def __new__(cls, ds):
                     raise ValueError("mismatch between the index length and the sum of the deduced trajectory lengths")
 
                 logger.debug(
-                    f"1D storage dataset; detected: {obs_dim = }, {timecoord = }, {trajectorycoord = }, {rowsizevar}"
+                    f"1D storage dataset; detected: {obs_dim = }, {timecoord = }, {trajectory_dim = }, {rowsizevar}"
                 )
 
-                return ocls(ds, obs_dim, timecoord, trajectorycoord, rowsizevar)
+                return ocls(ds, trajectory_dim, obs_dim, timecoord, rowsizevar)
 
             else:
                 logging.warning(f"{ds} has {tx.dims = } which is of dimension 1 but is not index; this is a bit unusual; try to parse with Traj1d or Traj2d")
@@ -137,4 +172,5 @@ def __new__(cls, ds):
                 f'Time variable has more than two dimensions: {ds[time_varname].shape}'
             )
 
-        return ocls(ds, obs_dim, time_varname)
+        # TODO: The provided attributes could perhaps be added here before returning
+        return ocls(ds, trajectory_dim, obs_dim, time_varname)

trajan/ragged.py

@@ -11,19 +11,17 @@
 class ContiguousRagged(Traj):
     """An unstructured dataset, where each trajectory may have observations at different times, and all the data for the different trajectories are stored in single arrays with one dimension, contiguously, one trajectory after the other. Typically from a collection of drifters. This class convert continous ragged datasets into 2d datasets, so that the Traj2d methods can be leveraged."""
 
-    trajdim: str
-    rowvar: str
+    rowvar: str  # TODO: Should have a more precise name than rowvar
 
-    def __init__(self, ds, obs_dim, time_varname, trajectorycoord, rowsizevar):
-        self.trajdim = trajectorycoord
+    def __init__(self, ds, trajectory_dim, obs_dim, time_varname, rowsizevar):
         self.rowvar = rowsizevar
-        super().__init__(ds, obs_dim, time_varname)
+        super().__init__(ds, trajectory_dim, obs_dim, time_varname)
 
     def to_2d(self, obs_dim='obs'):
         """This actually converts a contiguous ragged xarray Dataset into an xarray Dataset that follows the Traj2d conventions."""
         global_attrs = self.ds.attrs
 
-        nbr_trajectories = len(self.ds[self.trajdim])
+        nbr_trajectories = len(self.ds[self.trajectory_dim])
 
         # find the longest trajectory
         longest_trajectory = np.max(self.ds[self.rowvar].to_numpy())
@@ -32,7 +30,7 @@ def to_2d(self, obs_dim='obs'):
 
         # the trajectory dimension special case (as it is a different kind, and has a different dim than other variables)
 
-        array_instruments = self.ds[self.trajdim].to_numpy()
+        array_instruments = self.ds[self.trajectory_dim].to_numpy()
 
         # the time var (special case as it is of a different type)
 

trajan/traj.py

@@ -41,37 +41,40 @@ class Traj:
 
     __gcrs__: pyproj.CRS
 
-    def __init__(self, ds, obs_dim, time_varname):
+    def __init__(self, ds, trajectory_dim, obs_dim, time_varname):
         self.ds = ds
         self.__plot__ = None
         self.__animate__ = None
         self.__gcrs__ = pyproj.CRS.from_epsg(4326)
+        self.trajectory_dim = trajectory_dim  # name of trajectory dimension
         self.obs_dim = obs_dim  # dimension along which time increases
         self.time_varname = time_varname
 
     def __repr__(self):
         output = '=======================\n'
         output += 'TrajAn info:\n'
         output += '------------\n'
-        output += f'{self.ds.sizes["trajectory"]} trajectories\n'
-        if 'time' in self.ds.variables:
-            if self.time_varname in self.ds.sizes:
-                output += f'{self.ds.sizes[self.time_varname]} timesteps'
-                timevar = self.ds[self.time_varname]
-                output += f'    {timevar.name}{list(timevar.sizes)} ({len(timevar.sizes)}D)\n'
+        if self.trajectory_dim is None:
+            output += 'Single trajectory (no trajectory dimension)\n'
+        else:
+            output += f'{self.ds.sizes[self.trajectory_dim]} trajectories  [trajectory_dim: {self.trajectory_dim}]\n'
+        if self.time_varname is not None:
+            output += f'{self.ds.sizes[self.obs_dim]} timesteps      [obs_dim: {self.obs_dim}]\n'
+            timevar = self.ds[self.time_varname]
+            output += f'Time variable:    {timevar.name}{list(timevar.sizes)}   ({len(timevar.sizes)}D)\n'
             try:
                 timestep = self.timestep()
                 timestep = timedelta(seconds=int(timestep))
             except:
                 timestep = '[self.timestep returns error]'  # TODO
             output += f'Timestep:       {timestep}\n'
-            start_time = self.ds.time.min().data
-            end_time = self.ds.time.max().data
+            start_time = self.ds.time.min(skipna=True).data
+            end_time = self.ds.time.max(skipna=True).data
             output += f'Time coverage:  {start_time} - {end_time}\n'
         else:
             output += f'Dataset has no time variable'
-        output += f'Longitude span: {self.tx.min().data} to {self.tx.max().data}\n'
-        output += f'Latitude span:  {self.ty.min().data} to {self.ty.max().data}\n'
+        output += f'Longitude span: {self.tx.min(skipna=True).data} to {self.tx.max(skipna=True).data}\n'
+        output += f'Latitude span:  {self.ty.min(skipna=True).data} to {self.ty.max(skipna=True).data}\n'
         output += 'Variables:\n'
         for var in self.ds.variables:
             if var not in ['trajectory', self.obs_dim]:
@@ -352,31 +355,23 @@ def assign_cf_attrs(self,
         """
         ds = self.ds.copy(deep=True)
 
-        ds['trajectory'] = ds['trajectory'].astype(str)
-        ds['trajectory'].attrs = {
+        ds[self.trajectory_dim] = ds[self.trajectory_dim].astype(str)
+        ds[self.trajectory_dim].attrs = {
             'cf_role': 'trajectory_id',
             'long_name': 'trajectory name'
         }
 
         ds = ds.assign_attrs({
-            'Conventions':
-            'CF-1.10',
-            'featureType':
-            'trajectory',
-            'geospatial_lat_min':
-            np.nanmin(self.tlat),
-            'geospatial_lat_max':
-            np.nanmax(self.tlat),
-            'geospatial_lon_min':
-            np.nanmin(self.tlon),
-            'geospatial_lon_max':
-            np.nanmax(self.tlon),
-            'time_coverage_start':
-            pd.to_datetime(
+            'Conventions': 'CF-1.10',
+            'featureType': 'trajectory',
+            'geospatial_lat_min': np.nanmin(self.tlat),
+            'geospatial_lat_max': np.nanmax(self.tlat),
+            'geospatial_lon_min': np.nanmin(self.tlon),
+            'geospatial_lon_max': np.nanmax(self.tlon),
+            'time_coverage_start': pd.to_datetime(
                 np.nanmin(ds['time'].values[ds['time'].values != np.datetime64(
                     'NaT')])).isoformat(),
-            'time_coverage_end':
-            pd.to_datetime(
+            'time_coverage_end': pd.to_datetime(
                 np.nanmax(ds['time'].values[ds['time'].values != np.datetime64(
                     'NaT')])).isoformat(),
         })
@@ -479,24 +474,25 @@ def distance_to(self, other) -> xr.Dataset:
         """
 
         other = other.broadcast_like(self.ds)
+
         geod = pyproj.Geod(ellps='WGS84')
-        az_fwd, a2, distance = geod.inv(self.ds.traj.tlon, self.ds.traj.tlat,
+        az_fwd, a2, distance = geod.inv(self.tlon, self.tlat,
                                         other.traj.tlon, other.traj.tlat)
 
         ds = xr.Dataset()
         ds['distance'] = xr.DataArray(distance,
                                       name='distance',
-                                      coords=self.ds.traj.tlon.coords,
+                                      coords=self.tlon.coords,
                                       attrs={'units': 'm'})
 
         ds['az_fwd'] = xr.DataArray(az_fwd,
                                     name='forward azimuth',
-                                    coords=self.ds.traj.tlon.coords,
+                                    coords=self.tlon.coords,
                                     attrs={'units': 'degrees'})
 
         ds['az_bwd'] = xr.DataArray(a2,
                                     name='back azimuth',
-                                    coords=self.ds.traj.tlon.coords,
+                                    coords=self.tlon.coords,
                                     attrs={'units': 'degrees'})
 
         return ds

trajan/traj1d.py

@@ -14,8 +14,8 @@ class Traj1d(Traj):
     A structured dataset, where each trajectory is always given at the same times. Typically the output from a model or from a gridded dataset.
     """
 
-    def __init__(self, ds, obs_dim, time_varname):
-        super().__init__(ds, obs_dim, time_varname)
+    def __init__(self, ds, trajectory_dim, obs_dim, time_varname):
+        super().__init__(ds, trajectory_dim, obs_dim, time_varname)
 
     def timestep(self):
         """Time step between observations in seconds."""
@@ -32,7 +32,8 @@ def to_2d(self, obs_dim='obs'):
         ds = self.ds.copy()
         time = ds[self.time_varname].rename({
             self.time_varname: obs_dim
-        }).expand_dims(dim={'trajectory': ds.sizes['trajectory']})
+        }).expand_dims(dim={self.trajectory_dim: ds.sizes[self.trajectory_dim]})
+        # TODO should also add cf_role here
         ds = ds.rename({self.time_varname: obs_dim})
         ds[self.time_varname] = time
         ds[obs_dim] = np.arange(0, ds.sizes[obs_dim])
@@ -45,7 +46,7 @@ def time_to_next(self):
 
     def velocity_spectrum(self):
 
-        if self.ds.sizes['trajectory'] > 1:
+        if self.ds.sizes[self.trajectory_dim] > 1:
             raise ValueError(
                 'Spectrum can only be calculated for a single trajectory')
 
@@ -77,7 +78,7 @@ def rotary_spectrum(self):
         ### TODO unfinished method
 
         from .tools import rotary_spectra
-        if self.ds.sizes['trajectory'] > 1:
+        if self.ds.sizes[self.trajectory_dim] > 1:
             raise ValueError(
                 'Spectrum can only be calculated for a single trajectory')
 
@@ -95,9 +96,9 @@ def rotary_spectrum(self):
         plt.show()
 
     def skill(self, other, method='liu-weissberg', **kwargs):
-        if self.ds.sizes['trajectory'] != other.sizes['trajectory']:
+        if self.ds.sizes[self.trajectory_dim] != other.sizes[other.traj.trajectory_dim]:
             raise ValueError(
-                f"There must be the same number of trajectories in the two datasets that are compared. This dataset: {self.ds.sizes['trajectory']}, other: {other.sizes['trajectory']}."
+                f"There must be the same number of trajectories in the two datasets that are compared. This dataset: {self.ds.sizes[self.trajectory_dim]}, other: {other.sizes[other.traj.trajectory_dim]}."
             )
 
         diff = np.max(
@@ -110,31 +111,31 @@ def skill(self, other, method='liu-weissberg', **kwargs):
                 f"The two datasets must have approximately equal time coordinates, maximum difference: {diff} seconds. Consider using `gridtime` to interpolate one of the datasets."
             )
 
-        s = np.zeros((self.ds.sizes['trajectory']), dtype=np.float32)
+        s = np.zeros((self.ds.sizes[self.trajectory_dim]), dtype=np.float32)
 
         # ds = self.ds.dropna(dim=self.obs_dim)
         # other = other.dropna(dim=other.traj.obs_dim)
 
-        ds = self.ds.transpose('trajectory', self.obs_dim, ...)
-        other = other.transpose('trajectory', other.traj.obs_dim, ...)
+        ds = self.ds.transpose(self.trajectory_dim, self.obs_dim, ...)
+        other = other.transpose(other.traj.trajectory_dim, other.traj.obs_dim, ...)
 
-        lon0 = ds.traj.tlon
+        lon0 = ds.traj.tlon  # TODO should be self.tlon ?
         lat0 = ds.traj.tlat
         lon1 = other.traj.tlon
         lat1 = other.traj.tlat
 
         for ti in range(0, len(s)):
             if method == 'liu-weissberg':
-                s[ti] = skill.liu_weissberg(lon0.isel(trajectory=ti),
-                                            lat0.isel(trajectory=ti),
-                                            lon1.isel(trajectory=ti),
-                                            lat1.isel(trajectory=ti), **kwargs)
+                s[ti] = skill.liu_weissberg(lon0.isel({self.trajectory_dim: ti}),
+                                            lat0.isel({self.trajectory_dim: ti}),
+                                            lon1.isel({self.trajectory_dim: ti}),
+                                            lat1.isel({self.trajectory_dim: ti}), **kwargs)
             else:
                 raise ValueError(f"Unknown skill-score method: {method}.")
 
         return xr.DataArray(s,
                             name='Skillscore',
-                            coords={'trajectory': self.ds.trajectory},
+                            coords={self.trajectory_dim: self.ds.trajectory},
                             attrs={'method': method})
 
     def seltime(self, t0=None, t1=None):