def test( data1, data2, data1_name, data2_name, data1_price, data2_price, data1_price_name, data2_price_name, start_money, gap_days, trading_day, ): """ lundong model test with etf """ print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS], ) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS], ) print "Price1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_price_name, len(data1_price), data1_price[0][DATE_POS], data1_price[-1][DATE_POS], ) print "Price2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_price_name, len(data2_price), data2_price[0][DATE_POS], data2_price[-1][DATE_POS], ) # Preprocess datasets print "\npreprocessing datasets...\n" time_align(data1, data1_price) time_align(data2, data2_price) time_align(data1, data2) time_align(data1, data1_price) time_align(data2, data2_price) time_truncate(data1, "2013-02-22", "2015-10-30") time_truncate(data2, "2013-02-22", "2015-10-30") time_truncate(data1_price, "2013-02-22", "2015-10-30") time_truncate(data2_price, "2013-02-22", "2015-10-30") print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS], ) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS], ) print "Price1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_price_name, len(data1_price), data1_price[0][DATE_POS], data1_price[-1][DATE_POS], ) print "Price2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_price_name, len(data2_price), data2_price[0][DATE_POS], data2_price[-1][DATE_POS], ) # Get close price of two datasets and plot them data1_close_price = [item[CLOSE_PRICE_POS] for item in data1] data2_close_price = [item[CLOSE_PRICE_POS] for item in data2] data1_price_close_price = [item[CLOSE_PRICE_POS] for item in data1_price] data2_price_close_price = [item[CLOSE_PRICE_POS] for item in data2_price] analyser = data_analyser() analyser.draw_relevance(data1_close_price, data1_price_close_price, data1_name, data1_price_name) analyser.draw_relevance(data2_close_price, data2_price_close_price, data2_name, data2_price_name) plt.figure(1) plt.subplot(311) plt.plot(data1_close_price, "b-", data2_close_price, "r-") plt.legend([data1_name, data2_name], loc="upper left") plt.grid(True) # Initilize a lundong model, and get trading decisions by it lundong_model = lundong() # the test start with a blank account account = account_simulator(start_money) account_real = account_simulator(start_money) money = [] money_real = [] # record every trading plan result = [] last_valid_trading_plan = {} for item in data1: trading_plan = lundong_model.get_trading_plan( data1, data2, data1_name, data2_name, item[DATE_POS], gap_days, trading_day, last_valid_trading_plan ) # if there is a trading plan if trading_plan != {}: # choise == 0 means buy data1 and sell data2 if trading_plan["choise"] == 0: # fake # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # buy data1 as many as we can account.buy( data1_name, trading_plan["data1_close_price"], int(math.floor(account.money / trading_plan["data1_close_price"])), ) # real # if we have data2, sell it stock = account_real.get_stock(data2_price_name) if stock != None and stock.share != 0: real_data = get_data(data2_price, trading_plan["date_str"]) account_real.sell(data2_price_name, real_data[CLOSE_PRICE_POS], stock.share) # buy data1 as many as we can real_data = get_data(data1_price, trading_plan["date_str"]) account_real.buy( data1_price_name, real_data[CLOSE_PRICE_POS], int(math.floor(account_real.money / real_data[CLOSE_PRICE_POS])), ) # choise == 1 means sell data1 and buy data2 elif trading_plan["choise"] == 1: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # buy data2 as many as we can account.buy( data2_name, trading_plan["data2_close_price"], int(math.floor(account.money / trading_plan["data2_close_price"])), ) # real stock = account_real.get_stock(data1_price_name) if stock != None and stock.share != 0: real_data = get_data(data1_price, trading_plan["date_str"]) account_real.sell(data1_price_name, real_data[CLOSE_PRICE_POS], stock.share) # buy data1 as many as we can real_data = get_data(data2_price, trading_plan["date_str"]) account_real.buy( data2_price_name, real_data[CLOSE_PRICE_POS], int(math.floor(account_real.money / real_data[CLOSE_PRICE_POS])), ) # choise == 2 means sell all data2 elif trading_plan["choise"] == 2: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) stock = account_real.get_stock(data1_price_name) if stock != None and stock.share != 0: real_data = get_data(data1_price, trading_plan["date_str"]) account_real.sell(data1_price_name, real_data[CLOSE_PRICE_POS], stock.share) stock = account_real.get_stock(data2_price_name) if stock != None and stock.share != 0: real_data = get_data(data2_price, trading_plan["date_str"]) account_real.sell(data2_price_name, real_data[CLOSE_PRICE_POS], stock.share) # other choises means do not trade else: pass # if there is no trading plan, do not trade else: pass money.append(account.get_value()) money_real.append(account_real.get_value()) # money.append(account.money) result.append(trading_plan) if trading_plan != {} and trading_plan["choise"] != -1: last_valid_trading_plan = trading_plan # Get rising rates, choises and plot them data1_up = [] data2_up = [] choise_0 = [] choise_0_index = [] choise_1 = [] choise_1_index = [] choise_2 = [] choise_2_index = [] i = 0 for item in result: if item == {}: data1_up.append(0.0) data2_up.append(0.0) else: data1_up.append(item["data1_up"]) data2_up.append(item["data2_up"]) if item["choise"] == 0: choise_0.append(item["data1_up"]) choise_0_index.append(i) elif item["choise"] == 1: choise_1.append(item["data2_up"]) choise_1_index.append(i) elif item["choise"] == 2: choise_2.append(0) choise_2_index.append(i) i = i + 1 plt.subplot(312) plt.plot( range(len(data1_up)), data1_up, "b-", range(len(data2_up)), data2_up, "r-", choise_0_index, choise_0, "bo", choise_1_index, choise_1, "ro", choise_2_index, choise_2, "r^", ) # plt.plot(data1_up, 'b-', data2_up, 'r-', choise_0, 'bo', choise_1, 'r^') plt.legend( [data1_name, data2_name, "switch to:" + data1_name, "switch to:" + data2_name, "sell all"], loc="upper left" ) plt.grid(True) plt.subplot(313) plt.plot( [item / start_money for item in money], "g-", [item / start_money for item in money_real], "k-", [float(item[CLOSE_PRICE_POS]) / float(data1_price[0][CLOSE_PRICE_POS]) for item in data1_price], "b-", [float(item[CLOSE_PRICE_POS]) / float(data2_price[0][CLOSE_PRICE_POS]) for item in data2_price], "r-", ) plt.legend(["model", "real", data1_price_name, data2_price_name], loc="upper left") plt.grid(True) # plt.show() account.dump() return account.get_value() / start_money, account_real.get_value() / start_money
def test(data1, data2, data1_name, data2_name, start_money, gap_days, trading_day, show_pics): """ lundong model test """ print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) # Preprocess datasets print "preprocessing datasets..." pre_process(data1, data2) print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) # Get close price of two datasets and plot them data1_close_price = [ item[CLOSE_PRICE_POS] for item in data1 ] data2_close_price = [ item[CLOSE_PRICE_POS] for item in data2 ] dates = [ datetime.strptime(item[DATE_POS], "%Y-%m-%d") for item in data1] years = mdates.YearLocator() months = mdates.MonthLocator() yearsFmt = mdates.DateFormatter("%Y-%m") #plt.figure(1) #plt.subplot(311) if show_pics: fig, ax = plt.subplots() ax.set_title("股指变化图") ax.plot(dates, data1_close_price, 'b-', dates, data2_close_price, 'r-') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend([data1_name, data2_name], loc="upper left") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() # Initilize a lundong model, and get trading decisions by it lundong_model = lundong() # the test start with a blank account account = account_simulator(start_money) money = [] # record every trading plan result = [] last_valid_trading_plan = {} for item in data1: trading_plan = lundong_model.get_trading_plan(data1, data2, data1_name, data2_name, item[DATE_POS], gap_days, trading_day, last_valid_trading_plan) # if there is a trading plan if trading_plan != {}: # choise == 0 means buy data1 and sell data2 if trading_plan["choise"] == 0: # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # buy data1 as many as we can account.buy(data1_name, trading_plan["data1_close_price"], int(math.floor(account.money / trading_plan["data1_close_price"]))) # choise == 1 means sell data1 and buy data2 elif trading_plan["choise"] == 1: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # buy data2 as many as we can account.buy(data2_name, trading_plan["data2_close_price"], int(math.floor(account.money / trading_plan["data2_close_price"]))) # choise == 2 means sell all data2 elif trading_plan["choise"] == 2: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # other choises means do not trade else: pass # if there is no trading plan, do not trade else: pass money.append(account.get_value()) #money.append(account.money) result.append(trading_plan) if trading_plan != {} and trading_plan["choise"] != -1: last_valid_trading_plan = trading_plan # Get rising rates, choises and plot them data1_up = [] data2_up = [] choise_0 = [] choise_0_index = [] choise_1 = [] choise_1_index = [] choise_2 = [] choise_2_index = [] i = 0 for item in result: if item == {}: data1_up.append(0.0) data2_up.append(0.0) else: data1_up.append(item["data1_up"]) data2_up.append(item["data2_up"]) if item["choise"] == 0: choise_0.append(item["data1_up"]) choise_0_index.append(dates[i]) elif item["choise"] == 1: choise_1.append(item["data2_up"]) choise_1_index.append(dates[i]) elif item["choise"] == 2: choise_2.append(0) choise_2_index.append(dates[i]) i = i + 1 if show_pics: fig, ax = plt.subplots() ax.set_title("股指涨幅与模型交易方案") ax.plot(dates, data1_up, 'b-', dates, data2_up, 'r-',\ choise_0_index, choise_0, 'bo', choise_1_index, choise_1, 'ro',\ choise_2_index, choise_2, 'g^') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend([data1_name, data2_name, "轮动到:" + data1_name, "轮动到:" + data2_name, "清仓"], loc="upper center") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() if show_pics: fig, ax = plt.subplots() ax.set_title("投资收益对比") ax.plot(dates, [item / start_money for item in money], 'g-', \ dates, [ float(item[CLOSE_PRICE_POS]) / float(data1[0][CLOSE_PRICE_POS]) for item in data1], 'b-', \ dates, [ float(item[CLOSE_PRICE_POS]) / float(data2[0][CLOSE_PRICE_POS]) for item in data2], 'r-') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend(["模型收益率", "单独投资 " + data1_name + " 收益率", "单独投资 " + data2_name + " 收益率"], loc="upper left") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() account.dump() return account.get_value() / start_money, float(data1[-1][CLOSE_PRICE_POS]) / float(data1[0][CLOSE_PRICE_POS]), float(data2[-1][CLOSE_PRICE_POS]) / float(data2[0][CLOSE_PRICE_POS])
def test(data1, data2, data1_name, data2_name, data1_price, data2_price, data1_price_name, data2_price_name, start_money, gap_days, trading_day): """ lundong model test with etf """ print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) print "Price1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_price_name, len(data1_price), data1_price[0][DATE_POS], data1_price[-1][DATE_POS]) print "Price2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_price_name, len(data2_price), data2_price[0][DATE_POS], data2_price[-1][DATE_POS]) # Preprocess datasets print "\npreprocessing datasets...\n" time_align(data1, data1_price) time_align(data2, data2_price) time_align(data1, data2) time_align(data1, data1_price) time_align(data2, data2_price) time_truncate(data1, "2013-02-22", "2015-10-30") time_truncate(data2, "2013-02-22", "2015-10-30") time_truncate(data1_price, "2013-02-22", "2015-10-30") time_truncate(data2_price, "2013-02-22", "2015-10-30") print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) print "Price1: %s, data_len=%d, start_date=%s, end_date=%s" % (data1_price_name, len(data1_price), data1_price[0][DATE_POS], data1_price[-1][DATE_POS]) print "Price2: %s, data_len=%d, start_date=%s, end_date=%s" % (data2_price_name, len(data2_price), data2_price[0][DATE_POS], data2_price[-1][DATE_POS]) # Get close price of two datasets and plot them data1_close_price = [ item[CLOSE_PRICE_POS] for item in data1 ] data2_close_price = [ item[CLOSE_PRICE_POS] for item in data2 ] data1_price_close_price = [ item[CLOSE_PRICE_POS] for item in data1_price ] data2_price_close_price = [ item[CLOSE_PRICE_POS] for item in data2_price ] analyser = data_analyser() analyser.draw_relevance(data1_close_price, data1_price_close_price, data1_name, data1_price_name) analyser.draw_relevance(data2_close_price, data2_price_close_price, data2_name, data2_price_name) plt.figure(1) plt.subplot(311) plt.plot(data1_close_price, 'b-', data2_close_price, 'r-') plt.legend([data1_name, data2_name], loc="upper left") plt.grid(True) # Initilize a lundong model, and get trading decisions by it lundong_model = lundong() # the test start with a blank account account = account_simulator(start_money) account_real = account_simulator(start_money) money = [] money_real = [] # record every trading plan result = [] last_valid_trading_plan = {} for item in data1: trading_plan = lundong_model.get_trading_plan(data1, data2, data1_name, data2_name, item[DATE_POS], gap_days, trading_day, last_valid_trading_plan) # if there is a trading plan if trading_plan != {}: # choise == 0 means buy data1 and sell data2 if trading_plan["choise"] == 0: # fake # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # buy data1 as many as we can account.buy(data1_name, trading_plan["data1_close_price"], int(math.floor(account.money / trading_plan["data1_close_price"]))) # real # if we have data2, sell it stock = account_real.get_stock(data2_price_name) if stock != None and stock.share != 0: real_data = get_data(data2_price, trading_plan["date_str"]) account_real.sell(data2_price_name, real_data[CLOSE_PRICE_POS], stock.share) # buy data1 as many as we can real_data = get_data(data1_price, trading_plan["date_str"]) account_real.buy(data1_price_name, real_data[CLOSE_PRICE_POS], int(math.floor(account_real.money / real_data[CLOSE_PRICE_POS]))) # choise == 1 means sell data1 and buy data2 elif trading_plan["choise"] == 1: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # buy data2 as many as we can account.buy(data2_name, trading_plan["data2_close_price"], int(math.floor(account.money / trading_plan["data2_close_price"]))) # real stock = account_real.get_stock(data1_price_name) if stock != None and stock.share != 0: real_data = get_data(data1_price, trading_plan["date_str"]) account_real.sell(data1_price_name, real_data[CLOSE_PRICE_POS], stock.share) # buy data1 as many as we can real_data = get_data(data2_price, trading_plan["date_str"]) account_real.buy(data2_price_name, real_data[CLOSE_PRICE_POS], int(math.floor(account_real.money / real_data[CLOSE_PRICE_POS]))) # choise == 2 means sell all data2 elif trading_plan["choise"] == 2: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) stock = account_real.get_stock(data1_price_name) if stock != None and stock.share != 0: real_data = get_data(data1_price, trading_plan["date_str"]) account_real.sell(data1_price_name, real_data[CLOSE_PRICE_POS], stock.share) stock = account_real.get_stock(data2_price_name) if stock != None and stock.share != 0: real_data = get_data(data2_price, trading_plan["date_str"]) account_real.sell(data2_price_name, real_data[CLOSE_PRICE_POS], stock.share) # other choises means do not trade else: pass # if there is no trading plan, do not trade else: pass money.append(account.get_value()) money_real.append(account_real.get_value()) #money.append(account.money) result.append(trading_plan) if trading_plan != {} and trading_plan["choise"] != -1: last_valid_trading_plan = trading_plan # Get rising rates, choises and plot them data1_up = [] data2_up = [] choise_0 = [] choise_0_index = [] choise_1 = [] choise_1_index = [] choise_2 = [] choise_2_index = [] i = 0 for item in result: if item == {}: data1_up.append(0.0) data2_up.append(0.0) else: data1_up.append(item["data1_up"]) data2_up.append(item["data2_up"]) if item["choise"] == 0: choise_0.append(item["data1_up"]) choise_0_index.append(i) elif item["choise"] == 1: choise_1.append(item["data2_up"]) choise_1_index.append(i) elif item["choise"] == 2: choise_2.append(0) choise_2_index.append(i) i = i + 1 plt.subplot(312) plt.plot(range(len(data1_up)), data1_up, 'b-', range(len(data2_up)), data2_up, 'r-',\ choise_0_index, choise_0, 'bo', choise_1_index, choise_1, 'ro',\ choise_2_index, choise_2, 'r^') #plt.plot(data1_up, 'b-', data2_up, 'r-', choise_0, 'bo', choise_1, 'r^') plt.legend([data1_name, data2_name, "switch to:" + data1_name, "switch to:" + data2_name, "sell all"], loc="upper left") plt.grid(True) plt.subplot(313) plt.plot([item / start_money for item in money], 'g-', \ [item / start_money for item in money_real], 'k-', \ [ float(item[CLOSE_PRICE_POS]) / float(data1_price[0][CLOSE_PRICE_POS]) for item in data1_price], 'b-', \ [ float(item[CLOSE_PRICE_POS]) / float(data2_price[0][CLOSE_PRICE_POS]) for item in data2_price], 'r-') plt.legend(["model", "real", data1_price_name, data2_price_name], loc="upper left") plt.grid(True) #plt.show() account.dump() return account.get_value() / start_money, account_real.get_value() / start_money
def test(data1, data2, data1_name, data2_name, start_money, gap_days, trading_day, show_pics): """ lundong model test """ print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) # Preprocess datasets print "preprocessing datasets..." pre_process(data1, data2) print "Dataset1: %s, data_len=%d, start_date=%s, end_date=%s" % ( data1_name, len(data1), data1[0][DATE_POS], data1[-1][DATE_POS]) print "Dataset2: %s, data_len=%d, start_date=%s, end_date=%s" % ( data2_name, len(data2), data2[0][DATE_POS], data2[-1][DATE_POS]) # Get close price of two datasets and plot them data1_close_price = [item[CLOSE_PRICE_POS] for item in data1] data2_close_price = [item[CLOSE_PRICE_POS] for item in data2] dates = [datetime.strptime(item[DATE_POS], "%Y-%m-%d") for item in data1] years = mdates.YearLocator() months = mdates.MonthLocator() yearsFmt = mdates.DateFormatter("%Y-%m") #plt.figure(1) #plt.subplot(311) if show_pics: fig, ax = plt.subplots() ax.set_title("股指变化图") ax.plot(dates, data1_close_price, 'b-', dates, data2_close_price, 'r-') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend([data1_name, data2_name], loc="upper left") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() # Initilize a lundong model, and get trading decisions by it lundong_model = lundong() # the test start with a blank account account = account_simulator(start_money) money = [] # record every trading plan result = [] last_valid_trading_plan = {} for item in data1: trading_plan = lundong_model.get_trading_plan(data1, data2, data1_name, data2_name, item[DATE_POS], gap_days, trading_day, last_valid_trading_plan) # if there is a trading plan if trading_plan != {}: # choise == 0 means buy data1 and sell data2 if trading_plan["choise"] == 0: # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # buy data1 as many as we can account.buy( data1_name, trading_plan["data1_close_price"], int( math.floor(account.money / trading_plan["data1_close_price"]))) # choise == 1 means sell data1 and buy data2 elif trading_plan["choise"] == 1: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # buy data2 as many as we can account.buy( data2_name, trading_plan["data2_close_price"], int( math.floor(account.money / trading_plan["data2_close_price"]))) # choise == 2 means sell all data2 elif trading_plan["choise"] == 2: # if we have data1, sell it stock = account.get_stock(data1_name) if stock != None and stock.share != 0: account.sell(data1_name, trading_plan["data1_close_price"], stock.share) # if we have data2, sell it stock = account.get_stock(data2_name) if stock != None and stock.share != 0: account.sell(data2_name, trading_plan["data2_close_price"], stock.share) # other choises means do not trade else: pass # if there is no trading plan, do not trade else: pass money.append(account.get_value()) #money.append(account.money) result.append(trading_plan) if trading_plan != {} and trading_plan["choise"] != -1: last_valid_trading_plan = trading_plan # Get rising rates, choises and plot them data1_up = [] data2_up = [] choise_0 = [] choise_0_index = [] choise_1 = [] choise_1_index = [] choise_2 = [] choise_2_index = [] i = 0 for item in result: if item == {}: data1_up.append(0.0) data2_up.append(0.0) else: data1_up.append(item["data1_up"]) data2_up.append(item["data2_up"]) if item["choise"] == 0: choise_0.append(item["data1_up"]) choise_0_index.append(dates[i]) elif item["choise"] == 1: choise_1.append(item["data2_up"]) choise_1_index.append(dates[i]) elif item["choise"] == 2: choise_2.append(0) choise_2_index.append(dates[i]) i = i + 1 if show_pics: fig, ax = plt.subplots() ax.set_title("股指涨幅与模型交易方案") ax.plot(dates, data1_up, 'b-', dates, data2_up, 'r-',\ choise_0_index, choise_0, 'bo', choise_1_index, choise_1, 'ro',\ choise_2_index, choise_2, 'g^') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend([ data1_name, data2_name, "轮动到:" + data1_name, "轮动到:" + data2_name, "清仓" ], loc="upper center") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() if show_pics: fig, ax = plt.subplots() ax.set_title("投资收益对比") ax.plot(dates, [item / start_money for item in money], 'g-', \ dates, [ float(item[CLOSE_PRICE_POS]) / float(data1[0][CLOSE_PRICE_POS]) for item in data1], 'b-', \ dates, [ float(item[CLOSE_PRICE_POS]) / float(data2[0][CLOSE_PRICE_POS]) for item in data2], 'r-') ax.xaxis.set_major_locator(years) ax.xaxis.set_major_formatter(yearsFmt) ax.xaxis.set_minor_locator(months) ax.legend([ "模型收益率", "单独投资 " + data1_name + " 收益率", "单独投资 " + data2_name + " 收益率" ], loc="upper left") ax.grid(True) ax.set_xlim(dates[0], dates[-1] + timedelta(days=90)) #plt.xlim(0,len(data1_index)+50) plt.show() account.dump() return account.get_value() / start_money, float( data1[-1][CLOSE_PRICE_POS]) / float(data1[0][CLOSE_PRICE_POS]), float( data2[-1][CLOSE_PRICE_POS]) / float(data2[0][CLOSE_PRICE_POS])