SPY Direction Training Dataset

Build the first model-ready dataset from IBKR daily candle Parquets. The target is whether SPY closes above, below, or unchanged from today's close five trading days later.

Setup

The notebook is deliberately plain pandas so feature logic remains easy to inspect. VIX_INPUT_SYMBOL defaults to VIXY because that is the VIX-related Parquet currently present locally; switch it to VIX after fetching a VIX.parquet file.

In [1]:

from pathlib import Path

import numpy as np
import pandas as pd


def find_project_root(start: Path | None = None) -> Path:
    current = (start or Path.cwd()).resolve()
    for candidate in [current, *current.parents]:
        if (candidate / "pyproject.toml").exists():
            return candidate
    raise RuntimeError("Could not find project root containing pyproject.toml")


PROJECT_ROOT = find_project_root()
RAW_DATA_DIR = PROJECT_ROOT / "data" / "ibkr" / "daily"
OUTPUT_PATH = PROJECT_ROOT / "data" / "training" / "spy_direction_5d.parquet"

SPY_SYMBOL = "SPY"
VIX_INPUT_SYMBOL = "VIXY"
TLT_SYMBOL = "TLT"
USO_SYMBOL = "USO"

TRAIN_FRACTION = 0.70
VALIDATION_FRACTION = 0.15
TEST_FRACTION = 0.15

PROJECT_ROOT, RAW_DATA_DIR, OUTPUT_PATH

Out [1]:

(PosixPath('/home/jarno/repos/trading-bot'),
 PosixPath('/home/jarno/repos/trading-bot/data/ibkr/daily'),
 PosixPath('/home/jarno/repos/trading-bot/data/training/spy_direction_5d.parquet'))

Load Close Prices

Each symbol file is read, sorted by trading date, and reduced to a single close-price column. SPY defines the observation calendar through the inner join.

In [2]:

def load_close(symbol: str, alias: str | None = None) -> pd.DataFrame:
    alias = alias or symbol
    path = RAW_DATA_DIR / f"{symbol}.parquet"
    if not path.exists():
        raise FileNotFoundError(f"Missing raw data file: {path}")

    frame = pd.read_parquet(path)
    if "date" in frame.columns:
        frame = frame.set_index("date")
    if "close" not in frame.columns:
        raise ValueError(f"{path} does not contain a close column")

    frame = frame.copy()
    frame.index = pd.to_datetime(frame.index)
    frame.index.name = "date"
    frame = frame.sort_index()
    return frame[["close"]].rename(columns={"close": f"{alias}_close"})


prices = pd.concat(
    [
        load_close(SPY_SYMBOL, "SPY"),
        load_close(VIX_INPUT_SYMBOL, "VIX"),
        load_close(TLT_SYMBOL, "TLT"),
        load_close(USO_SYMBOL, "USO"),
    ],
    axis=1,
    join="inner",
)

prices.info()
prices.head()

Out [2]:

<class 'pandas.DataFrame'>
DatetimeIndex: 1251 entries, 2021-07-30 to 2026-07-24
Data columns (total 4 columns):
 #   Column     Non-Null Count  Dtype  
---  ------     --------------  -----  
 0   SPY_close  1251 non-null   float64
 1   VIX_close  1251 non-null   float64
 2   TLT_close  1251 non-null   float64
 3   USO_close  1251 non-null   float64
dtypes: float64(4)
memory usage: 48.9 KB

	SPY_close	VIX_close	TLT_close	USO_close
date
2021-07-30	438.51	495.4	149.52	50.66
2021-08-02	437.59	513.6	150.67	49.18
2021-08-03	441.15	488.0	150.75	48.85
2021-08-04	438.98	487.6	151.06	47.20
2021-08-05	441.76	474.8	150.29	48.10

Generate Features And Target

Feature values use only current and prior rows. The target looks five trading rows ahead in SPY_close; unchanged future prices are encoded as 0.5.

In [3]:

df = prices.copy()

