Refactor data handling and enhance functionality for ex-factor processing

- Updated `GenerateSplitBundle` to correct file naming and improve data writing logic.
- Introduced `GenerateExFactorBundle` class for handling ex-factor data, including methods for data retrieval and writing to HDF5.
- Enhanced `DataProxy` to support retrieval of ex-cum factors with appropriate error handling.
- Updated `BaseDataSource` to register ex-factor stores and adjust data retrieval methods accordingly.
- Improved `StockPosition` to incorporate unadjusted previous close and calculate position PnL accurately based on daily dividends and splits.
- Refactored `Position` class to enhance trading PnL calculations and streamline cost updates during trade applications.

Author: Cuizi7

rqalpha/data/base_data_source/data_source.py

@@ -82,7 +82,6 @@ def _p(name):
         # static registered storages
         self._future_info_store = FutureInfoStore(_p("future_info.json"), custom_future_info)
         self._yield_curve = YieldCurveStore(_p('yield_curve.h5'))
-        self._ex_cum_factor = SimpleFactorStore(_p('ex_cum_factor.h5'))
         self._share_transformation = ShareTransformationStore(_p('share_transformation.json'))
         self._suspend_days = [DateSet(_p('suspended_days.h5'))]  # type: List[AbstractDateSet]
         self._st_stock_days = DateSet(_p('st_stock_days.h5'))
@@ -94,6 +93,7 @@ def _p(name):
         self._dividend_stores: Dict[tuple[INSTRUMENT_TYPE, MARKET], AbstractDividendStore] = {}
         self._split_stores: Dict[tuple[INSTRUMENT_TYPE, MARKET], AbstractSimpleFactorStore] = {}
         self._calendar_stores: Dict[TRADING_CALENDAR_TYPE, AbstractCalendarStore] = {}
+        self._ex_factor_stores: Dict[tuple[INSTRUMENT_TYPE, MARKET], AbstractSimpleFactorStore] = {}
 
         # instruments
         self._id_instrument_map: Dict[str, Instrument] = {}
@@ -116,9 +116,11 @@ def _p(name):
         # register dividends and split factors stores
         dividend_store = DividendStore(_p('dividends.h5'))
         split_store = SimpleFactorStore(_p('split_factor.h5'))
+        ex_factor_store = SimpleFactorStore(_p('ex_cum_factor.h5'))
         for ins_type in [INSTRUMENT_TYPE.CS, INSTRUMENT_TYPE.ETF, INSTRUMENT_TYPE.LOF]:
             self.register_dividend_store(ins_type, dividend_store)
             self.register_split_store(ins_type, split_store)
+            self.register_ex_factor_store(ins_type, ex_factor_store)
 
         # register calendar stores
         self.register_calendar_store(TRADING_CALENDAR_TYPE.CN_STOCK, ExchangeTradingCalendarStore(_p("trading_dates.npy")))
