目录

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()

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