df["SPY_ret_5"] = df["SPY_close"].pct_change(5)
df["SPY_ret_20"] = df["SPY_close"].pct_change(20)

sma_50 = df["SPY_close"].rolling(50).mean()
df["SPY_dist_sma50"] = (df["SPY_close"] - sma_50) / sma_50

df["VIX_change_5"] = df["VIX_close"].pct_change(5)
df["VIX_rank_20"] = df["VIX_close"].rolling(20).rank(pct=True)

df["TLT_ret_10"] = df["TLT_close"].pct_change(10)
df["USO_ret_5"] = df["USO_close"].pct_change(5)
df["SPY_TLT_ratio_ret"] = (df["SPY_close"] / df["TLT_close"]).pct_change(5)

spy_forward_close = df["SPY_close"].shift(-5)
df["spy_up_5d"] = np.nan
df.loc[spy_forward_close > df["SPY_close"], "spy_up_5d"] = 1.0
df.loc[spy_forward_close < df["SPY_close"], "spy_up_5d"] = 0.0
df.loc[spy_forward_close == df["SPY_close"], "spy_up_5d"] = 0.5

FEATURE_COLUMNS = [
    "SPY_ret_5",
    "SPY_ret_20",
    "SPY_dist_sma50",
    "VIX_change_5",
    "VIX_rank_20",
    "TLT_ret_10",
    "USO_ret_5",
    "SPY_TLT_ratio_ret",
]
TARGET_COLUMN = "spy_up_5d"

df_model = df[FEATURE_COLUMNS + [TARGET_COLUMN]].dropna().copy()
df_model.head()

Out [3]:

	SPY_ret_5	SPY_ret_20	SPY_dist_sma50	VIX_change_5	VIX_rank_20	TLT_ret_10	USO_ret_5	SPY_TLT_ratio_ret	spy_up_5d
date
2021-10-08	0.008336	-0.017017	-0.011155	-0.075239	0.125	-0.034239	0.041495	0.032998	1.0
2021-10-11	0.014114	-0.026625	-0.018145	-0.090869	0.200	-0.033135	0.031015	0.038908	1.0
2021-10-12	0.001201	-0.023752	-0.020386	-0.074091	0.100	-0.001041	0.008628	-0.001303	1.0
2021-10-13	0.000644	-0.028356	-0.016597	-0.081006	0.050	0.006928	0.036928	-0.005897	1.0
2021-10-14	0.008754	-0.010443	-0.000214	-0.096759	0.050	0.010809	0.026192	-0.011991	1.0

Chronological Splits

Splits are assigned by row order after feature/target cleanup: oldest rows for training, newer rows for validation, newest rows for test. The split column is metadata and should not be used as a model input.

In [4]:

def assign_chronological_splits(
    frame: pd.DataFrame,
    train_fraction: float,
    validation_fraction: float,
    test_fraction: float,
) -> pd.Series:
    total_fraction = train_fraction + validation_fraction + test_fraction
    if not np.isclose(total_fraction, 1.0):
        raise ValueError(f"Split fractions must sum to 1.0, got {total_fraction}")

    n_rows = len(frame)
    train_end = int(n_rows * train_fraction)
    validation_end = train_end + int(n_rows * validation_fraction)

    split = pd.Series(index=frame.index, dtype="object")
    split.iloc[:train_end] = "train"
    split.iloc[train_end:validation_end] = "validation"
    split.iloc[validation_end:] = "test"
    return split


df_model["split"] = assign_chronological_splits(
    df_model,
    TRAIN_FRACTION,
    VALIDATION_FRACTION,
    TEST_FRACTION,
)

split_summary = df_model.groupby("split", sort=False).agg(
    rows=(TARGET_COLUMN, "size"),
    start_date=(TARGET_COLUMN, lambda values: values.index.min().date()),
    end_date=(TARGET_COLUMN, lambda values: values.index.max().date()),
    target_mean=(TARGET_COLUMN, "mean"),
)
split_summary

Out [4]:

	rows	start_date	end_date	target_mean
