use_cal added

README.md

@@ -17,3 +17,26 @@ mom = vbt.indicator("talib:MOM").run(t1data.data["BAC"].close, timeperiod=10, sk
 mom_anch_d = AnchoredIndicator("talib:MOM", anchor='D').run(t1data.data["BAC"].close, timeperiod=10, skipna=True) #anchored to D
 ```
 
+`create_mask_from_window` - creates mask of the same size AS INPUT, True values denotes that the window is open. Used to filter entry window or forced eod window. Range is denoted by pair (start, end) indicating minutes elapsed from the market start of that day.
+
+```python
+from ttools import create_mask_from_window
+
+entry_window_opens = 3 #in minutes from start of the market
+entry_window_closes = 388
+forced_exit_start = 387
+forced_exit_end = 390
+
+#create mask based on main session that day
+entry_window_opened = create_mask_from_window(entries, entry_window_opens, entry_window_closes)
+#limit entries to the window
+entries = entries & entry_window_opened
+
+#create forced exits mask
+forced_exits_window = create_mask_from_window(exits, forced_exit_start, forced_exit_end)
+
+#add forced_exits to exits
+exits = exits | forced_exits_window
+
+exits.tail(20)
+```

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='ttools',
-    version='0.1.6',
+    version='0.1.7',
     packages=find_packages(),
     install_requires=[
         'vectorbtpro',

ttools/vbtutils.py

@@ -2,61 +2,73 @@ import pandas as pd
 import vectorbtpro as vbt
 import pandas_market_calendars as mcal
 from typing import Any
+import datetime
 
-def create_mask_from_window(entries: Any, entry_window_opens:int, entry_window_closes:int):
+def create_mask_from_window(series: Any, entry_window_opens:int, entry_window_closes:int, use_cal: bool = True):
     """
-    Accepts entries and window range (number of minutes from market start) and returns boolean mask denoting 
+    Accepts series and window range (number of minutes from market start) and returns boolean mask denoting 
-     entries within the window.
+     series within the window.
 
     Parameters
     ----------
-    entries : pd.Series/pd:DataFrame
+    series : pd.Series/pd:DataFrame
-        Entries to be masked.
+        series to be masked.
     entry_window_opens : int
         Number of minutes from market start to open the window.
     entry_window_closes : int
         Number of minutes from market start to close the window.
+    use_cal : bool, default True
+        If True, uses NYSE calendar to determine market hours for each day. Otherwise uses 9:30 to 16:00 constant.
 
-    TODO: test with entries not covering whole main session
+    TODO: test with series not covering whole main session
 
     Returns
     -------
-    type of entries
+    type of series
     """
-    # Get the NYSE calendar
-    nyse = mcal.get_calendar("NYSE")
-    # Get the market hours data
-    market_hours = nyse.schedule(start_date=entries.index[0].to_pydatetime(), end_date=entries.index[-1].to_pydatetime(), tz=nyse.tz)
 
-    market_hours =market_hours.tz_localize(nyse.tz)
+    if use_cal:
+        # Get the NYSE calendar
+        nyse = mcal.get_calendar("NYSE")
+        # Get the market hours data
+        market_hours = nyse.schedule(start_date=series.index[0].to_pydatetime(), end_date=series.index[-1].to_pydatetime(), tz=nyse.tz)
 
-    # Ensure both entries and market_hours are timezone-aware and in the same timezone
+        market_hours =market_hours.tz_localize(nyse.tz)
-    if entries.index.tz is None:
-        entries.index = entries.index.tz_localize('America/New_York')
 
-    # Use merge_asof to align entries with the nearest market_open in market_hours
+        # Ensure both series and market_hours are timezone-aware and in the same timezone
-    merged = pd.merge_asof(
+        if series.index.tz is None:
-        entries.to_frame(), 
+            series.index = series.index.tz_localize('America/New_York')
-        market_hours[['market_open', 'market_close']], 
-        left_index=True, 
-        right_index=True, 
-        direction='backward'
-    )
 
-    # Calculate the time difference between each entry and its corresponding market_open
+        # Use merge_asof to align series with the nearest market_open in market_hours
-    elapsed_time = entries.index.to_series() - merged['market_open']
+        merged = pd.merge_asof(
+            series.to_frame(), 
+            market_hours[['market_open', 'market_close']], 
+            left_index=True, 
+            right_index=True, 
+            direction='backward'
+        )
 
-    # Convert the difference to minutes
+        # Calculate the time difference between each entry and its corresponding market_open
-    elapsed_minutes = elapsed_time.dt.total_seconds() / 60.0
+        elapsed_time = series.index.to_series() - merged['market_open']
 
-    #elapsed_minutes = pd.DataFrame(elapsed_minutes, index=entries.index)
+        # Convert the difference to minutes
+        elapsed_minutes = elapsed_time.dt.total_seconds() / 60.0
 
-    # Create a boolean mask for entries that are within the window
+        #elapsed_minutes = pd.DataFrame(elapsed_minutes, index=series.index)
-    window_opened = (elapsed_minutes >= entry_window_opens) & (elapsed_minutes < entry_window_closes)
 
-    # Return the mask as a series with the same index as entries
+        # Create a boolean mask for series that are within the window
-    return pd.Series(window_opened.values, index=entries.index)
+        window_opened = (elapsed_minutes >= entry_window_opens) & (elapsed_minutes < entry_window_closes)
 
+        # Return the mask as a series with the same index as series
+        return pd.Series(window_opened.values, index=series.index)
+    else:
+        # Calculate the time difference in minutes from market open for each timestamp
+        market_open = datetime.time(9, 30)
+        market_close = datetime.time(16, 0)
+        window_open= pd.Series(False, index=series.index)
+        elapsed_min_from_open = (series.index.hour - market_open.hour) * 60 + (series.index.minute - market_open.minute)
+        window_open[(elapsed_min_from_open >= entry_window_opens) & (elapsed_min_from_open < entry_window_closes)] = True
+        return window_open
 
 class AnchoredIndicator:
     """