Repository: microsoft/qlib
File: qlib/rl/order_execution/state.py (Python)

Code before:
    assert len(original_data) == len(found_time_list)
    tmp = dict(zip(found_time_list, original_data))
    fill_val = fill_method(original_data)
    return np.array([tmp.get(t, fill_val) for t in total_time_list])


class SAOEStateAdapter:
    """
    Maintain states of the environment. SAOEStateAdapter accepts execution results and update its internal state
    according to the execution results with additional information acquired from executors & exchange. For example,
    it gets the dealt order amount from execution results, and get the corresponding market price / volume from
    exchange.

    Example usage::

        adapter = SAOEStateAdapter(...)
        adapter.update(...)
        state = adapter.saoe_state
    """

    def __init__(
        self,
        order: Order,
        executor: BaseExecutor,
        exchange: Exchange,
        ticks_per_step: int,
        backtest_data: IntradayBacktestData,
    ) -> None:
        self.position = order.amount
        self.order = order
        self.executor = executor
        self.exchange = exchange
        self.backtest_data = backtest_data

        self.twap_price = self.backtest_data.get_deal_price().mean()

        metric_keys = list(SAOEMetrics.__annotations__.keys())  # pylint: disable=no-member
        self.history_exec = pd.DataFrame(columns=metric_keys).set_index("datetime")
        self.history_steps = pd.DataFrame(columns=metric_keys).set_index("datetime")
        self.metrics: Optional[SAOEMetrics] = None

        self.cur_time = max(backtest_data.ticks_for_order[0], order.start_time)
        self.ticks_per_step = ticks_per_step

    def _next_time(self) -> pd.Timestamp:
        current_loc = self.backtest_data.ticks_index.get_loc(self.cur_time)
        next_loc = current_loc + self.ticks_per_step
        next_loc = next_loc - next_loc % self.ticks_per_step
        if (
            next_loc < len(self.backtest_data.ticks_index)
            and self.backtest_data.ticks_index[next_loc] < self.order.end_time

Review comment (suggestion):
I suggest moving this adapter to `state_adapter.py` or `simulator_qlib.py`. If I'm using the simple simluator, there is no reason I will be interested in all these adapter-related logics.

Can be left as TODO. :)

Code after:
    """Fill missing data. We need this function to deal with data that have missing values in some minutes.

    TODO: making exchange return data without missing will make it more elegant. Fix this in the future.

    Parameters
    ----------
    original_data
        Original data without missing values.
    total_time_list
        All timestamps that required.
    found_time_list
        Timestamps found in the original data.
    fill_method
        Method used to fill the missing data.

    Returns
    -------
        The filled data.
    """
    assert len(original_data) == len(found_time_list)
    tmp = dict(zip(found_time_list, original_data))
    fill_val = fill_method(original_data)
    return np.array([tmp.get(t, fill_val) for t in total_time_list])


class SAOEStateAdapter:
    """
    Maintain states of the environment. SAOEStateAdapter accepts execution results and update its internal state
    according to the execution results with additional information acquired from executors & exchange. For example,
    it gets the dealt order amount from execution results, and get the corresponding market price / volume from
    exchange.

    Example usage::

        adapter = SAOEStateAdapter(...)
        adapter.update(...)
        state = adapter.saoe_state
    """

    def __init__(
        self,
        order: Order,
        executor: BaseExecutor,
        exchange: Exchange,
        ticks_per_step: int,
        backtest_data: IntradayBacktestData,
    ) -> None:
        self.position = order.amount
        self.order = order
        self.executor = executor
        self.exchange = exchange