Python 回测 Dataframe 形状错误问题

Question

我正在运行以下代码并收到此错误 - ChatGPT 无法解决此错误。请帮忙。 这是代码：

import pandas as pd
import numpy as np
import yfinance as yf
import matplotlib.pyplot as plt

# Fetch historical SPY data
def fetch_data(symbol, start_date, end_date):
    data = yf.download(symbol, start=start_date, end=end_date)
    return data

# Calculate RSI
def calculate_rsi(data, window=14):
    delta = data['Close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    data['RSI'] = rsi
    return data

# Create buy and sell signals based on RSI thresholds
def create_signals(data, buy_threshold=30, sell_threshold=70):
    data['Buy_Signal'] = np.where(data['RSI'] < buy_threshold, data['Close'], np.nan)
    data['Sell_Signal'] = np.where(data['RSI'] > sell_threshold, data['Close'], np.nan)
    return data

class RSIStrategy:
    def __init__(self, initial_investment):
        self.initial_investment = initial_investment
        self.portfolio_value = initial_investment
        self.shares_owned = 0
    
    def execute_trade(self, price, signal_type):
        if signal_type == 'buy':
            if self.shares_owned == 0:  # Buy only if no shares are owned
                self.shares_owned = self.portfolio_value / price
                self.portfolio_value = 0  # All investment goes to shares
                return "Bought"
        elif signal_type == 'sell':
            if self.shares_owned > 0:  # Sell only if shares are owned
                self.portfolio_value = self.shares_owned * price
                self.shares_owned = 0  # All shares sold
                return "Sold"
        return "No action"

def simulate_trading(data, initial_investment):
    strategy = RSIStrategy(initial_investment)
    trades = []

    for index, row in data.iterrows():
        if pd.notna(row['Buy_Signal']):
            action = strategy.execute_trade(row['Close'], 'buy')
            trades.append((index, action, row['Close']))
        elif pd.notna(row['Sell_Signal']):
            action = strategy.execute_trade(row['Close'], 'sell')
            trades.append((index, action, row['Close']))

    # Final portfolio value after all trades
    final_value = strategy.portfolio_value + (strategy.shares_owned * data['Close'].iloc[-1])
    return final_value, trades

# Define parameters and run the backtest
start_date = '2022-01-01'
end_date = '2023-12-31'
symbol = 'SPY'
initial_investment = 10000

# Fetch and prepare data
spy_data = fetch_data(symbol, start_date, end_date)
spy_data = calculate_rsi(spy_data)
spy_data = create_signals(spy_data)

# Simulate trading
final_value, trades = simulate_trading(spy_data, initial_investment)

# Output results
print(f"Final Portfolio Value: ${final_value:.2f}")
print("Trades executed:")
for trade in trades:
    print(trade)

# Plotting the results
plt.figure(figsize=(14, 7))
plt.plot(spy_data['Close'], label='SPY Close Price', alpha=0.5)
plt.scatter(spy_data.index, spy_data['Buy_Signal'], marker='^', color='g', label='Buy Signal', alpha=1)
plt.scatter(spy_data.index, spy_data['Sell_Signal'], marker='v', color='r', label='Sell Signal', alpha=1)
plt.title('SPY Trading Strategy with RSI Signals')
plt.xlabel('Date')
plt.ylabel('Price')
plt.legend()
plt.show()

我已经尝试了 ChatGPT 所要求的一切，但仍然不起作用。 chatGPT 是这么说的： 带有布尔检查的按位运算，特别是当您在 Pandas 上下文中工作时。但是，由于 self.shares_owned 是单个变量（不是 DataFrame 列），因此它本身不应导致此错误。”

> ---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/pandas/core/indexes/base.py in get_loc(self, key)
   3804         try:
-> 3805             return self._engine.get_loc(casted_key)
   3806         except KeyError as err:

index.pyx in pandas._libs.index.IndexEngine.get_loc()

index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'Buy_Signal'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
10 frames
KeyError: 'Buy_Signal'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/pandas/core/internals/managers.py in insert(self, loc, item, value, refs)
   1368             value = value.T
   1369             if len(value) > 1:
-> 1370                 raise ValueError(
   1371                     f"Expected a 1D array, got an array with shape {value.T.shape}"
   1372                 )

ValueError: Expected a 1D array, got an array with shape (501, 501)

Answer 1

yfinance 的数据有问题

有更新功能

# Fetch historical SPY data
def fetch_data(symbol, start_date, end_date):
    data = yf.download(symbol, start=start_date, end=end_date)
    if isinstance(data.columns, pd.MultiIndex) and data.columns.nlevels > 1:
        data.columns = data.columns.droplevel(1)
    return data