lightning_replot.py

import os
import torch
import argparse
import numpy as np
import pandas as pd
import pytorch_lightning as pl

from pytorch_lightning.callbacks import ModelCheckpoint
from torch.utils.data import DataLoader, Dataset
from vae_model import VAE, VAESmall
from util.plot_functions import *
from util.util_functions import *
from minisom import MiniSom


""" Custom Dataset for Pulsar Signals """
class PulsarDataset(Dataset):
    def __init__(self, dataset):
        self.data = dataset

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return torch.tensor(self.data[idx], dtype=torch.float32)


""" Lightning DataModule for Pulsar Dataset """
class PulsarDataModule(pl.LightningDataModule):
    def __init__(self, data_paths, batch_size):
        super().__init__()
        self.data_paths = data_paths
        self.batch_size = batch_size

    def prepare_data(self):
        self.dataset = None
        for month, day, year, antennae in self.data_paths:
            sub_dataset = np.load(f"data/{month}-{day}-{year}/{antennae}.npy", allow_pickle=True)
            shift = np.reshape(np.mean(sub_dataset, axis=1), [-1, 1])
            sub_dataset = sub_dataset - shift
            sub_dataset = get_window(sub_dataset)

            if self.dataset is None:
                self.dataset = sub_dataset
            else:
                self.dataset = np.vstack((self.dataset, sub_dataset))

    def setup(self, stage=None):
        self.train_dataset = PulsarDataset(self.dataset)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4, pin_memory=True)


""" Lightning Module for VAE Training """
class VAE_Lightning(pl.LightningModule):
    def __init__(self, args, x_dim, mse_lambda=0.5, kld_lambda=0.5, lr=1e-4):
        super(VAE_Lightning, self).__init__()
        self.args = args
        # self.model = VAE(x_dim=x_dim, h_dim=1280, device=args.device)
        self.model = VAE(x_dim=x_dim, h_dim=1280, device=args.device)
        self.mse_lambda = mse_lambda
        self.kld_lambda = kld_lambda
        self.lr = lr
        self.mse_loss = torch.nn.MSELoss(reduction='none')

    def loss_fn(self, recon_x, x, mu, logvar):
        # MSE = self.mse_lambda * self.mse_loss(recon_x, x).sum([1]).mean([0])
        
        # Build the upweighting vector
        upweight = torch.ones([1, recon_x.shape[1]], device=self.args.device)
        if x.shape[-1] == 200:
            upweight[:, 70:130] += 0.25
        else:
            upweight[:, 30:50] += 0.50
            upweight[:, :30] -= 0.50
            upweight[:, 50:] -= 0.50
        
        # Calculate the MSE loss
        MSE = (upweight * self.mse_loss(recon_x, x)).sum([1]).mean([0])
        
        # Calculate the KL divergence
        KLD = -self.kld_lambda * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
        return MSE + KLD, MSE, KLD

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        recon_x, z, h, mu, logvar = self(batch)
        loss, bce, kld = self.loss_fn(recon_x, batch, mu, logvar)
        self.log_dict({"train_loss": loss, "train_bce": bce, "train_kld": kld}, on_epoch=True, on_step=False, prog_bar=True)
        return loss

    def configure_optimizers(self):
        # Define optimizer and scheduler
        optim = torch.optim.AdamW(self.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=1, gamma=0.9, verbose=True)

        # Explicit dictionary to state how often to ping the scheduler
        scheduler = {
            'scheduler': scheduler,
            'frequency': 1,
            'interval': 'epoch'
        }

        return [optim], [scheduler]


