Python实例题:股票数据分析与可视化工具
·
目录
Python实例题
题目
股票数据分析与可视化工具
要求:
- 使用 Python 构建一个股票数据分析与可视化工具,支持以下功能:
- 从财经 API 获取股票数据(如收盘价、成交量等)
- 绘制股票价格走势图
- 计算并可视化技术指标(如 MA、MACD、RSI 等)
- 分析多只股票的相关性
- 基于历史数据生成简单预测
- 使用 Matplotlib 和 Seaborn 进行数据可视化。
- 添加命令行界面,支持用户输入股票代码和分析选项。
解题思路:
- 使用
yfinance库获取股票数据。 - 通过
pandas进行数据处理和分析。 - 使用
matplotlib和seaborn绘制图表。 - 实现技术指标计算和预测模型。
代码实现:
import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import argparse
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import warnings
# 忽略警告
warnings.filterwarnings('ignore')
# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
class StockAnalyzer:
def __init__(self):
self.data = {}
self.indicators = {}
def download_data(self, ticker, start_date=None, end_date=None, period='1y'):
"""下载股票数据"""
if not start_date:
end_date = end_date or datetime.now().strftime('%Y-%m-%d')
start_date = (datetime.strptime(end_date, '%Y-%m-%d') - timedelta(days=365)).strftime('%Y-%m-%d')
try:
stock = yf.Ticker(ticker)
df = stock.history(start=start_date, end=end_date, period=period)
if df.empty:
print(f"无法获取 {ticker} 的数据,请检查股票代码是否正确。")
return False
self.data[ticker] = df
self.indicators[ticker] = {}
return True
except Exception as e:
print(f"下载 {ticker} 数据时出错: {e}")
return False
def calculate_ma(self, ticker, periods=[5, 10, 20, 50, 200]):
"""计算移动平均线"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
self.indicators[ticker]['MA'] = {}
for period in periods:
self.indicators[ticker]['MA'][period] = df['Close'].rolling(window=period).mean()
def calculate_macd(self, ticker, fast=12, slow=26, signal=9):
"""计算MACD指标"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
# 计算EMA
ema_fast = df['Close'].ewm(span=fast, adjust=False).mean()
ema_slow = df['Close'].ewm(span=slow, adjust=False).mean()
# 计算MACD线和信号线
macd_line = ema_fast - ema_slow
signal_line = macd_line.ewm(span=signal, adjust=False).mean()
histogram = macd_line - signal_line
self.indicators[ticker]['MACD'] = {
'MACD': macd_line,
'Signal': signal_line,
'Histogram': histogram
}
def calculate_rsi(self, ticker, period=14):
"""计算RSI指标"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
# 计算价格变动
delta = df['Close'].diff()
# 分别计算上涨和下跌的变动
gain = delta.where(delta > 0, 0)
loss = -delta.where(delta < 0, 0)
# 计算平均收益和平均损失
avg_gain = gain.rolling(window=period).mean()
avg_loss = loss.rolling(window=period).mean()
# 计算RSI
rs = avg_gain / avg_loss
rsi = 100 - (100 / (1 + rs))
self.indicators[ticker]['RSI'] = rsi
def plot_price(self, ticker, show_ma=False, ma_periods=[5, 10, 20, 50]):
"""绘制价格走势图"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
plt.figure(figsize=(12, 6))
plt.plot(df.index, df['Close'], label='收盘价')
if show_ma:
if 'MA' not in self.indicators[ticker]:
self.calculate_ma(ticker, ma_periods)
for period in ma_periods:
plt.plot(df.index, self.indicators[ticker]['MA'][period],
label=f'{period}日均线')
plt.title(f'{ticker} 价格走势图')
plt.xlabel('日期')
plt.ylabel('价格')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
def plot_indicators(self, ticker, indicators=['MA', 'MACD', 'RSI']):
"""绘制技术指标图"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
# 根据选择的指标确定子图数量
nrows = len(indicators)
fig, axes = plt.subplots(nrows=nrows, ncols=1, figsize=(12, 4 * nrows), sharex=True)
if nrows == 1:
axes = [axes]
plot_idx = 0
if 'MA' in indicators:
if 'MA' not in self.indicators[ticker]:
self.calculate_ma(ticker)
ax = axes[plot_idx]
ax.plot(df.index, df['Close'], label='收盘价')
for period, ma in self.indicators[ticker]['MA'].items():
ax.plot(df.index, ma, label=f'{period}日均线')
ax.set_title(f'{ticker} 价格与均线')
ax.set_ylabel('价格')
ax.legend()
ax.grid(True)
plot_idx += 1
if 'MACD' in indicators:
if 'MACD' not in self.indicators[ticker]:
self.calculate_macd(ticker)
ax = axes[plot_idx]
macd_data = self.indicators[ticker]['MACD']
ax.plot(df.index, macd_data['MACD'], label='MACD线')
ax.plot(df.index, macd_data['Signal'], label='信号线')
# 绘制柱状图
for i in range(len(df.index)):
color = 'green' if macd_data['Histogram'][i] >= 0 else 'red'
ax.bar(df.index[i], macd_data['Histogram'][i], color=color, alpha=0.5)
ax.set_title(f'{ticker} MACD指标')
ax.set_ylabel('MACD值')
ax.legend()
ax.grid(True)
plot_idx += 1
if 'RSI' in indicators:
if 'RSI' not in self.indicators[ticker]:
self.calculate_rsi(ticker)
ax = axes[plot_idx]
rsi = self.indicators[ticker]['RSI']
ax.plot(df.index, rsi, label='RSI')
ax.axhline(y=70, color='r', linestyle='-', alpha=0.5)
ax.axhline(y=30, color='g', linestyle='-', alpha=0.5)
ax.set_title(f'{ticker} RSI指标')
ax.set_ylabel('RSI值')
ax.set_ylim(0, 100)
ax.legend()
ax.grid(True)
plot_idx += 1
plt.xlabel('日期')
plt.tight_layout()
plt.show()
def plot_volume(self, ticker):
"""绘制成交量图"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
df = self.data[ticker]
plt.figure(figsize=(12, 4))
plt.bar(df.index, df['Volume'])
plt.title(f'{ticker} 成交量图')
plt.xlabel('日期')
plt.ylabel('成交量')
plt.grid(True)
plt.tight_layout()
plt.show()
def analyze_correlation(self, tickers):
"""分析多只股票的相关性"""
# 确保所有股票数据都已下载
for ticker in tickers:
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return
# 创建收盘价DataFrame
close_df = pd.DataFrame()
for ticker in tickers:
close_df[ticker] = self.data[ticker]['Close']
# 计算相关性矩阵
corr_matrix = close_df.corr()
# 绘制热力图
plt.figure(figsize=(10, 8))
sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f', linewidths=0.5)
plt.title('股票相关性分析')
plt.tight_layout()
plt.show()
return corr_matrix
def predict_price(self, ticker, days=30, show_plot=True):
"""基于历史数据预测未来价格"""
if ticker not in self.data:
print(f"没有 {ticker} 的数据,请先下载数据。")
return None
df = self.data[ticker].copy()
# 创建预测特征
df['Date'] = df.index
df['Date'] = pd.to_datetime(df['Date'])
df['Date'] = df['Date'].map(datetime.toordinal)
# 准备训练数据
X = df[['Date']].values
y = df['Close'].values
# 训练线性回归模型
model = LinearRegression()
model.fit(X, y)
# 预测未来价格
last_date = df.index[-1]
future_dates = [last_date + timedelta(days=i) for i in range(1, days + 1)]
future_ordinals = [date.toordinal() for date in future_dates]
future_prices = model.predict(np.array(future_ordinals).reshape(-1, 1))
# 计算均方误差评估模型
y_pred = model.predict(X)
mse = mean_squared_error(y, y_pred)
rmse = np.sqrt(mse)
if show_plot:
plt.figure(figsize=(12, 6))
# 绘制历史价格
plt.plot(df.index, df['Close'], label='历史价格')
# 绘制预测价格
plt.plot(future_dates, future_prices, 'r--', label='预测价格')
plt.title(f'{ticker} 价格预测 (未来 {days} 天)')
plt.xlabel('日期')
plt.ylabel('价格')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
# 返回预测结果
prediction_df = pd.DataFrame({
'日期': future_dates,
'预测价格': future_prices
})
return prediction_df, rmse
def main():
parser = argparse.ArgumentParser(description='股票数据分析与可视化工具')
parser.add_argument('--ticker', type=str, help='股票代码 (例如: AAPL)')
parser.add_argument('--tickers', type=str, help='多个股票代码,用逗号分隔 (例如: AAPL,MSFT,GOOG)')
parser.add_argument('--start', type=str, help='开始日期 (YYYY-MM-DD)')
parser.add_argument('--end', type=str, help='结束日期 (YYYY-MM-DD)')
parser.add_argument('--period', type=str, default='1y', help='数据周期 (1d,5d,1mo,3mo,6mo,1y,2y,5y,10y,ytd,max)')
parser.add_argument('--price', action='store_true', help='绘制价格走势图')
parser.add_argument('--ma', action='store_true', help='在价格图中显示均线')
parser.add_argument('--ma_periods', type=str, default='5,10,20,50', help='均线周期,用逗号分隔')
parser.add_argument('--indicators', action='store_true', help='绘制技术指标图')
parser.add_argument('--volume', action='store_true', help='绘制成交量图')
parser.add_argument('--correlation', action='store_true', help='分析多只股票的相关性')
parser.add_argument('--predict', type=int, default=0, help='预测未来天数的价格')
args = parser.parse_args()
analyzer = StockAnalyzer()
if args.ticker:
# 下载单只股票数据
if analyzer.download_data(args.ticker, args.start, args.end, args.period):
# 绘制价格图
if args.price:
ma_periods = [int(p) for p in args.ma_periods.split(',')]
analyzer.plot_price(args.ticker, args.ma, ma_periods)
# 绘制技术指标图
if args.indicators:
analyzer.plot_indicators(args.ticker)
# 绘制成交量图
if args.volume:
analyzer.plot_volume(args.ticker)
# 预测价格
if args.predict > 0:
prediction, rmse = analyzer.predict_price(args.ticker, args.predict)
print(f"预测未来 {args.predict} 天的价格 (RMSE: {rmse:.2f}):")
print(prediction)
if args.tickers and args.correlation:
# 分析多只股票相关性
tickers = [t.strip() for t in args.tickers.split(',')]
# 下载所有股票数据
all_downloaded = True
for ticker in tickers:
if not analyzer.download_data(ticker, args.start, args.end, args.period):
all_downloaded = False
if all_downloaded:
corr_matrix = analyzer.analyze_correlation(tickers)
print("股票相关性矩阵:")
print(corr_matrix)
if __name__ == "__main__":
main()
更多推荐

所有评论(0)