Bug: 修正不满足持仓数量时最大卖出数量判断问题 (#16);

Bug: 简化买卖判断逻辑避免中间错误 (#15);

account.py

@@ -1,7 +1,7 @@
 import warnings
 import pandas as pd
 from rich import print as rprint
-from typing import Union, Iterable, Optional, Dict
+from typing import Union
 
 
 # 警告格式
@@ -39,8 +39,6 @@ class Account():
         if isinstance(leverage, (int,float)):
             if leverage > 1.5:
                 raise ValueError('leverage should less than 1.5')
-            # if leverage > 1:
-            #     warnings.warn('leverage(杠杆率)大于1时，优先满足杠杆，会出现高于指定换手率的情况！', category=CustomWarning)
         else:
             raise ValueError('leverage should be int or float')
         self.leverage = leverage

trader.py

@@ -5,7 +5,6 @@ import os
 import copy
 
 from typing import Union, Iterable, Dict
-from ordered_set import OrderedSet
 
 from account import Account
 from dataloader import DataLoader
@@ -64,7 +63,7 @@ class Trader(Account):
                  account: dict={},
                  **kwargs) -> None:
         # 初始化账户
-        super().__init__(account)
+        super().__init__(**account)
         if isinstance(signal, dict):
             self.signal = signal
             if 'close' in signal:
@@ -232,18 +231,16 @@ class Trader(Account):
         """
         # 计算持仓和最大可交易数量
         if len(self.position) > 0:
-            last_position = self.position['stock_code'].values
+            last_position = pd.Series(index=self.position['stock_code'])
-            last_position = factor.loc[last_position].sort_values(ascending=self.ascending)
             if len(self.position) <= self.num:
-                # 如果昨日持仓本身就不足持仓数量则不用足额换仓
+                # 如果昨日持仓本身就不足持仓数量则按照昨日持仓数量作为基准计算换仓数量
-                max_sell_num = min(int(self.interval.loc[date]*self.num), int(self.interval.loc[date]*self.num)+len(self.position)-self.num)
+                # 不足的数量通过买入列表自适应调整
-            else:
+                # 这样能实现在因子值不足时也正常换仓
-                # 如果昨日持仓本身就超额持仓数量则超额换仓
+                max_sell_num = self.interval.loc[date]*len(last_position)
-                max_sell_num = max(int(self.interval.loc[date]*self.num), int(self.interval.loc[date]*self.num)+len(self.position)-self.num)
         else:
             last_position = pd.Series()
             max_sell_num = self.num
-        target_list = []
+        
         # 获取用于筛选的数据
         stock_status = self.today_data['basic'].get(factor.index.values, 'opening_info').sort_index()
         stock_amt20 = self.today_data['basic'].get(factor.index.values, 'amount_20').sort_index()
@@ -255,8 +252,12 @@ class Trader(Account):
         stock_ipo_filter = stock_ipo_days.copy()
         stock_ipo_filter.loc[:] = 0
         stock_ipo_filter.loc[stock_ipo_days > self.ipo_days] = 1
-        # 剔除列表 = ipo筛选 + 成交量筛选 + 额外剔除
+        
-        # 其中额外提出会强制执行，因此单独保存一份强制执行的列表
+        # 剔除列表
+        # 包含强制列表和普通列表：
+        # 强制列表会将已经持仓的也强制剔除并且不算在换手率限制中
+        # 普通列表如果已经持有不会过滤只对新买入的过滤
+        # 强制过滤列表
         exclude_stock = []
         for exclude in self.exclude_list:
             if exclude == 'abnormal':
@@ -267,90 +268,80 @@ class Trader(Account):
                 exclude_stock += stock_recession.loc[stock_recession > 0].index.to_list()
         exclude_stock = list(set(exclude_stock))
         force_exclude = copy.deepcopy(exclude_stock)
-        exclude_stock += stock_ipo_filter.loc[stock_ipo_filter != 1].index.to_list()
+        # 普通过滤列表
-        exclude_stock += stock_amt_filter.loc[stock_amt_filter != 1].index.to_list()
+        normal_exclude = []
-        exclude_stock = list(set(exclude_stock))
+        normal_exclude += stock_ipo_filter.loc[stock_ipo_filter != 1].index.to_list()
+        normal_exclude += stock_amt_filter.loc[stock_amt_filter != 1].index.to_list()
+        normal_exclude = list(set(normal_exclude))
+        
         # 交易列表
         if self.leverage <= 1.0:
-            # 获取当前时间目标列表和冻结(无法交易)列表:
+            # 如果没有杠杆：
-            # 1 保留冻结的昨日持仓
-            # 2 目标持仓列表会对当日因子进行ST、成交量、上市时间滤后按因子排序方向选取
-            for stock in last_position.index:
-                # 如果停牌或者跌停继续持有
-                if stock_status.loc[stock] in [0,2]:
-                    target_list.append(stock)
-            # 剔除过滤条件后
-            after_filter_list = list(set(factor.index) - set(exclude_stock))
-            for stock in factor.loc[after_filter_list].dropna().sort_values(ascending=self.ascending).index.values:
-                # 目标持仓数量达到持仓数量则中止
-                if len(target_list) == self.num:
-                    break
-                # ST过滤
-                if self.with_st:
-                    if stock_status.loc[stock] in [0,2,5,7]:
-                        if stock in last_position.index:
-                            target_list.append(stock)
-                    else:
-                        target_list.append(stock)
-                else:
-                    # 非ST
-                    if stock_status.loc[stock] in [3,6]:
-                        target_list.append(stock)
-                    # 如果停牌或者跌停继续持有
-                    if stock_status.loc[stock] in [0,2,7]:
-                        if stock in last_position.index:
-                            target_list.append(stock)
-                target_list = list(OrderedSet(target_list))
-            
-            # 如果没有杠杆
             buy_list = []
             sell_list = []
+            untradable_list = []
+            target_list = []
             # ----- 卖出 -----
             # 异常强制卖出
             for stock in last_position.index:
-                if stock in force_exclude:
+                if stock_status.loc[stock] in [0,2,5,7]:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    untradable_list.append(stock)
-                        continue
+                else:
-                    sell_list.append(stock)
+                    if stock in force_exclude:
+                        sell_list.append(stock)
             force_sell_num = len(sell_list)
-            # 按照反向排名逐个卖出
+                
+            # 剔除无法交易列表后，按照当日因子反向排名逐个卖出
+            # 对无因子的异常股票进行最末位填充用于判断卖出，在判断买入列表时还是使用原始因子
             if self.ascending:
-                factor = factor.fillna(factor.max()+1)
+                factor_filled = factor.fillna(factor.max()+1)
             else:
-                factor = factor.fillna(factor.min()-1)
+                factor_filled = factor.fillna(factor.min()-1)
-            for stock in factor.loc[last_position.index].sort_values(ascending=self.ascending).index.values[::-1]:
+            for stock in factor_filled.loc[list(set(last_position.index)-set(untradable_list)-set(sell_list))].sort_values(ascending=self.ascending).index.values[::-1]:
                 if len(sell_list) >= max_sell_num + force_sell_num:
                     break
-                if stock in target_list:
+                if stock_status.loc[stock] in [0,2,5,7]:
                     continue
                 else:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    sell_list.append(stock)
-                        continue
-                    else:
-                        sell_list.append(stock)
                 sell_list = list(set(sell_list))
+                
             # ----- 买入 -----
-            # 卖出后持仓列表
+            # 剔除过滤条件后可买入列表
+            after_filter_list = list(set(factor.index) - set(normal_exclude) - set(force_exclude))
+            target_list = factor.loc[after_filter_list].dropna().sort_values(ascending=self.ascending).index.to_list()
+            
+            # 更新卖出后的持仓列表
             after_sell_list = set(last_position.index) - set(sell_list)
-            cant_buy_list = [] # 涨停股记录
+            limit_up_list = [] # 涨停股记录
             max_buy_num = max(0, self.num-len(last_position)+len(sell_list))
             for stock in target_list:
                 if stock in after_sell_list:
                     continue
                 else:
+                    if self.with_st:
+                        pass
+                    else:
+                        if stock_status.loc[stock] in [1,2]:
+                            continue
                     if stock_status.loc[stock] in [4,6]:
-                        cant_buy_list.append(stock)
+                        limit_up_list.append(stock)
                     buy_list.append(stock)
                 if len(buy_list) == max_buy_num:
                     break
+            # 剔除同时在sell_list和buy_list的股票
+            duplicate_stock = set(sell_list) & set(buy_list)
+            sell_list = list(set(sell_list) - duplicate_stock)
+            buy_list = list(set(buy_list) - duplicate_stock)
+            # 生成下一期持仓
             next_position = pd.DataFrame({'stock_code': list((set(last_position.index) - set(sell_list)) | set(buy_list))})
             next_position['date'] = date
             next_position['weight'] = self.get_weight(date, self.leverage, next_position)
             # 剔除无法买入的涨停股，这部分仓位空出
-            next_position = next_position[~next_position['stock_code'].isin(cant_buy_list)]
+            next_position = next_position[~next_position['stock_code'].isin(limit_up_list)]
             next_position['margin_trade'] = 0
         else:
-            # 如果有杠杆
+            # 如果有杠杆：
             def assign_stock(normal_list, margin_list, margin_needed, stock, status):
                 if status == 1:
                     if len(margin_list) < margin_needed:
@@ -363,12 +354,6 @@ class Trader(Account):
             margin_ratio = max(self.leverage-1, 0)
             margin_needed = round(self.num * margin_ratio)
             is_margin = self.today_data['basic'].get(factor.index.values, 'margin_list').sort_index()
-            # 获取当前时间目标列表和冻结(无法交易)列表:
-            # 1 保留冻结的昨日持仓
-            # 2 目标持仓列表会对当日因子进行ST、成交量、上市时间滤后按因子排序方向选取
-            # 3 根据融资比例将股票分为常规池子和融资池子
-            normal_list = []
-            margin_list = []
             
             # 获取历史融资融券池
             if len(last_position) > 0:
@@ -381,110 +366,97 @@ class Trader(Account):
             else:
                 last_normal_list = []
             
-            for stock in last_margin_list:
-                # 如果停牌或者跌停继续持有
-                if stock_status.loc[stock] in [0,2]:
-                    margin_list.append(stock)
-            for stock in last_normal_list:
-                # 如果停牌或者跌停继续持有
-                if stock_status.loc[stock] in [0,2]:
-                    normal_list.append(stock)
-            # 剔除过滤条件后
-            after_filter_list = list(set(factor.index) - set(exclude_stock))
-            for stock in factor.loc[after_filter_list].dropna().sort_values(ascending=self.ascending).index.values:
-                if len(normal_list + margin_list) == self.num:
-                    break
-                if self.with_st:
-                    if stock_status.loc[stock] in [0,2,5,7]:
-                        if stock in last_margin_list:
-                            margin_list.append(stock)
-                        else:
-                            normal_list.append(stock)
-                    else:
-                        normal_list, margin_list = assign_stock(normal_list, margin_list, margin_needed, stock, is_margin.loc[stock])
-                else:
-                    # 非ST
-                    if stock_status.loc[stock] in [3,6]:
-                        normal_list, margin_list = assign_stock(normal_list, margin_list, margin_needed, stock, is_margin.loc[stock])
-                    else:
-                        # 如果停牌或者跌停继续持有
-                        if stock_status.loc[stock] in [0,2,7]:
-                            if stock in last_margin_list:
-                                margin_list.append(stock)
-                            else:
-                                normal_list.append(stock)
-                margin_list = list(OrderedSet(margin_list))
-                normal_list = list(OrderedSet(normal_list))
-            target_list = normal_list + margin_list
             # ----- 卖出 -----            
             buy_list = []
             sell_list = []
-            # 融资融券池的和非融资融券池的分开更新
+            untradable_list = []
+            # 分别更新融资融券池的和非融资融券池
             # 更新融资融券池
             # 异常强制卖出
             for stock in last_margin_list:
-                if stock in force_exclude:
+                if stock_status.loc[stock] in [0,2,5,7]:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    untradable_list.append(stock)
-                        continue
-                    sell_list.append(stock)
-            force_sell_num = len(sell_list)
-            
-            for stock in factor.loc[last_margin_list].sort_values(ascending=self.ascending).index.values[::-1]:
-                if len(sell_list) >= int(max_sell_num * margin_ratio) + force_sell_num + 1:
-                    break
-                if stock in normal_list:
                     continue
                 else:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    if stock in force_exclude:
-                        continue
-                    else:
                         sell_list.append(stock)
+            force_sell_num = len(sell_list)
+            
+            # 对无因子的异常股票进行最末位填充用于判断卖出，在判断买入列表时还是使用原始因子
+            if self.ascending:
+                factor_filled = factor.fillna(factor.max()+1)
+            else:
+                factor_filled = factor.fillna(factor.min()-1)
+            
+            for stock in factor_filled.loc[list(set(last_margin_list)-set(untradable_list)-set(sell_list))].sort_values(ascending=self.ascending).index.values[::-1]:
+                if len(sell_list) >= int(max_sell_num * margin_ratio) + force_sell_num + 1:
+                    break
+                if stock_status.loc[stock] in [0,2,5,7]:
+                    continue
+                else:
+                    sell_list.append(stock)
                 sell_list = list(set(sell_list))
             next_margin_list = list(set(last_margin_list) - set(sell_list))
             # 更新非融资融券池
             # 异常强制卖出
+            untradable_list = []
             for stock in last_normal_list:
-                if stock in force_exclude:
+                if stock_status.loc[stock] in [0,2,5,7]:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    untradable_list.append(stock)
-                        continue
-                    sell_list.append(stock)
-                    force_sell_num += 1
-            for stock in factor.loc[last_normal_list].sort_values(ascending=self.ascending).index.values[::-1]:
-                if len(sell_list) >= max_sell_num + force_sell_num:
-                    break
-                if stock in normal_list:
                     continue
                 else:
-                    if stock_status.loc[stock] in [0,2,5,7]:
+                    if stock in force_exclude:
-                        continue
-                    else:
                         sell_list.append(stock)
+                        force_sell_num += 1
+            for stock in factor_filled.loc[list(set(last_normal_list)-set(untradable_list)-set(sell_list))].sort_values(ascending=self.ascending).index.values[::-1]:
+                if len(sell_list) >= max_sell_num + force_sell_num:
+                    break
+                if stock_status.loc[stock] in [0,2,5,7]:
+                    continue
+                else:
+                    sell_list.append(stock)
                 sell_list = list(set(sell_list))
             next_normal_list = list(set(last_normal_list) - set(sell_list))
+            
             # ----- 买入 -----
-            # 卖出后持仓列表
+            # 剔除过滤条件后可买入列表
+            after_filter_list = list(set(factor.index) - set(normal_exclude) - set(force_exclude))
+            target_list = factor.loc[after_filter_list].dropna().sort_values(ascending=self.ascending).index.to_list()
+            
+            # 更新卖出后的持仓列表
             after_sell_list = set(last_position.index) - set(sell_list)
-            cant_buy_list = [] # 涨停股记录
+            limit_up_list = [] # 涨停股记录
             # 融资融券池的和非融资融券池的分开更新
             # 更新融资融券池     
-            for stock in margin_list:
+            for stock in target_list:
                 if stock in after_sell_list:
                     continue
                 else:
-                    if stock_status.loc[stock] in [4,6]:
+                    if self.with_st:
-                        cant_buy_list.append(stock)
+                        pass
-                    next_margin_list.append(stock)
+                    else:
+                        if stock_status.loc[stock] in [1,2]:
+                            continue
+                    if is_margin.loc[stock]:
+                        if stock_status.loc[stock] in [4,6]:
+                            limit_up_list.append(stock)
+                        next_margin_list.append(stock)
                 if len(next_margin_list) >= margin_needed:
                     break
             # 更新非融资融券池 
-            for stock in normal_list:
+            for stock in target_list:
-                if stock in after_sell_list:
+                if stock in (set(after_sell_list) | set(next_margin_list)):
                     continue
                 else:
+                    if self.with_st:
+                        pass
+                    else:
+                        if stock_status.loc[stock] in [1,2]:
+                            continue
                     if stock_status.loc[stock] in [4,6]:
-                        cant_buy_list.append(stock)
+                        limit_up_list.append(stock)
                     next_normal_list.append(stock)
-                if len(next_normal_list) >= self.num - margin_needed:
+                if len(next_normal_list) >= self.num - len(next_margin_list):
                     break
             
             next_position = pd.DataFrame({'stock_code': next_margin_list + next_normal_list})
@@ -497,7 +469,7 @@ class Trader(Account):
             next_position.loc[next_margin_list, 'margin_trade'] = 1
             next_position = next_position.reset_index()
             # 剔除无法买入的涨停股，这部分仓位空出
-            next_position = next_position[~next_position['stock_code'].isin(cant_buy_list)]
+            next_position = next_position[~next_position['stock_code'].isin(limit_up_list)]
         # 检测当前持仓是否可以交易
         frozen_list = []
         if len(self.position) > 0: