dream_visualizer.py

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import os

# --- Configuration ---
LATENT_DIM = 32
HIDDEN_SIZE = 64
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")

# --- 1. Define Architectures (Must match training) ---
class VolatilityVAE(nn.Module):
    def __init__(self):
        super(VolatilityVAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 4, 2, 1), nn.BatchNorm2d(32), nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
            nn.Flatten()
        )
        self.fc_mu = nn.Linear(128 * 8 * 8, LATENT_DIM)
        self.fc_logvar = nn.Linear(128 * 8 * 8, LATENT_DIM)
        
        # Decoder: MUST match world_model_vae.py structure
        self.decoder_input = nn.Linear(LATENT_DIM, 128 * 8 * 8)  # Separate layer!
        self.decoder = nn.Sequential(
            nn.Unflatten(1, (128, 8, 8)),
            nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2),
            nn.ConvTranspose2d(64, 32, 4, 2, 1), nn.BatchNorm2d(32), nn.LeakyReLU(0.2),
            nn.ConvTranspose2d(32, 1, 4, 2, 1), nn.Sigmoid()
        )
        
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
        
    def forward(self, x):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        z = self.reparameterize(mu, logvar)
        z_expanded = self.decoder_input(z)  # Expand first!
        return self.decoder(z_expanded), mu, logvar
    
    def decode(self, z):
        """Helper method to decode from latent z"""
        z_expanded = self.decoder_input(z)
        return self.decoder(z_expanded)

class WorldModelLSTM(nn.Module):
    def __init__(self):
        super(WorldModelLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size=LATENT_DIM, hidden_size=HIDDEN_SIZE, batch_first=True)
        self.fc = nn.Linear(HIDDEN_SIZE, LATENT_DIM)
        
    def forward(self, x):
        out, _ = self.lstm(x)
        return self.fc(out[:, -1, :])

# --- 2. Visualization Logic ---
def visualize_dream():
    print("--> Loading Models...")
    vae = VolatilityVAE().to(DEVICE)
    vae.load_state_dict(torch.load("dataset_v1/vae_weights.pth", map_location=DEVICE))
    vae.eval()
    
    lstm = WorldModelLSTM().to(DEVICE)
    lstm.load_state_dict(torch.load("dataset_v1/dynamics_weights.pth", map_location=DEVICE))
    lstm.eval()
    
    # Load Real Data to seed the dream
    print("--> Loading Seed Data...")
    raw_data = np.load("dataset_v1/spy_iv_surface_dataset.npy") / 0.8
    # Take the LAST 10 frames from the dataset as the "Present Moment"
    seed_frames = raw_data[-10:] 
    
    # Encode Seed Frames into Latent Space
    seed_tensor = torch.Tensor(seed_frames).unsqueeze(1).to(DEVICE)
    with torch.no_grad():
        h = vae.encoder(seed_tensor)
        mu = vae.fc_mu(h)
        current_z_seq = mu.unsqueeze(0) # Shape: (1, 10, 32)

    # Dream 50 Steps into the Future
    print("--> Generating Hallucination...")
    dream_surfaces = []
    
    with torch.no_grad():
        for i in range(50):
            # 1. Predict Next Z
            next_z = lstm(current_z_seq) # Shape: (1, 32)
            
            # 2. Decode Z back to Surface Image
            decoded_surface = vae.decode(next_z)  # Use decode helper method
            dream_surfaces.append(decoded_surface.cpu().numpy()[0, 0])
            # 3. Update Sequence (Slide Window)
            next_z_reshaped = next_z.unsqueeze(1)
            current_z_seq = torch.cat((current_z_seq[:, 1:, :], next_z_reshaped), dim=1)

    # Create Animation
    print("--> Creating Animation...")
    fig, ax = plt.subplots(figsize=(8, 6))
    
    # Setup initial plot
    im = ax.imshow(dream_surfaces[0], aspect='auto', origin='lower', cmap='inferno', vmin=0, vmax=1.0)
    plt.colorbar(im, label="Implied Volatility (Scaled)")
    title = ax.set_title("World Model Dream: T+0")
    ax.set_xlabel("Time to Maturity")
    ax.set_ylabel("Moneyness")

    def update(frame_idx):
        im.set_array(dream_surfaces[frame_idx])
        title.set_text(f"World Model Dream: T+{frame_idx} (Future Hallucination)")
        return im, title

    ani = animation.FuncAnimation(fig, update, frames=len(dream_surfaces), interval=100, blit=True)
    
    # Save or Show
    try:
        ani.save('dataset_v1/market_dream.gif', writer='pillow')
        print("--> Saved 'market_dream.gif'")
    except:
        print("--> Could not save GIF (missing ffmpeg/pillow?), showing plot instead.")
        plt.show()

    plt.show()

if __name__ == "__main__":
    visualize_dream()