train_utils.py

from configs.global_cfg import DATASET_PATH
import os
import pickle
import pandas as pd
from torch.optim.lr_scheduler import _LRScheduler
import math
import torch
from safetensors.torch import save_file
from lifelines.utils import concordance_index
from torchmetrics import Metric, MetricCollection
from torchmetrics.utilities import dim_zero_cat


def get_splits_from_file(ds, all_configs, data_root, split=None):
    sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects
    if split is None:
        sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
    else:
        sample_dir = os.path.join(data_root, DATASET_PATH[ds], split, sample_modality)

    chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
    info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
    with open(info_path, 'rb') as f:
        sig_info = pickle.load(f)

    subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
    return subjects

def get_datasets(datasets, all_configs, data_root, splits=['train']):
    splits_set = dict()

    for split in splits:
        splits_set[split] = []

    for ds in datasets:
        # Use Physionet 2018 dataset        
        if ds == 'shhs1':
            def get_shhs_lists(shhs_split_file):
                ids = pd.read_csv(os.path.join("data","shhs_meta_data", shhs_split_file), header=None, names=['a'])
                ids = ids['a'].values.tolist()    
                list_files = ["shhs1-20"+id[11:15] for id in ids]
                return list_files
            list_files_val = get_shhs_lists("eval_list.txt")
            list_files_test = get_shhs_lists("test_list.txt")
            
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                missing_cases = ['shhs1-204709','shhs1-204806']

                train_subjects = set(train_subjects) - set(list_files_val) - set(list_files_test) - set(missing_cases)
                train_subjects = sorted(list(train_subjects)) 
                splits_set['train'].extend([('shhs1',f) for f in train_subjects])
            
            if 'val' in splits:
                val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
                val_subjects = set(list_files_val) & set(val_subjects)
                val_subjects = sorted(list(val_subjects))
                splits_set['val'].extend([('shhs1',f) for f in val_subjects])
            
            if 'test' in splits:
                test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
                test_subjects = set(list_files_test) & set(test_subjects)            
                test_subjects = sorted(list(test_subjects))
                splits_set['test'].extend([('shhs1',f) for f in test_subjects])
            
        elif ds == 'kiss' or 'cfs' in ds or 'isruc' in ds or 'sof' in ds:
            
            ds_main = 'kiss' if ds in 'kiss' else 'cfs' if 'cfs' in ds else 'isruc' if 'isruc' in ds else 'sof'            
            
            
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                splits_set['train'].extend([(ds_main,f) for f in train_subjects])
            
            if 'val' in splits:
                val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
                splits_set['val'].extend([(ds_main,f) for f in val_subjects])
            
            if 'test' in splits:
                test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
                splits_set['test'].extend([(ds_main,f) for f in test_subjects])
        elif ds in ['phy-tr', 'phy-te', 'mros1-val', 'mros1-te','mros2-val','mros2-te','mesa-val','mesa-te']:
            # sample_modality = all_configs.get('modalities', ['EEG'])[0]
            # sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
            # chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
            # info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
            # with open(info_path, 'rb') as f:
            #     sig_info = pickle.load(f)

            # all_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split=None)
                splits_set['train'].extend([(ds, f) for f in train_subjects])

        elif ds in ['shhs2', 'kiss','kvss','mros1','mros2','mesa']:
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                splits_set['train'].extend([(ds, f) for f in train_subjects])
            

            # splits_set['train'].extend([(ds, f) for f in all_subjects])

    return splits_set



def export_inference_weights(ckpt_path, model, out_path, metadata=None):
    """
    Strip a Lightning checkpoint down to the weights `model` loads and save them
    as safetensors (no optimizer / scheduler / callback state, no pickle).

    `model` is the inference backbone, e.g. `inference.build_micro_model()` or
    `inference.build_full_night_model()`. Keys are saved without the `model.`
    prefix; `pos_embed` is skipped since it is a fixed sin-cos table rebuilt at
    model construction. Keys absent from the checkpoint (e.g. `head` for a
    pretraining checkpoint) are not saved.

        from inference import build_micro_model
        export_inference_weights("experiments/micro.ckpt", build_micro_model(),
                                 "checkpoints/micro.safetensors", {"task": "micro"})
    """
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    src = ckpt["state_dict"] if "state_dict" in ckpt else ckpt

    weights = {}
    for k in model.state_dict():
        if "pos_embed" in k:
            continue
        if f"model.{k}" in src:
            weights[k] = src[f"model.{k}"].detach().clone().contiguous()
        elif k in src:
            weights[k] = src[k].detach().clone().contiguous()

    meta = {"source_epoch": str(ckpt.get("epoch")), "source_global_step": str(ckpt.get("global_step"))}
    meta.update({k: str(v) for k, v in (metadata or {}).items()})
    save_file(weights, out_path, metadata=meta)
    return weights


