Use xarray to output the combined multiunit

Eden-Kramer-Lab · Oct 7, 2018 · e4bedbc · e4bedbc
1 parent fd99db5
commit e4bedbc
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/loren_frank_data_processing/multiunit.py b/loren_frank_data_processing/multiunit.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+import xarray as xr
 from scipy.io import loadmat
 
 from loren_frank_data_processing.core import get_data_filename, logger
@@ -165,6 +166,9 @@ def get_all_multiunit_indicators(tetrode_keys, animals,
                                  time.total_seconds())
         time_index[time_index >= len(time)] = len(time) - 1
         multiunit_dfs.append(df.groupby(time[time_index]).mean()
-                             .reindex(index=time))
+                             .reindex(index=time)
+                             .to_xarray()
+                             .to_array('features'))
 
-    return pd.concat(multiunit_dfs, axis=1)
+    return (xr.concat(multiunit_dfs, dim='tetrodes')
+            .transpose('time', 'features', 'tetrodes'))