Merge pull request #139 from knutfrode/dev

Added new methods num_timesteps, num_trajectories and __repr__ to pri…

pyproject.toml

@@ -7,6 +7,7 @@ license = "GPLv2"
 
 [tool.poetry.scripts]
 trajanshow = "trajan.scripts.trajanshow:main"
+trajaninfo = "trajan.scripts.trajaninfo:main"
 
 [tool.poetry.dependencies]
 python = "^3.8"

tests/test_interpolate.py

@@ -77,8 +77,8 @@ def test_interpolate_non_floats(drifter_csv):
     dc = ta.read_csv(drifter_csv, name='Device', time='Time', lon='Longitude', lat='Latitude')
     dcg = dc.traj.gridtime('1h')
 
-    assert 'time' in dcg.dims
-    assert 'trajectory' in dcg.dims
+    assert 'time' in dcg.sizes
+    assert 'trajectory' in dcg.sizes
 
     print(dc, dcg)
 

tests/test_seltime.py

@@ -29,17 +29,17 @@ def test_iseltime(barents):
     assert ds.time.min(skipna=True) == barents.time.min(skipna=True)
     print(ds)
 
-    assert ds.dims['obs'] == 1
+    assert ds.sizes['obs'] == 1
     assert np.all(~pd.isna(ds.time))
 
     print(ds)
     ds = barents.traj.iseltime(-1)
     assert ds.time.max(skipna=True) == barents.time.max(skipna=True)
     print(ds)
-    assert ds.dims['obs'] == 1
+    assert ds.sizes['obs'] == 1
     assert np.all(~pd.isna(ds.time))
 
     ds = barents.traj.iseltime(slice(0, 2))
     print(ds)
-    assert ds.dims['obs'] == 2
+    assert ds.sizes['obs'] == 2
     assert np.all(~pd.isna(ds.time))

trajan/scripts/trajaninfo.py

@@ -0,0 +1,27 @@
+# Utility script to quickly print summary information about a drifter collection file
+
+import xarray as xr
+import trajan as _
+import click
+from pathlib import Path
+import lzma
+
+@click.command()
+@click.argument('tf')
+
+def main(tf):
+    tf = Path(tf)
+    if tf.suffix == '.xz':
+        with lzma.open(tf) as fd:
+            ds = xr.open_dataset(fd)
+            ds.load()
+    else:
+        ds = xr.open_dataset(tf)
+
+    if 'status' in ds:  # hack for OpenDrift files
+        ds = ds.where(ds.status>=0)
+
+    print(ds.traj)
+
+if __name__ == '__main__':
+    main()

trajan/traj.py

@@ -6,6 +6,7 @@
 """
 
 from abc import abstractmethod
+from datetime import timedelta
 import pyproj
 import numpy as np
 import xarray as xr
@@ -63,6 +64,37 @@ def __init__(self, ds, obsdim, timedim):
         self.obsdim = obsdim
         self.timedim = timedim
 
+    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.timedim in self.ds.sizes:
+                output += f'{self.ds.sizes[self.timedim]} timesteps\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
+            output += f'Time coverage:  {start_time} - {end_time}\n'
+        else:
+            output += f'Dataset has no time dimension'
+        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 += 'Variables:\n'
+        for var in self.ds.variables:
+            if var not in ['trajectory', 'obs']:
+                output += f'    {var}'
+                if 'standard_name' in self.ds[var].attrs:
+                    output += f'  [{self.ds[var].standard_name}]'
+                output += '\n'
+        output += '=======================\n'
+        return output
+
     @property
     def plot(self) -> Plot:
         """