1d: convert to 2d

Author: gauteh

tests/test_convert_datalayout.py

@@ -0,0 +1,15 @@
+from numpy.testing import assert_almost_equal
+import numpy as np
+import trajan as ta
+import xarray as xr
+import pandas as pd
+
+def test_to2d(barents):
+    # print(barents)
+    gr = barents.traj.gridtime('1H')
+    # print(gr)
+
+    assert gr.traj.is_1d()
+
+    b2d = gr.traj.to_2d()
+    assert b2d.traj.is_2d()

trajan/ragged.py

@@ -19,7 +19,7 @@ def __init__(self, ds, obsdim, timedim, trajectorycoord, rowsizevar):
         self.rowvar = rowsizevar
         super().__init__(ds, obsdim, timedim)
 
-    def to_2d(self):
+    def to_2d(self, obsdim='obs'):
         """This actually converts a contiguous ragged xarray Dataset into an xarray Dataset that follows the Traj2d conventions."""
         global_attrs = self.ds.attrs
 
@@ -70,7 +70,7 @@ def to_2d(self):
 
             # trajectory vars
             'time':
-            xr.DataArray(dims=["trajectory", "obs"],
+            xr.DataArray(dims=["trajectory", obsdim],
                          data=array_time,
                          attrs={
                              "standard_name": "time",
@@ -119,7 +119,7 @@ def to_2d(self):
                 crrt_data_var = "lat"
 
             ds_converted_to_traj2d[crrt_data_var] = \
-                xr.DataArray(dims=["trajectory", "obs"],
+                xr.DataArray(dims=["trajectory", obsdim],
                              data=crrt_var,
                              attrs=attrs)
 

trajan/traj.py

@@ -862,7 +862,7 @@ def condense_obs(self) -> xr.Dataset:
         """
 
     @abstractmethod
-    def to_2d(self) -> xr.Dataset:
+    def to_2d(self, obsdim='obs') -> xr.Dataset:
         """
         Convert dataset into a 2D dataset from.
         """

trajan/traj1d.py

@@ -28,6 +28,17 @@ def is_1d(self):
     def is_2d(self):
         return False
 
+    def to_2d(self, obsdim='obs'):
+        ds = self.ds.copy()
+        time = ds[self.timedim].rename({
+            self.timedim: obsdim
+        }).expand_dims(dim={'trajectory': ds.sizes['trajectory']})
+        ds = ds.rename({self.timedim: obsdim})
+        ds[self.timedim] = time
+        ds[obsdim] = np.arange(0, ds.sizes[obsdim])
+
+        return ds
+
     def time_to_next(self):
         time_step = self.ds.time[1] - self.ds.time[0]
         return time_step