train_withinuser.py

import os
import time
import argparse
import csv
import statistics
import pandas as pd
import numpy as np
import scipy.signal
from scipy.stats import pearsonr
from sklearn.model_selection import KFold
from sklearn.metrics.pairwise import cosine_similarity
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm

from model_orig import PCGtoSCG_Orig
from plot_npy import plot_eval_results
from dataset import MultiSignalSegmentDataset
from utils import compute_similarity_score, compute_spectral_overlap_index, collate_fn_processed

PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

def evaluate_model(model, dataloader, criterion, device, save_folder, epoch, fold):
    model.eval()
    total_loss = 0.0
    groundtruths = []  
    predictions = []
    os.makedirs(save_folder, exist_ok=True)
    
    with torch.no_grad():
        info_ls = []
        soi_ls = []
        sim_ls = []
        cos_ls = []
        for idx, (inearpcg_signal, target_signal, info) in enumerate(tqdm(dataloader, desc="Evaluating")):
            groundtruths = []
            predictions = []

            inearpcg_signal, target_signal = inearpcg_signal.to(device), target_signal.to(device)
            
            outputs = model(inearpcg_signal.unsqueeze(1))
            
            if outputs.size(2) > target_signal.size(1):
                outputs = outputs[:, :, :target_signal.size(1)]
            elif outputs.size(2) < target_signal.size(1):
                target_signal = target_signal[:, :outputs.size(2)]

            loss = criterion(outputs.squeeze(1), target_signal)
            total_loss += loss.item()

            # Save ground truth and predictions
            groundtruths = target_signal.cpu().numpy()
            predictions = outputs.squeeze(1).cpu().numpy()

            if idx>=0:
                input_file = os.path.join(save_folder, f"input_{fold+1}_{epoch}_{idx}.npy")
                groundtruth_file = os.path.join(save_folder, f"groundtruth_{fold+1}_{epoch}_{idx}.npy")
                prediction_file = os.path.join(save_folder, f"prediction_{fold+1}_{epoch}_{idx}.npy")
                # Save to files
                np.save(input_file, inearpcg_signal.cpu().numpy())
                np.save(groundtruth_file, target_signal.cpu().numpy())
                np.save(prediction_file, outputs.squeeze(1).cpu().numpy())
                info_ls.append(info)

            value_st = 20
            value_end = 200
            for gt, pred in zip(groundtruths, predictions):
                valuable_gt = gt[value_st:value_end]
                valuable_pred = pred[value_st:value_end]
                
                # Check for constant arrays before computing correlation
                if len(np.unique(valuable_gt)) == 1 or len(np.unique(valuable_pred)) == 1:
                    sim = 0.0  # or np.nan, depending on how you want to handle it
                else:
                    sim, _ = pearsonr(valuable_gt, valuable_pred)
                    sim = float(sim)  # Convert to Python float
                    if not np.isfinite(sim):  # Handle NaN or Inf
                        sim = 0.0
                
                soi = compute_spectral_overlap_index(valuable_gt.flatten(), valuable_pred.flatten())
                soi = float(soi) if np.isfinite(soi) else 0.0  # Convert and handle NaN
                
                # Compute cosine similarity
                valuable_gt_reshaped = valuable_gt.reshape(1, -1)
                valuable_pred_reshaped = valuable_pred.reshape(1, -1)
                cosine_sim = cosine_similarity(valuable_gt_reshaped, valuable_pred_reshaped)[0][0]
                cosine_sim = float(cosine_sim) if np.isfinite(cosine_sim) else 0.0
                
                sim_ls.append(sim)
                soi_ls.append(soi)
                cos_ls.append(cosine_sim)

            # plot_eval_results(eval_save_folder, fold+1, epoch, idx, similarity_score, idx_soi, loss, original_signal=inearpcg_signal.cpu().numpy(), save_plot=True)
        
    avg_sim = statistics.mean(sim_ls)
    med_sim = statistics.median(sim_ls)
    print(f"Average Similarity Score: {avg_sim:.4f}, Median Similarity Score: {med_sim:.4f}")

    avg_soi = statistics.mean(soi_ls)
    med_soi = statistics.median(soi_ls)
    print(f"Average Spectral Overlap Index (SOI): {avg_soi:.4f}, Median SOI: {med_soi:.4f}")
    
    avg_cos = statistics.mean(cos_ls)
    med_cos = statistics.median(cos_ls)
    print(f"Average Cosine Similarity: {avg_cos:.4f}, Median Cosine Similarity: {med_cos:.4f}")
    
    avg_loss = total_loss / len(dataloader)
    print(f"Evaluation Loss: {avg_loss:.4f}")
    
    for idx in range(1):
        plot_eval_results(eval_save_folder, fold+1, epoch, idx, avg_sim, avg_soi, avg_loss, save_plot=True, info=info_ls[idx])
        
        input_file = os.path.join(save_folder, f"input_{fold+1}_{epoch}_{idx}.npy")
        # groundtruth_file = os.path.join(save_folder, f"groundtruth_{fold+1}_{epoch}_{idx}.npy")
        # prediction_file = os.path.join(save_folder, f"prediction_{fold+1}_{epoch}_{idx}.npy")
        os.remove(input_file)
        # os.remove(groundtruth_file)
        # os.remove(prediction_file)

    return avg_loss, avg_sim, avg_soi, avg_cos, med_sim, med_soi, med_cos

