plot_npy.py

import numpy as np
import os
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
from sklearn.metrics.pairwise import cosine_similarity

from utils import compute_spectral_overlap_index

def compute_similarity_score(groundtruth, prediction):
    scores = []
    for gt, pred in zip(groundtruth, prediction):
        score, _ = pearsonr(gt.flatten(), pred.flatten())
        scores.append(score)
    return sum(scores) / len(scores)  # Average similarity score

def plot_eval_results(eval_folder, fold, epoch, batch_idx, avg_similarity_score, avg_soi, avg_loss, save_plot=True, info=None):
    """
    Plots the ground truth and prediction signals for a given epoch and batch index.

    Args:
        eval_folder (str): Path to the folder containing evaluation .npy files.
        epoch (int): The epoch number to load the files for.
        batch_idx (int): The batch index to load the files for.
        save_plot (bool): Whether to save the plot as a file. Default is True.

    Returns:
        None
    """
    # Define file paths
    groundtruth_file = os.path.join(eval_folder, f'groundtruth_{fold}_{epoch}_{batch_idx}.npy')
    prediction_file = os.path.join(eval_folder, f'prediction_{fold}_{epoch}_{batch_idx}.npy')
    original_signal_file = os.path.join(eval_folder, f'input_{fold}_{epoch}_{batch_idx}.npy')
    # import pdb; pdb.set_trace()

    # Load the .npy files
    groundtruth = np.load(groundtruth_file)  # Shape: (batch_size, data_length)
    prediction = np.load(prediction_file)    # Shape: (batch_size, data_length)
    original_signal = np.load(original_signal_file)
    # Compute the similarity score for the first piece of data
    # similarity_score, _ = pearsonr(groundtruth[0], prediction[0])
    # print(f"Similarity score for the first piece of data: {similarity_score:.4f}")

    batch_size = groundtruth.shape[0]

    # Create a figure with subplotsg
    fig, axes = plt.subplots(batch_size, 1, figsize=(10, 2 * batch_size), sharex=True)

    for i in range(batch_size):
        similarity_score, _ = pearsonr(groundtruth[i], prediction[i])
        soi = compute_spectral_overlap_index(groundtruth[i], prediction[i])
        cos = cosine_similarity(groundtruth[i].reshape(1, -1), prediction[i].reshape(1, -1))[0][0]
        similarity_score = round(similarity_score, 4)
        if np.any(original_signal) != None:
            axes[i].plot(original_signal[i], label='Input', linestyle='-', color='green', alpha=0.5)
        axes[i].plot(groundtruth[i], label='Ground Truth', linestyle='-', color='blue')
        axes[i].plot(prediction[i], label='Prediction', linestyle='-', color='red')
        if info!=None:
            axes[i].set_title(f'Sample {i} - {info[i]} - {similarity_score:.4f} - {soi:.4f} - {cos:.4f}')
        else:
            axes[i].set_title(f'Sample {i} Sim: {similarity_score:.4f} - SOI: {soi:.4f} - COS: {cos:.4f}')
        # axes[i].legend()
        axes[i].grid(True)

    # Add a common x-label and y-label
    fig.suptitle(f'Example Comparison of Ground Truth and Prediction (Fold {fold}, Epoch {epoch}, Batch {batch_idx})')
    fig.text(0.5, 0.97, f'Average Evaluation Loss: {avg_loss:.4f}', ha='center', fontsize=10)
    fig.text(0.5, 0.96, f'AVG Similarity Score (Pearson Correlation): {avg_similarity_score:.4f}', ha='center', fontsize=10)
    fig.text(0.5, 0.95, f'AVG SOI (Spectral Overlap Index): {avg_soi:.4f}', ha='center', fontsize=10)
    fig.supxlabel('Index')
    fig.supylabel('Value')

    # Adjust layout and optionally save the plot
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Leave space for the suptitle
    if save_plot:
        plot_path = os.path.join(eval_folder, f'comparison_{fold}_epoch_{epoch}_batch_{batch_idx}.png')
        plt.savefig(plot_path)
        print(f"Plot saved to {plot_path}")
    plt.show()
    
if __name__ == "__main__":
    plot_eval_results('eval/SCG', 1, 90, 0, 0, 0, 0, original_signal=None)