def get_layer_decay_param_groups(model, base_lr, weight_decay, layer_decay):
    """
    Create parameter groups with layer-wise learning rate decay.
    Args:
        model: The model whose parameters are grouped.
        base_lr: Base learning rate.
        weight_decay: Weight decay for regularization.
        layer_decay: Layer decay factor (e.g., 0.75).
    Returns:
        A list of parameter groups for the optimizer.
    """
    param_groups = []
    layers = list(model.children())
    num_layers = len(layers)

    # We use named_parameters to filter by name (e.g., A_log, dt, bias)
    # But we need to map them back to their layer index for LR decay.
    for i, layer in enumerate(layers):
        layer_scale = layer_decay ** (num_layers - i - 1)
        curr_lr = base_lr * layer_scale
        
        # We split parameters in each layer into 'decay' and 'no_decay'
        params_to_decay = []
        params_no_decay = []

        for name, param in layer.named_parameters():
            if not param.requires_grad:
                continue
            
            # Check if it's a Mamba-specific parameter or a standard "no-decay" candidate
            # Common Mamba names: A_log, D, dt_proj, x_proj, in_proj, conv1d
            # Standard no-decay: bias, norm weights
            is_ssm_param = any(ssm_name in name for ssm_name in ['A_log', 'D', 'dt', 'B', 'C'])
            is_bias_or_norm = name.endswith("bias") or "norm" in name
            
            if is_ssm_param or is_bias_or_norm:
                params_no_decay.append(param)
            else:
                params_to_decay.append(param)

        # Add groups for this specific layer
        if params_to_decay:
            param_groups.append({
                "params": params_to_decay,
                "lr": curr_lr,
                "weight_decay": weight_decay,
            })
        if params_no_decay:
            param_groups.append({
                "params": params_no_decay,
                "lr": curr_lr,
                "weight_decay": 0.0,
            })

    return param_groups

class WarmupCosineAnnealingLR(_LRScheduler):
    def __init__(self, optimizer, warmup_epochs, max_epochs, eta_min=1e-8, last_epoch=-1):
        """
        Step-based Warmup Cosine Annealing Learning Rate Scheduler.

        Args:
            optimizer: The optimizer for which to schedule the learning rate.
            warmup_epochs: Number of steps for the warmup phase.
            max_epochs: Total number of steps for the scheduler.
            eta_min: Minimum learning rate at the end of the schedule.
            last_epoch: The index of the last step (default: -1).
        """
        self.warmup_epochs = warmup_epochs
        self.max_epochs = max_epochs
        self.eta_min = eta_min
        super(WarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.warmup_epochs:
            # Linear warmup
            return [base_lr * (self.last_epoch + 1) / self.warmup_epochs for base_lr in self.base_lrs]
        else:
            # Cosine annealing
            progress = (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)
            return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2 for base_lr in self.base_lrs]



def cox_ph_loss(hazards, event_times, is_event):
    # Sort event times and get corresponding indices for sorting other tensors
    event_times, sorted_idx = event_times.sort(dim=0, descending=True)
   

    # Make gathering robust: if sorted_idx is 1D and hazards/is_event are 2D,
    # use advanced indexing; otherwise fall back to gather which requires
    # the index tensor to have the same number of dimensions as input.
    if sorted_idx.dim() == 1 and hazards.dim() == 2:
        hazards = hazards[sorted_idx]
        is_event = is_event[sorted_idx]
    else:
        hazards = hazards.gather(0, sorted_idx)
        is_event = is_event.gather(0, sorted_idx)

    log_cumulative_hazard = torch.logcumsumexp(hazards.float(), dim=0)

    # Calculate losses for all labels simultaneously
    losses = (hazards - log_cumulative_hazard) * is_event
    losses = -losses  # Negative for maximization

    # Average loss per label
    label_loss = losses.sum(dim=0) / (is_event.sum(dim=0) + 1e-9)  # Avoid division by zero

    # Average across labels
    total_loss = label_loss.mean()

    return total_loss


