File size: 3,179 Bytes
bbf5d55 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import os
import sys
import torch
import pandas as pd
import numpy as np
from torch.utils.data import DataLoader
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from src.model import LSTMAutoEncoder, QuarterlyStockDataset # pylint: disable=import-error
from src.logger import get_logger
logger = get_logger(__name__)
if __name__ == "__main__":
try:
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
processed_data_path = os.path.join(base_dir, "data", "processed", "stock_data.parquet")
model_path = os.path.join(base_dir, "models", "lstm_autoencoder.pth")
latent_vectors_path = os.path.join(base_dir, "..", "GAN", "data", "processed", "latent_vectors.npy")
ticker_mapping_path = os.path.join(base_dir, "..", "GAN", "data", "processed", "ticker_mapping.npy")
device = "cuda" if torch.cuda.is_available() else "cpu"
df = pd.read_parquet(processed_data_path)
tickers = df["Ticker"].unique()
num_tickers = df["Ticker_Encoded"].nunique()
model = LSTMAutoEncoder(
input_dim=5,
num_tickers=num_tickers,
embed_dim=16,
hidden_size=64,
latent_dim=32,
num_layers=2
).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
def encode(model, x, ticker_id):
ticker_emb = model.ticker_embed(ticker_id).unsqueeze(1).repeat(1, x.size(1), 1)
x_in = torch.cat([x, ticker_emb], dim=2)
enc_out, _ = model.encoder(x_in)
latent = model.fc_enc(enc_out[:, -1, :])
return latent
all_latents = []
all_tickers = []
for ticker in tickers:
ticker_df = df[df["Ticker"] == ticker].copy()
if len(ticker_df) < 90:
continue
dataset = QuarterlyStockDataset(ticker_df, sequence_length=90)
loader = DataLoader(dataset, batch_size=64, shuffle=False)
ticker_latents = []
with torch.no_grad():
for batch_x, batch_ticker in loader:
batch_x, batch_ticker = batch_x.to(device), batch_ticker.to(device)
latent = encode(model, batch_x, batch_ticker)
ticker_latents.append(latent.cpu().numpy())
if ticker_latents:
ticker_latents = np.concatenate(ticker_latents, axis=0)
all_latents.append(ticker_latents)
all_tickers.extend([ticker] * len(ticker_latents))
logger.info(f"Extracted {len(ticker_latents)} latent vectors for {ticker}.")
all_latents = np.concatenate(all_latents, axis=0)
all_tickers = np.array(all_tickers)
np.save(latent_vectors_path, all_latents)
np.save(ticker_mapping_path, all_tickers)
logger.info(f"Saved {len(all_latents)} latent vectors to {latent_vectors_path}")
logger.info(f"Saved ticker mapping to {ticker_mapping_path}")
except Exception as e:
logger.error("Error extracting latent space vectors: %s", e)
raise
|