@@ -141,6 +143,9 @@ def register_split_store(self, instrument_type: INSTRUMENT_TYPE, split_store: Ab
     def register_calendar_store(self, calendar_type: TRADING_CALENDAR_TYPE, calendar_store: AbstractCalendarStore):
         self._calendar_stores[calendar_type] = calendar_store
 
+    def register_ex_factor_store(self, instrument_type: INSTRUMENT_TYPE, ex_factor_store: AbstractSimpleFactorStore, market: MARKET = MARKET.CN):
+        self._ex_factor_stores[instrument_type, market] = ex_factor_store
+
     def append_suspend_date_set(self, date_set):
         # type: (AbstractDateSet) -> None
         self._suspend_days.append(date_set)
@@ -241,8 +246,13 @@ def _are_fields_valid(fields, valid_fields):
                 return False
         return True
 
-    def get_ex_cum_factor(self, order_book_id):
-        return self._ex_cum_factor.get_factors(order_book_id)
+    def get_ex_cum_factor(self, instrument: Instrument):
+        try:
+            ex_factor_store = self._ex_factor_stores[instrument.type, instrument.market]
+        except KeyError:
+            return None
+
+        return ex_factor_store.get_factors(instrument.order_book_id)
 
     def _update_weekly_trading_date_index(self, idx):
         env = Environment.get_instance()
@@ -320,7 +330,7 @@ def history_bars(
                 week_bars = self.resample_week_bars(bars, bar_count, fields)
                 return week_bars if fields is None else week_bars[fields]
 
-            adjust_bars_date = adjust_bars(bars, self.get_ex_cum_factor(instrument.order_book_id),
+            adjust_bars_date = adjust_bars(bars, self.get_ex_cum_factor(instrument),
                                            fields, adjust_type, adjust_orig)
             adjust_week_bars = self.resample_week_bars(adjust_bars_date, bar_count, fields)
             return adjust_week_bars if fields is None else adjust_week_bars[fields]
@@ -337,7 +347,7 @@ def history_bars(
         if isinstance(fields, str) and fields not in FIELDS_REQUIRE_ADJUSTMENT:
             return bars if fields is None else bars[fields]
 
-        bars = adjust_bars(bars, self.get_ex_cum_factor(instrument.order_book_id),
+        bars = adjust_bars(bars, self.get_ex_cum_factor(instrument),
                            fields, adjust_type, adjust_orig)
 
         return bars if fields is None else bars[fields]

rqalpha/data/bundle.py

@@ -119,7 +119,7 @@ def _get_split(self):
         return rqdatac.get_split(stocks)
 
     def _write(self, data_iter: Iterable[tuple[str, np.ndarray]]):
-        with h5py.File(os.path.join(self.d, 'split_factors.h5'), "w") as h5:
+        with h5py.File(os.path.join(self.d, 'split_factor.h5'), "w") as h5:
             for order_book_id, data in data_iter:
                 h5.create_dataset(order_book_id, data=data)
 
@@ -130,32 +130,45 @@ def __call__(self):
         split['split_factor'] = split['split_coefficient_to'] / split['split_coefficient_from']
         split = split[['split_factor']]
         split.reset_index(inplace=True)
-        split.rename(columns={'ex_dividend_date': 'ex_date'}, inplace=True)
+        split.rename(columns={'ex_dividend_date': 'ex_date'}, inplace=True)  # type: ignore
         split['ex_date'] = [convert_date_to_int(d) for d in split['ex_date']]
         split.set_index(['order_book_id', 'ex_date'], inplace=True)
         self._write([(
             order_book_id, split.loc[order_book_id].to_records()
         ) for order_book_id in split.index.levels[0]])  # type: ignore
 
 
+class GenerateExFactorBundle:
+    def __init__(self, d: str):
+        self.d = d
+    
+    def _get_ex_factor(self):
+        stocks = rqdatac.all_instruments().order_book_id.tolist()
+        return rqdatac.get_ex_factor(stocks)
 
-def gen_ex_factor(d):
-    stocks = rqdatac.all_instruments().order_book_id.tolist()
-    ex_factor = rqdatac.get_ex_factor(stocks)
-    ex_factor.reset_index(inplace=True)
-    ex_factor['ex_date'] = [convert_date_to_int(d) for d in ex_factor['ex_date']]
-    ex_factor.rename(columns={'ex_date': 'start_date'}, inplace=True)
-    ex_factor.set_index(['order_book_id', 'start_date'], inplace=True)
-    ex_factor = ex_factor[['ex_cum_factor']]
+    def _write(self, data_iter: Iterable[tuple[str, np.ndarray]]):
+        with h5py.File(os.path.join(self.d, 'ex_cum_factor.h5'), "w") as h5:
+            for order_book_id, data in data_iter:
+                h5.create_dataset(order_book_id, data=data)
+    
+    def __call__(self):
+        ex_factor = self._get_ex_factor()
+        if ex_factor is None:
+            raise RuntimeError("Got no ex factor data")
+        ex_factor.reset_index(inplace=True)
+        ex_factor['ex_date'] = [convert_date_to_int(d) for d in ex_factor['ex_date']]
+        ex_factor.rename(columns={'ex_date': 'start_date'}, inplace=True)
+        ex_factor.set_index(['order_book_id', 'start_date'], inplace=True)
+        ex_factor = ex_factor[['ex_cum_factor']]
 
-    dtype = ex_factor.loc[ex_factor.index.levels[0][0]].to_records().dtype
-    initial = np.empty((1,), dtype=dtype)
-    initial['start_date'] = 0
-    initial['ex_cum_factor'] = 1.0
+        dtype = ex_factor.loc[ex_factor.index.levels[0][0]].to_records().dtype  # type: ignore
+        initial = np.empty((1,), dtype=dtype)
+        initial['start_date'] = 0
+        initial['ex_cum_factor'] = 1.0
 
-    with h5py.File(os.path.join(d, 'ex_cum_factor.h5'), 'w') as h5:
-        for order_book_id in ex_factor.index.levels[0]:
-            h5[order_book_id] = np.concatenate([initial, ex_factor.loc[order_book_id].to_records()])
+        self._write(((
+            order_book_id, np.concatenate([initial, ex_factor.loc[order_book_id].to_records()])
+        ) for order_book_id in ex_factor.index.levels[0]))  # type: ignore
 
 
 def gen_share_transformation(d):
@@ -484,7 +497,7 @@ def gather_tasks(path: str, create: bool, enable_compression: bool, **h5_kwargs)
     )
 
     gen_file_funcs = (
-        gen_instruments, gen_trading_dates, gen_ex_factor, gen_st_days,
+        gen_instruments, gen_trading_dates, gen_st_days,
         gen_suspended_days, gen_yield_curve, gen_share_transformation, gen_future_info
     )
     kwargs = {}
@@ -497,6 +510,7 @@ def gather_tasks(path: str, create: bool, enable_compression: bool, **h5_kwargs)
         tasks.append(GenerateFileTask(func, path))
     tasks.append(GenerateFileTask(GenerateDividendBundle(path)))
     tasks.append(GenerateFileTask(GenerateSplitBundle(path)))
+    tasks.append(GenerateFileTask(GenerateExFactorBundle(path)))
     return tasks
 
 

rqalpha/data/data_proxy.py

@@ -118,16 +118,17 @@ def get_split(self, order_book_id: str) -> np.ndarray | None:
         return self._data_source.get_split(instrument)
 
     @lru_cache(10240)
-    def _get_prev_close(self, order_book_id, dt):
+    def _get_prev_close(self, order_book_id, dt, adjust_type: str = "pre"):
         instrument = self.instrument_not_none(order_book_id)
         prev_trading_date = self.get_previous_trading_date(dt)
         bar = self._data_source.history_bars(instrument, 1, '1d', 'close', prev_trading_date,
-                                             skip_suspended=False, include_now=False, adjust_orig=dt)
+                                             skip_suspended=False, include_now=False, adjust_type=adjust_type, adjust_orig=dt)
         if bar is None or len(bar) < 1:
             return np.nan
         return bar[0]
 
-    def get_prev_close(self, order_book_id, dt):
+    def get_prev_close(self, order_book_id, dt, adjust_type: str = "pre"):
+        # 获取（基于当日前复权过的）昨收价
         return self._get_prev_close(order_book_id, dt.replace(hour=0, minute=0, second=0))
 
     @lru_cache(10240)

rqalpha/mod/rqalpha_mod_sys_accounts/position_model.py

@@ -15,6 +15,7 @@
 #         在此前提下，对本软件的使用同样需要遵守 Apache 2.0 许可，Apache 2.0 许可与本许可冲突之处，以本许可为准。
 #         详细的授权流程，请联系 public@ricequant.com 获取。
 from datetime import date
+from typing import Optional
 from functools import cached_property
 
 from decimal import Decimal
@@ -61,6 +62,17 @@ def __init__(self, order_book_id, direction, init_quantity=0, init_price=None):
         self._pending_transform = None
         self._non_closable = 0
 
+        # 当日发生的拆分和分红，用于 position_pnl 的计算
+        self._daily_dividend: float = 0.
+        self._daily_split: float = 1.
+        self._unadjusted_prev_close = None
+
+    @property
+    def unadjusted_prev_close(self) -> float:
+        if self._unadjusted_prev_close is None:
+            self._unadjusted_prev_close = self._env.data_proxy.get_prev_close(self._order_book_id, self._env.trading_dt, "none")
+        return self._unadjusted_prev_close
+
     @property
     def dividend_receivable(self):
         # type: () -> float
@@ -83,6 +95,15 @@ def equity(self):
     def market_value_local(self):
         return self.market_value
 
+    @property
+    def position_pnl(self) -> float:
+        if not self._logical_old_quantity:
+            # 新股第一天，没有 prev_close
+            return 0
+        return (self._logical_old_quantity * self._daily_split * (
+            self.last_price - self.unadjusted_prev_close
+        ) + self._daily_dividend) * self._direction_factor
+
     @property
     def closable(self):
         # type: () -> int
@@ -111,14 +132,15 @@ def get_state(self):
     def before_trading(self, trading_date):
         # type: (date) -> float
         delta_cash = super(StockPosition, self).before_trading(trading_date)
+        self._unadjusted_prev_close = self.last_price
         if self._quantity == 0 and not self._dividend_receivable:
             return delta_cash
         if self.direction != POSITION_DIRECTION.LONG:
             raise RuntimeError("direction of stock position {} is not supposed to be short".format(self._order_book_id))
         data_proxy = self._env.data_proxy
-        self._handle_dividend_book_closure(trading_date, data_proxy)
+        self._daily_dividend = self._handle_dividend_book_closure(trading_date, data_proxy)
         delta_cash += self._handle_dividend_payable(trading_date)
-        self._handle_split(trading_date, data_proxy)
+        self._daily_split = self._handle_split(trading_date, data_proxy)
         return delta_cash
 
     def apply_trade(self, trade):
@@ -198,11 +220,10 @@ def _get_dividends_or_splits(self, events: ndarray | None, trading_date: date, d
         events = events[left_pos: right_pos]
         return events
 
-    def _handle_dividend_book_closure(self, trading_date, data_proxy):
-        # type: (date, DataProxy) -> None
+    def _handle_dividend_book_closure(self, trading_date: date, data_proxy: DataProxy) -> float:
         dividends = self._get_dividends_or_splits(self._all_dividends, trading_date, "ex_dividend_date")  # type: ignore[reportIncompatibleVariableOverride]
         if dividends is None or len(dividends) == 0:
-            return
+            return 0
         dividend_per_share: float = (dividends["dividend_cash_before_tax"] / dividends["round_lot"]).sum() * (1 - self.dividend_tax_rate)
         self._avg_price -= dividend_per_share
         # 前一天结算发生了除息, 此时 last_price 还是前一个交易日的收盘价，需要改为 除息后收盘价, 否则影响在before_trading中查看盈亏
@@ -211,6 +232,7 @@ def _handle_dividend_book_closure(self, trading_date, data_proxy):
         # FIXME: 这里隐含了获取的多条 dividend 的 payable_date 都相同的假设
         payable_date = _int_to_date(dividends["payable_date"][-1])
         self._dividend_receivable = (payable_date, self._quantity * dividend_per_share)
+        return self._quantity * dividend_per_share
 
     def _handle_dividend_payable(self, trading_date):
         # type: (date) -> float
@@ -236,18 +258,18 @@ def _handle_dividend_payable(self, trading_date):
         else:
             return dividend_value
 
-    def _handle_split(self, trading_date, data_proxy):
+    def _handle_split(self, trading_date, data_proxy) -> float:
         splits = self._get_dividends_or_splits(self._all_splits, trading_date, "ex_date")  # type: ignore[reportIncompatibleVariableOverride]
         if splits is None or len(splits) == 0:
-            return
+            return 1.
         ratio: float = splits["split_factor"].cumprod()[-1]
         self._avg_price /= ratio
         self._last_price /= ratio
         ratio_decimal = Decimal(ratio)
         # int(6000 * 1.15) -> 6899
         self._old_quantity = self._quantity = round(Decimal(self._quantity) * ratio_decimal)
-        self._logical_old_quantity = round(Decimal(self._logical_old_quantity) * ratio_decimal)
         self._queue.handle_split(ratio_decimal, self._quantity)
+        return ratio
 
 
 class FuturePosition(Position):

rqalpha/portfolio/position.py

@@ -21,7 +21,7 @@
 from typing import Dict, Iterable, Tuple, Optional, Deque, List
 from functools import cached_property
 
-from rqalpha.const import POSITION_DIRECTION, POSITION_EFFECT, MARKET
+from rqalpha.const import POSITION_DIRECTION, POSITION_EFFECT, MARKET, SIDE
 from rqalpha.environment import Environment
 from rqalpha.interface import AbstractPosition
 from rqalpha.model.instrument import Instrument
@@ -123,14 +123,28 @@ def avg_price(self):
         return self._avg_price
 
     @property
-    def trading_pnl(self):
-        # type: () -> float
+    def trading_pnl(self) -> float:
+        """
+        交易盈亏
+        即当日所做的交易相比使用最新价交易产生的盈亏
+
+        买方向：sum(t.quantity * (last_price - t.price) for t in trades)
+        卖方向：sum(t.quantity * (t.price - last_price) for t in trades)
+
+        direction = 1 if buy else -1
+        trading_pnl 
+            = sum(t.quantity * (last_price - t.price) * direction for t in trades)
+            = sum(t.quantity * last_price * direction for t in trades) - sum(t.quantity * t.price * direction for t in trades)
+            = net_trade_quantity * last_price * direction - sum(t.quantity * t.price * direction for t in trades)
+        
+        trade_cost = sum(t.quantity * t.price * direction for t in trades)
+        trading_pnl = net_trade_quantity * last_price * direction - trade_cost
+        """
         trade_quantity = self._quantity - self._logical_old_quantity
         return (trade_quantity * self.last_price - self._trade_cost) * self._direction_factor
 
     @property
-    def position_pnl(self):
-        # type: () -> float
+    def position_pnl(self) -> float:
         if self._logical_old_quantity:
             return self._logical_old_quantity * (self.last_price - self.prev_close) * self._direction_factor
         else:
@@ -155,8 +169,8 @@ def equity(self):
         return self.last_price * self._quantity if self._quantity else 0
 
     @property
-    def prev_close(self):
-        # type: () -> float
+    def prev_close(self) -> float:
+        # （前复权过的）昨收价
         if not is_valid_price(self._prev_close):
             self._prev_close = self._env.data_proxy.get_prev_close(self._order_book_id, self._env.trading_dt)
             if not is_valid_price(self._prev_close):
@@ -232,10 +246,17 @@ def before_trading(self, trading_date):
         self._prev_close = None
         return 0
 
+    def _update_costs(self, trade: Trade):
+        self._transaction_cost += trade.transaction_cost
+        if trade.side == SIDE.BUY:
+            self._trade_cost += trade.last_price * trade.last_quantity
+        else:
+            self._trade_cost -= trade.last_price * trade.last_quantity
+
     def apply_trade(self, trade):
         # type: (Trade) -> float
         # 返回总资金的变化量
-        self._transaction_cost += trade.transaction_cost
+        self._update_costs(trade)
         if trade.position_effect == POSITION_EFFECT.OPEN:
             self._queue.handle_trade(trade.last_quantity, self._env.trading_dt.date())
             if self._quantity < 0:
@@ -244,14 +265,12 @@ def apply_trade(self, trade):
                 cost = self._quantity * self._avg_price + trade.last_quantity * trade.last_price
                 self._avg_price = cost / (self._quantity + trade.last_quantity)
             self._quantity += trade.last_quantity
-            self._trade_cost += trade.last_price * trade.last_quantity
             return (-1 * trade.last_price * trade.last_quantity) - trade.transaction_cost
         elif trade.position_effect == POSITION_EFFECT.CLOSE:
             # 先平昨，后平今
             self._queue.handle_trade(-trade.last_quantity, self._env.trading_dt.date())
             self._old_quantity -= min(trade.last_quantity, self._old_quantity)
             self._quantity -= trade.last_quantity
-            self._trade_cost -= trade.last_price * trade.last_quantity
             return trade.last_price * trade.last_quantity - trade.transaction_cost
         else:
             raise NotImplementedError("{} does not support position effect {}".format(