src/gan_model.py

# conditional_wgangp_train.py
import os
import sys
import json
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
try:
    from src.logger import get_logger
    logger = get_logger(__name__)
except Exception:
    import logging
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

class LatentTickerDataset(Dataset):
    def __init__(self, latent_path, ticker_path):
        self.latents = np.load(latent_path)
        self.tickers = np.load(ticker_path)
        assert self.latents.shape[0] == self.tickers.shape[0], "Latents and tickers length mismatch"

    def __len__(self):
        return self.latents.shape[0]

    def __getitem__(self, idx):
        x = self.latents[idx].astype(np.float32)
        y = int(self.tickers[idx])
        return x, y

class ConditionalGenerator(nn.Module):
    def __init__(self, noise_dim, embed_dim, num_tickers, latent_dim, hidden_dim=128):
        super().__init__()
        self.ticker_emb = nn.Embedding(num_tickers, embed_dim)
        input_dim = noise_dim + embed_dim
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, latent_dim)
        )

    def forward(self, z, ticker_ids):
        emb = self.ticker_emb(ticker_ids) 
        x = torch.cat([z, emb], dim=1)
        return self.net(x)

class ConditionalDiscriminator(nn.Module):
    def __init__(self, latent_dim, embed_dim, num_tickers, hidden_dim=128):
        super().__init__()
        self.ticker_emb = nn.Embedding(num_tickers, embed_dim)
        input_dim = latent_dim + embed_dim
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, x, ticker_ids):
        emb = self.ticker_emb(ticker_ids)
        x_cat = torch.cat([x, emb], dim=1)
        return self.net(x_cat)

def gradient_penalty_cond(D, real, fake, ticker_ids, device):
    """Compute gradient penalty for conditional discriminator D(x, ticker_ids)."""
    batch_size = real.size(0)
    alpha = torch.rand(batch_size, 1).to(device)
    interpolates = (alpha * real + (1 - alpha) * fake).requires_grad_(True)
    d_interpolates = D(interpolates, ticker_ids)
    grad_outputs = torch.ones_like(d_interpolates).to(device)
    gradients = torch.autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=grad_outputs,
        create_graph=True,
        retain_graph=True,
        only_inputs=True
    )[0]
    gradients = gradients.view(batch_size, -1)
    gp = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gp

if __name__ == "__main__":
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    latent_path = os.path.join(base_dir, "data", "processed", "latent_vectors.npy")
    ticker_path = os.path.join(base_dir, "data", "processed", "sequence_tickers.npy")
    models_dir = os.path.join(base_dir, "models")
    resources_dir = os.path.join(base_dir, "resources")
    os.makedirs(models_dir, exist_ok=True)
    os.makedirs(resources_dir, exist_ok=True)

    logger.info("Loading latent vectors from: %s", latent_path)
    latent_vectors = np.load(latent_path)
    logger.info("Loaded latent vectors shape: %s", latent_vectors.shape)

    logger.info("Loading sequence ticker IDs from: %s", ticker_path)
    sequence_tickers = np.load(ticker_path)
    logger.info("Loaded ticker IDs shape: %s", sequence_tickers.shape)

    scaler = StandardScaler()
    latent_scaled = scaler.fit_transform(latent_vectors)
    scaler_save = {"mean": scaler.mean_.tolist(), "scale": scaler.scale_.tolist()}
    np.save(os.path.join(resources_dir, "latent_scaler.npy"), scaler_save)
    logger.info("Saved latent scaler params to resources.")

    dataset = LatentTickerDataset(latent_path, ticker_path)
    dataset.latents = latent_scaled 
    batch_size = 256
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, drop_last=True, num_workers=2)

    noise_dim = 64
    hidden_dim = 128
    n_epochs = 300
    lr = 1e-4
    lambda_gp = 10
    n_critic = 5
    embed_dim = 16 

    latent_dim = latent_scaled.shape[1]
    num_tickers = int(sequence_tickers.max()) + 1

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    G = ConditionalGenerator(noise_dim=noise_dim, embed_dim=embed_dim,
                             num_tickers=num_tickers, latent_dim=latent_dim,
                             hidden_dim=hidden_dim).to(device)
    D = ConditionalDiscriminator(latent_dim=latent_dim, embed_dim=embed_dim,
                                 num_tickers=num_tickers, hidden_dim=hidden_dim).to(device)

    opt_G = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.9))
    opt_D = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.9))

    losses = {"epoch": [], "D_loss": [], "G_loss": []}

    logger.info("Starting Conditional WGAN-GP training...")
    for epoch in range(n_epochs):
        D_losses_epoch = []
        G_losses_epoch = []

        for real_batch, tickers_batch in tqdm(loader, desc=f"Epoch {epoch+1}/{n_epochs}", leave=False):
            real = real_batch.to(device)
            tickers = tickers_batch.to(device).long()
            bsize = real.size(0)
            
            for _ in range(n_critic):
                z = torch.randn(bsize, noise_dim).to(device)
                fake = G(z, tickers)

                d_real = D(real, tickers)
                d_fake = D(fake.detach(), tickers)

                gp = gradient_penalty_cond(D, real, fake.detach(), tickers, device)
                d_loss = -(d_real.mean() - d_fake.mean()) + lambda_gp * gp

                opt_D.zero_grad()
                d_loss.backward()
                opt_D.step()

            z = torch.randn(bsize, noise_dim).to(device)
            fake = G(z, tickers)
            g_loss = -D(fake, tickers).mean()

            opt_G.zero_grad()
            g_loss.backward()
            opt_G.step()

            D_losses_epoch.append(d_loss.item())
            G_losses_epoch.append(g_loss.item())

        mean_D = float(np.mean(D_losses_epoch)) if len(D_losses_epoch) else 0.0
        mean_G = float(np.mean(G_losses_epoch)) if len(G_losses_epoch) else 0.0

        losses["epoch"].append(epoch + 1)
        losses["D_loss"].append(mean_D)
        losses["G_loss"].append(mean_G)

        logger.info(f"[{epoch+1}/{n_epochs}] D_loss={mean_D:.4f}, G_loss={mean_G:.4f}")


    losses_df = pd.DataFrame(losses)
    losses_csv_path = os.path.join(resources_dir, "latent_gan_losses.csv")
    losses_df.to_csv(losses_csv_path, index=False)
    logger.info("Saved training losses to %s", losses_csv_path)

    torch.save(G.state_dict(), os.path.join(models_dir, "latent_gan_generator_conditional.pth"))
    torch.save(D.state_dict(), os.path.join(models_dir, "latent_gan_discriminator_conditional.pth"))
    logger.info("Saved GAN models to models/")

    with open(os.path.join(resources_dir, "gan_config.json"), "w") as f:
        json.dump({
            "model": "WGAN-GP-conditional",
            "noise_dim": noise_dim,
            "latent_dim": latent_dim,
            "hidden_dim": hidden_dim,
            "epochs": n_epochs,
            "batch_size": batch_size,
            "lr": lr,
            "lambda_gp": lambda_gp,
            "n_critic": n_critic,
            "embed_dim": embed_dim,
            "num_tickers": num_tickers
        }, f, indent=4)

    logger.info("Saved GAN config to resources/gan_config.json")
    logger.info("Training completed successfully.")