vae_pipeline.py

"""
@file vae_train.py
@author Ryan Missel
@source https://github.com/sksq96/pytorch-vae/blob/master/vae-cnn.ipynb

Handles the full pipeline of training a VAE on the given pulsar data before running it through an SOM
to get all of its cluster graphs
"""
import os
import torch
import argparse
import numpy as np
import torch.nn.functional as F
import pandas as pd
import matplotlib.pyplot as plt

from tqdm import tqdm
from vae_model import VAE
from minisom import MiniSom
from util.plot_functions import *
from util.util_functions import *

""" 
Arg parsing and Data setup
"""
# Arg parsing
parser = argparse.ArgumentParser()

parser.add_argument('--seed', type=int, default=123, help='random seed')
parser.add_argument('--device', '-g', type=int, default=0, help='which GPU to run on')
parser.add_argument('--checkpt', type=str, default='None', help='checkpoint to resume training from')

parser.add_argument('--epochs', '-e', type=int, default=25, help='how many epochs to run')
parser.add_argument('--batch', '-b', type=int, default=32, help='size of batch')
parser.add_argument('--resume', '-r', type=bool, default=False, help='whether to resume training')

parser.add_argument('--month', '-m', type=str, default='april', help='month of obs')
parser.add_argument('--day', '-d', type=str, default='17', help='day of obs')
parser.add_argument('--year', '-y', type=str, default='2024', help='year of obs')
parser.add_argument('--antenna', '-a', type=str, default='A2', help='antenna of obs')

args = parser.parse_args()

# Get vars
month, day, year, antenna = args.month, args.day, args.year, args.antenna

# Build path
if not os.path.exists(f"models/{month}-{day}-{year}-{antenna}/"):
    os.mkdir(f"models/{month}-{day}-{year}-{antenna}/")

# Set seed if given
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)

# Device configuration
device = torch.device('cuda:2' if torch.cuda.is_available() else 'cpu')
torch.cuda.set_device(args.device)

# Load in the given dataset
dataset = np.load(f"data/{month}-{day}-{year}/{antenna}.npy", allow_pickle=True)
print("Loaded in cache dataset: ", "data/{}{}/{}.npy".format(month, day, antenna), " shape {}".format(dataset.shape))

# Sort the dataset for the largest pulse baseline indices
maxes = dataset[:, 0].argsort()[-15:][::-1]

# Rebase the dataset to 0
shift = np.reshape(np.mean(dataset, axis=1), [-1, 1])
dataset = dataset - shift

# Get the centered window of the signal
dataset = get_window(dataset)

# Transform datasets into loaders
kwargs = {'num_workers': 0, 'pin_memory': True}
dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch, shuffle=True, **kwargs)
testloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch, shuffle=False, **kwargs)

""" 
Model setup 
"""
mse_lambda = 0.5
kld_lambda = 0.5


def loss_fn(recon_x, x, mu, logvar):
    # Reconstruction loss
    BCE = mse_lambda * F.mse_loss(recon_x, x)

    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -kld_lambda * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD, BCE, KLD


# Build the model and load in a checkpoint if given
model = VAE(x_dim=dataset.shape[1], h_dim=1280, device=device)
if args.resume is True:
    model.load_state_dict(torch.load(f'models/{month}-{day}-{year}-{antenna}.torch'))
model = model.to(device)

# Define optimizer to use
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# Train the model for the given epochs
print("=> Starting training...")
for epoch in range(args.epochs):
    epochloss = 0
    bceloss = 0
    kldloss = 0

    for idx, signals in enumerate(dataloader):
        signals = signals.to(device).float()

        recon_signals, z, h, mu, logvar = model(signals)
        loss, bce, kld = loss_fn(recon_signals, signals, mu, logvar)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epochloss += loss.item()
        bceloss += bce.item()
        kldloss += kld.item()

    torch.save(model.state_dict(), f'models/{month}-{day}-{year}-{antenna}/{month}-{day}-{year}-{antenna}.torch')
    print(f"Epoch[{epoch + 1}/{args.epochs}] Loss: {epochloss / args.batch:0.3f} {bceloss / args.batch:0.3f} {kldloss / args.batch:0.3f}")

# Save model checkpoint out
print("=> Complete!")
torch.save(model.state_dict(), f'models/{month}-{day}-{year}-{antenna}/{month}-{day}-{year}-{antenna}.torch')
print(f"=> Model saved at models/{month}-{day}-{year}-{antenna}/{month}-{day}-{year}-{antenna}.torch")


""" Evaluation """
# Reconstruct the entire dataset
recons = None
for idx, signals in tqdm(enumerate(testloader)):
    signals = signals.to(device).float()
    recon_signals, _, _, _, _ = model(signals)
    recon_signals = recon_signals.detach().cpu().numpy()

    if recons is None:
        recons = recon_signals
    else:
        recons = np.vstack((recons, recon_signals))


# Iterate over the SOM shapes
for som_shape in [(2, 2), (2, 3), (3, 3)]:
    # Full path for graph output
    fullpath = f'models/{month}-{day}-{year}-{antenna}/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] < 20:
            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"{month}-{day}-{year}-{antenna}\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.antenna}" + ".}\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.antenna}", fontdict={'fontsize': 14, 'fontweight': 3})
            plt.savefig(f"{fullpath}/examples/{idxs}signalCluster{clus + 1}.png")
            plt.close()


    # Plotting the centroid signals of each node
    plot = True
    if plot is True:
        plot_centroids(fullpath, som, som_shape, dataset, cluster_index)

    # 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.antenna}')

    # 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.antenna}', vertical_fix=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)

    # Grid graph of the physical features using colormap
    plot = False
    if plot is True:
        plot_gridgraph(fullpath, som_shape, xs, ys, maxes, argmaxes, widths, skews, mses)