|
|
import akshare as ak
import pandas as pd
import time
import talib
import pandas_ta as ta
# 获取上证指数日线数据
start_time = time.time()
stock_data = ak.stock_zh_a_hist(symbol="000001", period="daily",
start_date="20250102", end_date="20250908",
adjust="hfq") # 前复权
data_fetch_time = time.time() - start_time
# 重命名列名,方便后续处理
stock_data = stock_data[['日期', '收盘', '最高', '最低', '开盘', '成交量']]
stock_data.columns = ['date', 'close', 'high', 'low', 'open', 'volume']
# 将日期列设为索引并转换为datetime格式
stock_data['date'] = pd.to_datetime(stock_data['date'])
stock_data.set_index('date', inplace=True)
# 查看数据前5行
print("上证指数数据前5行:")
print(stock_data.head())
print("\n数据获取时间:{:.6f}秒".format(data_fetch_time))
# ========== 手动roll计算SMA ==========
print("\n========== 手动roll计算SMA ==========")
# 单独计算5日SMA
start_time = time.time()
stock_data['sma_manual'] = stock_data['close'].rolling(window=5).mean()
manual_sma5_time = time.time() - start_time
# 单独计算20日SMA
start_time = time.time()
stock_data['sma20_manual'] = stock_data['close'].rolling(window=20).mean()
manual_sma20_time = time.time() - start_time
# 同时计算5日和20日SMA
start_time = time.time()
stock_data['sma_manual_both'] = stock_data['close'].rolling(window=5).mean()
stock_data['sma20_manual_both'] = stock_data['close'].rolling(window=20).mean()
manual_both_time = time.time() - start_time
# 查看结果(后10行)
print("\n手动计算SMA结果(后10行):")
print(stock_data[['close', 'sma_manual', 'sma20_manual']].tail(10))
print("\n手动计算时间统计:")
print("单独计算5日SMA时间:{:.6f}秒".format(manual_sma5_time))
print("单独计算20日SMA时间:{:.6f}秒".format(manual_sma20_time))
print("同时计算5日和20日SMA时间:{:.6f}秒".format(manual_both_time))
# ========== TA-Lib计算SMA ==========
print("\n========== TA-Lib计算SMA ==========")
# 准备数据
close_prices = stock_data['close'].values
# 单独计算5日SMA
start_time = time.time()
stock_data['sma_talib'] = talib.SMA(close_prices, timeperiod=5)
talib_sma5_time = time.time() - start_time
# 单独计算20日SMA
start_time = time.time()
stock_data['sma20_talib'] = talib.SMA(close_prices, timeperiod=20)
talib_sma20_time = time.time() - start_time
# 同时计算5日和20日SMA
start_time = time.time()
stock_data['sma_talib_both'] = talib.SMA(close_prices, timeperiod=5)
stock_data['sma20_talib_both'] = talib.SMA(close_prices, timeperiod=20)
talib_both_time = time.time() - start_time
# 查看结果(后10行)
print("\nTA-Lib计算SMA结果(后10行):")
print(stock_data[['close', 'sma_talib', 'sma20_talib']].tail(10))
print("\nTA-Lib计算时间统计:")
print("单独计算5日SMA时间:{:.6f}秒".format(talib_sma5_time))
print("单独计算20日SMA时间:{:.6f}秒".format(talib_sma20_time))
print("同时计算5日和20日SMA时间:{:.6f}秒".format(talib_both_time))
# ========== pandas-ta计算SMA ==========
print("\n========== pandas-ta计算SMA ==========")
# 复制数据
data_ta = stock_data.copy()
# 单独计算5日SMA
start_time = time.time()
data_ta.ta.sma(length=5, append=True)
pandas_ta_sma5_time = time.time() - start_time
# 单独计算20日SMA
start_time = time.time()
data_ta.ta.sma(length=20, append=True)
pandas_ta_sma20_time = time.time() - start_time
# 同时计算5日和20日SMA
start_time = time.time()
data_ta.ta.sma(length=5, append=True)
data_ta.ta.sma(length=20, append=True)
pandas_ta_both_time = time.time() - start_time
# 查看结果(后10行)
print("\npandas-ta计算SMA结果(后10行):")
print(data_ta[['close', 'SMA_5', 'SMA_20']].tail(10))
print("\npandas-ta计算时间统计:")
print("单独计算5日SMA时间:{:.6f}秒".format(pandas_ta_sma5_time))
print("单独计算20日SMA时间:{:.6f}秒".format(pandas_ta_sma20_time))
print("同时计算5日和20日SMA时间:{:.6f}秒".format(pandas_ta_both_time))
# ========== 性能对比总结 ==========
print("\n========== 性能对比总结 ==========")
# 创建对比表格
comparison_data = {
'方法': ['手动roll计算', 'TA-Lib', 'pandas-ta'],
'5日SMA(秒)': [manual_sma5_time, talib_sma5_time, pandas_ta_sma5_time],
'20日SMA(秒)': [manual_sma20_time, talib_sma20_time, pandas_ta_sma20_time],
'同时计算(秒)': [manual_both_time, talib_both_time, pandas_ta_both_time]
}
comparison_df = pd.DataFrame(comparison_data)
print(comparison_df)
# 计算相对性能(以手动计算为基准)
comparison_df['5日SMA相对性能'] = comparison_df['5日SMA(秒)'] / manual_sma5_time
comparison_df['20日SMA相对性能'] = comparison_df['20日SMA(秒)'] / manual_sma20_time
comparison_df['同时计算相对性能'] = comparison_df['同时计算(秒)'] / manual_both_time
print("\n相对性能对比(以手动计算为基准1.0):")
print(comparison_df[['方法', '5日SMA相对性能', '20日SMA相对性能', '同时计算相对性能']])
上面的代码分别用 手动roll计算、 talib计算、 pandas-ta进行了计算,还有其他的开源库, 可以自行探索。 其中pandas-ta需要Python3.12以上版本, 如果python版本低 可以选择版本升级或删掉对应代码。
这里说下 这几个库的应用场景。
TA-Lib
在计算SMA指标时性能最佳,比手动计算快约5倍,这是因为其底层使用C语言实现。 缺点是 安装ta-lib 相对比较麻烦, 初学者可以参考官方文档安装。
手动roll计算
性能居中,代码简单直观,适合学习和理解指标原理。 初学的时候,想了解均线、 MACD、KDJ、RSI、BOLL等代码实现可以手动撸。 我之前写过几篇对应的文章,也可以作为参考。
pandas-ta
虽然计算速度较慢,但提供了最简洁的API和与Pandas的无缝集成,适合快速开发和策略原型验证。
对于大多数应用场景,三种方法的计算时间都在毫秒级别,差异不大。但在处理大量数据(如分钟级数据)或需要高频计算时,TA-Lib的性能优势会更加明显。 |
|