Enrich test cases

qlib/rl/order_execution/simulator_qlib.py

@@ -98,7 +98,7 @@ class StateMaintainer:
         self.history_steps = pd.DataFrame(columns=metric_keys).set_index("datetime")
         self.metrics = None
 
-    def update(self, inner_executor: BaseExecutor, inner_strategy: DecomposedStrategy) -> None:
+    def update(self, inner_executor: BaseExecutor, inner_strategy: DecomposedStrategy, done: bool) -> None:
         execute_order = inner_strategy.execute_order
         execute_result = inner_strategy.execute_result
         exec_vol = np.array([e[0].deal_amount for e in execute_result])
@@ -117,7 +117,7 @@ class StateMaintainer:
             market_volume = np.array([exchange.get_volume(execute_order.stock_id, t, t) for t in minutes])
 
             datetime_list = _get_ticks_slice(
-                self._tick_index, execute_result[0][0].start_time, execute_result[-1][0].start_time, include_end=True
+                self._tick_index, execute_result[0][0].start_time, execute_result[-1][0].start_time, include_end=True,
             )
         else:
             market_price = np.array([])
@@ -157,6 +157,16 @@ class StateMaintainer:
             ],
         )
 
+        if done:
+            self.metrics = self._metrics_collect(
+                self._order,
+                self._tick_index[0],  # start time
+                self.history_exec["market_volume"],
+                self.history_exec["market_price"],
+                self.history_steps["amount"].sum(),
+                self.history_exec["deal_amount"],
+            )
+
     def _metrics_collect(
         self,
         order: Order,
@@ -248,13 +258,19 @@ class QlibSimulator(Simulator[Order, SAOEState, float]):
 
         exchange = self._inner_executor.trade_exchange
         self._ticks_index = pd.DatetimeIndex([e[1] for e in list(exchange.quote_df.index)])
-        self._ticks_for_order = _get_ticks_slice(self._ticks_index, self._order.start_time, self._order.end_time)
+        self._ticks_for_order = _get_ticks_slice(
+            self._ticks_index,
+            self._order.start_time,
+            self._order.end_time,
+            include_end=True,
+        )
 
-        twap_price = exchange.get_deal_price(
+        self.twap_price = exchange.get_deal_price(
             order.stock_id,
             pd.Timestamp(self._ticks_for_order[0]),
-            pd.Timestamp(self._ticks_for_order[1]),
+            pd.Timestamp(self._ticks_for_order[-1]),
             direction=order.direction,
+            method="mean",
         )
 
         top_strategy = SingleOrderStrategy(common_infra, order, self._trade_range, instrument)
@@ -270,7 +286,7 @@ class QlibSimulator(Simulator[Order, SAOEState, float]):
         self._maintainer = StateMaintainer(
             order=self._order,
             tick_index=self._ticks_index,
-            twap_price=twap_price,
+            twap_price=self.twap_price,
         )
 
     def _iter_strategy(self, action: float = None) -> DecomposedStrategy:
@@ -281,6 +297,8 @@ class QlibSimulator(Simulator[Order, SAOEState, float]):
         return strategy
 
     def step(self, action: float) -> None:
+        assert not self._done, "Simulator has already done!"
+
         try:
             self._iter_strategy(action=action)
         except StopIteration:
@@ -289,6 +307,7 @@ class QlibSimulator(Simulator[Order, SAOEState, float]):
         self._maintainer.update(
             inner_executor=self._inner_executor,
             inner_strategy=self._inner_strategy,
+            done=self._done,
         )
 
     def get_state(self) -> SAOEState:

qlib/rl/order_execution/tests/common.py

@@ -0,0 +1,69 @@
+from pathlib import Path
+
+from qlib.backtest.decision import Order
+from qlib.backtest.executor import NestedExecutor, SimulatorExecutor
+from qlib.backtest.utils import CommonInfrastructure
+from qlib.config import QlibConfig
+from qlib.contrib.strategy import TWAPStrategy
+from qlib.rl.order_execution.simulator_qlib import ExchangeConfig, QlibSimulator
+
+# fmt: off
+qlib_config = QlibConfig(
+    {
+        "provider_uri_day": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_1d"),
+        "provider_uri_1min": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_1min"),
+        "feature_root_dir": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_handler_stock"),
+        "feature_columns_today": [
+            "$open", "$high", "$low", "$close", "$vwap", "$bid", "$ask", "$volume",
+            "$bidV", "$bidV1", "$bidV3", "$bidV5", "$askV", "$askV1", "$askV3", "$askV5",
+        ],
+        "feature_columns_yesterday": [
+            "$open_1", "$high_1", "$low_1", "$close_1", "$vwap_1", "$bid_1", "$ask_1", "$volume_1",
+            "$bidV_1", "$bidV1_1", "$bidV3_1", "$bidV5_1", "$askV_1", "$askV1_1", "$askV3_1", "$askV5_1",
+        ],
+    }
+)
+# fmt: on
+
+exchange_config = ExchangeConfig(
+    limit_threshold=("$ask == 0", "$bid == 0"),
+    deal_price=("If($ask == 0, $bid, $ask)", "If($bid == 0, $ask, $bid)"),
+    volume_threshold={
+        "all": ("cum", "0.2 * DayCumsum($volume, '9:30', '14:29')"),
+        "buy": ("current", "$askV1"),
+        "sell": ("current", "$bidV1"),
+    },
+    open_cost=0.0005,
+    close_cost=0.0015,
+    min_cost=5.0,
+    trade_unit=None,
+    cash_limit=None,
+    generate_report=False,
+)
+
+
+def _inner_executor_fn(time_per_step: str, common_infra: CommonInfrastructure) -> NestedExecutor:
+    return NestedExecutor(
+        time_per_step=time_per_step,
+        inner_strategy=TWAPStrategy(),
+        inner_executor=SimulatorExecutor(
+            time_per_step="1min",
+            verbose=False,
+            trade_type=SimulatorExecutor.TT_SERIAL,
+            generate_report=False,
+            common_infra=common_infra,
+            track_data=True,
+        ),
+        common_infra=common_infra,
+        track_data=True,
+    )
+
+
+def get_simulator(order: Order) -> QlibSimulator:
+    return QlibSimulator(
+        order=order,
+        time_per_step="30min",
+        qlib_config=qlib_config,
+        inner_executor_fn=_inner_executor_fn,
+        exchange_config=exchange_config,
+    )

qlib/rl/order_execution/tests/test_simulator_qlib.py

@@ -1,98 +1,89 @@
-from pathlib import Path
-
 import pandas as pd
 
 from qlib.backtest.decision import Order, OrderDir
-from qlib.backtest.executor import NestedExecutor, SimulatorExecutor
-from qlib.backtest.utils import CommonInfrastructure
-from qlib.config import QlibConfig
-from qlib.contrib.strategy import TWAPStrategy
 from qlib.rl.order_execution import CategoricalActionInterpreter
-from qlib.rl.order_execution.simulator_qlib import ExchangeConfig, QlibSimulator
+from qlib.rl.order_execution.tests.common import get_simulator
-
-# fmt: off
-qlib_config = QlibConfig(
-    {
-        "provider_uri_day": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_1d"),
-        "provider_uri_1min": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_1min"),
-        "feature_root_dir": Path("C:/workspace/NeuTrader/data_sample/cn/qlib_amc_handler_stock"),
-        "feature_columns_today": [
-            "$open", "$high", "$low", "$close", "$vwap", "$bid", "$ask", "$volume",
-            "$bidV", "$bidV1", "$bidV3", "$bidV5", "$askV", "$askV1", "$askV3", "$askV5",
-        ],
-        "feature_columns_yesterday": [
-            "$open_1", "$high_1", "$low_1", "$close_1", "$vwap_1", "$bid_1", "$ask_1", "$volume_1",
-            "$bidV_1", "$bidV1_1", "$bidV3_1", "$bidV5_1", "$askV_1", "$askV1_1", "$askV3_1", "$askV5_1",
-        ],
-    }
-)
-# fmt: on
-
-exchange_config = ExchangeConfig(
-    limit_threshold=("$ask == 0", "$bid == 0"),
-    deal_price=("If($ask == 0, $bid, $ask)", "If($bid == 0, $ask, $bid)"),
-    volume_threshold={
-        "all": ("cum", "0.2 * DayCumsum($volume, '9:45', '14:44')"),
-        "buy": ("current", "$askV1"),
-        "sell": ("current", "$bidV1"),
-    },
-    open_cost=0.0005,
-    close_cost=0.0015,
-    min_cost=5.0,
-    trade_unit=None,
-    cash_limit=None,
-    generate_report=False,
-)
 
 
-def _inner_executor_fn(time_per_step: str, common_infra: CommonInfrastructure) -> NestedExecutor:
+def is_close(a: float, b: float, epsilon: float = 1e-4) -> bool:
-    return NestedExecutor(
+    return abs(a - b) <= epsilon
-        time_per_step=time_per_step,
-        inner_strategy=TWAPStrategy(),
-        inner_executor=SimulatorExecutor(
-            time_per_step="1min",
-            verbose=False,
-            trade_type=SimulatorExecutor.TT_SERIAL,
-            generate_report=False,
-            common_infra=common_infra,
-            track_data=True,
-        ),
-        common_infra=common_infra,
-        track_data=True,
-    )
 
 
-def test():
+def test_simulator_first_step():
+    TOTAL_POSITION = 2100.0
+
     order = Order(
         stock_id="SH600000",
-        amount=1078.644160270691,
+        amount=TOTAL_POSITION,
-        direction=OrderDir(1),
+        direction=OrderDir.BUY,
-        start_time=pd.Timestamp("2019-03-04 09:45:00"),
+        start_time=pd.Timestamp("2019-03-04 09:30:00"),
-        end_time=pd.Timestamp("2019-03-04 14:44:00"),
+        end_time=pd.Timestamp("2019-03-04 14:29:00"),
     )
 
-    simulator = QlibSimulator(
+    simulator = get_simulator(order)
-        order=order,
-        time_per_step="30min",
-        qlib_config=qlib_config,
-        inner_executor_fn=_inner_executor_fn,
-        exchange_config=exchange_config,
-    )
-
-    interpreter_action = CategoricalActionInterpreter(values=4)
-
     state = simulator.get_state()
-    print(state.position)
+    assert state.cur_time == pd.Timestamp('2019-03-04 09:30:00')
-    for i in range(10):
+    assert state.position == TOTAL_POSITION
-        print(f"Step {i}")
-        simulator.step(interpreter_action(state, 1))
 
-        state = simulator.get_state()
+    AMOUNT = 300.0
-        print(state.position)
+    simulator.step(AMOUNT)
+    state = simulator.get_state()
+    assert state.cur_time == pd.Timestamp('2019-03-04 10:00:00')
+    assert state.position == TOTAL_POSITION - AMOUNT
+    assert len(state.history_exec) == 30
+    assert state.history_exec.index[0] == pd.Timestamp('2019-03-04 09:30:00')
 
-        if simulator.done():
+    assert is_close(state.history_exec["market_volume"].iloc[0], 109382.382812)
-            break
+    assert is_close(state.history_exec["market_price"].iloc[0], 149.566483)
+    assert (state.history_exec["amount"] == AMOUNT / 30).all()
+    assert (state.history_exec["deal_amount"] == AMOUNT / 30).all()
+    assert is_close(state.history_exec["trade_price"].iloc[0], 149.566483)
+    assert is_close(state.history_exec["trade_value"].iloc[0], 1495.664825)
+    assert is_close(state.history_exec["position"].iloc[0], TOTAL_POSITION - AMOUNT / 30)
+    # assert state.history_exec["ffr"].iloc[0] == 1 / 60  # FIXME
+
+    assert is_close(state.history_steps["market_volume"].iloc[0], 1254848.5756835938)
+    assert state.history_steps["amount"].iloc[0] == AMOUNT
+    assert state.history_steps["deal_amount"].iloc[0] == AMOUNT
+    assert state.history_steps["ffr"].iloc[0] == 1.0
+    assert is_close(
+        state.history_steps["pa"].iloc[0] * (1.0 if order.direction == OrderDir.SELL else -1.0),
+        (state.history_steps["trade_price"].iloc[0] / simulator.twap_price - 1) * 10000,
+    )
+
+
+def test_simulator_stop_twap() -> None:
+    TOTAL_POSITION = 2100.0
+
+    order = Order(
+        stock_id="SH600000",
+        amount=TOTAL_POSITION,
+        direction=OrderDir.BUY,
+        start_time=pd.Timestamp("2019-03-04 09:30:00"),
+        end_time=pd.Timestamp("2019-03-04 14:29:00"),
+    )
+
+    simulator = get_simulator(order)
+    NUM_STEPS = 7
+    for i in range(NUM_STEPS):
+        simulator.step(TOTAL_POSITION / NUM_STEPS)
+
+    HISTORY_STEP_LENGTH = 30 * NUM_STEPS
+    state = simulator.get_state()
+    assert len(state.history_exec) == HISTORY_STEP_LENGTH
+
+    assert (state.history_exec["deal_amount"] == TOTAL_POSITION / HISTORY_STEP_LENGTH).all()
+    assert is_close(state.history_steps["position"].iloc[0], TOTAL_POSITION * (NUM_STEPS - 1) / NUM_STEPS)
+    assert is_close(state.history_steps["position"].iloc[-1], 0.0)
+    assert is_close(state.position, 0.0)
+    assert is_close(state.metrics["ffr"], 1.0)
+
+    # assert abs(state.metrics["market_price"] - state.backtest_data.get_deal_price().mean()) < 1e-4
+    # assert np.isclose(state.metrics["market_volume"], state.backtest_data.get_volume().sum())
+    assert is_close(state.metrics["trade_price"], state.metrics["market_price"])
+    assert is_close(state.metrics["pa"], 0.0)
 
 
 if __name__ == "__main__":
-    test()
+    test_simulator_first_step()
+    test_simulator_stop_twap()