聚宽平台——双均线策略实现

import tushare as ts
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

#1.获取股票代码为002400的股票数据
#df = ts.get_k_data('002400', start='1988')
#df.to_csv('002400.csv')
df = pd.read_csv('002400.csv', index_col='date', parse_dates=['date'])[['open', 'close', 'high', 'low']]
print(df)

#2.计算其5日和30日均线(收盘价)
df['ma5'] = np.nan
df['ma30'] = np.nan
df['ma5'] = df['close'].rolling(5).mean()
df['ma30'] = df['close'].rolling(30).mean()
#print(df)

#3.使用matplotlib模块实现收盘价、ma5、ma30的可视化
# plt.plot(df.index, df['close'].values, 'r', label='close')
# plt.plot(df.index, df['ma5'].values, 'b', label='ma5')
# plt.plot(df.index, df['ma30'].values, 'k', label='ma30')
# plt.show()

#4.分析输出所有金叉和死叉日期
golden_cross = []
death_cross = []
df = df['2011':]
for i in range(1, len(df)):
    if (df['ma5'][i]>df['ma30'][i] and df['ma5'][i-1]<=df['ma30'][i-1]):
        golden_cross.append(df.index[i].to_pydatetime())
    if (df['ma5'][i] < df['ma30'][i] and df['ma5'][i - 1] >= df['ma30'][i - 1]):
        death_cross.append(df.index[i].to_pydatetime())

sr1 = df['ma5'] >= df['ma30']
sr2 = df['ma5'] < df['ma30']
sr2 = sr2.shift(1)
golden_cross = df[sr1 & sr2].index
death_cross = df[~(sr1 | sr2)].index

#5.从2011/01/01开始，初始资金为100000元，金叉尽量买入，死叉全部卖出，那么最终我的收益为多少？
goldsr = pd.Series(1, index=golden_cross)
deathsr = pd.Series(0, index=death_cross)
sr = goldsr.append(deathsr).sort_index()

first_money = 100000
cash = first_money
hold = 0
for i in range(len(sr)):
    p = df['open'][sr.index[i]]
    print('p:'+str(p))
    if sr.iloc[i] == 1:#金叉，买入
        buy = (cash // (100*p))
        hold += buy
        cash -= buy*100*p
    else:
        cash += hold*100*p
        print('cash:'+str(cash))
        hold = 0

now_money = hold*100*df['open'][-1] + cash
print(now_money)