dataProcess.py

import numpy as np
import torch
import yfinance as yf
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import DataLoader, Dataset


def dataFetch():
    return yf.download("^GSPC", start="2010-01-01", multi_level_index=False, progress=False)


def dataFeatures(df):
    df = df.copy()

    # Moving avgs | Window sizes explicitly specified due to readability.
    df["SMA_20"] = df["Close"].rolling(window=20).mean()
    df["SMA_50"] = df["Close"].rolling(window=50).mean()
    df["EMA_12"] = df["Close"].ewm(span=12).mean()

    # RSI
    change = df["Close"].diff()
    avgGain = change.clip(lower=0).rolling(14).mean()
    avgLoss = (-change.clip(upper=0)).rolling(14).mean()
    rs = avgGain / avgLoss
    df["RSI"] = 100 - (100 / (1 + rs))

    # ATR
    prev_close = df["Close"].shift(1)
    tr = np.maximum(df["High"], prev_close) - np.minimum(df["Low"], prev_close)
    df["ATR"] = tr.rolling(14).mean()
    df["ATR_pct"] = (df["ATR"] / df["Close"])

    # BB
    df["BB_upper"] = df["SMA_20"] + (2 * df["Close"].rolling(window=20).std())
    df["BB_lower"] = df["SMA_20"] - (2 * df["Close"].rolling(window=20).std())
    df["BB_width"] = (df["BB_upper"] - df["BB_lower"]) / df["SMA_20"]
    df["BB_pct"] = (df["Close"] - df["BB_lower"]) / (df["BB_upper"] - df["BB_lower"])

    # Price-level ratios (stationary alternatives to raw prices)
    df["EMA_ratio"] = df["EMA_12"] / df["SMA_20"]
    df["SMA_ratio"] = df["SMA_20"] / df["SMA_50"]

    # Volume signal
    df["Volume_MA"] = df["Volume"].rolling(window=20).mean()
    df["Volume_relative"] = df["Volume"] / df["Volume_MA"]

    # Target log return
    df["Log_Return"] = np.log(df["Close"]).diff()

    # Drop NaN rows
    df.dropna(inplace=True)
    return df


def dataScaler(df):
    feature_columns = [
        "RSI",
        "ATR_pct",
        "BB_width",
        "BB_pct",
        "EMA_ratio",
        "SMA_ratio",
        "Volume_relative",
    ]

    target_column = "Log_Return"

    feature_data = df[feature_columns].values
    target_data = df[[target_column]].values  # Double brackets for preserving 2D array

    feature_scaler = MinMaxScaler()
    target_scaler = MinMaxScaler()

    split_idx = int(len(feature_data) * 0.8)

    feature_scaled = np.vstack(
        [
            feature_scaler.fit_transform(feature_data[:split_idx]),
            feature_scaler.transform(feature_data[split_idx:]),
        ]
    )
    target_scaled = np.vstack(
        [
            target_scaler.fit_transform(target_data[:split_idx]),
            target_scaler.transform(target_data[split_idx:]),
        ]
    )
    return feature_scaled, target_scaled, target_scaler


def create_sequences(features, targets, seq_len=60):
    X, y = [], []
    for i in range(seq_len, len(features)):
        X.append(features[i - seq_len : i])
        y.append(targets[i, 0])
    return np.array(X), np.array(y)


class StockDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)  # type: ignore
        self.y = torch.tensor(y, dtype=torch.float32)  # type: ignore

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


def runPipeline(seq_len=60, batch_size=32):
    df = dataFetch()
    df = dataFeatures(df)

    features, targets, target_scaler = dataScaler(df)
    X, y = create_sequences(features, targets, seq_len=seq_len)

    split_idx = int(len(X) * 0.8)
    X_train, X_test = X[:split_idx], X[split_idx:]
    y_train, y_test = y[:split_idx], y[split_idx:]

    train_dataset = StockDataset(X_train, y_train)
    test_dataset = StockDataset(X_test, y_test)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return train_loader, test_loader, target_scaler