split
train	837	2021-10-08	2025-02-07	0.583035
validation	179	2025-02-10	2025-10-24	0.653631
test	181	2025-10-27	2026-07-17	0.558011

Inspect And Save

model_input_columns is the source of truth for columns that are allowed into the model. The saved Parquet includes date, features, target, and split metadata.

In [5]:

model_input_columns = FEATURE_COLUMNS
model_target_column = TARGET_COLUMN

X_train = df_model.loc[df_model["split"] == "train", model_input_columns]
y_train = df_model.loc[df_model["split"] == "train", model_target_column]

print(f"Rows: {len(df_model):,}")
print(f"Feature columns: {model_input_columns}")
print(f"Target column: {model_target_column}")
print(f"Training matrix shape: {X_train.shape}")

df_model.describe(include="all")

Out [5]:

Rows: 1,197
Feature columns: ['SPY_ret_5', 'SPY_ret_20', 'SPY_dist_sma50', 'VIX_change_5', 'VIX_rank_20', 'TLT_ret_10', 'USO_ret_5', 'SPY_TLT_ratio_ret']
Target column: spy_up_5d
Training matrix shape: (837, 8)

	SPY_ret_5	SPY_ret_20	SPY_dist_sma50	VIX_change_5	VIX_rank_20	TLT_ret_10	USO_ret_5	SPY_TLT_ratio_ret	spy_up_5d	split
count	1197.000000	1197.000000	1197.000000	1197.000000	1197.000000	1197.000000	1197.000000	1197.000000	1197.000000	1197
unique	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	3
top	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	train
freq	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	837
mean	0.002556	0.009841	0.011067	-0.008646	0.387009	-0.004083	0.004662	0.004927	0.589808	NaN
std	0.023081	0.043634	0.037371	0.092413	0.333697	0.028557	0.052319	0.027656	0.492074	NaN
min	-0.114962	-0.123975	-0.141995	-0.387709	0.050000	-0.092880	-0.196652	-0.117208	0.000000	NaN
25%	-0.009535	-0.017017	-0.008847	-0.058376	0.075000	-0.023517	-0.026200	-0.010013	0.000000	NaN
50%	0.003869	0.015797	0.017906	-0.020655	0.250000	-0.004339	0.004817	0.005710	1.000000	NaN
75%	0.015957	0.038385	0.038332	0.027778	0.700000	0.014272	0.031883	0.021406	1.000000	NaN
max	0.082843	0.157566	0.088103	0.937828	1.000000	0.091322	0.327273	0.129204	1.000000	NaN

In [6]:

OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
dataset_to_save = df_model.reset_index()
dataset_to_save.to_parquet(OUTPUT_PATH, index=False)

OUTPUT_PATH, dataset_to_save.shape, dataset_to_save.head()

Out [6]:

(PosixPath('/home/jarno/repos/trading-bot/data/training/spy_direction_5d.parquet'),
 (1197, 11),
         date  SPY_ret_5  SPY_ret_20  SPY_dist_sma50  VIX_change_5  \
 0 2021-10-08   0.008336   -0.017017       -0.011155     -0.075239   
 1 2021-10-11   0.014114   -0.026625       -0.018145     -0.090869   
 2 2021-10-12   0.001201   -0.023752       -0.020386     -0.074091   
 3 2021-10-13   0.000644   -0.028356       -0.016597     -0.081006   
 4 2021-10-14   0.008754   -0.010443       -0.000214     -0.096759   
 
    VIX_rank_20  TLT_ret_10  USO_ret_5  SPY_TLT_ratio_ret  spy_up_5d  split  
 0        0.125   -0.034239   0.041495           0.032998        1.0  train  
 1        0.200   -0.033135   0.031015           0.038908        1.0  train  
 2        0.100   -0.001041   0.008628          -0.001303        1.0  train  
 3        0.050    0.006928   0.036928          -0.005897        1.0  train  
 4        0.050    0.010809   0.026192          -0.011991        1.0  train  )

29 KiB Raw Blame History