added close position and trade session

lib/pt_trading/trading_pair.py

@@ -23,11 +23,17 @@ class TradingPair:
     predicted_df_: Optional[pd.DataFrame]
 
     def __init__(
-        self, market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
+        self, config: Dict[str, Any], market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
     ):
         self.symbol_a_ = symbol_a
         self.symbol_b_ = symbol_b
         self.price_column_ = price_column
+        self.set_market_data(market_data)
+        self.user_data_ = {}
+        self.predicted_df_ = None
+        self.config_ = config
+
+    def set_market_data(self, market_data: pd.DataFrame) -> None:
         self.market_data_ = pd.DataFrame(
             self._transform_dataframe(market_data)[["tstamp"] + self.colnames()]
         )
@@ -36,9 +42,19 @@ class TradingPair:
         self.market_data_['tstamp'] = pd.to_datetime(self.market_data_['tstamp'])
         self.market_data_ = self.market_data_.sort_values('tstamp')
 
- 
-        self.user_data_ = {}
-        self.predicted_df_ = None
+    def get_begin_index(self) -> int:
+        if "trading_hours" not in self.config_:
+            return 0
+        start_time = pd.to_datetime(self.config_["trading_hours"]["begin_session"]).time()
+        mask = self.market_data_['tstamp'].dt.time >= start_time
+        return int(self.market_data_.index[mask].min())
+
+    def get_end_index(self) -> int:
+        if "trading_hours" not in self.config_:
+            return 0
+        end_time = pd.to_datetime(self.config_["trading_hours"]["end_session"]).time()
+        mask = self.market_data_['tstamp'].dt.time <= end_time
+        return int(self.market_data_.index[mask].max())
 
     def _transform_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
         # Select only the columns we need