多空组合构建代码

#R-1是未来一期的return,permno是股票代码，pre是预测的未来一期return
def long_short(result):
    result["R-1"] = result["R-1"].astype(float)
    result = result*100
    result_shift = result.groupby('PERMNO')['R-1',"pre"].shift(-1)
    real_ret = pd.DataFrame(result_shift["R-1"])
    real_ret["PERMNO"] = result["PERMNO"]
    real_ret =  real_ret.pivot_table(real_ret,index=[u'date',u'PERMNO'])
    pf_ret = real_ret.unstack(level=-1)
    pf_ret.columns = pf_ret.columns.droplevel()
    
    pre_ret = pd.DataFrame(result_shift["pre"])
    pre_ret["PERMNO"] = result["PERMNO"]
    pre_ret = pre_ret.pivot_table(pre_ret,index=[u'date',u'PERMNO'])
    pre_pf_ret = pre_ret.unstack(level=-1)
    pre_pf_ret.columns = pre_pf_ret.columns.droplevel()
    
    ranks = pre_pf_ret.rank(axis=1, method="min")
    stock_num = ranks.shape[1] - ranks.isna().sum(axis=1)
    ranks = ranks.div(stock_num,axis=0)
    long = ranks > 0.9
    long = long*pf_ret
    long = long.sum(axis=1).div(stock_num,axis=0)*10
    long.name = "long"
    dflong = long.to_frame()
    short = ranks < 0.1
    short = short*pf_ret
    short = short.sum(axis=1).div(stock_num,axis=0)*10
    short.name = "short"
    dfshort = short.to_frame()
    pf = long-short
    pf.name = "pf"
    
    strategy = pd.concat([pf,dflong ,dfshort],axis = 1)
    return strategy



figure["long"]= strategy['dflong']
figure["short"]= strategy['dfshort']

figure = figure/100 +1 
value = figure.cumprod()
lgvalue = np.log(value)

lgvalue.plot()
plt.ylabel("log($ value of investment)")
plt.legend(["risk-free,end value=$9.10", "market, end value=$187.81","long, end value=$1,391,449.90", "short, end value=$3.26"])