model_predict.py

import pandas as pd
from sqlalchemy import create_engine
import torch
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt


device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device: ", device )

# Database connection info
DATABASE_URL = "postgresql://postgres:123456789@localhost/trader_master"

engine = create_engine(DATABASE_URL)

df = pd.read_sql("SELECT * FROM candle_tsm_minute_5;", engine)

df['time'] = pd.to_datetime(df['ts'], unit='s')

df = df.sort_values('time')

data = df[["o", "h", "l", "c", "v"]].values # shape(N, 5)

scaler = MinMaxScaler()
data_scaled = scaler.fit_transform(data)

# INPUT_LEN = 60 * 24 * 15 # 1440 * 15 = 21600
INPUT_LEN = 60 * 5
PRED_INDEX = list(range(1, 13))

def create_sequence(data, input_len, pred_index):
    X, y = [], []
    max_h = max(pred_index)

    for i in range(len(data) - input_len - max_h):
        X.append(data[i : i + input_len])

        y_sub = []
        for h in pred_index:
            close_future = data[i + input_len + h -1][3]
            y_sub.append(close_future)

        y.append(y_sub)

    return np.array(X), np.array(y)

X, y = create_sequence(data_scaled, INPUT_LEN, PRED_INDEX)
print("X:", X.shape, "Y:", y.shape)

# Train/test split
split = int(len(X) * 0.95)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]

X_test = torch.tensor(X_test, dtype=torch.float32).to(device)
y_test = torch.tensor(y_test, dtype=torch.float32).to(device)

batch_size = 512

class LSTMModel(nn.Module):
    def __init__(self):
        super(LSTMModel, self).__init__()

        self.lstm1 = nn.LSTM(
            input_size = 5,
            hidden_size = 64,
            num_layers = 1,
            batch_first = True
        )

        self.dropout1 = nn.Dropout(0.2)

        self.lstm2 = nn.LSTM(
            input_size = 64,
            hidden_size = 64,
            num_layers = 1,
            batch_first = True
        )

        self.dropout2 = nn.Dropout(0.2)

        self.fc1 = nn.Linear(64, 32)
        self.fc2 = nn.Linear(32, 12)

    def forward(self, x):
        x, _ = self.lstm1(x)
        x = self.dropout1(x)

        x, _ = self.lstm2(x)
        x = self.dropout2(x)

        x = x[:, -1, :]
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    
model = LSTMModel().to(device)

model = LSTMModel()
model.load_state_dict(torch.load("./model_train/model/tsm_model.pt", map_location=device))
model.eval()

with torch.no_grad():
    y_pred = model(X_test)

print("Pred shape:", y_pred.shape)
print("True shape:", y_test.shape)

mse = torch.mean((y_pred - y_test) ** 2).item()

print("MSE:", mse)

mse_per_index = {}

for i, h in enumerate(PRED_INDEX):
    mse_h = torch.mean((y_pred[:, i] - y_test[:, i]) ** 2).item()
    mse_per_index[h] = mse_h

print("MSE per index")
for h, val in mse_per_index.items():
    print(f" {h * 5} min: {val}")

y_test_np = y_test.cpu().numpy()
y_pred_np = y_pred.cpu().numpy()


CLOSE_IDX = 3

def inverse_transform_close(y_np, sclear):
    n_samples, n_preds = y_np.shape
    dummy = np.zeros((n_samples, scaler.n_features_in_))
    y_inv_list = []

    for i in range(n_preds):
        dummy[:, CLOSE_IDX] = y_np[:, i]
        y_inv = scaler.inverse_transform(dummy)[:, CLOSE_IDX]
        y_inv_list.append(y_inv)

    return np.array(y_inv_list).T

y_test_inv = inverse_transform_close(y_test_np, scaler)
y_pred_inv = inverse_transform_close(y_pred_np, scaler)

plt.figure(figsize=(14, 12))

for i, h in enumerate(PRED_INDEX):
    plt.subplot(len(PRED_INDEX), 1, i+1)
    plt.plot(y_test_inv[:, i], label="True")
    plt.plot(y_pred_inv[:, i], label="Predicted")
    plt.title(f"{h*5}-minute Prediction")
    plt.legend()

plt.tight_layout()
plt.show()