def c_index_torch(scores, event_times, is_event):    
    return float(concordance_index(event_times, -scores, is_event))



class CIndex(Metric):
    def __init__(self, num_classes=6, **kwargs):
        super().__init__(**kwargs)
        self.add_state("scores", default=[], dist_reduce_fx="cat")
        self.add_state("event_time", default=[], dist_reduce_fx="cat")
        self.add_state("is_event", default=[], dist_reduce_fx="cat")
        self.num_classes = num_classes

    def update(self, score, t, e) -> None:
        self.scores.append(score)
        self.event_time.append(t)
        self.is_event.append(e)

    def compute(self):

        # parse inputs
        # Convert to float32 before numpy() since numpy doesn't support BFloat16
        scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
        event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
        is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()
        
        from lifelines.utils import concordance_index
        cindex = []
        for i in range(self.num_classes):
            try:
                c = concordance_index(event_time[:,i], -scores[:,i], is_event[:,i])
                cindex.append(c)
            except (ZeroDivisionError, ValueError):
                # Not enough admissible pairs (common during sanity check or with small batches)
                # Return 0.5 (random predictor) as default
                cindex.append(0.5)
        return torch.tensor(cindex)
    
class AUROC(Metric):
    def __init__(self, num_classes=6, **kwargs):
        super().__init__(**kwargs)
        self.add_state("scores", default=[], dist_reduce_fx="cat")
        self.add_state("event_time", default=[], dist_reduce_fx="cat")
        self.add_state("is_event", default=[], dist_reduce_fx="cat")
        self.num_classes = num_classes
        self.measure_date = 365 * 6 # at 6 years

    def update(self, score, t, e) -> None:
        self.scores.append(score)
        self.event_time.append(t)
        self.is_event.append(e)

    def compute(self):

        # parse inputs
        # Convert to float32 before numpy() since numpy doesn't support BFloat16
        scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
        event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
        is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()
        
        from sklearn.metrics import roc_auc_score, roc_curve
        # auroc = []
        # # For each disease, calculate AUROC
        # for i in range(self.num_classes):
        #     y_true = (event_time[:,i] <= self.measure_date).astype(int)
        #     y_scores = scores[:,i]
        #     try:
        #         c = roc_auc_score(y_true, y_scores)
        #         auroc.append(c)
        #     except (ZeroDivisionError, ValueError):
        #         # Not enough admissible pairs (common during sanity check or with small batches)
        #         # Return 0.5 (random predictor) as default
        #         auroc.append(0.5)
        # return torch.tensor(auroc)
        auroc = []
        # For each disease, calculate AUROC at 6-year follow-up (time-dependent)
        for i in range(self.num_classes):
            # Filter out zero times (invalid entries)
            valid_mask = event_time[:, i] > 0
            if valid_mask.sum() < 2:
                auroc.append(0.5)
                continue
            
            valid_scores = scores[valid_mask, i]
            valid_times = event_time[valid_mask, i]
            valid_events = is_event[valid_mask, i]
            
            # Time-dependent AUROC at measure_date:
            # - Samples with event before measure_date: y=1
            # - Samples with no event and observed until after measure_date: y=0
            # - Samples with censoring before measure_date: excluded (potentially informative censoring)
            
            # Consider only samples observed until measure_date
            observed_mask = valid_times >= self.measure_date
            # Or event occurred before measure_date
            event_before_mask = valid_events == 1
            
            # Use all samples with valid time info for AUROC
            # y_true: 1 if event before measure_date, 0 otherwise
            y_true = (valid_events == 1) & (valid_times <= self.measure_date)
            y_true = y_true.astype(int)
            y_scores = valid_scores
            
            try:
                # Check if there's variation in labels
                if len(np.unique(y_true)) < 2:
                    auroc.append(0.5)
                else:
                    c = roc_auc_score(y_true, y_scores)
                    auroc.append(c)
            except (ZeroDivisionError, ValueError) as e:
                # Not enough admissible pairs
                auroc.append(0.5)
        return torch.tensor(auroc)