# Main function
if __name__ == "__main__":
    time_start = time.time()
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_folder', type=str, default=os.path.join(PROJECT_ROOT, 'data'), help='Path to the data folder')
    parser.add_argument('--config_folder', type=str, default=os.path.join(PROJECT_ROOT, 'data_label', 'label_output'), help='Path to the config folder')
    parser.add_argument('--label_folder', type=str, default=os.path.join(PROJECT_ROOT, 'data_label', 'label_cycle'), help='Path to the label folder')
    parser.add_argument('--target_fs', type=int, default=500, help='Target sampling frequency')
    parser.add_argument('--segment_length', type=int, default=400, help='Segment length in samples')
    parser.add_argument('--batch_size', type=int, default=16, help='Batch size for training')
    parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate for the optimizer')
    parser.add_argument('--num_epochs', type=int, default=40, help='Number of epochs for training')
    parser.add_argument('--signal_type', type=str, default='scg', help='Target signal type: scg or gcg')
    parser.add_argument('--groups', nargs='+', default=['mic_wireless'], help='List of groups to process')
    parser.add_argument('--identifiers', nargs='+', required=True, help='Participant folder names under data/ (one per within-user run)')
    parser.add_argument('--mode', type=str, default='within', help='Mode of operation (e.g., cross, intra)')
    args = parser.parse_args()

    data_folder = args.data_folder
    config_folder = args.config_folder
    label_folder = args.label_folder
    target_fs = args.target_fs
    segment_length = args.segment_length
    batch_size = args.batch_size
    learning_rate = args.learning_rate
    num_epochs = args.num_epochs
    identifiers = args.identifiers

    signal_type = args.signal_type
    groups = args.groups
    mode = args.mode

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print('DEVICE:', device)
    
    for identifier in identifiers:
        print(f"\nProcessing identifier: {identifier}")
        identifier = [identifier]  # Make identifier a single-element list
        eval_save_folder = os.path.join(PROJECT_ROOT, 'eval', mode, signal_type, groups[0], f"{signal_type}_{identifier[0]}_{groups[0]}")
        save_path = os.path.join(PROJECT_ROOT, 'pth', mode, signal_type, groups[0], f"{signal_type}_{identifier[0]}_{groups[0]}")
        processed_dataset_path = os.path.join(PROJECT_ROOT, 'data', 'saved_dataset', f"{identifier[0]}.pt")
        print("Save path:", eval_save_folder)
        
        # Parameters for cross-validation
        k_folds = 5  # Number of folds
        
        # Set a random seed for reproducibility
        random_seed = 42
        generator = torch.Generator().manual_seed(random_seed)

        # Load dataset and split into train and evaluation sets
        if os.path.exists(processed_dataset_path):
            print("Loading processed dataset from file...")
            dataset = torch.load(processed_dataset_path)
        else:
            print("\nProcessing DATASET...")
            dataset = MultiSignalSegmentDataset(identifiers=identifier, groups=groups, signal_type=signal_type, data_path=data_folder, config_path=config_folder, label_path=label_folder)
            if len(dataset) == 0:
                print(f"Dataset is empty for {identifier}")
                continue
            # print("Saving processed dataset to file...")
            # torch.save(dataset, processed_dataset_path)
            # print("Dataset saved.")
        
        kf = KFold(n_splits=k_folds, shuffle=True)
        dataset_indices = list(range(len(dataset)))
        
        crossfold_sim = {}
        crossfold_soi = {}
        crossfold_cos = {}
        # Track best sum (sim + soi) seen so far across evaluations to keep only best outputs
        best_sum_global = float('-inf')
        best_model_dir = os.path.join(PROJECT_ROOT, 'pth', mode, identifier[0])
        os.makedirs(best_model_dir, exist_ok=True)
        best_model_path = None
        # Cross-validation loop
        for fold, (train_indices, eval_indices) in enumerate(kf.split(dataset_indices)):
            print(f"\nStarting Fold {fold + 1}/{k_folds}")
            
            # Initialize models
            model = PCGtoSCG_Orig()
            model.to(device)
            
            # Loss and optimizers
            # criterion = SimilarityLoss()
            criterion = nn.MSELoss()
            model_optimizer = optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999))
            
            # Initialize the learning rate scheduler
            model_scheduler = ReduceLROnPlateau(model_optimizer, mode='min', factor=0.5, patience=5, verbose=True)

            # Create train and evaluation datasets for this fold
            train_subset = torch.utils.data.Subset(dataset, train_indices)
            eval_subset = torch.utils.data.Subset(dataset, eval_indices)
            
            train_dataloader = DataLoader(train_subset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn_processed)
            eval_dataloader = DataLoader(eval_subset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn_processed)
            
            # Training loop for this fold
            for epoch in range(num_epochs):
                model.train()
                running_loss = 0.0
                for inearpcg_signal, target_signal, info in tqdm(train_dataloader, desc=f"Fold {fold + 1}, Epoch {epoch + 1}/{num_epochs}"):

                    inearpcg_signal, target_signal = inearpcg_signal.to(device), target_signal.to(device)
                                    
                    # Forward pass through main model
                    model_optimizer.zero_grad()
                    outputs = model(inearpcg_signal.unsqueeze(1))

                    if outputs.size(2) > target_signal.size(1):
                        outputs = outputs[:, :, :target_signal.size(1)]
                    elif outputs.size(2) < target_signal.size(1):
                        target_signal = target_signal[:, :outputs.size(2)]

                    # Calculate loss based on final task performance
                    reconstruction_loss = criterion(outputs.squeeze(1), target_signal)
                    loss = reconstruction_loss 

                    # Backward pass and optimization
                    loss.backward()
                    model_optimizer.step()

                    running_loss += loss.item()

                print(f"Fold {fold + 1}, Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_dataloader):.4f}")

                # Evaluate the model periodically
                if (epoch + 1) % 5 == 0 and epoch > 10:
                    os.makedirs(eval_save_folder, exist_ok=True)
                    # Perform evaluation (this will save per-sample .npy files into eval_save_folder)
                    eval_loss, avg_sim, avg_soi, avg_cos, med_sim, med_soi, med_cos = evaluate_model(model, eval_dataloader, criterion, device, eval_save_folder, epoch=epoch + 1, fold=fold)

                    # Step the scheduler based on evaluation loss
                    model_scheduler.step(eval_loss)

                    # Decide whether to keep the saved files for this (fold, epoch)
                    current_sum = float(avg_sim) + float(avg_soi)
                    epoch_eval = epoch + 1
                    prefixes_current = [f"input_{fold+1}_{epoch_eval}_", f"groundtruth_{fold+1}_{epoch_eval}_", f"prediction_{fold+1}_{epoch_eval}_"]
                    # graph file prefix for this epoch
                    graph_prefix_current = f"comparison_{fold+1}_epoch_{epoch_eval}_"
                    if current_sum > best_sum_global:
                        # New global best: remove all previously saved .npy files and comparison graphs except the ones for this (fold,epoch)
                        best_sum_global = current_sum
                        print(f'New best (fold {fold+1}, epoch {epoch_eval}) with sim+soi = {current_sum:.4f} -> keeping files and removing previous ones')
                        for root, dirs, files in os.walk(eval_save_folder):
                            for fname in files:
                                # remove previous .npy files following naming convention
                                if fname.endswith('.npy') and (fname.startswith('input_') or fname.startswith('groundtruth_') or fname.startswith('prediction_')):
                                    if not any(fname.startswith(p) for p in prefixes_current):
                                        fp = os.path.join(root, fname)
                                        try:
                                            os.remove(fp)
                                            print(f'Deleted previous {fp}')
                                        except Exception as e:
                                            print(f'Failed to delete {fp}: {e}')
                                # remove previous comparison graphs (png) not for current best
                                if fname.endswith('.png') and fname.startswith('comparison_'):
                                    if not fname.startswith(graph_prefix_current):
                                        fp = os.path.join(root, fname)
                                        try:
                                            os.remove(fp)
                                            print(f'Deleted previous graph {fp}')
                                        except Exception as e:
                                            print(f'Failed to delete graph {fp}: {e}')
                        # Remove empty directories if any (except top folder)
                        for root, dirs, files in os.walk(eval_save_folder, topdown=False):
                            if root == eval_save_folder:
                                continue
                            try:
                                if not os.listdir(root):
                                    os.rmdir(root)
                                    print(f'Removed empty directory {root}')
                            except Exception:
                                pass
                        # Model saving logic for best epoch
                        new_model_path = os.path.join(best_model_dir, f"{mode}_{identifier[0]}_{groups[0]}_fold{fold+1}_epoch{epoch_eval}.pth")
                        # Delete previous best model file if it exists and is different
                        if best_model_path and os.path.exists(best_model_path) and best_model_path != new_model_path:
                            try:
                                os.remove(best_model_path)
                                print(f"Deleted previous best model: {best_model_path}")
                            except Exception as e:
                                print(f"Failed to delete previous best model: {best_model_path}, error: {e}")
                        torch.save(model.state_dict(), new_model_path)
                        print(f"Best model saved to {new_model_path}")
                        best_model_path = new_model_path
                    else:
                        # Not a new best: delete files and graphs for this epoch to save space
                        print(f'Current (fold {fold+1}, epoch {epoch_eval}) sim+soi = {current_sum:.4f} <= best {best_sum_global:.4f} -> deleting files for this epoch to save space')
                        prefixes = prefixes_current
                        for root, dirs, files in os.walk(eval_save_folder):
                            for fname in files:
                                for p in prefixes:
                                    if fname.startswith(p) and fname.endswith('.npy'):
                                        fp = os.path.join(root, fname)
                                        try:
                                            os.remove(fp)
                                            print(f'Deleted {fp}')
                                        except Exception as e:
                                            print(f'Failed to delete {fp}: {e}')
                                # delete graphs for this epoch
                                if fname.startswith(graph_prefix_current) and fname.endswith('.png'):
                                    fp = os.path.join(root, fname)
                                    try:
                                        os.remove(fp)
                                        print(f'Deleted graph {fp}')
                                    except Exception as e:
                                        print(f'Failed to delete graph {fp}: {e}')
                        # Remove empty directories if any
                        for root, dirs, files in os.walk(eval_save_folder, topdown=False):
                            if root == eval_save_folder:
                                continue
                            try:
                                if not os.listdir(root):
                                    os.rmdir(root)
                                    print(f'Removed empty directory {root}')
                            except Exception:
                                pass

                    # # Save the model periodically
                    # if (epoch + 1) % 20 == 0:
                    # model_save_path = os.path.join(save_path, f"model_fold_{fold + 1}_epoch_{epoch + 1}.pth")
                    # torch.save(model.state_dict(), model_save_path)
                    # print(f"Models saved for Fold {fold + 1} at epoch {epoch + 1} \n")

                    # Update crossfold dictionaries to include 'avg' and 'med' levels
                    if str(fold+1) not in crossfold_sim:
                        crossfold_sim[str(fold+1)] = {}
                    if str(epoch+1) not in crossfold_sim[str(fold+1)]:
                        crossfold_sim[str(fold+1)][str(epoch+1)] = {'avg': round(avg_sim, 4), 'med': round(med_sim, 4)}

                    if str(fold+1) not in crossfold_soi:
                        crossfold_soi[str(fold+1)] = {}
                    if str(epoch+1) not in crossfold_soi[str(fold+1)]:
                        crossfold_soi[str(fold+1)][str(epoch+1)] = {'avg': round(avg_soi, 4), 'med': round(med_soi, 4)}

                    if str(fold+1) not in crossfold_cos:
                        crossfold_cos[str(fold+1)] = {}
                    if str(epoch+1) not in crossfold_cos[str(fold+1)]:
                        crossfold_cos[str(fold+1)][str(epoch+1)] = {'avg': round(avg_cos, 4), 'med': round(med_cos, 4)}

        output_file = os.path.join(eval_save_folder, f'{identifier[0]}_number.csv')
        with open(output_file, mode='w', newline='') as file:
            writer = csv.writer(file)
            # Write the header
            writer.writerow(['Fold', 'Epoch', 'Metric', 'Average', 'Median'])
            # Write the data
            for fold in crossfold_sim:
                for epoch in crossfold_sim[fold]:
                    sim_avg = round(crossfold_sim[fold][epoch]['avg'], 4)
                    sim_med = round(crossfold_sim[fold][epoch]['med'], 4)
                    writer.writerow([fold, epoch, 'Similarity Score', sim_avg, sim_med])

                    soi_avg = round(crossfold_soi[fold][epoch]['avg'], 4)
                    soi_med = round(crossfold_soi[fold][epoch]['med'], 4)
                    writer.writerow([fold, epoch, 'SOI', soi_avg, soi_med])

                    cos_avg = round(crossfold_cos[fold][epoch]['avg'], 4)
                    cos_med = round(crossfold_cos[fold][epoch]['med'], 4)
                    writer.writerow([fold, epoch, 'Cosine Similarity', cos_avg, cos_med])
        print(f"\n . . . Results saved to {output_file} \n")  
        
        # Collect metrics for std dev calculation
        avg_sim_list = []
        med_sim_list = []
        avg_soi_list = []
        med_soi_list = []
        avg_cos_list = []
        med_cos_list = []

        for fold in crossfold_sim:
            # Find the epoch with the highest median similarity score in the current fold
            max_epoch = max(crossfold_sim[fold], key=lambda epoch: crossfold_soi[fold][epoch]['med'])
            
            max_med_sim = crossfold_sim[fold][max_epoch]['med']
            corr_avg_sim = crossfold_sim[fold][max_epoch]['avg']
            corr_avg_soi = crossfold_soi[fold][max_epoch]['avg']
            corr_med_soi = crossfold_soi[fold][max_epoch]['med']
            corr_avg_cos = crossfold_cos[fold][max_epoch]['avg']
            corr_med_cos = crossfold_cos[fold][max_epoch]['med']

            avg_sim_list.append(corr_avg_sim)
            med_sim_list.append(max_med_sim)
            avg_soi_list.append(corr_avg_soi)
            med_soi_list.append(corr_med_soi)
            avg_cos_list.append(corr_avg_cos)
            med_cos_list.append(corr_med_cos)

            # Print the fold-specific results
            print(f"Fold {fold}: Med Sim = {max_med_sim}, Avg Sim = {corr_avg_sim} ; Avg SOI = {corr_avg_soi}, Med SOI = {corr_med_soi} ; Avg COS = {corr_avg_cos}, Med COS = {corr_med_cos}")

        # Calculate the averages and std devs across all folds
        average_avg_sim = sum(avg_sim_list) / k_folds
        average_med_sim = sum(med_sim_list) / k_folds
        average_avg_soi = sum(avg_soi_list) / k_folds
        average_med_soi = sum(med_soi_list) / k_folds
        average_avg_cos = sum(avg_cos_list) / k_folds
        average_med_cos = sum(med_cos_list) / k_folds

        std_avg_sim = statistics.stdev(avg_sim_list) if k_folds > 1 else 0.0
        std_med_sim = statistics.stdev(med_sim_list) if k_folds > 1 else 0.0
        std_avg_soi = statistics.stdev(avg_soi_list) if k_folds > 1 else 0.0
        std_med_soi = statistics.stdev(med_soi_list) if k_folds > 1 else 0.0
        std_avg_cos = statistics.stdev(avg_cos_list) if k_folds > 1 else 0.0
        std_med_cos = statistics.stdev(med_cos_list) if k_folds > 1 else 0.0

        # Print the averages and std devs
        print(f'{signal_type} {groups[0]} {identifier[0]}')
        print(f"Average of Avg Similarity Score across all folds: {average_avg_sim:.4f} ± {std_avg_sim:.4f}")
        print(f"Average of Med Similarity Score across all folds: {average_med_sim:.4f} ± {std_med_sim:.4f}")
        print(f"Average of Avg SOI across all folds: {average_avg_soi:.4f} ± {std_avg_soi:.4f}")
        print(f"Average of Med SOI across all folds: {average_med_soi:.4f} ± {std_med_soi:.4f}")
        print(f"Average of Avg COS across all folds: {average_avg_cos:.4f} ± {std_avg_cos:.4f}")
        print(f"Average of Med COS across all folds: {average_med_cos:.4f} ± {std_med_cos:.4f}")

        # Fix: set time_end here before using it below
        time_end = time.time()

        print(f"\nTotal time used: {round((time_end - time_start)/60, 2)} minutes")

        # Write summary to file
        summary_file = os.path.join(eval_save_folder, f'{identifier[0]}_summary.txt')
        with open(summary_file, 'w') as f:
            f.write(f'{signal_type} {groups[0]} {identifier[0]}\n')
            f.write(f"Average of Avg Similarity Score across all folds: {average_avg_sim:.4f} ± {std_avg_sim:.4f}\n")
            f.write(f"Average of Med Similarity Score across all folds: {average_med_sim:.4f} ± {std_med_sim:.4f}\n")
            f.write(f"Average of Avg SOI across all folds: {average_avg_soi:.4f} ± {std_avg_soi:.4f}\n")
            f.write(f"Average of Med SOI across all folds: {average_med_soi:.4f} ± {std_med_soi:.4f}\n")
            f.write(f"Average of Avg COS across all folds: {average_avg_cos:.4f} ± {std_avg_cos:.4f}\n")
            f.write(f"Average of Med COS across all folds: {average_med_cos:.4f} ± {std_med_cos:.4f}\n")
            f.write(f"\nTotal time used: {round((time_end - time_start)/60, 2)} minutes\n")
        print(f"Summary written to {summary_file}")