Flow Processor. Get datasets data

.idea/csv-plugin.xml

@@ -17,6 +17,13 @@
             </Attribute>
           </value>
         </entry>
+        <entry key="/cfatools/tests/test_processor/test_flow.py">
+          <value>
+            <Attribute>
+              <option name="separator" value="," />
+            </Attribute>
+          </value>
+        </entry>
       </map>
     </option>
   </component>

cfatools/processor/flow.py

@@ -4,6 +4,7 @@ The flow processor module provides functions to determine where each ice core wa
 import datetime
 import pandas as pd
 import numpy as np
+import typing
 from pandarallel import pandarallel
 
 from cfatools.logreader.dataset import DatasetReader
@@ -233,6 +234,62 @@ def get_bubble_density_from_conduct(dataset: DatasetReader):
     return compump_df
 
 
+def get_datasets_data(dataset: DatasetReader)\
+        -> typing.Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame,]:
+    """Get the dataframes useful for flow analysis
+
+    Parameters
+    ----------
+    dataset: DatasetReader
+        Dataset whose data should be retrieved.
+
+    Returns
+    -------
+    tuple of pd.DataFrame
+        encoder_df, iceblock_df, melting_df, conduct_df, pump_df, picarro_df, collector_df
+    """
+    # Encoder
+    encoder_df = dataset.get_timeseries("ENCODER_periodic", "speed")
+    encoder_df = encoder_df.rename(columns={"value": "speed"})
+
+    # Iceblock
+    iceblock_df = dataset.get_timeseries("ICBKCTRL_instant")
+    iceblock_df = iceblock.get_melting_timeseries(iceblock_df)
+    iceblock_df["icbk_datetime"] = iceblock_df.index
+    iceblock_df = iceblock_df.rename(columns={"value_int": "icbk_code",
+                                              "value": "icbk_name"})
+
+    # Melting
+    melting_df = iceblock_df.copy()
+    melting_df = melting_df.drop(columns="icbk_datetime")
+    melting_df = pd.merge_asof(encoder_df, melting_df, left_index=True, right_index=True)
+    melting_df = melting_df.groupby("icbk_code").apply(__compute_mm__)
+
+    # Conductivity
+    conduct_df = dataset.get_timeseries("CONDUCTI_periodic")
+
+    # Pumps' flow rate.
+    pump_df = get_bubble_density_from_pump(dataset)
+    pump_df["colpic"] = pump_df["col"] + pump_df["pic"]
+
+    # Picarro
+    picarro_df = dataset.get_timeseries("PICARRO_periodic")
+
+    # Collector
+    collector_df = dataset.get_timeseries("COLLECTR_instant", "flask")
+    collector_df = collector_df.rename(columns={"value": "flask"})
+
+    return encoder_df, iceblock_df, melting_df, conduct_df, pump_df, picarro_df, collector_df
+
+
+def __compute_mm__(df: pd.DataFrame) -> pd.DataFrame:
+    df["datetime"] = df.index
+    df["mm_step"] = df["speed"].abs() * df["datetime"].diff().dt.total_seconds().fillna(0)
+    df["mm_tot"] = df["mm_step"].cumsum()
+    df = df.drop(columns=["datetime", "speed", "mm_step"])
+    return df
+
+
 def get_tubing_volume_dict(filepath: str,
                            max_datetime: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)) -> dict:
     """Get the characteristic tubing volumes

cfatools/tests/test_flow/test_flow.py

-from unittest import TestCase
-import datetime
-
-from cfatools.processor import flow
-
-
-class TestFlow(TestCase):
-    def test_get_tubing_volume_dict(self):
-        vol_dict = flow.get_tubing_volume_dict("tubing_volumes.csv")
-        self.assertEqual(vol_dict["melt_debub"], 3.5)
-
-        vol_dict = flow.get_tubing_volume_dict("tubing_volumes.csv", datetime.datetime(2021, 1, 1))
-        self.assertEqual(vol_dict["T1_needle"], 7)
-
-        with self.assertRaises(ValueError):
-            flow.get_tubing_volume_dict("tubing_volumes.csv", datetime.datetime(2010, 1, 1))

cfatools/tests/test_processor/test_flow.py

@@ -103,4 +103,24 @@ class TestFlow(TestCase):
     def test_get_bubble_density(self):
         dataset_name = "20201223_mock_core_bronkhorst"
         dataset = DatasetReader(self.base_path, dataset_name)
-        flow.get_bubble_density(dataset)
\ No newline at end of file
+        flow.get_bubble_density(dataset)
+
+    def test_get_datasets_data(self):
+        dataset_name = "20210202_mock_core_flasks"
+        dataset = DatasetReader(self.base_path, dataset_name)
+
+        try:
+            encoder_df, iceblock_df, melting_df, conduct_df, pump_df, picarro_df, collector_df = \
+                flow.get_datasets_data(dataset)
+        except:
+            self.fail("Exception raised!")
+
+    def test_get_tubing_volume_dict(self):
+        vol_dict = flow.get_tubing_volume_dict("../test_processor/tubing_volumes.csv")
+        self.assertEqual(vol_dict["melt_debub"], 3.5)
+
+        vol_dict = flow.get_tubing_volume_dict("../test_processor/tubing_volumes.csv", datetime.datetime(2021, 1, 1))
+        self.assertEqual(vol_dict["T1_needle"], 7)
+
+        with self.assertRaises(ValueError):
+            flow.get_tubing_volume_dict("../test_processor/tubing_volumes.csv", datetime.datetime(2010, 1, 1))
\ No newline at end of file