volume limit update

qlib/contrib/ops/high_freq.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 import numpy as np
 import pandas as pd
+from datetime import datetime
 
 import qlib
 from qlib.data import D
@@ -12,7 +13,9 @@ from qlib.data.ops import ElemOperator
 
 
 def get_calendar_day(freq="1min", future=False):
-    """Load High-Freq Calendar Date Using Memcache.
+    """
+    Load High-Freq Calendar Date Using Memcache.
+    !!!NOTE: Loading the calendar is quite slow. So loading calendar before start multiprocessing will make it faster.
 
     Parameters
     ----------
@@ -36,20 +39,57 @@ def get_calendar_day(freq="1min", future=False):
 
 
 class DayCumsum(ElemOperator):
-    """DayLast Operator
+    """DayCumsum Operator during start time and end time.
 
     Parameters
     ----------
     feature : Expression
         feature instance
+    start : str
+        the start time of backtest in one day.
+        !!!NOTE: "9:30" means the time period of (9:30, 9:31) is in transaction.
+    end : str
+        the end time of backtest in one day.
+        !!!NOTE: "14:59" means the time period of (14:59, 15:00) is in transaction,
+                but (15:00, 15:01) is not.
+        So start="9:30" and end="14:59" means trading all day.
 
     Returns
     ----------
     feature:
-        a series of that each value equals the last value of its day
+        a series of that each value equals the cumsum value during start time and end time.
+        Otherwise, the value is zero.
     """
 
+    def __init__(self, feature, start: str = "9:30", end: str = "14:59"):
+        self.feature = feature
+        self.start = datetime.strptime(start, "%H:%M")
+        self.end = datetime.strptime(end, "%H:%M")
+
+        self.morning_open = datetime.strptime("9:30", "%H:%M")
+        self.morning_close = datetime.strptime("11:30", "%H:%M")
+        self.noon_open = datetime.strptime("13:00", "%H:%M")
+        self.noon_close = datetime.strptime("15:00", "%H:%M")
+
+        self.start_id = self.time_to_index(self.start)
+        self.end_id = self.time_to_index(self.end)
+
+    def time_to_index(self, t):
+        if t >= self.morning_open and t < self.morning_close:
+            return int((t - self.morning_open).total_seconds() / 60)
+        elif t >= self.noon_open and t < self.noon_close:
+            return int((t - self.noon_open).total_seconds() / 60) + 120
+        else:
+            raise ValueError(f"{t} is not the opening time of the stock market")
+
+    def period_cusum(self, df):
+        assert len(df) == 240
+        df.iloc[0 : self.start_id] = 0
+        df = df.cumsum()
+        df.iloc[self.end_id + 1 : 240] = 0
+        return df
+
     def _load_internal(self, instrument, start_index, end_index, freq):
         _calendar = get_calendar_day(freq=freq)
         series = self.feature.load(instrument, start_index, end_index, freq)
-        return series.groupby(_calendar[series.index]).cumsum()
+        return series.groupby(_calendar[series.index]).transform(self.period_cusum)