### Clone Repository and Set Up Conda Environment Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md Clone the SleepFM repository and create a Conda environment with all necessary dependencies. This is the initial setup step for running the project. ```bash git clone https://github.com/zou-group/sleepfm-clinical.git cd sleepfm-clinical conda env create -f env.yml conda activate sleepfm_env ``` -------------------------------- ### Import Core Libraries Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Imports essential Python libraries for data manipulation, machine learning, and deep learning. Ensure these libraries are installed. ```python import torch from torch import nn import numpy as np from sklearn.metrics import f1_score, confusion_matrix import os import tqdm import random import matplotlib.pyplot as plt import seaborn as sns import sys from collections import Counter import pandas as pd sys.path.append("..") sys.path.append("../sleepfm") from preprocessing.preprocessing import EDFToHDF5Converter from models.dataset import SetTransformerDataset, collate_fn from models.models import SetTransformer, SleepEventLSTMClassifier, DiagnosisFinetuneFullLSTMCOXPHWithDemo import h5py from utils import load_config, load_data, save_data, count_parameters from torch.utils.data import Dataset, DataLoader ``` -------------------------------- ### Output Directory Setup for Embeddings Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Creates output directories for storing generated embeddings and aggregated embeddings. Ensure the base save path is correctly defined. ```python output = os.path.join(base_save_path, "demo_emb") output_5min_agg = os.path.join(base_save_path, "demo_5min_agg_emb") os.makedirs(output, exist_ok=True) os.makedirs(output_5min_agg, exist_ok=True) ``` -------------------------------- ### Initialize and Print Sleep Staging Model Info Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Initializes the sleep staging model and prints its trainable parameters and number of layers. This is useful for verifying model setup. ```python print(f"Model initialized: {sleep_staging_model_name}") total_layers, total_params = count_parameters(sleep_staging_model) print(f'Trainable parameters: {total_params / 1e6:.2f} million') print(f'Number of layers: {total_layers}') ``` -------------------------------- ### Example Sleep Stage Labels Format Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md Illustrates the expected format for sleep stage labels, which are used for training and evaluation. Ensure your label files match this structure. ```csv Start,Stop,StageName,StageNumber 0.0,5190.0,Wake,0 0.0,5190.0,Wake,0, 0.0,5190.0,Wake,0 ``` -------------------------------- ### Get Item from Dataset Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Retrieves data for a given index from the dataset. It reads HDF5 files, extracts relevant features, converts them to PyTorch tensors, and returns them along with labels and metadata. ```python def __getitem__(self, idx): hdf5_path, tte_event = self.index_map[idx] event_time = tte_event["event_time"] is_event = tte_event["is_event"] demo_feats = tte_event["demo_feats"] x_data = [] with h5py.File(hdf5_path, 'r') as hf: dset_names = [] for dset_name in hf.keys(): if isinstance(hf[dset_name], h5py.Dataset) and dset_name in self.config["modality_types"]: dset_names.append(dset_name) random.shuffle(dset_names) for dataset_name in dset_names: x_data.append(hf[dataset_name][:]) if not x_data: # Skip this data point if x_data is empty return self.__getitem__((idx + 1) % self.total_len) # Convert x_data list to a single numpy array x_data = np.array(x_data) # Convert x_data to tensor x_data = torch.tensor(x_data, dtype=torch.float32) event_time = torch.tensor(event_time, dtype=torch.float32) is_event = torch.tensor(is_event) demo_feats = torch.tensor(demo_feats, dtype=torch.float32) return x_data, event_time, is_event, demo_feats, self.max_channels, self.max_seq_len, hdf5_path ``` -------------------------------- ### Get Shape of Flattened Arrays Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Retrieves the shapes of the flattened data arrays after concatenation. Useful for verifying data dimensions. ```python all_logits_flat.shape, all_outputs_flat.shape, all_targets_flat.shape, all_masks_flat.shape ``` -------------------------------- ### Initialize Output Directory Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Creates a directory to store demo data. Ensure the path is writable. ```python base_save_path = "demo_data" os.makedirs(base_save_path, exist_ok=True) ``` -------------------------------- ### Initialize EDF to HDF5 Converter Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Sets up the converter for EDF files. Dummy root and target directories are provided as they are not used for single file conversion in this demo. ```python base_save_path = "demo_data" os.makedirs(base_save_path, exist_ok=True) root_dir = "/edf_root" # dummy root not used for a single file conversion target_dir = "/note" # dummy target not used for a single file conversion edf_path = "demo_data/demo_psg.edf" hdf5_path = os.path.join(base_save_path, "demo_psg.hdf5") converter = EDFToHDF5Converter( root_dir=root_dir, target_dir=target_dir, resample_rate=128 ) # run for single file conversion converter.convert(edf_path, hdf5_path) ``` -------------------------------- ### Initialize Diagnosis Dataset Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Sets up the dataset for disease prediction fine-tuning. Specifies configuration, channel groups, split, HDF5 file paths, and demo labels. ```python hdf5_paths = [os.path.join(base_save_path, "demo_emb/demo_psg.hdf5")] demo_labels_path = os.path.join(base_save_path, "demo_age_gender.csv") config["labels_path"] = base_save_path test_dataset = DiagnosisFinetuneFullCOXPHWithDemoDataset(config, channel_groups, split="test", hdf5_paths=hdf5_paths, demo_labels_path=demo_labels_path) ``` -------------------------------- ### Initialize and Configure SleepFM Model Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Instantiates the SleepFM model based on the extracted parameters and moves it to the appropriate device (GPU if available). It also prints the total trainable parameters and layers. ```python model_class = getattr(sys.modules[__name__], config['model']) model = model_class(in_channels, patch_size, embed_dim, num_heads, num_layers, pooling_head=pooling_head, dropout=dropout) device = torch.device("cuda") if device.type == "cuda": model = torch.nn.DataParallel(model) model.to(device) total_layers, total_params = count_parameters(model) print(f'Trainable parameters: {total_params / 1e6:.2f} million') print(f'Number of layers: {total_layers}') ``` -------------------------------- ### Parallelize Model and Print Info Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Wraps the initialized model with `nn.DataParallel` for multi-GPU training and prints model initialization status, trainable parameters, and layer count. Requires a `count_parameters` function. ```python model = nn.DataParallel(model) print(f"Model initialized: {model_name}") total_layers, total_params = count_parameters(model) print(f'Trainable parameters: {total_params / 1e6:.2f} million') print(f'Number of layers: {total_layers}') ``` -------------------------------- ### Load Finetuned Sleep Staging Model Configuration Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the configuration for a finetuned sleep staging model from a JSON file. Initializes the model class and parameters based on the loaded configuration. ```python sleep_staging_model_path = "../sleepfm/checkpoints/model_sleep_staging" sleep_staging_config = load_data(os.path.join(sleep_staging_model_path, "config.json")) sleep_staging_model_params = sleep_staging_config['model_params'] sleep_staging_model_class = getattr(sys.modules[__name__], sleep_staging_config['model']) sleep_staging_model = sleep_staging_model_class(**sleep_staging_model_params).to(device) sleep_staging_model_name = type(sleep_staging_model).__name__ ``` -------------------------------- ### Load Model Configuration and Channels Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the configuration and channel group data for the pretrained SleepFM model. Ensure the paths to checkpoints and configs are correct. ```python model_path = "../sleepfm/checkpoints/model_base" channel_groups_path = "../sleepfm/configs/channel_groups.json" config_path = os.path.join(model_path, "config.json") config = load_config(config_path) channel_groups = load_data(channel_groups_path) ``` -------------------------------- ### Configuration Object Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Displays the configuration object used in the notebook. This is useful for understanding the parameters and settings applied to the model or process. ```python config ``` -------------------------------- ### Create Save Directory Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Creates a directory to save prediction outputs. Ensures the directory exists, creating it if necessary. ```python save_path = os.path.join(base_save_path, "demo_diagnosis") os.makedirs(save_path, exist_ok=True) ``` -------------------------------- ### Define Class Labels and Mapping Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Sets up a list of class labels and creates a dictionary mapping these labels to their corresponding integer indices. This is crucial for consistent interpretation of model outputs and targets. ```python class_labels = ["Wake", "Stage 1", "Stage 2", "Stage 3", "REM"] # class_labels = ["No-Apnea", "Apnea"] class_mapping = {label: idx for idx, label in enumerate(class_labels)} ``` -------------------------------- ### Initialize Disease Prediction Model Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Configures and initializes the disease prediction model based on loaded parameters. It dynamically retrieves the model class and moves the model to the appropriate device. ```python config["model_params"]["dropout"] = 0.0 model_params = config['model_params'] model_class = getattr(sys.modules[__name__], config['model']) model = model_class(**model_params).to(device) model_name = type(model).__name__ ``` -------------------------------- ### Specify Checkpoint Path Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Defines the file path to the pre-trained model checkpoint. Ensure the path is correct. ```python checkpoint_path = os.path.join(disease_model_path, "best.pth") ``` -------------------------------- ### Create DataLoader Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Initializes a DataLoader for the test dataset. Configures batch size, shuffling, number of workers, and a custom collate function. ```python test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=1, collate_fn=diagnosis_finetune_full_coxph_with_demo_collate_fn) ``` -------------------------------- ### Load Model State Dictionary Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the saved state dictionary from a checkpoint file into the initialized model. This step is essential for resuming training or performing inference with pre-trained weights. ```python checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint) ``` -------------------------------- ### Load Pretrained Model Weights Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the best performing weights from a checkpoint file into the initialized SleepFM model and sets the model to evaluation mode. This prepares the model for inference. ```python checkpoint = torch.load(os.path.join(model_path, "best.pt")) model.load_state_dict(checkpoint["state_dict"]) model.eval() ``` -------------------------------- ### Model Configuration Parameters Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Configuration parameters for the SetTransformer model, including dataset paths, model dimensions, and training settings. ```python Result: {'seed': 42, 'model': 'SetTransformer', 'in_channels': 1, 'batch_size': 128, 'epochs': 1, 'lr': 0.001, 'lr_step_period': 2, 'gamma': 0.1, 'temperature': 0.0, 'momentum': 0.9, 'num_workers': 16, 'embed_dim': 128, 'num_heads': 8, 'num_layers': 6, 'pooling_head': 8, 'dropout': 0.3, 'split_path': 'path_to_/dataset_split.json', 'save_path': 'path_to_/models', 'weight_decay': 0.0, 'mode': 'leave_one_out', 'save_iter': 5000, 'eval_iter': 5000, 'log_interval': 100, 'use_wandb': True, 'BAS_CHANNELS': 10, 'RESP_CHANNELS': 7, 'EKG_CHANNELS': 2, 'EMG_CHANNELS': 4, 'max_files': None, 'val_size': 100, 'sampling_duration': 5, 'sampling_freq': 128, 'patch_size': 640, 'modality_types': ['BAS', 'RESP', 'EKG', 'EMG']} ``` -------------------------------- ### Load Sleep Staging Model Checkpoint Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the best performing checkpoint for the sleep staging model. Ensure the model architecture is defined before loading weights. ```python sleep_staging_checkpoint_path = os.path.join(sleep_staging_model_path, "best.pth") sleep_staging_checkpoint = torch.load(sleep_staging_checkpoint_path) sleep_staging_model.load_state_dict(sleep_staging_checkpoint) ``` -------------------------------- ### Initialize Sleep Event Classification Dataset Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb This snippet shows how to initialize a `SleepEventClassificationDataset` for testing. It specifies HDF5 paths and label files, and sets the split to 'test'. Ensure `sleep_staging_config` and `channel_groups` are defined prior to use. ```python hdf5_paths = [os.path.join(base_save_path, "demo_emb/demo_psg.hdf5")] label_files = [os.path.join(base_save_path, "demo_psg.csv")] test_dataset = SleepEventClassificationDataset(sleep_staging_config, channel_groups, split="test", hdf5_paths=hdf5_paths, label_files=label_files) ``` -------------------------------- ### Create DataLoader for Sleep Staging Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb This code creates a PyTorch DataLoader for the sleep staging dataset. It uses a batch size of 8, disables shuffling, and specifies the custom `sleep_event_finetune_full_collate_fn` for batch collation. Set `num_workers` according to your system's capabilities. ```python test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=1, collate_fn=sleep_event_finetune_full_collate_fn) ``` -------------------------------- ### Parallelize Sleep Staging Model Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Wraps the sleep staging model with DataParallel for multi-GPU training and prints the number of GPUs being used. ```python sleep_staging_model = nn.DataParallel(sleep_staging_model) print(f"Using {torch.cuda.device_count()} GPUs") ``` -------------------------------- ### Display Mapped Labels Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Shows the first few rows of the labels DataFrame, now augmented with model predictions, event status, and event times. ```python labels_df.head() ``` -------------------------------- ### Load SleepFM Dataset and DataLoader Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the dataset using SetTransformerDataset and creates a DataLoader for batch processing. Ensure the base save path and dataset configuration are correctly set. ```python hdf5_paths = [os.path.join(base_save_path, "demo_psg.hdf5")] dataset = SetTransformerDataset(config, channel_groups, hdf5_paths=hdf5_paths, split="test") dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, num_workers=1, shuffle=False, collate_fn=collate_fn) ``` -------------------------------- ### Load Model Configuration Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the configuration file for the disease prediction model. Assumes the config file is in JSON format. ```python disease_model_path = "../sleepfm/checkpoints/model_diagnosis" config = load_data(os.path.join(disease_model_path, "config.json")) ``` -------------------------------- ### Extract Model Parameters from Config Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Retrieves model architecture parameters from the loaded configuration. These parameters define the structure of the SleepFM model. ```python modality_types = config["modality_types"] in_channels = config["in_channels"] patch_size = config["patch_size"] embed_dim = config["embed_dim"] num_heads = config["num_heads"] num_layers = config["num_layers"] pooling_head = config["pooling_head"] dropout = 0.0 ``` -------------------------------- ### Collate Function for Diagnosis Finetuning Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Pads and stacks data from a batch of samples for diagnosis finetuning. It handles variable sequence lengths and channel numbers, creating masks for padded regions. ```python def diagnosis_finetune_full_coxph_with_demo_collate_fn(batch): x_data, event_time, is_event, demo_feats, max_channels_list, max_seq_len_list, hdf5_path_list = zip(*batch) num_channels = max(max_channels_list) if max_seq_len_list[0] == None: max_seq_len = max([item.size(1) for item in x_data]) else: max_seq_len = max_seq_len_list[0] padded_x_data = [] padded_mask = [] for item in x_data: c, s, e = item.size() c = min(c, num_channels) s = min(s, max_seq_len) # Ensure the sequence length doesn't exceed max_seq_len # Create a padded tensor and a mask tensor padded_item = torch.zeros((num_channels, max_seq_len, e)) mask = torch.ones((num_channels, max_seq_len)) # Copy the actual data to the padded tensor and set the mask for real data padded_item[:c, :s, :e] = item[:c, :s, :e] mask[:c, :s] = 0 # 0 for real data, 1 for padding padded_x_data.append(padded_item) padded_mask.append(mask) # Stack all tensors into a batch x_data = torch.stack(padded_x_data) event_time = torch.stack(event_time) is_event = torch.stack(is_event) demo_feats = torch.stack(demo_feats) padded_mask = torch.stack(padded_mask) return x_data, event_time, is_event, demo_feats, padded_mask, hdf5_path_list ``` -------------------------------- ### Load Autoreload Extension Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the autoreload extension for interactive development. This is typically used in Jupyter environments. ```python %load_ext autoreload %autoreload 2 ``` -------------------------------- ### SleepFM Model Architecture Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Displays the detailed architecture of the SleepFM model, including its various layers for patch embedding, spatial pooling, and transformer encoding. This is useful for understanding the model's structure and how it processes input data. ```python Result: DataParallel( (module): SetTransformer( (patch_embedding): Tokenizer( (tokenizer): Sequential( (0): Conv1d(1, 4, kernel_size=(5,), stride=(2,), padding=(2,)) (1): BatchNorm1d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (2): ELU(alpha=1.0) (3): LayerNorm((4, 320), eps=1e-05, elementwise_affine=True) (4): Conv1d(4, 8, kernel_size=(5,), stride=(2,), padding=(2,)) (5): BatchNorm1d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (6): ELU(alpha=1.0) (7): LayerNorm((8, 160), eps=1e-05, elementwise_affine=True) (8): Conv1d(8, 16, kernel_size=(5,), stride=(2,), padding=(2,)) (9): BatchNorm1d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (10): ELU(alpha=1.0) (11): LayerNorm((16, 80), eps=1e-05, elementwise_affine=True) (12): Conv1d(16, 32, kernel_size=(5,), stride=(2,), padding=(2,)) (13): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (14): ELU(alpha=1.0) (15): LayerNorm((32, 40), eps=1e-05, elementwise_affine=True) (16): Conv1d(32, 64, kernel_size=(5,), stride=(2,), padding=(2,)) (17): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (18): ELU(alpha=1.0) (19): LayerNorm((64, 20), eps=1e-05, elementwise_affine=True) (20): Conv1d(64, 128, kernel_size=(5,), stride=(2,), padding=(2,)) (21): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (22): ELU(alpha=1.0) (23): LayerNorm((128, 10), eps=1e-05, elementwise_affine=True) (24): AdaptiveAvgPool1d(output_size=1) (25): Flatten(start_dim=1, end_dim=-1) (26): Linear(in_features=128, out_features=128, bias=True) ) ) (spatial_pooling): AttentionPooling( (transformer_layer): TransformerEncoderLayer( (self_attn): MultiheadAttention( (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True) ) (linear1): Linear(in_features=128, out_features=2048, bias=True) (dropout): Dropout(p=0.0, inplace=False) (linear2): Linear(in_features=2048, out_features=128, bias=True) (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (dropout1): Dropout(p=0.0, inplace=False) (dropout2): Dropout(p=0.0, inplace=False) ) ) (positional_encoding): PositionalEncoding() (layer_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (transformer_encoder): TransformerEncoder( (layers): ModuleList( (0-5): 6 x TransformerEncoderLayer( (self_attn): MultiheadAttention( (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True) ) (linear1): Linear(in_features=128, out_features=2048, bias=True) (dropout): Dropout(p=0.0, inplace=False) (linear2): Linear(in_features=2048, out_features=128, bias=True) (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (dropout1): Dropout(p=0.0, inplace=False) (dropout2): Dropout(p=0.0, inplace=False) ) ) ) (temporal_pooling): AttentionPooling( (transformer_layer): TransformerEncoderLayer( (self_attn): MultiheadAttention( (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True) ) (linear1): Linear(in_features=128, out_features=2048, bias=True) (dropout): Dropout(p=0.0, inplace=False) (linear2): Linear(in_features=2048, out_features=128, bias=True) (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True) (dropout1): Dropout(p=0.0, inplace=False) (dropout2): Dropout(p=0.0, inplace=False) ) ) ) ) ``` -------------------------------- ### Sleep Staging Model Validation Loop Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb This snippet outlines the validation loop for a sleep staging model. It iterates through the test DataLoader, performs inference, and collects targets, logits, outputs, masks, and file paths. The results are then saved using a `save_data` function. Ensure the model, device, and `tqdm` are imported and initialized. ```python # Validation loop at the end of each epoch model.eval() all_targets = [] all_logits = [] all_outputs = [] all_masks = [] all_paths = [] count = 0 with torch.no_grad(): for (x_data, y_data, padded_matrix, hdf5_path_list) in tqdm.tqdm(test_loader, desc="Evaluating"): x_data, y_data, padded_matrix, hdf5_path_list = x_data.to(device), y_data.to(device), padded_matrix.to(device), list(hdf5_path_list) outputs, mask = sleep_staging_model(x_data, padded_matrix) all_targets.append(y_data.cpu().numpy()) all_outputs.append(torch.softmax(outputs, dim=-1).cpu().numpy()) all_logits.append(outputs.cpu().numpy()) all_masks.append(mask.cpu().numpy()) all_paths.append(hdf5_path_list) save_path = os.path.join(base_save_path, "demo_sleep_staging") os.makedirs(save_path, exist_ok=True) targets_path = os.path.join(save_path, "all_targets.pickle") outputs_path = os.path.join(save_path, "all_outputs.pickle") logits_path = os.path.join(save_path, "all_logits.pickle") mask_path = os.path.join(save_path, "all_masks.pickle") file_paths = os.path.join(save_path, "all_paths.pickle") save_data(all_targets, targets_path) save_data(all_outputs, outputs_path) save_data(all_logits, logits_path) save_data(all_masks, mask_path) save_data(all_paths, file_paths) ``` -------------------------------- ### Compute and Visualize F1 Scores and Confusion Matrix Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Calculates F1 scores for each class and generates a normalized confusion matrix with percentages. Requires `f1_score` and `confusion_matrix` from scikit-learn, and plotting libraries like matplotlib and seaborn. Ensure `class_labels` are defined correctly. ```python # Step 1: Get predicted labels (argmax on probabilities) predicted_labels = np.argmax(all_outputs_filtered, axis=1) fontsize = 12 # Step 2: Compute F1 score for each class f1_scores = f1_score(all_targets_filtered, predicted_labels, average=None, labels=range(len(class_labels))) for idx, label in enumerate(class_labels): print(f"F1 Score for {label}: {f1_scores[idx]:.3f}") # Step 3: Create a confusion matrix and normalize it by row to get percentages conf_matrix = confusion_matrix(all_targets_filtered, predicted_labels, labels=range(len(class_labels))) conf_matrix_percent = conf_matrix / conf_matrix.sum(axis=1, keepdims=True) * 100 # Plotting the confusion matrix with percentages plt.figure(figsize=(6, 4)) sns.heatmap( conf_matrix_percent, annot=True, fmt=".1f", cmap="Blues", xticklabels=class_labels, yticklabels=class_labels, annot_kws={\"size\": fontsize}, # Font size for numbers inside the heatmap cbar_kws={\"shrink\": 1}, # Adjust colorbar size ) # Customizing axis labels and ticks plt.xlabel("Predicted Labels", fontsize=fontsize) plt.ylabel("True Labels", fontsize=fontsize) plt.xticks(fontsize=12, ha="center") # Font size for x-axis tick labels with rotation plt.yticks(fontsize=12) # Font size for y-axis tick labels # Adjust layout and save the figure plt.tight_layout() plt.show() ``` -------------------------------- ### Load Disease Labels Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads the mapping between disease indices and their corresponding phecodes and phenotypes from a CSV file. ```python labels_df = pd.read_csv("../sleepfm/configs/label_mapping.csv") ``` -------------------------------- ### Detailed Shapes of Model Outputs and Targets Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb This snippet provides the detailed shapes of all collected outputs from the validation loop: logits, softmax outputs, targets, and masks. This is useful for debugging and ensuring data consistency before further processing. ```python all_logits[0].shape, all_outputs[0].shape, all_targets[0].shape, all_masks[0].shape ``` -------------------------------- ### SleepFM Project BibTeX Entry Source: https://github.com/zou-group/sleepfm-clinical/blob/main/README.md The BibTeX entry for the SleepFM paper. Use this for citing the project in academic work. ```bibtex @article{thapa2026multimodal, title={A multimodal sleep foundation model for disease prediction}, author={Thapa, Rahul and Kjaer, Magnus Ruud and He, Bryan and Covert, Ian and Moore IV, Hyatt and Hanif, Umaer and Ganjoo, Gauri and Westover, M Brandon and Jennum, Poul and Brink-Kjaer, Andreas and others}, journal={Nature Medicine}, pages={1--11}, year={2026}, publisher={Nature Publishing Group US New York} } ``` -------------------------------- ### Map Predictions to Labels Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Integrates the model's prediction outputs, event indicators, and event times into the labels DataFrame for further analysis. ```python labels_df["output"] = all_outputs[0] labels_df["is_event"] = all_is_event[0] labels_df["event_time"] = all_event_times[0] ``` -------------------------------- ### Load and Prepare Clinical Data for Disease Prediction Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Loads disease labels from CSV files and resolves HDF5 paths for a given split. Filters studies based on available labels and optionally truncates the dataset. Builds an index map linking HDF5 file paths to their corresponding labels. ```python is_event_df = pd.read_csv(os.path.join(self.config["labels_path"], "is_event.csv")) event_time_df = pd.read_csv(os.path.join(self.config["labels_path"], "time_to_event.csv")) is_event_df = is_event_df.set_index('Study ID') event_time_df = event_time_df.set_index('Study ID') # --- Resolve HDF5 paths (explicit precedence) --- if hdf5_paths: # Use provided paths directly hdf5_paths = [f for f in hdf5_paths if os.path.exists(f)] else: # Load from split file split_paths = load_data(config["split_path"])[split] hdf5_paths = [f for f in split_paths if os.path.exists(f)] # Filter by available demo labels hdf5_paths = [ f for f in hdf5_paths if os.path.basename(f).split(".")[0] in study_ids ] # Optional truncation if config.get("max_files"): hdf5_paths = hdf5_paths[:config["max_files"]] labels_dict = {} # Loop over each study_id for study_id in tqdm.tqdm(study_ids): # Extract the row as a whole for both dataframes (faster than iterating over columns) is_event_row = list(is_event_df.loc[study_id].values) event_time_row = list(event_time_df.loc[study_id].values) demo_feats = list(demo_labels_df.loc[study_id].values) # values = [[event_time, is_event] for is_event, event_time in zip(is_event_row, event_time_row)] labels_dict[study_id] = { "is_event": is_event_row, "event_time": event_time_row, "demo_feats": demo_feats } # --- Build index map --- self.index_map = [ (path, labels_dict[os.path.basename(path).split(".")[0]]) for path in hdf5_paths ] print(f"Number of files in {split} set: {len(hdf5_paths)}") print(f"Number of files to be processed in {split} set: {len(self.index_map)}") self.total_len = len(self.index_map) self.max_seq_len = config["model_params"]["max_seq_length"] if self.total_len == 0: raise ValueError(f"No valid HDF5 files found for split='{split}'.") ``` -------------------------------- ### Model Evaluation Loop Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Evaluates the trained model on the test dataset. It iterates through the DataLoader, performs inference, and collects outputs, event times, and event indicators. ```python model.eval() all_event_times = [] all_is_event = [] all_outputs = [] all_paths = [] with torch.no_grad(): for item in tqdm.tqdm(test_loader, desc="Evaluating"): x_data, event_times, is_event, demo_feats, padded_matrix, hdf5_path_list = item x_data, event_times, is_event, demo_feats, padded_matrix, hdf5_path_list = x_data.to(device), event_times.to(device), is_event.to(device), demo_feats.to(device), padded_matrix.to(device), list(hdf5_path_list) outputs = model(x_data, padded_matrix, demo_feats) logits = outputs.cpu().numpy() all_outputs.append(logits) all_event_times.append(event_times.cpu().numpy()) all_is_event.append(is_event.cpu().numpy()) all_paths.append(hdf5_path_list) all_outputs = np.concatenate(all_outputs, axis=0) all_event_times = np.concatenate(all_event_times, axis=0) all_is_event = np.concatenate(all_is_event, axis=0) all_paths = np.concatenate(all_paths) outputs_path = os.path.join(save_path, "all_outputs.pickle") event_times_path = os.path.join(save_path, "all_event_times.pickle") is_event_path = os.path.join(save_path, "all_is_event.pickle") file_paths = os.path.join(save_path, "all_paths.pickle") save_data(all_outputs, outputs_path) save_data(all_event_times, event_times_path) save_data(all_is_event, is_event_path) save_data(all_paths, file_paths) ``` -------------------------------- ### Shape of Model Outputs Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb This snippet displays the shapes of the collected model outputs and targets after the validation loop. It confirms the dimensions for logits, outputs, targets, and masks, which are crucial for subsequent analysis or evaluation. ```python all_outputs[0].shape, all_targets[0].shape ``` -------------------------------- ### Custom Sleep Event Classification Dataset Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb A custom PyTorch Dataset class for loading and preprocessing sleep event classification data from HDF5 files and label CSVs. It handles data indexing, context windowing, and channel selection. ```python class SleepEventClassificationDataset(Dataset): def __init__( self, config, channel_groups, hdf5_paths, label_files, split="train", ): self.config = config self.max_channels = self.config["max_channels"] self.context = int(self.config["context"]) self.channel_like = self.config["channel_like"] self.max_seq_len = config["model_params"]["max_seq_length"] # --- Build label lookup: {study_id: label_csv_path} --- # study_id = filename without extension, e.g. "SSC_12345" labels_dict = { os.path.basename(p).rsplit(".", 1)[0]: p for p in label_files if os.path.exists(p) } # --- Filter to HDF5s that exist and have a matching label file --- hdf5_paths = [p for p in hdf5_paths if os.path.exists(p)] hdf5_paths = [ p for p in hdf5_paths if os.path.basename(p).rsplit(".", 1)[0] in labels_dict ] if config.get("max_files"): hdf5_paths = hdf5_paths[: config["max_files"]] self.hdf5_paths = hdf5_paths self.labels_dict = labels_dict # --- Build index map --- # Each item is (hdf5_path, label_path, start_index) if self.context == -1: self.index_map = [ (p, labels_dict[os.path.basename(p).rsplit(".", 1)[0]], -1) for p in self.hdf5_paths ] else: self.index_map = [] loop = tqdm(self.hdf5_paths, total=len(self.hdf5_paths), desc=f"Indexing {split} data") for hdf5_file_path in loop: file_prefix = os.path.basename(hdf5_file_path).rsplit(".", 1)[0] label_path = labels_dict[file_prefix] with h5py.File(hdf5_file_path, "r") as hf: dset_names = list(hf.keys()) if len(dset_names) == 0: continue # Use first dataset to define length (same as your original behavior) first_name = dset_names[0] dataset_length = hf[first_name].shape[0] for i in range(0, dataset_length, self.context): self.index_map.append((hdf5_file_path, label_path, i)) # If you have logger, keep; otherwise you can remove these. # logger.info(f"Number of files in {split} set: {len(self.hdf5_paths)}") # logger.info(f"Number of files to be processed in {split} set: {len(self.index_map)}") self.total_len = len(self.index_map) def __len__(self): return self.total_len def get_index_map(self): return self.index_map def __getitem__(self, idx): hdf5_path, label_path, start_index = self.index_map[idx] labels_df = pd.read_csv(label_path) labels_df["StageNumber"] = labels_df["StageNumber"].replace(-1, 0) y_data = labels_df["StageNumber"].to_numpy() if self.context != -1: y_data = y_data[start_index : start_index + self.context] x_data = [] with h5py.File(hdf5_path, "r") as hf: dset_names = list(hf.keys()) for dataset_name in dset_names: if dataset_name in self.channel_like: if self.context == -1: x_data.append(hf[dataset_name][:]) else: x_data.append(hf[dataset_name][start_index : start_index + self.context]) if not x_data: # Skip this data point if x_data is empty return self.__getitem__((idx + 1) % self.total_len) x_data = np.array(x_data) # (C, T, F) assuming each channel returns (T, F) x_data = torch.tensor(x_data, dtype=torch.float32) y_data = torch.tensor(y_data, dtype=torch.float32) min_length = min(x_data.shape[1], len(y_data)) x_data = x_data[:, :min_length, :] y_data = y_data[:min_length] return x_data, y_data, self.max_channels, self.max_seq_len, hdf5_path ``` -------------------------------- ### Extract and Save Embeddings from Clinical Data Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Processes batches of clinical data to extract and save both 5-second level and 5-minute aggregated embeddings into HDF5 files. Requires model, dataloader, device, and configuration to be set up. ```python with torch.no_grad(): with tqdm.tqdm(total=len(dataloader)) as pbar: for batch in dataloader: batch_data, mask_list, file_paths, dset_names_list, chunk_starts = batch (bas, resp, ekg, emg) = batch_data (mask_bas, mask_resp, mask_ekg, mask_emg) = mask_list bas = bas.to(device, dtype=torch.float) resp = resp.to(device, dtype=torch.float) ekg = ekg.to(device, dtype=torch.float) emg = emg.to(device, dtype=torch.float) mask_bas = mask_bas.to(device, dtype=torch.bool) mask_resp = mask_resp.to(device, dtype=torch.bool) mask_ekg = mask_ekg.to(device, dtype=torch.bool) mask_emg = mask_emg.to(device, dtype=torch.bool) embeddings = [ model(bas, mask_bas), model(resp, mask_resp), model(ekg, mask_ekg), model(emg, mask_emg), ] # Model gives two kinds of embeddings. Granular 5 second-level embeddings and aggregated 5 minute-level embeddings. We save both of them below. embeddings_new = [e[0].unsqueeze(1) for e in embeddings] for i in range(len(file_paths)): file_path = file_paths[i] chunk_start = chunk_starts[i] subject_id = os.path.basename(file_path).split('.')[0] output_path = os.path.join(output_5min_agg, f"{subject_id}.hdf5") with h5py.File(output_path, 'a') as hdf5_file: for modality_idx, modality_type in enumerate(config["modality_types"]): if modality_type in hdf5_file: dset = hdf5_file[modality_type] chunk_start_correct = chunk_start // (embed_dim * 5 * 60) chunk_end = chunk_start_correct + embeddings_new[modality_idx][i].shape[0] if dset.shape[0] < chunk_end: dset.resize((chunk_end,) + embeddings_new[modality_idx][i].shape[1:]) dset[chunk_start_correct:chunk_end] = embeddings_new[modality_idx][i].cpu().numpy() else: hdf5_file.create_dataset(modality_type, data=embeddings_new[modality_idx][i].cpu().numpy(), chunks=(embed_dim,) + embeddings_new[modality_idx][i].shape[1:], maxshape=(None,) + embeddings_new[modality_idx][i].shape[1:]) embeddings_new = [e[1] for e in embeddings] for i in range(len(file_paths)): file_path = file_paths[i] chunk_start = chunk_starts[i] subject_id = os.path.basename(file_path).split('.')[0] output_path = os.path.join(output, f"{subject_id}.hdf5") with h5py.File(output_path, 'a') as hdf5_file: for modality_idx, modality_type in enumerate(config["modality_types"]): if modality_type in hdf5_file: dset = hdf5_file[modality_type] chunk_start_correct = chunk_start // (embed_dim * 5) chunk_end = chunk_start_correct + embeddings_new[modality_idx][i].shape[0] if dset.shape[0] < chunk_end: dset.resize((chunk_end,) + embeddings_new[modality_idx][i].shape[1:]) dset[chunk_start_correct:chunk_end] = embeddings_new[modality_idx][i].cpu().numpy() else: hdf5_file.create_dataset(modality_type, data=embeddings_new[modality_idx][i].cpu().numpy(), chunks=(embed_dim,) + embeddings_new[modality_idx][i].shape[1:], maxshape=(None,) + embeddings_new[modality_idx][i].shape[1:]) pbar.update() ``` -------------------------------- ### Define Custom Dataset Class Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Defines a custom PyTorch Dataset class for loading and preprocessing data for disease prediction. It handles loading demographic features from a CSV file and setting the 'Study ID' as the index. ```python class DiagnosisFinetuneFullCOXPHWithDemoDataset(Dataset): def __init__(self, config, channel_groups, hdf5_paths=None, demo_labels_path=None, split="train"): self.config = config self.channel_groups = channel_groups self.max_channels = self.config["max_channels"] # --- Load demographic features --- if not demo_labels_path: demo_labels_path = config["demo_labels_path"] demo_labels_df = pd.read_csv(demo_labels_path) demo_labels_df = demo_labels_df.set_index("Study ID") study_ids = set(demo_labels_df.index) ``` -------------------------------- ### Inspect Prediction Shapes Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Checks the shapes of the collected prediction outputs, event times, and event indicators to ensure data integrity after evaluation. ```python all_outputs.shape, all_event_times.shape, all_is_event.shape ``` -------------------------------- ### Filter Data Using Mask Source: https://github.com/zou-group/sleepfm-clinical/blob/main/notebooks/demo.ipynb Applies a binary mask to filter specific elements from flattened data arrays. The mask should have a boolean dtype. ```python mask_filter = all_masks_flat == 0 # Apply the mask to each flattened array all_logits_filtered = all_logits_flat[mask_filter] all_outputs_filtered = all_outputs_flat[mask_filter] all_targets_filtered = all_targets_flat[mask_filter] ```