""" Training and Evaluation Script """
def main():
    # Argument parsing
    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', type=int, default=1234, help='random seed')
    parser.add_argument('--epochs', '-e', type=int, default=10, help='number of epochs')
    parser.add_argument('--batch', '-b', type=int, default=32, help='batch size')
    parser.add_argument('--device', '-g', type=str, default="cuda:0", help='GPU device')
    
    parser.add_argument('--month', '-m', type=str, default="july")
    parser.add_argument('--day', '-d', type=str, default="21")
    parser.add_argument('--year', '-y', type=str, default="2024")
    parser.add_argument('--antennae', '-a', type=str, default="A2")
    args = parser.parse_args()

    # Set random seed
    pl.seed_everything(args.seed)

    # Initialize DataModule
    data_module = PulsarDataModule(data_paths=[(args.month, args.day, args.year, args.antennae)], batch_size=args.batch)
    data_module.prepare_data()
    x_dim = data_module.dataset.shape[1]
    print(f"=> Dataset Shape: {data_module.dataset.shape}")

    # Define VAE model
    vae = VAE_Lightning(args, x_dim=x_dim)

    # Load in checkpoint
    vae = vae.load_from_checkpoint(f"models/{args.month}-{args.day}-{args.year}-{args.antennae}/vae_model.ckpt", args=args, x_dim=x_dim, map_location="cuda:0")

    # Evaluation
    vae_model = vae.model.eval().cuda()
    dataset = np.load(f"data/{args.month}-{args.day}-{args.year}/{args.antennae}.npy", allow_pickle=True)
    shift = np.reshape(np.mean(dataset, axis=1), [-1, 1])
    dataset = dataset - shift
    dataset = get_window(dataset)

    testloader = DataLoader(PulsarDataset(dataset), batch_size=args.batch, shuffle=False)

    # Reconstruct data
    recons = []
    for signals in testloader:
        signals = signals.cuda()
        with torch.no_grad():
            recon_signals, _, _, _, _ = vae_model(signals)
            recons.append(recon_signals.cpu().numpy())
    recons = np.vstack(recons)

    # Iterate over the SOM shapes
    for som_shape in [(2, 2), (2, 3), (3, 3)]:
        # Full path for graph output
        fullpath = f'models/{args.month}-{args.day}-{args.year}-{args.antennae}/vae_{som_shape[0]}_{som_shape[1]}/'
        print(f"=> Graph output path: {fullpath}")

        # Setting up folders for graph saving
        if not os.path.exists(fullpath):
            os.makedirs(fullpath)

        # Handles plotting the mean signals of the raw data and the reconstruction
        plot_set_mean(fullpath, recons, dataset)

        # Train SOM
        som = MiniSom(x=som_shape[0], y=som_shape[1], input_len=dataset.shape[1], sigma=0.3, learning_rate=0.5, random_seed=args.seed)
        som.train(recons, 10000, verbose=False)

        # Extract coordinates for each cluster and assignment
        winner_coordinates = np.array([som.winner(x) for x in recons]).T
        cluster_index = np.ravel_multi_index(winner_coordinates, som_shape)

        # Get maxes across clusters
        maxes = get_maxes(dataset, cluster_index, som_shape)
        resort = list(np.array(maxes).argsort()[::-1])

        # Remap to sort via peak height
        map_dict = {float(k): float(v) for k, v in zip(resort, range(som_shape[0] * som_shape[1]))}
        cluster_index = vector_map(cluster_index, map_dict)
        cluster_index = cluster_index.astype(np.int64)
        print(np.unique(cluster_index, return_counts=True))
        
        # Get cluster-specific metrics
        retrain = False
        for idx in range(som_shape[0] * som_shape[1]):
            subidxs = np.where(cluster_index == idx)[0]
            if subidxs.shape[0] < 100:
                dataset = np.delete(dataset, subidxs, axis=0)
                recons = np.delete(recons, subidxs, axis=0)
                retrain = True
        
        # Re-train SOM
        if retrain is True:
            som = MiniSom(x=som_shape[0], y=som_shape[1], input_len=dataset.shape[1], sigma=0.3, learning_rate=0.5, random_seed=args.seed)
            som.train(recons, 10000, verbose=False)

            # Extract coordinates for each cluster and assignment
            winner_coordinates = np.array([som.winner(x) for x in recons]).T
            cluster_index = np.ravel_multi_index(winner_coordinates, som_shape)

            # Get maxes across clusters
            maxes = get_maxes(dataset, cluster_index, som_shape)
            resort = list(np.array(maxes).argsort()[::-1])

            # Remap to sort via peak height
            map_dict = {float(k): float(v) for k, v in zip(resort, range(som_shape[0] * som_shape[1]))}
            cluster_index = vector_map(cluster_index, map_dict)
            cluster_index = cluster_index.astype(np.int64)
            print(np.unique(cluster_index, return_counts=True))
        
        # Save cluster IDs to file
        np.savetxt(f'{fullpath}/clusterIDs.csv', cluster_index + 1, delimiter=',')

        """ Statistics """
        # Gets statistics to be used in multiple plots
        shapes, indices, xs, ys, (maxes, maxes_std), (argmaxes, argmaxes_std), (widths, widths_std), (skews, skews_std), (mses, mses_std) \
            = cluster_statistics(fullpath, som_shape, recons, dataset, cluster_index, save=True)

        # Combined and save as a CSV over all
        combined = np.array([maxes, maxes_std, argmaxes, argmaxes_std, widths, widths_std, skews, skews_std, mses, mses_std]).T
        df = pd.DataFrame(np.array([['total'] + [i for i in range(som_shape[0] * som_shape[1])], maxes, maxes_std, argmaxes, argmaxes_std, widths, widths_std, skews, skews_std, mses, mses_std]).T,
                            columns=['cluster', 'maxes', 'maxes_std', 'argmaxes', 'argmaxes_std', 'widths', 'widths_std', 'skews', 'skews_std', 'mses', 'mses_std'])
        df.to_csv(f'{fullpath}/som{som_shape[0]}{som_shape[1]}_cluster_values.csv')

        # Extract stats to latex table
        metrics = [
            ['peak loc', argmaxes, argmaxes_std],
            ['peak height', maxes, maxes_std],
            ['peak width', widths, widths_std],
            ['peak skew', skews, skews_std],
            ['MSE', mses, mses_std]
        ]

        f = open("{}/statistics_as_latex.txt".format(fullpath), 'w')
        f.write(f"{args.month}-{args.day}-{args.year}-{args.antennae}\n")
        f.write("\\begin{table*}\n")
        f.write("\t \\centering\n")
        f.write("\t \\caption{SOM Clustering for " + f"{args.month.title()} {args.day} " + "with Antenna " + f"{args.antennae}" + ".}\n")
        f.write("\t \\label{tab:[]}\n")
        f.write("\t \\begin{tabular}{cclllll}\n")
        f.write("\t \t \\hline\n")
        f.write("\t \t Cluster \\# & \\# Pulses & Peak Loc & Peak Height & Peak Width & Peak Skew & MSE \\\\ \n")
        f.write("\t \t \\hline \n")

        for cidx, pulses in zip(range(1 + (som_shape[0] * som_shape[1])), shapes):
            num_pulses = pulses[0]
            peak_loc_mean, peak_loc_std = metrics[0][1][cidx], metrics[0][2][cidx]
            peak_height_mean, peak_height_std = metrics[1][1][cidx], metrics[1][2][cidx]
            peak_width_mean, peak_width_std = metrics[2][1][cidx], metrics[2][2][cidx]
            peak_skew_mean, peak_skew_std = metrics[3][1][cidx], metrics[3][2][cidx]
            mse_mean, mse_std = metrics[4][1][cidx], metrics[4][2][cidx]

            f.write(
                f"\t \t {cidx} & {num_pulses} & "
                f"${peak_loc_mean:0.2f} \\pm {peak_loc_std:0.2f}$ & "
                f"${peak_height_mean:0.2f} \\pm {peak_height_std:0.2f}$ & "
                f"${peak_width_mean:0.2f} \\pm {peak_width_std:0.2f}$ & "
                f"${peak_skew_mean:0.2f} \\pm {peak_skew_std:0.2f}$ & "
                f"${mse_mean:0.5f} \\pm {mse_std:0.5f}$  \\\\ \n"
            )

        f.write("\t \t \\hline \n")
        f.write("\t \\end{tabular} \n")
        f.write("\end{table*} \n")


        """ Plotting """
        if not os.path.exists(f'{fullpath}/examples/'):
            os.mkdir(f'{fullpath}/examples/')

        # Plot reconstruction examples per cluster
        for _ in range(4):
            for clus in range(som_shape[0] * som_shape[1]):
                if len(np.where(cluster_index == clus)[0]) == 0:
                    continue

                idxs = np.random.choice(np.where(cluster_index == clus)[0])
                signal = dataset[idxs]
                reconsig = recons[idxs]

                plt.figure(1)
                ax = plt.gca()
                ax.tick_params(axis='both', which='major', labelsize=12)

                plt.plot(signal)
                plt.plot(reconsig)
                plt.legend(('Raw', 'Reconstructed'), fontsize=11)
                plt.title(f"Cluster #{clus + 1} Single Reconstruction on {args.month.title()} {args.day} {args.antennae}", fontdict={'fontsize': 14, 'fontweight': 3})
                plt.savefig(f"{fullpath}/examples/{idxs}signalCluster{clus + 1}.png")
                plt.close()

        # Mean plots of each cluster plotted on each other
        plot = True
        if plot is True:
            plot_means(fullpath, som_shape, recons, dataset, cluster_index, f'{args.month.title()} {args.day} {args.antennae}')

        # Mean plots of each cluster plotted on each other *scaled to a common vertical axis*
        plot = True
        if plot is True:
            plot_means(fullpath, som_shape, recons, dataset, cluster_index, f'{args.month.title()} {args.day} {args.antennae}', vertical_fix=True)

        # Mean plots of each cluster plotted on each other *scaled to a common vertical axis*
        plot = True
        if plot is True:
            plot_means_paper(fullpath, som_shape, recons, dataset, cluster_index, f'{args.month.title()} {args.day} {args.antennae}', narrow=False)

        # Mean plots of each cluster plotted on each other *scaled to a common vertical axis*
        plot = True
        if plot is True:
            plot_means_paper(fullpath, som_shape, recons, dataset, cluster_index, f'{args.month.title()} {args.day} {args.antennae}', narrow=True)

        # Plotting the raw and recon mean of each cluster over each other
        plot = True
        if plot is True:
            plot_mean_comparison(fullpath, som_shape, recons, dataset, cluster_index)


if __name__ == "__main__":
    main()