adding mth5_path property to kds

Author: kujaku11

mtpy/processing/kernel_dataset.py

@@ -60,6 +60,7 @@
 # =============================================================================
 # Imports
 # =============================================================================
+from pathlib import Path
 import copy
 from typing import Optional, Union
 
@@ -197,6 +198,38 @@ def df(self, value):
             self._set_datetime_columns(self._add_columns(value)), inplace=False
         )
 
+    def _has_df(self):
+        """
+        check to see if dataframe is set
+        """
+        if self._df is not None:
+            if not self._df.empty:
+                return True
+            return False
+        return False
+
+    def _df_has_local_station_id(self):
+        """
+        Check to make sure the dataframe has the local station id
+
+        :return: DESCRIPTION
+        :rtype: bool
+
+        """
+        if self._has_df():
+            return (self._df.station == self.local_station_id).any()
+
+    def _df_has_remote_station_id(self):
+        """
+        Check to make sure the dataframe has the local station id
+
+        :return: DESCRIPTION
+        :rtype: bool
+
+        """
+        if self._has_df():
+            return (self._df.station == self.remote_station_id).any()
+
     def _set_datetime_columns(self, df):
         """
         be sure to set start and end to be date time objects
@@ -241,6 +274,119 @@ def _add_columns(self, df):
                 )
         return df
 
+    @property
+    def local_station_id(self):
+        return self._local_station_id
+
+    @local_station_id.setter
+    def local_station_id(self, value):
+        if value is None:
+            self._local_station_id = None
+        else:
+            try:
+                self._local_station_id = str(value)
+            except ValueError:
+                raise ValueError(
+                    f"Bad type {type(value)}. "
+                    "Cannot convert local_station_id value to string."
+                )
+            if self._has_df():
+                if not self._df_has_local_station_id():
+                    raise NameError(
+                        f"Could not find {self._local_station_id} in dataframe"
+                    )
+
+    @property
+    def local_mth5_path(self):
+        """
+
+        :return: Local station MTH5 path, a property extracted from the dataframe
+        :rtype: Path
+
+        """
+        if self._has_df():
+            return Path(
+                self._df.loc[
+                    self._df.station == self.local_station_id, "mth5_path"
+                ].unique()[0]
+            )
+        else:
+            return None
+
+    # @local_mth5_path.setter
+    # def local_mth5_path(self, value):
+    #     self._local_mth5_path = self.set_path(value)
+
+    def has_local_mth5(self):
+        """test if local mth5 exists"""
+        if self.local_mth5_path is None:
+            return False
+        else:
+            return self.local_mth5_path.exists()
+
+    @property
+    def remote_station_id(self):
+        return self._remote_station_id
+
+    @remote_station_id.setter
+    def remote_station_id(self, value):
+        if value is None:
+            self._remote_station_id = None
+        else:
+            try:
+                self._remote_station_id = str(value)
+            except ValueError:
+                raise ValueError(
+                    f"Bad type {type(value)}. "
+                    "Cannot convert remote_station_id value to string."
+                )
+            if self._has_df():
+                if not self._df_has_remote_station_id():
+                    raise NameError(
+                        f"Could not find {self._remote_station_id} in dataframe"
+                    )
+
+    @property
+    def remote_mth5_path(self):
+        """
+
+        :return: remote station MTH5 path, a property extracted from the dataframe
+        :rtype: Path
+
+        """
+        if self._has_df() and self.remote_station_id is not None:
+            return Path(
+                self._df.loc[
+                    self._df.station == self.remote_station_id, "mth5_path"
+                ].unique()[0]
+            )
+        else:
+            return None
+
+    # @remote_mth5_path.setter
+    # def remote_mth5_path(self, value):
+    #     self._remote_mth5_path = self.set_path(value)
+
+    def has_remote_mth5(self):
+        """test if remote mth5 exists"""
+        if self.remote_mth5_path is None:
+            return False
+        else:
+            return self.remote_mth5_path.exists()
+
+    @classmethod
+    def set_path(self, value):
+        return_path = None
+        if value is not None:
+            if isinstance(value, (str, Path)):
+                return_path = Path(value)
+                if not return_path.exists():
+                    raise IOError(f"Cannot find file: {return_path}")
+            else:
+                raise ValueError(f"Cannot convert type{type(value)} to Path")
+
+        return return_path
+
     def from_run_summary(
         self,
         run_summary: RunSummary,

mtpy/processing/run_summary.py

@@ -189,20 +189,6 @@ def set_sample_rate(self, sample_rate: float, inplace: bool = False):
             new_rs.df = new_rs.df[new_rs.df.sample_rate == sample_rate]
             return new_rs
 
-    # BELOW FUNCTION CAN BE COPIED FROM METHOD IN KernelDataset()
-    # def drop_runs_shorter_than(self, duration, units="s"):
-    #     if units != "s":
-    #         raise NotImplementedError
-    #     if "duration" not in self.df.columns:
-    #         self.add_duration()
-    #     drop_cond = self.df.duration < duration
-    #     # df = self.df[drop_cond]
-    #     self.df.drop(self.df[drop_cond].index, inplace=True)
-    #     df = df.reset_index()
-    #
-    #     self.df = df
-    #     return df
-
 
 def extract_run_summary_from_mth5(mth5_obj, summary_type="run"):
     """

