feature(get_time_series): returns dataloader IDs

src/wela/dataloader.py

@@ -80,6 +80,7 @@ class dataloader:
     """
 
     def __init__(self, h5dir=None, wdir=".", ls=True, projection="max"):
+        """Initialise by defining possible signals."""
         # method to collapse Z stacks to one image
         self.projection = projection
         # from grouper.signals to abbreviations
@@ -96,7 +97,7 @@ class dataloader:
             self.ls
 
     def define_g2a_dict(self, fl_channels=["GFP", "mCherry"]):
-        """Dict to convert names of grouper.signals to short forms."""
+        """Define dict to convert names of grouper.signals to short forms."""
         projection = self.projection
         g2a_dict = {
             f"extraction/cy5/{projection}/median": "cy5",
@@ -729,7 +730,7 @@ class dataloader:
             sdf = sdf[sdf.time <= tmax]
         return sdf
 
-    def get_time_series(self, signal, group=None, df=None):
+    def get_time_series(self, signal, group=None, df=None, return_ids=False):
         """
         Extract a signal as a 2D array with each row a time series.
 
@@ -737,6 +738,13 @@ class dataloader:
         ----------
         signal: string
             The signal to be extracted
+        group: string (optional)
+            Selected data only for a particular group
+        df: pd.DataFrame (optional)
+            If specified, use this dataframe not dataloader's.
+        return_ids: boolean (optional)
+            If True, return too the dataloader IDs for the cells.
+
 
         Returns
         -------
@@ -745,10 +753,8 @@ class dataloader:
         values: 2D array
             The values of the signal with each row a separate time
             series
-        group: string (optional)
-            Selected data only for a particular group
-        df: pd.DataFrame (optional)
-            If specified, use this dataframe not dataloader's.
+        ids: 1D array (optional)
+            The dataloader IDs for the cells in each row.
         """
         if df is None:
             if group is None:
@@ -770,7 +776,11 @@ class dataloader:
         time = wdf.columns.to_numpy()
         # each column is a time series
         values = wdf.to_numpy()
-        return time, values
+        if return_ids:
+            ids = np.array(wdf.index)
+            return time, values, ids
+        else:
+            return time, values
 
     def put_time_series(self, values, signal):
         """