StockSenseSpace / testing.py
amitke
test
09fbd2c
Raw
History Blame Contribute Delete
4.21 kB
import os, json, math, pickle
from datetime import datetime, timedelta
import numpy as np
import yfinance as yf
from sklearn.metrics import mean_absolute_error, mean_squared_error
import torch
from models import StockLSTM
os.environ["CUDA_VISIBLE_DEVICES"] = ""
ARTIFACTS_DIR = "artifacts"
@torch.no_grad()
def evaluate(symbol: str):
base = os.path.join(ARTIFACTS_DIR, symbol.upper())
model = StockLSTM(input_dim=1, hidden_dim=64, num_layers=2, dropout=0.2)
model.load_state_dict(torch.load(os.path.join(base, "model.pt"), map_location="cpu"))
model.eval()
with open(os.path.join(base, "scaler.pkl"), "rb") as f:
scaler = pickle.load(f)
with open(os.path.join(base, "meta.json"), "r") as f:
meta = json.load(f)
seq_len = meta["seq_len"]
end = datetime.utcnow().date()
start = end - timedelta(days=5*365)
df = yf.download(symbol, start=start.isoformat(), end=end.isoformat(), progress=False, auto_adjust=True)
data = df[["Close"]].dropna()
# Compute returns
data['LogReturn'] = np.log(data['Close'] / data['Close'].shift(1))
data = data.dropna()
returns = data['LogReturn'].values.reshape(-1, 1)
scaled = scaler.transform(returns)
split_idx = int(len(scaled) * 0.8)
test_scaled = scaled[split_idx - seq_len:] # include tail of train for continuity
# build sequences
X, y = [], []
for i in range(seq_len, len(test_scaled)):
X.append(test_scaled[i-seq_len:i])
y.append(test_scaled[i])
X = np.array(X, dtype=np.float32)
y = np.array(y, dtype=np.float32)
X_t = torch.from_numpy(X) # [N, T, 1]
pred_scaled = model(X_t).numpy()
# inverse returns
pred_returns = scaler.inverse_transform(pred_scaled).flatten()
# Reconstruct prices
# We need the price before the first test prediction
# The test set in 'data' starts at split_idx
# The first prediction corresponds to return at split_idx
# So base is price at split_idx - 1
# Note: 'data' here is the return-df (shifted).
# We need indices from the original df.
# It's cleaner to just align by length.
# Get original prices aligned with returns
# df['Close'] has N+1 items if returns has N items.
# data indices are a subset of df indices
# Let's match by index
test_indices = data.index[split_idx:]
# Price predecessors (bases)
# If a return is at time t, it depends on Price[t-1]
# Simple reconstruction:
# Get the price immediately preceding the test set
base_price_idx = split_idx - 1
if base_price_idx < 0:
# Fallback if split is at 0 (unlikely)
base_price = df['Close'].iloc[0]
else:
# The return at data.iloc[base_price_idx] is NOT the price
# data only has returns.
# We need to look at the original DF
# The 'data' was created by dropping first row of df.
# So data.iloc[0] corresponds to df.iloc[1].
# data.iloc[split_idx] is roughly df.iloc[split_idx+1]
# Exact alignment:
# data index i matches df index i (if we kept index)
pass
# Let's rely on the original df
# The returns in "y" (targets) correspond to `data.iloc[split_idx:]`
# The Prices we want to compare against are `df['Close'][data.index[split_idx:]]`
y_true_prices = df['Close'].loc[data.index[split_idx:]].values
# Base price for the FIRST prediction:
# The first return predicted is for data.index[split_idx]
# So we need Price at data.index[split_idx-1] (previous day)
# OR simpler: df['Close'].loc[data.index[split_idx-1]]
first_test_idx_pos = df.index.get_loc(data.index[split_idx])
base_price = df['Close'].iloc[first_test_idx_pos - 1]
reconstructed = []
curr = base_price
for r in pred_returns:
curr = curr * np.exp(r)
reconstructed.append(curr)
pred = np.array(reconstructed)
y_true = y_true_prices[:len(pred)] # sync lengths
rmse = math.sqrt(mean_squared_error(y_true, pred))
mae = mean_absolute_error(y_true, pred)
return {"symbol": symbol.upper(), "rmse": rmse, "mae": mae, "n": len(y_true)}