tests/processing/test_kernel_dataset.py

@@ -1,6 +1,7 @@
 # =============================================================================
 # Imports
 # =============================================================================
+from pathlib import Path
 import pandas as pd
 import unittest
 
@@ -166,7 +167,7 @@ def setUp(self):
                 ],
                 "has_data": [True, True, True, True],
                 "input_channels": ["hx, hy"] * 4,
-                "mth5_path": ["path"] * 4,
+                "mth5_path": ["path"] * 2 + ["remote_path"] * 2,
                 "n_samples": [86400, 86400, 79200, 7200],
                 "output_channels": ["hz, ex, ey"] * 4,
                 "run": ["01", "02", "03", "04"],
@@ -187,23 +188,25 @@ def setUp(self):
         self.kd = KernelDataset()
         self.kd.from_run_summary(self.run_summary, self.local, self.remote)
 
-    def test_from_run_summary(self):
+    def test_from_run_summary_local_station_id(self):
+        self.assertEqual(self.kd.local_station_id, self.local)
 
-        with self.subTest("local_station_id"):
-            self.assertEqual(self.kd.local_station_id, self.local)
-        with self.subTest("remote_station_id"):
-            self.assertEqual(self.kd.remote_station_id, self.remote)
-        with self.subTest("has remote column"):
-            self.assertIn("remote", self.kd.df.columns)
-        with self.subTest("has fc column"):
-            self.assertIn("fc", self.kd.df.columns)
-        with self.subTest("has_duration"):
-            self.assertFalse((self.kd.df.duration == 0).all())
-
-        with self.subTest("has_all_columns"):
-            self.assertListEqual(
-                sorted(self.kd.df.columns), sorted(KERNEL_DATASET_COLUMNS)
-            )
+    def test_from_run_summary_remote_station_id(self):
+        self.assertEqual(self.kd.remote_station_id, self.remote)
+
+    def test_from_run_summary_has_duration(self):
+        self.assertFalse((self.kd.df.duration == 0).all())
+
+    def test_from_run_summary_has_all_columns(self):
+        self.assertListEqual(
+            sorted(self.kd.df.columns), sorted(KERNEL_DATASET_COLUMNS)
+        )
+
+    def test_from_run_summary_local_mth5_path(self):
+        self.assertEqual(self.kd.local_mth5_path, Path("path"))
+
+    def test_from_run_summary_remote_mth5_path(self):
+        self.assertEqual(self.kd.remote_mth5_path, Path("remote_path"))
 
     def test_num_sample_rates(self):
         self.assertEqual(self.kd.num_sample_rates, 1)
@@ -223,6 +226,18 @@ def test_update_duration_column_not_inplace(self):
 
         self.assertTrue((new_df.duration == self.kd.df.duration).all())
 
+    def test_set_local_station_id_fail(self):
+        def set_station(value):
+            self.kd.local_station_id = value
+
+        self.assertRaises(NameError, set_station, "mt03")
+
+    def test_set_remote_station_id_fail(self):
+        def set_station(value):
+            self.kd.remote_station_id = value
+
+        self.assertRaises(NameError, set_station, "mt03")
+
 
 class TestKernelDatasetMethodsFail(unittest.TestCase):
     def setUp(self):