From 2f9644a08af2df85ec5ca07edda236ddd1b4d44d Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 24 Mar 2026 00:36:54 +0100 Subject: [PATCH 01/37] Enable training a new head --- src/asap/__init__.py | 2 +- src/asap/task.py | 64 +++++++++++++++++++++++++++++++++++++ src/asap/trainer/trainer.py | 21 +++++++++--- 3 files changed, 81 insertions(+), 6 deletions(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index b412d77..c0c4bdf 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,2 +1,2 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head diff --git a/src/asap/task.py b/src/asap/task.py index 67d2d6c..362acd1 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -75,6 +75,69 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) +def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False): + ''' + Evaluate the model on the given dataset. + Args: + base_experiment_name (str): The name of the model whose weights will be loaded. + new_experiment_name (str): The name of the experiment for the new head. + model (str): The model to evaluate. + eval_dataset: The evaluation dataset. + logs_dir (str): The directory to load model checkpoints from. + batch_size (int): The batch size for evaluation. + use_map (bool): If mappability information was used during training. + ''' + if n_gpus > 0 and not torch.cuda.is_available(): + n_gpus = 0 + print("No GPU available, using CPU instead.") + + # Count the number of GPUs available + if n_gpus > torch.cuda.device_count(): + n_gpus = torch.cuda.device_count() + print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") + + # Initialize the model + model = _get_model(model, use_map=use_map) + + # Initialize the trainer with the model and datasets + trainer = Trainer( + filename=new_experiment_name, + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + linear_probe=True, + ) + + # train the new head based on the previous model + print(f'Loading best model weights from {base_experiment_name}') + checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth' + trainer.load_weights(checkpoint_path) + + # replace head + in_features = trainer.model.core.linear_out.in_features + out_features = trainer.model.core.linear_out.out_features + trainer.model.core.linear_out = nn.Linear(in_features, out_features) + + # freeze everything + for param in trainer.model.parameters(): + param.requires_grad = False + + # unfreeze head + for param in trainer.model.core.linear_out.parameters(): + param.requires_grad = True + + # set modes such that there is no dropout for the core + trainer.model.eval() + trainer.model.core.linear_out.train() + + + # Start training + trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) + print("trained a new model") + def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False): ''' Evaluate the model on the given dataset. @@ -121,6 +184,7 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs scores[chrom] = {key: result_metrics[key] for key in ['pearson_r', 'mse', 'poisson_nll', 'spearman_r', 'kendall_tau']} return scores + def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17): ''' Evaluate the robustness of the model on the given dataset. diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 4875843..32dbf79 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -28,6 +28,7 @@ def __init__(self, batch_size: int = None, logger: Logger = None, n_gpus: int = None, + linear_probe=False, ): self.filename = filename self.model = model @@ -40,6 +41,7 @@ def __init__(self, self.nr_tracks = 1 self.nr_devices = n_gpus self.batch_size = batch_size + self.linear_probe = linear_probe if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' @@ -71,6 +73,7 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.logger, self.filename, self.nr_devices, + self.linear_probe, port ), nprocs=self.nr_devices @@ -100,7 +103,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.unmap_criterion, self.logger, self.filename, - ddp_enabled=False + ddp_enabled=False, + linear_probe=self.linear_probe, ) def predict(self, gen): @@ -276,7 +280,8 @@ def _ddp_and_fit( logger, filename, world_size, - port=12355 + linear_probe, + port=12355, ): #model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = setup_ddp(rank, world_size, model, port) @@ -303,7 +308,8 @@ def _ddp_and_fit( unmap_criterion=unmap_criterion, logger=logger, filename=filename, - ddp_enabled=True + ddp_enabled=True, + linear_probe=linear_probe, ) dist.destroy_process_group() @@ -319,9 +325,14 @@ def _fit( unmap_criterion, logger: Logger, filename: str, - ddp_enabled: bool + ddp_enabled: bool, + linear_probe=False, ): - optimizer = configure_adamw(model, lr=learning_rate) + + if linear_probe: + optimizer = configure_adamw(model.core.linear_out, lr=learning_rate) + else: + optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( optimizer=optimizer, warmup_steps=int(len(train_gen) * 0.25), From 39ca964b183bb00aa12821ae93a43474cf20b679 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 24 Mar 2026 00:40:32 +0100 Subject: [PATCH 02/37] add running instructions --- tutorials/train.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tutorials/train.py b/tutorials/train.py index 3836023..f561661 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -45,5 +45,16 @@ def main(): n_gpus=n_gpus, ) + asap.train_new_head( + base_experiment_name=experiment_name, + new_experiment_name=f"{experiment_name}-new", + model=model_name, + train_dataset=train, + val_dataset=val, + logs_dir=logs_dir, + n_gpus=n_gpus, + max_epochs=15, + ) + if __name__ == "__main__": main() From 7843b0885aa79b59bff59c30f266fc4468eff26d Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Mon, 30 Mar 2026 23:08:24 +0200 Subject: [PATCH 03/37] Two headed model template --- src/asap/__init__.py | 2 +- src/asap/dataloader/base.py | 8 +- src/asap/models/convnext_dcnn.py | 6 +- src/asap/task.py | 100 +++++++++++++++++++++- src/asap/trainer/trainer.py | 96 +++++++++++++++------ tutorials/train.py | 141 +++++++++++++++++++++++++++---- 6 files changed, 303 insertions(+), 50 deletions(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index b412d77..f8aa669 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,2 +1,2 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_multiheaded_model, eval_multihead_model diff --git a/src/asap/dataloader/base.py b/src/asap/dataloader/base.py index 473eba7..e6545c6 100644 --- a/src/asap/dataloader/base.py +++ b/src/asap/dataloader/base.py @@ -31,10 +31,12 @@ def __init__( ): super().__init__() self.genome = genome - if signal_files is not None: - self.signal_files = [signal_files] - else: + if signal_files is None: self.signal_files = None + elif type(signal_files) is list: + self.signal_files = signal_files + else: + self.signal_files = [signal_files] # if we randomly shift to augment the data self.random_shift = random_shift diff --git a/src/asap/models/convnext_dcnn.py b/src/asap/models/convnext_dcnn.py index 37dbbad..a26c8f5 100644 --- a/src/asap/models/convnext_dcnn.py +++ b/src/asap/models/convnext_dcnn.py @@ -20,7 +20,8 @@ 'kernel2': 1, 'dropout': 0.3, 'final_dropout': 0.05, - 'use_map': False + 'use_map': False, + 'num_heads': 1, } @@ -41,7 +42,7 @@ def __init__(self, **kwargs) -> None: if self.use_map: self.unmap_predictor = UnmapPredictor(channels_in=config['filters0']) - self.core = BasenjiCoreBlock(nr_tracks=1, window=window, filters_in=config['filters0'], + self.core = BasenjiCoreBlock(nr_tracks=config['num_heads'], window=window, filters_in=config['filters0'], nr_res_blocks=config['residual_blocks'], rate_mult=config['dilation_mult'], bin_size=config['bin_size'], @@ -93,6 +94,7 @@ def __init__(self, nr_tracks: int, window: int, filters_in, pool_size = bin_size // 2 self.pool = nn.AvgPool1d(kernel_size=pool_size, stride=pool_size) + print("creating output layer with tracks:", nr_tracks) self.linear_out = nn.Linear( in_features=filters3, out_features=nr_tracks, diff --git a/src/asap/task.py b/src/asap/task.py index 67d2d6c..1621a2f 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -13,7 +13,7 @@ from .models import VanillaCNN, CNN_LSTM, DilatedCNN, ConvNextTransformer, ConvNeXtCNN, ConvNeXtLSTM, ConvNeXtDCNN from .snv import make_pcawg_df, add_predictions -def _get_model(model_name: str, use_map: bool = False): +def _get_model(model_name: str, use_map: bool = False, num_heads=1): if model_name == 'cnn': return VanillaCNN(use_map=use_map) elif model_name == 'lstm': @@ -27,7 +27,7 @@ def _get_model(model_name: str, use_map: bool = False): elif model_name == 'convnext_lstm': return ConvNeXtLSTM(use_map=use_map) elif model_name == 'convnext_dcnn': - return ConvNeXtDCNN(use_map=use_map) + return ConvNeXtDCNN(use_map=use_map, num_heads=num_heads) else: raise ValueError(f'Unknown model name: {model_name}') @@ -75,7 +75,52 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False): +def train_multiheaded_model(experiment_name : str, model: str, train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): + """ + Train the model with the given datasets and parameters. + + Args: + experiment_name (str): The name of the experiment. + model (str): The model to train. + train_dataset: The training dataset. + val_dataset: The validation dataset. + logs_dir (str): The directory to save logs. + n_gpus (int): The number of GPUs to use for training. + max_epochs (int): The maximum number of epochs to train. + learning_rate (float): The learning rate for the optimizer. + batch_size (int): The batch size for training. + use_map (bool): Whether to use mappability for training. + num_heads (int): The number of heads (=num of prediction signals) + """ + + # Check if gpu is available + if n_gpus > 0 and not torch.cuda.is_available(): + n_gpus = 0 + print("No GPU available, using CPU instead.") + + # Count the number of GPUs available + if n_gpus > torch.cuda.device_count(): + n_gpus = torch.cuda.device_count() + print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") + + # Initialize the model + model = _get_model(model, use_map=use_map, num_heads=num_heads) + + # Initialize the trainer with the model and datasets + trainer = Trainer( + filename=experiment_name, + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + ) + + # Start training + trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) + +def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, target_head:int = 0): ''' Evaluate the model on the given dataset. Args: @@ -89,7 +134,54 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs n_gpus = 1 if torch.cuda.is_available() else 0 # Initialize the model - model = _get_model(model, use_map=use_map) + model = _get_model(model, use_map=use_map, num_heads=num_heads) + + trainer = Trainer( + filename=experiment_name, + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + ) + + # for evaluation use the checkpoint of the best model + print(f'Loading best model weights from {trainer.filename}') + checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' + trainer.load_weights(checkpoint_path) + + test_chroms = eval_dataset.chroms + scores = {} + for chrom in test_chroms: + eval_dataset.set_chroms([chrom]) + test_gen = make_dataloader( + ddp_enabled=False, + dataset=eval_dataset, + batch_size=batch_size, + is_train=False + ) + + _, _, result_metrics = trainer.predict_and_evaluate_multihead(test_gen, target_head=target_head) + scores[chrom] = {key: result_metrics[key] for key in ['pearson_r', 'mse', 'poisson_nll', 'spearman_r', 'kendall_tau']} + return scores + + +def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1): + ''' + Evaluate the model on the given dataset. + Args: + experiment_name (str): The name of the experiment. + model (str): The model to evaluate. + eval_dataset: The evaluation dataset. + logs_dir (str): The directory to load model checkpoints from. + batch_size (int): The batch size for evaluation. + use_map (bool): If mappability information was used during training. + ''' + n_gpus = 1 if torch.cuda.is_available() else 0 + + # Initialize the model + model = _get_model(model, use_map=use_map, num_heads=num_heads) trainer = Trainer( filename=experiment_name, diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 4875843..121caf2 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -28,6 +28,7 @@ def __init__(self, batch_size: int = None, logger: Logger = None, n_gpus: int = None, + nr_tracks: int = 1, ): self.filename = filename self.model = model @@ -37,7 +38,7 @@ def __init__(self, self.logger: Logger = logger self.logspace = True - self.nr_tracks = 1 + self.nr_tracks = nr_tracks self.nr_devices = n_gpus self.batch_size = batch_size if self.nr_devices > 1: @@ -115,6 +116,43 @@ def evaluate(self, test_gen) -> dict: print(result_metrics) return result_metrics + def predict_and_evaluate_multihead(self, gen, metrics_for_track=None, no_eval=False, target_head=0) -> Tuple[np.ndarray, np.ndarray, dict]: + self.model.eval() + + try: + gen.dataset.margin_size + except AttributeError: + print("Generator has no margin size -- assuming full prediction.") + + predictions, true = self.predict(gen) + num_heads = predictions.shape[-1] + predictions = predictions[..., target_head: target_head+1] + predictions = predictions.reshape((-1, 1)).detach().cpu().numpy() + true = true.reshape((-1, 1)).detach().cpu().numpy() + + if no_eval: + return true, predictions, None + + metric_results = {} + if metrics_for_track is None: + # compute metrics for all tracks + metrics_for_track = range(self.nr_tracks) + metrics_ = compute_metrics( + predictions.flatten(), + true.flatten(), + logspace_input=self.logspace + ) + metric_results.update(metrics_) + else: + for track in metrics_for_track: + metrics_ = compute_metrics( + predictions[:, track], + true[:, track], + logspace_input=self.logspace + ) + metric_results.update(prepend_to_keys(metrics_, f"_{track}_")) + return true, predictions, metric_results + def predict_and_evaluate(self, gen, metrics_for_track=None, no_eval=False) -> Tuple[np.ndarray, np.ndarray, dict]: self.model.eval() @@ -356,28 +394,31 @@ def _fit( logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() - predictions, true = predictions[..., 0].flatten().numpy(), true[..., 0].flatten().numpy() - val_log_payload = compute_metrics( - predictions, - true, - logspace_input=val_gen.dataset.logspace - ) - val_log_payload = prepend_to_keys(val_log_payload, 'val/') - logger.log(val_log_payload, step=epoch) - print(f'Epoch {epoch} - {datetime.now()}') - if train_log_payload is not None: - print(f'\tTrain loss: {train_log_payload["train/loss"]}') - print(f'\tVal pearson r: {val_log_payload["val/pearson_r"]}') - print('-----------------------------------------') - if val_log_payload['val/pearson_r'] > best_val_score: - best_val_score = val_log_payload['val/pearson_r'] - logger.save_model(model, filename) - # handle early stopping - no_improvement_for = 0 - else: - no_improvement_for += 1 - if (epoch != nr_epochs -1) and early_stopping_after_no_improvement and no_improvement_for >= early_stopping_after_no_improvement: - stop_early += 1 + num_tracks = predictions.shape[-1] + for track in range(num_tracks): + print("Head: ", track+1) + predictions_head, true_head = predictions[..., track].flatten().numpy(), true[..., track].flatten().numpy() + val_log_payload = compute_metrics( + predictions_head, + true_head, + logspace_input=val_gen.dataset.logspace + ) + val_log_payload = prepend_to_keys(val_log_payload, 'val/') + logger.log(val_log_payload, step=epoch) + print(f'Epoch {epoch} - {datetime.now()}') + if train_log_payload is not None: + print(f'\tTrain loss: {train_log_payload["train/loss"]}') + print(f'\tVal pearson r: {val_log_payload["val/pearson_r"]}') + print('-----------------------------------------') + if val_log_payload['val/pearson_r'] > best_val_score: + best_val_score = val_log_payload['val/pearson_r'] + logger.save_model(model, filename) + # handle early stopping + no_improvement_for = 0 + else: + no_improvement_for += 1 + if (epoch != nr_epochs -1) and early_stopping_after_no_improvement and no_improvement_for >= early_stopping_after_no_improvement: + stop_early += 1 if ddp_enabled: dist.all_reduce(stop_early) @@ -407,17 +448,22 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ m_i = m_i.to(rank) y_i = y_i.to(rank) optimizer.zero_grad() + nr_tracks = y_i.shape[-1] if train_unmap: output, output_m_i = model(X_i, return_unmap=True) # trim m_i in case of unpadded conv in stem m_len = output_m_i.shape[1] unmap_loss = unmap_criterion(output_m_i, m_i[:, :m_len]) - base_loss = criterion(output, y_i) + base_loss = 0 + for track in range(nr_tracks): + base_loss += criterion(output[..., track], y_i[..., track]) loss = base_loss + unmap_loss else: output = model(X_i) - loss = criterion(output, y_i) + loss = 0 + for track in range(nr_tracks): + loss += criterion(output[..., track], y_i[..., track]) loss.backward() optimizer.step() diff --git a/tutorials/train.py b/tutorials/train.py index 3836023..3e5824c 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -1,32 +1,46 @@ +import sys +import os +import numpy as np +import pyBigWig + +sys.path.append(os.path.abspath("src")) + import asap # The below code is a complete script that sets up the training of a model using the ASAP library. # It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. # The script then trains the model using the specified parameters. + def main(): # Data paths - signal_file = "../data/TCGA-A6-A567/TCGA-A6-A567.nodup.no_chrM_MT.tn5.pval.signal.bigwig" - genome = "../data/hg38.fa" - blacklist_file = ["../data/basenji_blacklist.bed", "../data/example_snv.vcf"] - unmap_file = "../data/basenji_unmappable.bed" - generated = "../tmp" - logs_dir = "../tmp/logs" + signal_file_1 = "data/GM12878/ENCFF667MDI-signal.bigWig" + peak_file_1 = "data/GM12878/ENCFF748UZH-peak.bed" + + signal_file_2 = "data/K562/ENCFF357GNC-signal.bigWig" + peak_file_2 = "data/K562/ENCFF333TAT-peak.bed.gz" + + genome = "data/hg38.fa" + blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] + unmap_file = "data/basenji_unmappable.bed" + generated = "tmp" + logs_dir = "tmp/logs" # Model parameters model_name = "convnext_dcnn" - experiment_name = "TCGA-A6-A567_convnext_dcnn" + experiment_name = "GM12878_K562" # Training parameters test_chroms = [2, 10, 14, 19, 21] - val_chroms = [1, 11, 20, 13] - train_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in val_chroms] - n_gpus = 4 + train_chroms = [1, 11, 20, 13] + val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in val_chroms] + n_gpus = 2 # Create the training and validation datasets - train, val = asap.training_datasets( - signal_file=signal_file, + print("create the dataset") + train_comb, val_comb = asap.training_datasets( + signal_file=[signal_file_1, signal_file_2], genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -35,15 +49,112 @@ def main(): unmap_file=unmap_file, ) + + + + # print(train_comb.X[0].shape) + + # print(train_comb.y[0].shape) + + # print(train_comb.y[0][0]) + + + print("start to train!!!\n") + # Train the model - asap.train_model( + asap.train_multiheaded_model( experiment_name=experiment_name, model=model_name, - train_dataset=train, - val_dataset=val, + num_heads=2, + train_dataset=train_comb, + val_dataset=val_comb, logs_dir=logs_dir, n_gpus=n_gpus, ) + print("Training done.") + print("Create eval ds") + print() + + + peak1 = asap.peak_dataset( + signal_file=signal_file_1, + peak_file=peak_file_1, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + peak2 = asap.peak_dataset( + signal_file=signal_file_2, + peak_file=peak_file_2, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + + # print(len(peak1.X), len(peak1.y)) + # print("y") + # print(peak.y[0][0][100]) + # print(peak1.y[0].shape) + # print(peak2.y[0][0][100]) + + # Evaluate the model + peak_scores_head1 = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak1, + logs_dir=logs_dir, + num_heads=2, + target_head=0, + ) + print("Peak scores head1:", peak_scores_head1) + print() + print() + + + peak_scores_head2 = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak2, + logs_dir=logs_dir, + num_heads=2, + target_head=1, + ) + print("Peak scores head2:", peak_scores_head2) + print() + print() + + peak_scores_bad = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak1, + logs_dir=logs_dir, + num_heads=2, + target_head=1, + ) + print("Peak scores bad:", peak_scores_bad) + print() + + print("Finished") + + if __name__ == "__main__": + print("hi") + # sudo mount -a main() + + # bigwig (signl) + + # import pyBigWig + # bw = pyBigWig.open(signal_file) + # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] + # print(bw.chroms()) [chr1:463772] {chrom:len} + + # fasta (genome): AGGGCAAAA... + + # bed (blacklist): chr1 10468 11447 + # (unmapabble region): if 65% overlap, remove 2046 window From 912cf78574ec0fd70100fe635291fd0872339438 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 31 Mar 2026 01:01:55 +0200 Subject: [PATCH 04/37] Replace tracks by heads --- src/asap/models/convnext_dcnn.py | 26 +++++++++++--------- src/asap/task.py | 2 ++ src/asap/trainer/trainer.py | 41 +++++++++++++++++++------------- 3 files changed, 41 insertions(+), 28 deletions(-) diff --git a/src/asap/models/convnext_dcnn.py b/src/asap/models/convnext_dcnn.py index a26c8f5..f3f01b4 100644 --- a/src/asap/models/convnext_dcnn.py +++ b/src/asap/models/convnext_dcnn.py @@ -42,7 +42,7 @@ def __init__(self, **kwargs) -> None: if self.use_map: self.unmap_predictor = UnmapPredictor(channels_in=config['filters0']) - self.core = BasenjiCoreBlock(nr_tracks=config['num_heads'], window=window, filters_in=config['filters0'], + self.core = BasenjiCoreBlock(nr_tracks=1, window=window, filters_in=config['filters0'], nr_res_blocks=config['residual_blocks'], rate_mult=config['dilation_mult'], bin_size=config['bin_size'], @@ -51,7 +51,8 @@ def __init__(self, **kwargs) -> None: kernel1=config['kernel1'], kernel2=config['kernel2'], dropout=config['dropout'], - final_dropout=config['final_dropout']) + final_dropout=config['final_dropout'], + num_heads=config['num_heads']) def forward(self, x:torch.Tensor, return_unmap=False) -> torch.Tensor: x = torch.transpose(x, dim0=-1, dim1=-2) @@ -71,7 +72,7 @@ class BasenjiCoreBlock(nn.Module): def __init__(self, nr_tracks: int, window: int, filters_in, nr_res_blocks: int = 11, rate_mult: float = 1.5, bin_size: int = 100, filters1: int = 128, filters3: int = None, kernel1: int = 3, kernel2: int = 1, dropout: float = 0.3, - final_dropout: float = 0.05): + final_dropout: float = 0.05, num_heads=1): super().__init__() if not filters3: filters3 = window @@ -94,11 +95,11 @@ def __init__(self, nr_tracks: int, window: int, filters_in, pool_size = bin_size // 2 self.pool = nn.AvgPool1d(kernel_size=pool_size, stride=pool_size) - print("creating output layer with tracks:", nr_tracks) - self.linear_out = nn.Linear( - in_features=filters3, - out_features=nr_tracks, - ) + print("creating output layer with tracks:", num_heads) + self.heads = nn.ModuleList([ + nn.Linear(filters3, nr_tracks) + for _ in range(num_heads) + ]) self.activation = nn.Softplus() def forward(self, x:torch.Tensor) -> torch.Tensor: @@ -117,9 +118,12 @@ def forward(self, x:torch.Tensor) -> torch.Tensor: x = self.final_dropout(x) x = self.pool(x) x = torch.transpose(x, -2, -1) - x = self.linear_out(x) - x = self.activation(x) - return x + outputs = [] + for head in self.heads: + out = head(x) + out = self.activation(out) + outputs.append(out) + return outputs class ConvBlock(nn.Module): diff --git a/src/asap/task.py b/src/asap/task.py index 1621a2f..9283d83 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -115,6 +115,7 @@ def train_multiheaded_model(experiment_name : str, model: str, train_dataset: L batch_size=batch_size, logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, + num_heads=num_heads ) # Start training @@ -144,6 +145,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat batch_size=batch_size, logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, + num_heads=num_heads ) # for evaluation use the checkpoint of the best model diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 121caf2..e15efc9 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -29,6 +29,7 @@ def __init__(self, logger: Logger = None, n_gpus: int = None, nr_tracks: int = 1, + num_heads: int = 1, ): self.filename = filename self.model = model @@ -41,6 +42,7 @@ def __init__(self, self.nr_tracks = nr_tracks self.nr_devices = n_gpus self.batch_size = batch_size + self.num_heads = num_heads if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' @@ -72,7 +74,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.logger, self.filename, self.nr_devices, - port + port, + self.num_heads ), nprocs=self.nr_devices ) @@ -101,7 +104,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.unmap_criterion, self.logger, self.filename, - ddp_enabled=False + ddp_enabled=False, + num_heads=self.num_heads ) def predict(self, gen): @@ -125,7 +129,6 @@ def predict_and_evaluate_multihead(self, gen, metrics_for_track=None, no_eval=Fa print("Generator has no margin size -- assuming full prediction.") predictions, true = self.predict(gen) - num_heads = predictions.shape[-1] predictions = predictions[..., target_head: target_head+1] predictions = predictions.reshape((-1, 1)).detach().cpu().numpy() true = true.reshape((-1, 1)).detach().cpu().numpy() @@ -314,7 +317,8 @@ def _ddp_and_fit( logger, filename, world_size, - port=12355 + port=12355, + num_heads=1, ): #model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = setup_ddp(rank, world_size, model, port) @@ -341,7 +345,8 @@ def _ddp_and_fit( unmap_criterion=unmap_criterion, logger=logger, filename=filename, - ddp_enabled=True + ddp_enabled=True, + num_heads=num_heads, ) dist.destroy_process_group() @@ -357,7 +362,8 @@ def _fit( unmap_criterion, logger: Logger, filename: str, - ddp_enabled: bool + ddp_enabled: bool, + num_heads: int, ): optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( @@ -380,7 +386,8 @@ def _fit( optimizer, scheduler, criterion, - unmap_criterion) + unmap_criterion, + num_heads) if train_log_payload is not None and (not ddp_enabled or rank == 0): logger.log(train_log_payload, step=epoch) @@ -394,10 +401,10 @@ def _fit( logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() - num_tracks = predictions.shape[-1] - for track in range(num_tracks): - print("Head: ", track+1) - predictions_head, true_head = predictions[..., track].flatten().numpy(), true[..., track].flatten().numpy() + for head in range(num_heads): + print("Head: ", head) + print("Debug", predictions.shape, true.shape) + predictions_head, true_head = predictions[..., head].flatten().numpy(), true[..., head].flatten().numpy() val_log_payload = compute_metrics( predictions_head, true_head, @@ -434,7 +441,7 @@ def _fit( print('Completed training!') -def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion): +def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, num_heads=1): model.train() train_unmap = unmap_criterion is not None @@ -448,7 +455,6 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ m_i = m_i.to(rank) y_i = y_i.to(rank) optimizer.zero_grad() - nr_tracks = y_i.shape[-1] if train_unmap: output, output_m_i = model(X_i, return_unmap=True) # trim m_i in case of unpadded conv in stem @@ -456,14 +462,14 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ unmap_loss = unmap_criterion(output_m_i, m_i[:, :m_len]) base_loss = 0 - for track in range(nr_tracks): - base_loss += criterion(output[..., track], y_i[..., track]) + for head in range(num_heads): + base_loss += criterion(output[head], y_i[..., head:head+1]) loss = base_loss + unmap_loss else: output = model(X_i) loss = 0 - for track in range(nr_tracks): - loss += criterion(output[..., track], y_i[..., track]) + for head in range(num_heads): + loss += criterion(output[head], y_i[..., head:head+1]) loss.backward() optimizer.step() @@ -509,6 +515,7 @@ def _predict(model, gen, rank, ddp_enabled): with torch.no_grad(): p_i = model(X_i) + p_i = torch.cat(p_i, dim=-1) if margin_size is not None: y_i.flatten() From e88a7b39f6221baf4d6d238e778feda67edcbc86 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Wed, 15 Apr 2026 00:08:03 +0200 Subject: [PATCH 05/37] ASAP 4 Cell lines --- src/asap/trainer/trainer.py | 6 +-- tutorials/train.py | 82 ++++++++++++++++++++++++++----------- 2 files changed, 60 insertions(+), 28 deletions(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index e15efc9..2db481c 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -39,7 +39,7 @@ def __init__(self, self.logger: Logger = logger self.logspace = True - self.nr_tracks = nr_tracks + self.nr_tracks = 1 self.nr_devices = n_gpus self.batch_size = batch_size self.num_heads = num_heads @@ -401,9 +401,7 @@ def _fit( logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() - for head in range(num_heads): - print("Head: ", head) - print("Debug", predictions.shape, true.shape) + for head in range(num_heads): predictions_head, true_head = predictions[..., head].flatten().numpy(), true[..., head].flatten().numpy() val_log_payload = compute_metrics( predictions_head, diff --git a/tutorials/train.py b/tutorials/train.py index 3e5824c..5fa4e32 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -20,6 +20,12 @@ def main(): signal_file_2 = "data/K562/ENCFF357GNC-signal.bigWig" peak_file_2 = "data/K562/ENCFF333TAT-peak.bed.gz" + + signal_file_3 = "data/HepG2/ENCFF262URW-signal.bigWig" + peak_file_3 = "data/HepG2/ENCFF439EIO-peak.bed.gz" + + signal_file_4 = "data/IMR90/ENCFF770EAV-signal.bigWig" + peak_file_4 = "data/IMR90/ENCFF243NTP-peak.bed.gz" genome = "data/hg38.fa" blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] @@ -29,18 +35,18 @@ def main(): # Model parameters model_name = "convnext_dcnn" - experiment_name = "GM12878_K562" + experiment_name = "all4" # Training parameters - test_chroms = [2, 10, 14, 19, 21] - train_chroms = [1, 11, 20, 13] - val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in val_chroms] + test_chroms = [1, 11, 20, 13] + train_chroms = [2, 10, 14, 19, 21] + val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] n_gpus = 2 # Create the training and validation datasets print("create the dataset") train_comb, val_comb = asap.training_datasets( - signal_file=[signal_file_1, signal_file_2], + signal_file=[signal_file_1, signal_file_2, signal_file_3, signal_file_4], genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -49,15 +55,6 @@ def main(): unmap_file=unmap_file, ) - - - - # print(train_comb.X[0].shape) - - # print(train_comb.y[0].shape) - - # print(train_comb.y[0][0]) - print("start to train!!!\n") @@ -65,7 +62,7 @@ def main(): asap.train_multiheaded_model( experiment_name=experiment_name, model=model_name, - num_heads=2, + num_heads=4, train_dataset=train_comb, val_dataset=val_comb, logs_dir=logs_dir, @@ -95,12 +92,24 @@ def main(): blacklist_file=blacklist_file, unmap_file=unmap_file, ) - - # print(len(peak1.X), len(peak1.y)) - # print("y") - # print(peak.y[0][0][100]) - # print(peak1.y[0].shape) - # print(peak2.y[0][0][100]) + peak3 = asap.peak_dataset( + signal_file=signal_file_3, + peak_file=peak_file_3, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + peak4 = asap.peak_dataset( + signal_file=signal_file_4, + peak_file=peak_file_4, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) # Evaluate the model peak_scores_head1 = asap.eval_multihead_model( @@ -108,7 +117,7 @@ def main(): model=model_name, eval_dataset=peak1, logs_dir=logs_dir, - num_heads=2, + num_heads=4, target_head=0, ) print("Peak scores head1:", peak_scores_head1) @@ -121,19 +130,44 @@ def main(): model=model_name, eval_dataset=peak2, logs_dir=logs_dir, - num_heads=2, + num_heads=4, target_head=1, ) print("Peak scores head2:", peak_scores_head2) print() print() + peak_scores_head3 = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak3, + logs_dir=logs_dir, + num_heads=4, + target_head=2, + ) + print("Peak scores head3:", peak_scores_head3) + print() + print() + + peak_scores_head4 = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak4, + logs_dir=logs_dir, + num_heads=4, + target_head=3, + ) + print("Peak scores head4:", peak_scores_head4) + print() + print() + + peak_scores_bad = asap.eval_multihead_model( experiment_name=experiment_name, model=model_name, eval_dataset=peak1, logs_dir=logs_dir, - num_heads=2, + num_heads=4, target_head=1, ) print("Peak scores bad:", peak_scores_bad) From 3cfce94e210ea2b8a880641eb03361549afb532b Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 21 Apr 2026 00:56:21 +0200 Subject: [PATCH 06/37] plots --- plots.py | 98 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 plots.py diff --git a/plots.py b/plots.py new file mode 100644 index 0000000..eb27d19 --- /dev/null +++ b/plots.py @@ -0,0 +1,98 @@ +# remove this file if the branch is merged +import numpy as np +import matplotlib.pyplot as plt +from scipy import stats + +test_chroms = [1, 11, 20, 13] +datasets = ["GM12878", "K562", "HepG2", "IMR90"] + +# 4 models +GM12878_single = [0.701, 0.695, 0.692, 0.687] # 11 epochs +K562_single = [0.711, 0.713, 0.723, 0.698] # 11 epochs +HepG2_single = [0.704, 0.704, 0.682, 0.701] # 12 epochs +IMR90_single = [0.739, 0.734, 0.731, 0.714] # 11 epochs + +# one model, 2 heads +# 8 epochs +GM12878_2head = [0.706, 0.699, 0.695, 0.688] +K562_2head = [0.715, 0.718, 0.722, 0.699] + +# one model, 4 heads +# 7 epochs +GM12878_4head = [0.703, 0.702, 0.692, 0.685] +K562_4head = [0.720, 0.723, 0.730, 0.705] +HepG2_4head = [0.716, 0.718, 0.69, 0.713] +IMR90_4head = [0.742, 0.733, 0.729, 0.708] + + + + +data_single = [GM12878_single, K562_single, HepG2_single, IMR90_single] +data_4head = [GM12878_4head, K562_4head, HepG2_4head, IMR90_4head] + +# Compute means and stds +means_single = [np.mean(d) for d in data_single] +stds_single = [np.std(d) for d in data_single] +means_4head = [np.mean(d) for d in data_4head] +stds_4head = [np.std(d) for d in data_4head] +print(means_4head) +print(means_single) + +# 1. Calculate P-values and Stars +def get_sig_star(p): + if p < 0.001: return '***' + elif p < 0.01: return '**' + elif p < 0.05: return '*' + else: return 'ns' + +p_values = [] +for s, h in zip(data_single, data_4head): + # Using paired t-test because observations are on the same chromosomes + _, p = stats.ttest_rel(s, h) + p_values.append(p) + +# Plotting +x = np.arange(len(datasets)) +width = 0.35 + +plt.figure(figsize=(8, 6)) + +# Plot the points with error bars +plt.errorbar(x - width/4, means_4head, yerr=stds_4head, + fmt='o', capsize=5, label='1 model (4 heads)', color='#1f77b4') + +plt.errorbar(x + width/4, means_single, yerr=stds_single, + fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') + +# 2. Add Significance Brackets +for i in range(len(datasets)): + # Determine height of the bracket + y_max = max(means_4head[i] + stds_4head[i], means_single[i] + stds_single[i]) + y_line = y_max + 0.005 # Bracket baseline + h = 0.003 # Bracket tick height + + # Draw bracket line + # (x_start, x_end), (y_start, y_end) + plt.plot([x[i] - width/4, x[i] - width/4, x[i] + width/4, x[i] + width/4], + [y_line, y_line + h, y_line + h, y_line], lw=1, c='black') + + # Add star/ns text + star = get_sig_star(p_values[i]) + plt.text(x[i], y_line + h, star, ha='center', va='bottom', fontsize=12) + +# Labels and formatting +plt.xticks(x, datasets, fontsize=12) +plt.yticks(fontsize=12) +plt.xlabel("Dataset", fontsize=13) +plt.ylabel("Pearson's R", fontsize=13) +plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) +plt.legend(frameon=False) +plt.ylim(0.6, 0.8) # Adjusted to fit brackets +plt.tight_layout() + +plt.savefig("4head-comparison-with-stats.png") +plt.show() + +# Print values for verification +for i, ds in enumerate(datasets): + print(f"{ds}: p-value = {p_values[i]:.4f} ({get_sig_star(p_values[i])})") \ No newline at end of file From 08d37247b109c18207f22ffc923ce8da805d2484 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sat, 2 May 2026 23:00:33 +0200 Subject: [PATCH 07/37] 13 head train.py --- .gitignore | 2 + plots.py | 29 ++++---- tutorials/train.py | 165 +++++++++++++++++---------------------------- 3 files changed, 75 insertions(+), 121 deletions(-) diff --git a/.gitignore b/.gitignore index 501eeae..17d9625 100644 --- a/.gitignore +++ b/.gitignore @@ -3,12 +3,14 @@ generated/** wandb/** */.DS_Store .idea +*.png /data /checkpoint /logs /plots /archive/imgs/ +/tmp # Created by https://www.gitignore.io/api/macos,python,R,linux,vim,emacs,visualstudiocode,intellij diff --git a/plots.py b/plots.py index eb27d19..6b80c5f 100644 --- a/plots.py +++ b/plots.py @@ -1,10 +1,9 @@ -# remove this file if the branch is merged import numpy as np import matplotlib.pyplot as plt from scipy import stats test_chroms = [1, 11, 20, 13] -datasets = ["GM12878", "K562", "HepG2", "IMR90"] +datasets = ["GM12878", "K562"] # 4 models GM12878_single = [0.701, 0.695, 0.692, 0.687] # 11 epochs @@ -27,18 +26,15 @@ -data_single = [GM12878_single, K562_single, HepG2_single, IMR90_single] -data_4head = [GM12878_4head, K562_4head, HepG2_4head, IMR90_4head] +data_single = [GM12878_single, K562_single] +data_2head = [GM12878_2head, K562_2head] # Compute means and stds means_single = [np.mean(d) for d in data_single] stds_single = [np.std(d) for d in data_single] -means_4head = [np.mean(d) for d in data_4head] -stds_4head = [np.std(d) for d in data_4head] -print(means_4head) -print(means_single) +means_2head = [np.mean(d) for d in data_2head] +stds_2head = [np.std(d) for d in data_2head] -# 1. Calculate P-values and Stars def get_sig_star(p): if p < 0.001: return '***' elif p < 0.01: return '**' @@ -46,9 +42,8 @@ def get_sig_star(p): else: return 'ns' p_values = [] -for s, h in zip(data_single, data_4head): - # Using paired t-test because observations are on the same chromosomes - _, p = stats.ttest_rel(s, h) +for s, h in zip(data_single, data_2head): + _, p = stats.wilcoxon(s, h) # Wilcox p_values.append(p) # Plotting @@ -58,8 +53,8 @@ def get_sig_star(p): plt.figure(figsize=(8, 6)) # Plot the points with error bars -plt.errorbar(x - width/4, means_4head, yerr=stds_4head, - fmt='o', capsize=5, label='1 model (4 heads)', color='#1f77b4') +plt.errorbar(x - width/4, means_2head, yerr=stds_2head, + fmt='o', capsize=5, label='1 model (2 heads)', color='#1f77b4') plt.errorbar(x + width/4, means_single, yerr=stds_single, fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') @@ -67,7 +62,7 @@ def get_sig_star(p): # 2. Add Significance Brackets for i in range(len(datasets)): # Determine height of the bracket - y_max = max(means_4head[i] + stds_4head[i], means_single[i] + stds_single[i]) + y_max = max(means_2head[i] + stds_2head[i], means_single[i] + stds_single[i]) y_line = y_max + 0.005 # Bracket baseline h = 0.003 # Bracket tick height @@ -87,10 +82,10 @@ def get_sig_star(p): plt.ylabel("Pearson's R", fontsize=13) plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) plt.legend(frameon=False) -plt.ylim(0.6, 0.8) # Adjusted to fit brackets +plt.ylim(0.67, 0.74) # Adjusted to fit brackets plt.tight_layout() -plt.savefig("4head-comparison-with-stats.png") +plt.savefig("2head-comparison-with-stats.png") plt.show() # Print values for verification diff --git a/tutorials/train.py b/tutorials/train.py index 5fa4e32..1ebd538 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -14,19 +14,32 @@ def main(): - # Data paths - signal_file_1 = "data/GM12878/ENCFF667MDI-signal.bigWig" - peak_file_1 = "data/GM12878/ENCFF748UZH-peak.bed" - - signal_file_2 = "data/K562/ENCFF357GNC-signal.bigWig" - peak_file_2 = "data/K562/ENCFF333TAT-peak.bed.gz" - - signal_file_3 = "data/HepG2/ENCFF262URW-signal.bigWig" - peak_file_3 = "data/HepG2/ENCFF439EIO-peak.bed.gz" - - signal_file_4 = "data/IMR90/ENCFF770EAV-signal.bigWig" - peak_file_4 = "data/IMR90/ENCFF243NTP-peak.bed.gz" + leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" + + datasets = [ + ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], + ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], + ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], + ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], + ["HG03432", "ENCFF993BIL.bigWig", "ENCFF831FGS.bed"], + ["MCF-7", "ENCFF976UNK.bigWig", "ENCFF821OEF.bed"], + ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], + ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], + ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], + ["GM12878", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], + ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], + ["K562", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], + ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] + ] + + print(len(datasets)) + signal_files = [f"{leomed_path}/{dataset[0]}.bigWig" for dataset in datasets] + signal_files[0] = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/HCT116.bigwig" + peak_files = [f"{leomed_path}/{dataset[0]}.bed" for dataset in datasets] + + print(signal_files, peak_files) + genome = "data/hg38.fa" blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] unmap_file = "data/basenji_unmappable.bed" @@ -35,7 +48,7 @@ def main(): # Model parameters model_name = "convnext_dcnn" - experiment_name = "all4" + experiment_name = "allCellLines" # Training parameters test_chroms = [1, 11, 20, 13] @@ -46,7 +59,7 @@ def main(): # Create the training and validation datasets print("create the dataset") train_comb, val_comb = asap.training_datasets( - signal_file=[signal_file_1, signal_file_2, signal_file_3, signal_file_4], + signal_file=signal_files, genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -62,7 +75,7 @@ def main(): asap.train_multiheaded_model( experiment_name=experiment_name, model=model_name, - num_heads=4, + num_heads=len(signal_files), train_dataset=train_comb, val_dataset=val_comb, logs_dir=logs_dir, @@ -73,101 +86,40 @@ def main(): print("Create eval ds") print() - - peak1 = asap.peak_dataset( - signal_file=signal_file_1, - peak_file=peak_file_1, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak2 = asap.peak_dataset( - signal_file=signal_file_2, - peak_file=peak_file_2, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak3 = asap.peak_dataset( - signal_file=signal_file_3, - peak_file=peak_file_3, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak4 = asap.peak_dataset( - signal_file=signal_file_4, - peak_file=peak_file_4, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - - # Evaluate the model - peak_scores_head1 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak1, - logs_dir=logs_dir, - num_heads=4, - target_head=0, - ) - print("Peak scores head1:", peak_scores_head1) - print() - print() - - - peak_scores_head2 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak2, - logs_dir=logs_dir, - num_heads=4, - target_head=1, - ) - print("Peak scores head2:", peak_scores_head2) - print() - print() - - peak_scores_head3 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak3, - logs_dir=logs_dir, - num_heads=4, - target_head=2, - ) - print("Peak scores head3:", peak_scores_head3) - print() - print() - - peak_scores_head4 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak4, - logs_dir=logs_dir, - num_heads=4, - target_head=3, - ) - print("Peak scores head4:", peak_scores_head4) - print() - print() + peak = [] + for i in range(len(signal_files)): + print("Evaluating ", datasets[i][0]) + peak.append(asap.peak_dataset( + signal_file=signal_files[i], + peak_file=peak_files[i], + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + ) + + # Evaluate the model + peak_scores_head = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak[i], + logs_dir=logs_dir, + num_heads=len(signal_files), + target_head=i, + ) + print(f"Peak scores head {i}:", peak_scores_head) + print() + print() peak_scores_bad = asap.eval_multihead_model( experiment_name=experiment_name, model=model_name, - eval_dataset=peak1, + eval_dataset=peak, logs_dir=logs_dir, - num_heads=4, + num_heads=len(signal_files), target_head=1, ) print("Peak scores bad:", peak_scores_bad) @@ -192,3 +144,8 @@ def main(): # bed (blacklist): chr1 10468 11447 # (unmapabble region): if 65% overlap, remove 2046 window + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? From ffd98b08c51ef79d820c65fa5c16adbe6a4b0a03 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sun, 3 May 2026 00:18:52 +0200 Subject: [PATCH 08/37] Change task.py for linear probing on multihead --- src/asap/task.py | 16 ++++++++------- src/asap/trainer/trainer.py | 14 +++++++++---- tutorials/train.py | 41 +++++++++++++++++++++++++++++++++---- 3 files changed, 56 insertions(+), 15 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index 22f3084..ca696d9 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -75,7 +75,7 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False): +def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): ''' Evaluate the model on the given dataset. Args: @@ -109,6 +109,7 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, linear_probe=True, + num_heads=num_heads, ) # train the new head based on the previous model @@ -116,22 +117,23 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth' trainer.load_weights(checkpoint_path) - # replace head - in_features = trainer.model.core.linear_out.in_features - out_features = trainer.model.core.linear_out.out_features - trainer.model.core.linear_out = nn.Linear(in_features, out_features) + # add new head + in_features = trainer.model.core.heads[0].in_features + out_features = trainer.model.core.heads[0].out_features + new_head = nn.Linear(in_features, out_features) + trainer.model.core.heads.append(new_head) # freeze everything for param in trainer.model.parameters(): param.requires_grad = False # unfreeze head - for param in trainer.model.core.linear_out.parameters(): + for param in trainer.model.core.heads[-1].parameters(): param.requires_grad = True # set modes such that there is no dropout for the core trainer.model.eval() - trainer.model.core.linear_out.train() + trainer.model.core.heads[-1].train() # Start training diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index b356160..1c09590 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -374,7 +374,7 @@ def _fit( ): if linear_probe: - optimizer = configure_adamw(model.core.linear_out, lr=learning_rate) + optimizer = configure_adamw(model.core.heads[-1], lr=learning_rate) else: optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( @@ -398,6 +398,7 @@ def _fit( scheduler, criterion, unmap_criterion, + linear_probe, num_heads) if train_log_payload is not None and (not ddp_enabled or rank == 0): @@ -412,7 +413,8 @@ def _fit( logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() - for head in range(num_heads): + start_head = num_heads-1 if linear_probe else 0 + for head in range(start_head, num_heads): predictions_head, true_head = predictions[..., head].flatten().numpy(), true[..., head].flatten().numpy() val_log_payload = compute_metrics( predictions_head, @@ -450,8 +452,12 @@ def _fit( print('Completed training!') -def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, num_heads=1): - model.train() +def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, linear_probe, num_heads=1): + if linear_probe: + model.eval() + model.core.heads[-1].train() + else: + model.train() train_unmap = unmap_criterion is not None diff --git a/tutorials/train.py b/tutorials/train.py index e1f22e1..a36e880 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -86,6 +86,7 @@ def main(): print("Create eval ds") print() + # Eval the multihead model peak = [] for i in range(len(signal_files)): print("Evaluating ", datasets[i][0]) @@ -125,21 +126,53 @@ def main(): print("Peak scores bad:", peak_scores_bad) print() - print("Finished") - + print("Finished evaluating the model") + # linear probing for a new head print("training a new head") + experiment_name_new_head = f"{experiment_name}-new-head" asap.train_new_head( base_experiment_name=experiment_name, - new_experiment_name=f"{experiment_name}-new-head", + new_experiment_name=experiment_name_new_head, model=model_name, train_dataset=train_lp, val_dataset=val_lp, logs_dir=logs_dir, n_gpus=n_gpus, - max_epochs=15, ) + # Re-evaluate to make sure the body was not changed + for i in range(len(signal_files)): + # Evaluate the model + peak_scores_head = asap.eval_multihead_model( + experiment_name=experiment_name_new_head, + model=model_name, + eval_dataset=peak[i], + logs_dir=logs_dir, + num_heads=len(signal_files)+1, + target_head=i, + ) + print(f"Peak scores head {i}:", peak_scores_head) + + peak_new_head = asap.peak_dataset( + signal_file=signal_file_new_head, + peak_file=peak_file_new_head, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + peak_scores_last_head = asap.eval_multihead_model( + experiment_name=experiment_name_new_head, + model=model_name, + eval_dataset=peak_new_head, + logs_dir=logs_dir, + num_heads=len(signal_files)+1, + target_head=len(signal_files), + ) + print(f"Peak scores new head {i}:", peak_scores_last_head) + if __name__ == "__main__": print("hi") # sudo mount -a From 58036c196b036c6370e2852fb55363a9db0f30f9 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sat, 2 May 2026 23:00:33 +0200 Subject: [PATCH 09/37] 13 head train.py --- .gitignore | 2 + plots.py | 29 ++++---- tutorials/train.py | 165 +++++++++++++++++---------------------------- 3 files changed, 75 insertions(+), 121 deletions(-) diff --git a/.gitignore b/.gitignore index 501eeae..17d9625 100644 --- a/.gitignore +++ b/.gitignore @@ -3,12 +3,14 @@ generated/** wandb/** */.DS_Store .idea +*.png /data /checkpoint /logs /plots /archive/imgs/ +/tmp # Created by https://www.gitignore.io/api/macos,python,R,linux,vim,emacs,visualstudiocode,intellij diff --git a/plots.py b/plots.py index eb27d19..6b80c5f 100644 --- a/plots.py +++ b/plots.py @@ -1,10 +1,9 @@ -# remove this file if the branch is merged import numpy as np import matplotlib.pyplot as plt from scipy import stats test_chroms = [1, 11, 20, 13] -datasets = ["GM12878", "K562", "HepG2", "IMR90"] +datasets = ["GM12878", "K562"] # 4 models GM12878_single = [0.701, 0.695, 0.692, 0.687] # 11 epochs @@ -27,18 +26,15 @@ -data_single = [GM12878_single, K562_single, HepG2_single, IMR90_single] -data_4head = [GM12878_4head, K562_4head, HepG2_4head, IMR90_4head] +data_single = [GM12878_single, K562_single] +data_2head = [GM12878_2head, K562_2head] # Compute means and stds means_single = [np.mean(d) for d in data_single] stds_single = [np.std(d) for d in data_single] -means_4head = [np.mean(d) for d in data_4head] -stds_4head = [np.std(d) for d in data_4head] -print(means_4head) -print(means_single) +means_2head = [np.mean(d) for d in data_2head] +stds_2head = [np.std(d) for d in data_2head] -# 1. Calculate P-values and Stars def get_sig_star(p): if p < 0.001: return '***' elif p < 0.01: return '**' @@ -46,9 +42,8 @@ def get_sig_star(p): else: return 'ns' p_values = [] -for s, h in zip(data_single, data_4head): - # Using paired t-test because observations are on the same chromosomes - _, p = stats.ttest_rel(s, h) +for s, h in zip(data_single, data_2head): + _, p = stats.wilcoxon(s, h) # Wilcox p_values.append(p) # Plotting @@ -58,8 +53,8 @@ def get_sig_star(p): plt.figure(figsize=(8, 6)) # Plot the points with error bars -plt.errorbar(x - width/4, means_4head, yerr=stds_4head, - fmt='o', capsize=5, label='1 model (4 heads)', color='#1f77b4') +plt.errorbar(x - width/4, means_2head, yerr=stds_2head, + fmt='o', capsize=5, label='1 model (2 heads)', color='#1f77b4') plt.errorbar(x + width/4, means_single, yerr=stds_single, fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') @@ -67,7 +62,7 @@ def get_sig_star(p): # 2. Add Significance Brackets for i in range(len(datasets)): # Determine height of the bracket - y_max = max(means_4head[i] + stds_4head[i], means_single[i] + stds_single[i]) + y_max = max(means_2head[i] + stds_2head[i], means_single[i] + stds_single[i]) y_line = y_max + 0.005 # Bracket baseline h = 0.003 # Bracket tick height @@ -87,10 +82,10 @@ def get_sig_star(p): plt.ylabel("Pearson's R", fontsize=13) plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) plt.legend(frameon=False) -plt.ylim(0.6, 0.8) # Adjusted to fit brackets +plt.ylim(0.67, 0.74) # Adjusted to fit brackets plt.tight_layout() -plt.savefig("4head-comparison-with-stats.png") +plt.savefig("2head-comparison-with-stats.png") plt.show() # Print values for verification diff --git a/tutorials/train.py b/tutorials/train.py index 5fa4e32..1ebd538 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -14,19 +14,32 @@ def main(): - # Data paths - signal_file_1 = "data/GM12878/ENCFF667MDI-signal.bigWig" - peak_file_1 = "data/GM12878/ENCFF748UZH-peak.bed" - - signal_file_2 = "data/K562/ENCFF357GNC-signal.bigWig" - peak_file_2 = "data/K562/ENCFF333TAT-peak.bed.gz" - - signal_file_3 = "data/HepG2/ENCFF262URW-signal.bigWig" - peak_file_3 = "data/HepG2/ENCFF439EIO-peak.bed.gz" - - signal_file_4 = "data/IMR90/ENCFF770EAV-signal.bigWig" - peak_file_4 = "data/IMR90/ENCFF243NTP-peak.bed.gz" + leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" + + datasets = [ + ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], + ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], + ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], + ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], + ["HG03432", "ENCFF993BIL.bigWig", "ENCFF831FGS.bed"], + ["MCF-7", "ENCFF976UNK.bigWig", "ENCFF821OEF.bed"], + ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], + ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], + ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], + ["GM12878", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], + ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], + ["K562", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], + ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] + ] + + print(len(datasets)) + signal_files = [f"{leomed_path}/{dataset[0]}.bigWig" for dataset in datasets] + signal_files[0] = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/HCT116.bigwig" + peak_files = [f"{leomed_path}/{dataset[0]}.bed" for dataset in datasets] + + print(signal_files, peak_files) + genome = "data/hg38.fa" blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] unmap_file = "data/basenji_unmappable.bed" @@ -35,7 +48,7 @@ def main(): # Model parameters model_name = "convnext_dcnn" - experiment_name = "all4" + experiment_name = "allCellLines" # Training parameters test_chroms = [1, 11, 20, 13] @@ -46,7 +59,7 @@ def main(): # Create the training and validation datasets print("create the dataset") train_comb, val_comb = asap.training_datasets( - signal_file=[signal_file_1, signal_file_2, signal_file_3, signal_file_4], + signal_file=signal_files, genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -62,7 +75,7 @@ def main(): asap.train_multiheaded_model( experiment_name=experiment_name, model=model_name, - num_heads=4, + num_heads=len(signal_files), train_dataset=train_comb, val_dataset=val_comb, logs_dir=logs_dir, @@ -73,101 +86,40 @@ def main(): print("Create eval ds") print() - - peak1 = asap.peak_dataset( - signal_file=signal_file_1, - peak_file=peak_file_1, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak2 = asap.peak_dataset( - signal_file=signal_file_2, - peak_file=peak_file_2, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak3 = asap.peak_dataset( - signal_file=signal_file_3, - peak_file=peak_file_3, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak4 = asap.peak_dataset( - signal_file=signal_file_4, - peak_file=peak_file_4, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - - # Evaluate the model - peak_scores_head1 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak1, - logs_dir=logs_dir, - num_heads=4, - target_head=0, - ) - print("Peak scores head1:", peak_scores_head1) - print() - print() - - - peak_scores_head2 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak2, - logs_dir=logs_dir, - num_heads=4, - target_head=1, - ) - print("Peak scores head2:", peak_scores_head2) - print() - print() - - peak_scores_head3 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak3, - logs_dir=logs_dir, - num_heads=4, - target_head=2, - ) - print("Peak scores head3:", peak_scores_head3) - print() - print() - - peak_scores_head4 = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak4, - logs_dir=logs_dir, - num_heads=4, - target_head=3, - ) - print("Peak scores head4:", peak_scores_head4) - print() - print() + peak = [] + for i in range(len(signal_files)): + print("Evaluating ", datasets[i][0]) + peak.append(asap.peak_dataset( + signal_file=signal_files[i], + peak_file=peak_files[i], + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + ) + + # Evaluate the model + peak_scores_head = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak[i], + logs_dir=logs_dir, + num_heads=len(signal_files), + target_head=i, + ) + print(f"Peak scores head {i}:", peak_scores_head) + print() + print() peak_scores_bad = asap.eval_multihead_model( experiment_name=experiment_name, model=model_name, - eval_dataset=peak1, + eval_dataset=peak, logs_dir=logs_dir, - num_heads=4, + num_heads=len(signal_files), target_head=1, ) print("Peak scores bad:", peak_scores_bad) @@ -192,3 +144,8 @@ def main(): # bed (blacklist): chr1 10468 11447 # (unmapabble region): if 65% overlap, remove 2046 window + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? From 4a8c793fb92ea75bc5c8aba7a37ccf7cb29417f3 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Thu, 7 May 2026 23:43:01 +0200 Subject: [PATCH 10/37] Fix Cuda OOM on eval --- src/asap/trainer/trainer.py | 56 +++++++++++++++++++------------------ 1 file changed, 29 insertions(+), 27 deletions(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 2db481c..7d59c4a 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -401,6 +401,7 @@ def _fit( logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() + del val_res for head in range(num_heads): predictions_head, true_head = predictions[..., head].flatten().numpy(), true[..., head].flatten().numpy() val_log_payload = compute_metrics( @@ -505,33 +506,34 @@ def _predict(model, gen, rank, ddp_enabled): predictions = [] true = [] - for i, (X_i, _, y_i) in enumerate(gen): - X_i = X_i.to(rank) - y_i = y_i.to(rank) - if i == 2040: - print(X_i[0,1005:-1005,:]) - - with torch.no_grad(): - p_i = model(X_i) - p_i = torch.cat(p_i, dim=-1) - - if margin_size is not None: - y_i.flatten() - p_i = p_i[..., trim:-trim, :] - - y_i, p_i = y_i.contiguous(), p_i.contiguous() - if ddp_enabled: - all_predictions = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())] - all_true = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())] - dist.all_gather(all_predictions, p_i) - dist.all_gather(all_true, y_i) - - if rank == 0: - predictions.extend(all_predictions) - true.extend(all_true) - else: - predictions.append(p_i) - true.append(y_i) + with torch.no_grad(): + for i, (X_i, _, y_i) in enumerate(gen): + X_i = X_i.to(rank) + y_i = y_i.to(rank) + if i == 2040: + print(X_i[0,1005:-1005,:]) + + with torch.no_grad(): + p_i = model(X_i) + p_i = torch.cat(p_i, dim=-1) + + if margin_size is not None: + y_i.flatten() + p_i = p_i[..., trim:-trim, :] + + y_i, p_i = y_i.contiguous(), p_i.contiguous() + if ddp_enabled: + all_predictions = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())] + all_true = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())] + dist.all_gather(all_predictions, p_i) + dist.all_gather(all_true, y_i) + + if rank == 0: + predictions.extend([t.detach().cpu() for t in all_predictions]) + true.extend([t.detach().cpu() for t in all_true]) + else: + predictions.append(p_i.detach().cpu()) + true.append(y_i.detach().cpu()) print(i) return predictions, true From 935c2e06480c3e35e256616971abafd44b9aa975 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Fri, 8 May 2026 16:04:07 +0200 Subject: [PATCH 11/37] 13head experiment name + fix datasets --- tutorials/train.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tutorials/train.py b/tutorials/train.py index 1ebd538..66a7a38 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -17,7 +17,7 @@ def main(): leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" datasets = [ - ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], + ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], # missing ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], @@ -26,9 +26,9 @@ def main(): ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], - ["GM12878", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], + ["GM12878_XSC", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], # missing, but I have run this ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], - ["K562", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], + ["K562_FGK", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] ] @@ -48,13 +48,13 @@ def main(): # Model parameters model_name = "convnext_dcnn" - experiment_name = "allCellLines" + experiment_name = "allCellLines13_1gpu" # Training parameters test_chroms = [1, 11, 20, 13] train_chroms = [2, 10, 14, 19, 21] val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] - n_gpus = 2 + n_gpus = 1 # Create the training and validation datasets print("create the dataset") @@ -80,6 +80,7 @@ def main(): val_dataset=val_comb, logs_dir=logs_dir, n_gpus=n_gpus, + batch_size=32 ) print("Training done.") @@ -113,11 +114,11 @@ def main(): print() print() - + print("Peak scores bad") peak_scores_bad = asap.eval_multihead_model( experiment_name=experiment_name, model=model_name, - eval_dataset=peak, + eval_dataset=peak[0], logs_dir=logs_dir, num_heads=len(signal_files), target_head=1, From 8532ce4ca5532e78d1a5c4a13559301039f11fc8 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sat, 9 May 2026 01:25:45 +0200 Subject: [PATCH 12/37] Fix num_heads --- src/asap/task.py | 6 +- src/asap/trainer/trainer.py | 3 +- tutorials/train.py | 125 ++++++++++++------------------------ 3 files changed, 47 insertions(+), 87 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index ca696d9..719864b 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -75,7 +75,7 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): +def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): ''' Evaluate the model on the given dataset. Args: @@ -97,7 +97,7 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") # Initialize the model - model = _get_model(model, use_map=use_map) + model = _get_model(model, use_map=use_map, num_heads=num_original_heads) # Initialize the trainer with the model and datasets trainer = Trainer( @@ -109,7 +109,7 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, linear_probe=True, - num_heads=num_heads, + num_heads=num_original_heads, ) # train the new head based on the previous model diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 1c09590..1c4b2da 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -455,7 +455,8 @@ def _fit( def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, linear_probe, num_heads=1): if linear_probe: model.eval() - model.core.heads[-1].train() + model.core.heads[-1].train() + num_heads = num_heads+1 else: model.train() diff --git a/tutorials/train.py b/tutorials/train.py index a36e880..87b709d 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -26,9 +26,9 @@ def main(): ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], - ["GM12878", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], + ["GM12878_XSC", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], - ["K562", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], + ["K562_FGK", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] ] @@ -48,18 +48,24 @@ def main(): # Model parameters model_name = "convnext_dcnn" - experiment_name = "allCellLines" + experiment_name = "allCellLines13_1gpu" # Training parameters test_chroms = [1, 11, 20, 13] train_chroms = [2, 10, 14, 19, 21] val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] - n_gpus = 2 + n_gpus = 1 - # Create the training and validation datasets - print("create the dataset") - train_comb, val_comb = asap.training_datasets( - signal_file=signal_files, + + # linear probing for a new head + print("create a new dataset") + experiment_name_new_head = f"{experiment_name}-new-head" + + signal_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bigWig" + peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw./T_cell_f_21.bed" + + train_lp, val_lp = asap.training_datasets( + signal_file=signal_file_new_head, genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -68,69 +74,7 @@ def main(): unmap_file=unmap_file, ) - - print("start to train!!!\n") - - # Train the model - asap.train_multiheaded_model( - experiment_name=experiment_name, - model=model_name, - num_heads=len(signal_files), - train_dataset=train_comb, - val_dataset=val_comb, - logs_dir=logs_dir, - n_gpus=n_gpus, - ) - - print("Training done.") - print("Create eval ds") - print() - - # Eval the multihead model - peak = [] - for i in range(len(signal_files)): - print("Evaluating ", datasets[i][0]) - peak.append(asap.peak_dataset( - signal_file=signal_files[i], - peak_file=peak_files[i], - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - ) - - # Evaluate the model - peak_scores_head = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak[i], - logs_dir=logs_dir, - num_heads=len(signal_files), - target_head=i, - ) - print(f"Peak scores head {i}:", peak_scores_head) - print() - print() - - - peak_scores_bad = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak, - logs_dir=logs_dir, - num_heads=len(signal_files), - target_head=1, - ) - print("Peak scores bad:", peak_scores_bad) - print() - - print("Finished evaluating the model") - - # linear probing for a new head print("training a new head") - experiment_name_new_head = f"{experiment_name}-new-head" asap.train_new_head( base_experiment_name=experiment_name, new_experiment_name=experiment_name_new_head, @@ -139,21 +83,10 @@ def main(): val_dataset=val_lp, logs_dir=logs_dir, n_gpus=n_gpus, + num_original_heads=len(datasets) ) - # Re-evaluate to make sure the body was not changed - for i in range(len(signal_files)): - # Evaluate the model - peak_scores_head = asap.eval_multihead_model( - experiment_name=experiment_name_new_head, - model=model_name, - eval_dataset=peak[i], - logs_dir=logs_dir, - num_heads=len(signal_files)+1, - target_head=i, - ) - print(f"Peak scores head {i}:", peak_scores_head) - + print("Create new eval dataset") peak_new_head = asap.peak_dataset( signal_file=signal_file_new_head, peak_file=peak_file_new_head, @@ -163,6 +96,8 @@ def main(): blacklist_file=blacklist_file, unmap_file=unmap_file, ) + + print("Eval new head") peak_scores_last_head = asap.eval_multihead_model( experiment_name=experiment_name_new_head, model=model_name, @@ -173,6 +108,30 @@ def main(): ) print(f"Peak scores new head {i}:", peak_scores_last_head) + # Re-evaluate to make sure the body was not modified + for i in range(len(signal_files)): + print("Reevaluating ", datasets[i][0]) + peak_dataset = asap.peak_dataset( + signal_file=signal_files[i], + peak_file=peak_files[i], + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + peak_scores_head = asap.eval_multihead_model( + experiment_name=experiment_name_new_head, + model=model_name, + eval_dataset=peak_dataset, + logs_dir=logs_dir, + num_heads=len(signal_files)+1, + target_head=i, + ) + print(f"Peak scores head {i}:", peak_scores_head) + + + if __name__ == "__main__": print("hi") # sudo mount -a From 389b049accd160ee4b6ddf8ee087e26cafe40860 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sun, 10 May 2026 14:10:47 +0200 Subject: [PATCH 13/37] Fix num heads --- src/asap/task.py | 2 +- src/asap/trainer/trainer.py | 10 +++++++--- tutorials/train.py | 3 ++- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index 719864b..1fea17e 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -109,7 +109,7 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, linear_probe=True, - num_heads=num_original_heads, + num_heads=num_original_heads+1, ) # train the new head based on the previous model diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 183e55a..d37def1 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -374,7 +374,8 @@ def _fit( ): if linear_probe: - optimizer = configure_adamw(model.core.heads[-1], lr=learning_rate) + base_model = model.module if hasattr(model, "module") else model + optimizer = configure_adamw(base_model.core.heads[-1], lr=learning_rate) else: optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( @@ -416,7 +417,11 @@ def _fit( predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() start_head = num_heads-1 if linear_probe else 0 for head in range(start_head, num_heads): - predictions_head, true_head = predictions[..., head].flatten().numpy(), true[..., head].flatten().numpy() + predictions_head = predictions[head].flatten().numpy(), + if linear_probe: + true_head = true[..., 0].flatten().numpy() + else: + true_head = true[..., head].flatten().numpy() val_log_payload = compute_metrics( predictions_head, true_head, @@ -457,7 +462,6 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ if linear_probe: model.eval() model.core.heads[-1].train() - num_heads = num_heads+1 else: model.train() diff --git a/tutorials/train.py b/tutorials/train.py index e713432..1f92710 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -83,7 +83,8 @@ def main(): val_dataset=val_lp, logs_dir=logs_dir, n_gpus=n_gpus, - num_original_heads=len(datasets) + num_original_heads=len(datasets), + max_epochs=2 ) print("Create new eval dataset") From cfdf723afef36fd4f22c60931ab84e5c664f299f Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sun, 10 May 2026 17:07:39 +0200 Subject: [PATCH 14/37] Fix types --- src/asap/trainer/trainer.py | 9 ++++++--- tutorials/train.py | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index d37def1..293168b 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -417,7 +417,7 @@ def _fit( predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() start_head = num_heads-1 if linear_probe else 0 for head in range(start_head, num_heads): - predictions_head = predictions[head].flatten().numpy(), + predictions_head = predictions[..., head].flatten().numpy() if linear_probe: true_head = true[..., 0].flatten().numpy() else: @@ -470,6 +470,8 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ if rank == 0: # pbar if on rank 0 train_gen = tqdm(train_gen) + + start_head = num_heads-1 if linear_probe else 0 for X_i, m_i, y_i in train_gen: X_i = X_i.to(rank) @@ -482,14 +484,15 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ m_len = output_m_i.shape[1] unmap_loss = unmap_criterion(output_m_i, m_i[:, :m_len]) + # only use the last head base_loss = 0 - for head in range(num_heads): + for head in range(start_head, num_heads): base_loss += criterion(output[head], y_i[..., head:head+1]) loss = base_loss + unmap_loss else: output = model(X_i) loss = 0 - for head in range(num_heads): + for head in range(start_head, num_heads): loss += criterion(output[head], y_i[..., head:head+1]) loss.backward() diff --git a/tutorials/train.py b/tutorials/train.py index 1f92710..94256d4 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -62,7 +62,7 @@ def main(): experiment_name_new_head = f"{experiment_name}-new-head" signal_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bigWig" - peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw./T_cell_f_21.bed" + peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bed" train_lp, val_lp = asap.training_datasets( signal_file=signal_file_new_head, @@ -107,7 +107,7 @@ def main(): num_heads=len(signal_files)+1, target_head=len(signal_files), ) - print(f"Peak scores new head {i}:", peak_scores_last_head) + print(f"Peak scores new head:", peak_scores_last_head) # Re-evaluate to make sure the body was not modified for i in range(len(signal_files)): From d500e733970f711410feafcf234a6a698d6e84ac Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Mon, 11 May 2026 01:33:56 +0200 Subject: [PATCH 15/37] Fix criterion, model initialisation --- src/asap/task.py | 36 ++++++++++++++++-------------------- src/asap/trainer/trainer.py | 11 +++++------ tutorials/train.py | 2 +- 3 files changed, 22 insertions(+), 27 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index 1fea17e..4b3c0c1 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -97,7 +97,19 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") # Initialize the model - model = _get_model(model, use_map=use_map, num_heads=num_original_heads) + model = _get_model(model, use_map=use_map, num_heads=num_original_heads+1) + + # freeze everything + for param in model.parameters(): + param.requires_grad = False + + # unfreeze head + for param in model.core.heads[-1].parameters(): + param.requires_grad = True + + # set modes such that there is no dropout for the core + model.eval() + model.core.heads[-1].train() # Initialize the trainer with the model and datasets trainer = Trainer( @@ -112,28 +124,12 @@ def train_new_head(base_experiment_name: str, new_experiment_name: str, model: s num_heads=num_original_heads+1, ) - # train the new head based on the previous model print(f'Loading best model weights from {base_experiment_name}') checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth' - trainer.load_weights(checkpoint_path) - - # add new head - in_features = trainer.model.core.heads[0].in_features - out_features = trainer.model.core.heads[0].out_features - new_head = nn.Linear(in_features, out_features) - trainer.model.core.heads.append(new_head) - - # freeze everything - for param in trainer.model.parameters(): - param.requires_grad = False - # unfreeze head - for param in trainer.model.core.heads[-1].parameters(): - param.requires_grad = True - - # set modes such that there is no dropout for the core - trainer.model.eval() - trainer.model.core.heads[-1].train() + # Load the file manually to use strict=False + checkpoint = torch.load(checkpoint_path, map_location='cpu') # Load to CPU first to avoid OOM + trainer.model.load_state_dict(checkpoint, strict=False) # Start training diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 293168b..5df7654 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -374,8 +374,7 @@ def _fit( ): if linear_probe: - base_model = model.module if hasattr(model, "module") else model - optimizer = configure_adamw(base_model.core.heads[-1], lr=learning_rate) + optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) else: optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( @@ -493,8 +492,10 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ output = model(X_i) loss = 0 for head in range(start_head, num_heads): - loss += criterion(output[head], y_i[..., head:head+1]) - + if linear_probe: + loss += criterion(output[head], y_i) + else: + loss += criterion(output[head], y_i[..., head:head+1]) loss.backward() optimizer.step() scheduler.step() @@ -535,8 +536,6 @@ def _predict(model, gen, rank, ddp_enabled): for i, (X_i, _, y_i) in enumerate(gen): X_i = X_i.to(rank) y_i = y_i.to(rank) - if i == 2040: - print(X_i[0,1005:-1005,:]) with torch.no_grad(): p_i = model(X_i) diff --git a/tutorials/train.py b/tutorials/train.py index 94256d4..eaab828 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -84,7 +84,7 @@ def main(): logs_dir=logs_dir, n_gpus=n_gpus, num_original_heads=len(datasets), - max_epochs=2 + max_epochs=20 ) print("Create new eval dataset") From 307dba303fbfda37d19831bd799f484a32643508 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 12 May 2026 16:48:16 +0200 Subject: [PATCH 16/37] helper functions, remove them --- plot_13eval.py | 94 ++++++++++++++++ plot_lp_13.py | 18 +++ plot_validation.py | 213 ++++++++++++++++++++++++++++++++++++ src/asap/trainer/trainer.py | 2 + tutorials/train.py | 50 +++++---- tutorials/train20.py | 170 ++++++++++++++++++++++++++++ tutorials/train_t.py | 105 ++++++++++++++++++ 7 files changed, 628 insertions(+), 24 deletions(-) create mode 100644 plot_13eval.py create mode 100644 plot_lp_13.py create mode 100644 plot_validation.py create mode 100644 tutorials/train20.py create mode 100644 tutorials/train_t.py diff --git a/plot_13eval.py b/plot_13eval.py new file mode 100644 index 0000000..518fdc0 --- /dev/null +++ b/plot_13eval.py @@ -0,0 +1,94 @@ +import numpy as np +import matplotlib.pyplot as plt + +datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90'] +chroms = [1, 11, 20, 13] + +# combined, 13 epochs (time limit) +HCT116_13head = [0.748, 0.747, 0.731, 0.742] +A549_RGS_13head = [0.736, 0.733, 0.707, 0.718] +WTC11_13head = [0.791, 0.783, 0.777, 0.783] +GM23338_13head = [0.796, 0.786, 0.782, 0.799] +HG03432_13head = [0.74, 0.735, 0.738, 0.74] +MCF7_13head = [0.738, 0.734, 0.695, 0.742] +PC3_13head = [0.714, 0.721, 0.699, 0.714] +Panc1_13head = [0.766, 0.759, 0.746, 0.745] +RWPE2_13head = [0.717, 0.716, 0.692, 0.705] +GM12878_XSC_13head = [0.703, 0.699, 0.692, 0.69] +HEPG2_GJU_13head = [0.712, 0.711, 0.685, 0.708] +K562_FGK_13head = [0.714, 0.715, 0.722, 0.701] +IMR90_13head = [0.738, 0.726, 0.73, 0.709] + +data_13heads = [ + HCT116_13head, + A549_RGS_13head, + WTC11_13head, + GM23338_13head, + HG03432_13head, + MCF7_13head, + PC3_13head, + Panc1_13head, + RWPE2_13head, + GM12878_XSC_13head, + HEPG2_GJU_13head, + K562_FGK_13head, + IMR90_13head +] + +# 13 separate models +HCT116_single = [0.758] +A549_RGS_single = [0.737] +WTC11_single = [0.8] +GM23338_single = [0.821] +HG03432_single = [0.715] +MCF7_single = [0.738] +PC3_single = [0.7] +Panc1_single = [0.744] +RWPE2_single = [0.745] +GM12878_XSC_single = [0.701, 0.695, 0.692, 0.687] +HEPG2_GJU_single = [0.704, 0.704, 0.682, 0.701] +K562_FGK_single = [0.711, 0.713, 0.723, 0.698] +IMR90_single = [0.739, 0.734, 0.731, 0.714] + +data_single = [ + HCT116_single, + A549_RGS_single, + WTC11_single, + GM23338_single, + HG03432_single, + MCF7_single, + PC3_single, + Panc1_single, + RWPE2_single, + GM12878_XSC_single, + HEPG2_GJU_single, + K562_FGK_single, + IMR90_single +] + + +# Compute mean and std across chroms +means_13head = [np.mean(x) for x in data_13heads] +stds_13head = [np.std(x) for x in data_13heads] +means_single = [np.mean(x) for x in data_single] +stds_single = [np.std(x) for x in data_single] + +# Plot +plt.figure(figsize=(12, 6)) + +width = 0.35 +x = np.arange(len(datasets)) +plt.errorbar(x - width/4, means_13head, yerr=stds_13head, + fmt='o', capsize=5, label='1 model (13 heads)', color='#1f77b4') + +plt.errorbar((x + width/4)[-4:], means_single[-4:], yerr=stds_single[-4:], + fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') + +plt.xlabel("Dataset", fontsize=13) +plt.ylabel("Pearson's R", fontsize=13) +plt.title("Model Comparison Across Datasets (13 headmodel vs single models)") +plt.xticks(x, datasets, fontsize=8) +plt.legend(frameon=False) + +plt.tight_layout() +plt.savefig("13heads_eval.png", dpi=300, bbox_inches='tight') \ No newline at end of file diff --git a/plot_lp_13.py b/plot_lp_13.py new file mode 100644 index 0000000..46188f3 --- /dev/null +++ b/plot_lp_13.py @@ -0,0 +1,18 @@ +T_cell_f_21_lp = [0.681, 0.675, 0.666, 0.656] #0.669 +T_cell_f_21_separate = [0.7057, 0.701, 0.699, 0.692] #0.699 + +a_B_cell_m_22_treated_lp = [0.723, 0.712, 0.706, 0.730] #0.71775 +a_B_cell_m_22_treated_separated = [0.702, 0.706, 0.703, 0.707] #0.704 + +# request rtx4090 +# change order and rerun 19heads +# lp on all new heads +# record runtimes +# keep finetuning 13 to 14 + + +print(sum(T_cell_f_21_lp)/4) +print(sum(T_cell_f_21_separate)/4) + +print(sum(a_B_cell_m_22_treated_lp)/4) +print(sum(a_B_cell_m_22_treated_separated)/4) diff --git a/plot_validation.py b/plot_validation.py new file mode 100644 index 0000000..90cbc26 --- /dev/null +++ b/plot_validation.py @@ -0,0 +1,213 @@ +import re +import matplotlib.pyplot as plt + +raw_text = """ Val pearson r: 0.6964387753816834 + Val pearson r: 0.6870994218353818 + Val pearson r: 0.7291391529590228 + Val pearson r: 0.7444582223654782 + Val pearson r: 0.6241408496209283 + Val pearson r: 0.6560767789907609 + Val pearson r: 0.6151724151653086 + Val pearson r: 0.7034725605829238 + Val pearson r: 0.6497986103344944 + Val pearson r: 0.6155472763204469 + Val pearson r: 0.6552607604094044 + Val pearson r: 0.630054931594626 + Val pearson r: 0.6185279754623397 + Val pearson r: 0.7216549256354842 + Val pearson r: 0.7055036708491718 + Val pearson r: 0.7511119742668912 + Val pearson r: 0.7635181946177819 + Val pearson r: 0.6537835362569802 + Val pearson r: 0.6835558226295194 + Val pearson r: 0.6388907881539803 + Val pearson r: 0.7190023912652485 + Val pearson r: 0.6743860834015354 + Val pearson r: 0.6469024054589176 + Val pearson r: 0.6794421250988107 + Val pearson r: 0.6472221303262831 + Val pearson r: 0.640377184784484 + Val pearson r: 0.7295442255629667 + Val pearson r: 0.7136852801487562 + Val pearson r: 0.7589122433393578 + Val pearson r: 0.7707581572647464 + Val pearson r: 0.6622760109618123 + Val pearson r: 0.6936889889396246 + Val pearson r: 0.6511044945298297 + Val pearson r: 0.7263597983993362 + Val pearson r: 0.6833798111148207 + Val pearson r: 0.6560634099565105 + Val pearson r: 0.6966597691502155 + Val pearson r: 0.6540557195808048 + Val pearson r: 0.6567890556048219 + Val pearson r: 0.7355667790321698 + Val pearson r: 0.7181773762583913 + Val pearson r: 0.762306945735433 + Val pearson r: 0.7735634438674696 + Val pearson r: 0.6709062488003141 + Val pearson r: 0.7021438188342437 + Val pearson r: 0.6612441596265453 + Val pearson r: 0.7299128731387278 + Val pearson r: 0.6938203863166237 + Val pearson r: 0.6645958744627727 + Val pearson r: 0.7062905065624483 + Val pearson r: 0.6594368702854894 + Val pearson r: 0.6751533285462077 + Val pearson r: 0.7374728764598621 + Val pearson r: 0.7196044648882977 + Val pearson r: 0.7670996660386562 + Val pearson r: 0.7780080149901397 + Val pearson r: 0.6742898696492363 + Val pearson r: 0.7058934639201252 + Val pearson r: 0.6672949258734164 + Val pearson r: 0.7312364038403028 + Val pearson r: 0.6955165547211278 + Val pearson r: 0.6682926530595027 + Val pearson r: 0.7098759765598139 + Val pearson r: 0.6632259104204266 + Val pearson r: 0.6801936203807368 + Val pearson r: 0.7396180166522474 + Val pearson r: 0.7219697031367717 + Val pearson r: 0.7700792666119827 + Val pearson r: 0.7812709213412251 + Val pearson r: 0.6768472576134353 + Val pearson r: 0.7089772560114744 + Val pearson r: 0.6724248145994438 + Val pearson r: 0.7328345447485672 + Val pearson r: 0.70032733213575 + Val pearson r: 0.6708278062936569 + Val pearson r: 0.712862836202867 + Val pearson r: 0.6696105337131105 + Val pearson r: 0.6866884996048938 + Val pearson r: 0.7411891055219084 + Val pearson r: 0.7245902689326984 + Val pearson r: 0.7708178484437126 + Val pearson r: 0.7812205607572208 + Val pearson r: 0.6805039766603462 + Val pearson r: 0.7134462216944767 + Val pearson r: 0.6765413814193819 + Val pearson r: 0.7341053857499662 + Val pearson r: 0.7045731751043073 + Val pearson r: 0.6743634373727768 + Val pearson r: 0.716458337219881 + Val pearson r: 0.6729699794543645 + Val pearson r: 0.6892055798526996 + Val pearson r: 0.7405559298361586 + Val pearson r: 0.7237466444041802 + Val pearson r: 0.7730430125437816 + Val pearson r: 0.7841607376771775 + Val pearson r: 0.679896613150622 + Val pearson r: 0.7170980883605286 + Val pearson r: 0.679135088372655 + Val pearson r: 0.7333611224807183 + Val pearson r: 0.704946036120285 + Val pearson r: 0.6750229864727569 + Val pearson r: 0.7170640526514216 + Val pearson r: 0.6754801310366083 + Val pearson r: 0.6892097632370298 + Val pearson r: 0.7457581148412622 + Val pearson r: 0.7281806164610638 + Val pearson r: 0.775695809322687 + Val pearson r: 0.7861311517116034 + Val pearson r: 0.6837907068494478 + Val pearson r: 0.7254445193158302 + Val pearson r: 0.6826638561237102 + Val pearson r: 0.7371976392755624 + Val pearson r: 0.7129005819104313 + Val pearson r: 0.6785930010014146 + Val pearson r: 0.7203582832818095 + Val pearson r: 0.6786407513936481 + Val pearson r: 0.6946855095794354 + Val pearson r: 0.7439361343698927 + Val pearson r: 0.7271142408196352 + Val pearson r: 0.7758318372488294 + Val pearson r: 0.7857087072185104 + Val pearson r: 0.6817283255573829 + Val pearson r: 0.72494311855335 + Val pearson r: 0.6845077675512952 + Val pearson r: 0.7362547679616372 + Val pearson r: 0.7108034558528062 + Val pearson r: 0.6770417456291605 + Val pearson r: 0.7195329931669482 + Val pearson r: 0.6786599228741661 + Val pearson r: 0.6946550478293225 + Val pearson r: 0.7412733773670532 + Val pearson r: 0.7245576204965146 + Val pearson r: 0.7746532810950112 + Val pearson r: 0.7848534129186068 + Val pearson r: 0.6793816030317934 + Val pearson r: 0.7236412771844641 + Val pearson r: 0.6818868719419771 + Val pearson r: 0.7336746683379142 + Val pearson r: 0.7099762496923312 + Val pearson r: 0.6744342750812531 + Val pearson r: 0.7184549823385983 + Val pearson r: 0.6760434079584795 + Val pearson r: 0.6925774202111611 + Val pearson r: 0.7424062650662471 + Val pearson r: 0.7266368686876203 + Val pearson r: 0.7755794476134695 + Val pearson r: 0.7842987925516092 + Val pearson r: 0.68165292486155 + Val pearson r: 0.7258095444959102 + Val pearson r: 0.6834873311576279 + Val pearson r: 0.7354560756083379 + Val pearson r: 0.71199184196435 + Val pearson r: 0.6763740993440086 + Val pearson r: 0.7189250457504082 + Val pearson r: 0.6788830698904805 + Val pearson r: 0.6943106983132857 + Val pearson r: 0.7450365912906849 + Val pearson r: 0.7288023282895445 + Val pearson r: 0.7765684548470875 + Val pearson r: 0.7856408276313993 + Val pearson r: 0.6818362862528512 + Val pearson r: 0.7281351907168422 + Val pearson r: 0.6848795112165097 + Val pearson r: 0.736865135898467 + Val pearson r: 0.7149251550514664 + Val pearson r: 0.6773219620250895 + Val pearson r: 0.7227276417064363 + Val pearson r: 0.6809720803882103 + Val pearson r: 0.695572395715137 + Val pearson r: 0.7416012986516195 + Val pearson r: 0.7273640296381129 + Val pearson r: 0.7773683142944958 + Val pearson r: 0.7862939212395638 + Val pearson r: 0.6804398988843954 + Val pearson r: 0.7260217080921356 + Val pearson r: 0.6836087222999384 + Val pearson r: 0.7352362622176543 + Val pearson r: 0.7121195524652735 + Val pearson r: 0.674739671521881 + Val pearson r: 0.7216534767164122 + Val pearson r: 0.679829676185234 + Val pearson r: 0.6943985773305369""" + +datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90'] + +# Extract all floating point numbers +values = [float(x) for x in re.findall(r"\d+\.\d+", raw_text)] + +heads = 13 +epochs = len(values) // heads + +# Reshape into (epochs x heads) +data = [values[i*heads:(i+1)*heads] for i in range(epochs)] + +# Transpose to get per-head time series +per_head = list(zip(*data)) + +# Plot +plt.figure() + +for i, head_vals in enumerate(per_head): + plt.plot(range(1, epochs + 1), head_vals, label=f"{datasets[i]}") + +plt.xlabel("Epoch") +plt.ylabel("Validation Pearson r") +plt.title("Validation Scores per Head over Epochs") +plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left") +plt.tight_layout() + +plt.savefig("validation_scores_13head.png") \ No newline at end of file diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 5df7654..cfb2377 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -388,6 +388,7 @@ def _fit( best_val_score = -1 for epoch in range(nr_epochs): + print(f'\nEpoch {epoch}: start training at {datetime.now()}') if ddp_enabled: train_gen.sampler.set_epoch(epoch) train_log_payload = _train_epoch( @@ -400,6 +401,7 @@ def _fit( unmap_criterion, linear_probe, num_heads) + print(f'Epoch {epoch}: stop training at {datetime.now()}') if train_log_payload is not None and (not ddp_enabled or rank == 0): logger.log(train_log_payload, step=epoch) diff --git a/tutorials/train.py b/tutorials/train.py index eaab828..1eeb6d7 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -56,13 +56,15 @@ def main(): val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] n_gpus = 1 + # change test for val for the next + # linear probing for a new head print("create a new dataset") - experiment_name_new_head = f"{experiment_name}-new-head" + experiment_name_new_head = f"{experiment_name}-linProbe-a_B_cell_m_22_treated" - signal_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bigWig" - peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bed" + signal_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig" + peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed" train_lp, val_lp = asap.training_datasets( signal_file=signal_file_new_head, @@ -109,27 +111,27 @@ def main(): ) print(f"Peak scores new head:", peak_scores_last_head) - # Re-evaluate to make sure the body was not modified - for i in range(len(signal_files)): - print("Reevaluating ", datasets[i][0]) - peak_dataset = asap.peak_dataset( - signal_file=signal_files[i], - peak_file=peak_files[i], - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - peak_scores_head = asap.eval_multihead_model( - experiment_name=experiment_name_new_head, - model=model_name, - eval_dataset=peak_dataset, - logs_dir=logs_dir, - num_heads=len(signal_files)+1, - target_head=i, - ) - print(f"Peak scores head {i}:", peak_scores_head) + # # Re-evaluate to make sure the body was not modified + # for i in range(len(signal_files)): + # print("Reevaluating ", datasets[i][0]) + # peak_dataset = asap.peak_dataset( + # signal_file=signal_files[i], + # peak_file=peak_files[i], + # genome=genome, + # chroms=test_chroms, + # generated=generated, + # blacklist_file=blacklist_file, + # unmap_file=unmap_file, + # ) + # peak_scores_head = asap.eval_multihead_model( + # experiment_name=experiment_name_new_head, + # model=model_name, + # eval_dataset=peak_dataset, + # logs_dir=logs_dir, + # num_heads=len(signal_files)+1, + # target_head=i, + # ) + # print(f"Peak scores head {i}:", peak_scores_head) diff --git a/tutorials/train20.py b/tutorials/train20.py new file mode 100644 index 0000000..55b577e --- /dev/null +++ b/tutorials/train20.py @@ -0,0 +1,170 @@ +import sys +import os +import numpy as np +import pyBigWig + +sys.path.append(os.path.abspath("src")) + +import asap + +# The below code is a complete script that sets up the training of a model using the ASAP library. +# It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. +# The script then trains the model using the specified parameters. + + +def main(): + + leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" + + cell_lines = [ + ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], + ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], + ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], + ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], + ["HG03432", "ENCFF993BIL.bigWig", "ENCFF831FGS.bed"], + ["MCF-7", "ENCFF976UNK.bigWig", "ENCFF821OEF.bed"], + ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], + ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], + ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], + ["GM12878_XSC", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], + ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], + ["K562_FGK", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], + ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] + ] + + signal_files_cell_lines = [f"{leomed_path}/{dataset[0]}.bigWig" for dataset in cell_lines] + signal_files_cell_lines[0] = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/HCT116.bigwig" + peak_files_cell_lines = [f"{leomed_path}/{dataset[0]}.bed" for dataset in cell_lines] + + signal_files_primary = ["/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bigWig", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/nk_cell_f_41.bigWig", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/t_helper_17_m_50.bigWig", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/td_CD8_ab_T_m_30.bigWig", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/foreski_ker_m.bigWig"] + peak_files_primary = ["/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bed", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/nk_cell_f_41.bed", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/t_helper_17_m_50.bed", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/td_CD8_ab_T_m_30.bed", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed", + "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/foreski_ker_m.bed"] + signal_files = signal_files_cell_lines + signal_files_primary + peak_files = peak_files_cell_lines + peak_files_primary + + genome = "data/hg38.fa" + blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] + unmap_file = "data/basenji_unmappable.bed" + generated = "tmp" + logs_dir = "tmp/logs" + + # Model parameters + model_name = "convnext_dcnn" + experiment_name = "all19" + + # Training parameters + val_chroms = [1, 11, 20, 13] + train_chroms = [2, 10, 14, 19, 21] + test_chroms = [x for x in range(1, 23) if x not in val_chroms and x not in train_chroms] + n_gpus = 1 + + + # Create the training and validation datasets + print("create the dataset") + train_comb, val_comb = asap.training_datasets( + signal_file=signal_files, + genome=genome, + train_chroms=train_chroms, + val_chroms=val_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + + + print("start to train!!!\n") + + # Train the model + asap.train_multiheaded_model( + experiment_name=experiment_name, + model=model_name, + num_heads=len(signal_files), + train_dataset=train_comb, + val_dataset=val_comb, + logs_dir=logs_dir, + n_gpus=n_gpus, + ) + + print("Training done.") + print("Create eval ds") + print() + + peak = [] + for i in range(len(signal_files)): + print("Evaluating ", peak_files[i]) + peak.append(asap.peak_dataset( + signal_file=signal_files[i], + peak_file=peak_files[i], + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + ) + + # Evaluate the model + peak_scores_head = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak[i], + logs_dir=logs_dir, + num_heads=len(signal_files), + target_head=i, + ) + print(f"Peak scores head {i}:", peak_scores_head) + print() + print() + + print("Peak scores bad") + peak_scores_bad = asap.eval_multihead_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak[0], + logs_dir=logs_dir, + num_heads=len(signal_files), + target_head=1, + ) + print("Peak scores bad:", peak_scores_bad) + print() + + print("Finished") + + + + +if __name__ == "__main__": + print("hi") + # sudo mount -a + main() + + # bigwig (signl) + + # import pyBigWig + # bw = pyBigWig.open(signal_file) + # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] + # print(bw.chroms()) [chr1:463772] {chrom:len} + + # fasta (genome): AGGGCAAAA... + + # bed (blacklist): chr1 10468 11447 + # (unmapabble region): if 65% overlap, remove 2046 window + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? diff --git a/tutorials/train_t.py b/tutorials/train_t.py new file mode 100644 index 0000000..5d13cf0 --- /dev/null +++ b/tutorials/train_t.py @@ -0,0 +1,105 @@ +import sys +import os +import numpy as np +import pyBigWig + +sys.path.append(os.path.abspath("src")) + +import asap + +# The below code is a complete script that sets up the training of a model using the ASAP library. +# It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. +# The script then trains the model using the specified parameters. + + +def main(): + genome = "data/hg38.fa" + blacklist_file = ["data/basenji_blacklist.bed"] + unmap_file = "data/basenji_unmappable.bed" + generated = "tmp" + logs_dir = "tmp/logs" + + # Model parameters + model_name = "convnext_dcnn" + experiment_name = "a_B_cell_m_22_treated" + + # Training parameters + val_chroms = [1, 11, 20, 13] + train_chroms = [2, 10, 14, 19, 21] + test_chroms = [x for x in range(1, 23) if x not in val_chroms and x not in train_chroms] + n_gpus = 1 + + + signal_file = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig" + peak_file = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed" + + train, val = asap.training_datasets( + signal_file=signal_file, + genome=genome, + train_chroms=train_chroms, + val_chroms=val_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + + print("training a new head") + asap.train_model( + experiment_name=experiment_name, + model=model_name, + train_dataset=train, + val_dataset=val, + logs_dir=logs_dir, + n_gpus=n_gpus, + max_epochs=20 + ) + print() + print("Create new eval dataset") + peak = asap.peak_dataset( + signal_file=signal_file, + peak_file=peak_file, + genome=genome, + chroms=test_chroms, + generated=generated, + blacklist_file=blacklist_file, + unmap_file=unmap_file, + ) + print() + print("Eval") + peak_scores = asap.eval_model( + experiment_name=experiment_name, + model=model_name, + eval_dataset=peak, + logs_dir=logs_dir, + ) + print(f"Peak scores: ", peak_scores) + + + + +if __name__ == "__main__": + print("hi") + # sudo mount -a + main() + + # bigwig (signl) + + # import pyBigWig + # bw = pyBigWig.open(signal_file) + # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] + # print(bw.chroms()) [chr1:463772] {chrom:len} + + # fasta (genome): AGGGCAAAA... + + # bed (blacklist): chr1 10468 11447 + # (unmapabble region): if 65% overlap, remove 2046 window + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? + + +# send means for 2 and 4 +# Train 13 to validate +# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? From d5652ea8dc3d9cfba7b3d5e82ab81f81182106db Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Sun, 17 May 2026 20:50:01 +0200 Subject: [PATCH 17/37] Time epochs, fix early stopping, plot lp --- plot_14heads.py | 128 ++++++++++++++++++++++++++++++++++++ plot_lp_13.py | 111 ++++++++++++++++++++++++++++--- src/asap/trainer/trainer.py | 28 ++++---- tutorials/train_t.py | 4 +- 4 files changed, 249 insertions(+), 22 deletions(-) create mode 100644 plot_14heads.py diff --git a/plot_14heads.py b/plot_14heads.py new file mode 100644 index 0000000..b1deedc --- /dev/null +++ b/plot_14heads.py @@ -0,0 +1,128 @@ +import numpy as np +import matplotlib.pyplot as plt + +datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90', "T_cell_f_21_lp"] +chroms = [1, 11, 20, 13] + +# combined, 13 epochs (time limit) +HCT116_13head = [0.748, 0.747, 0.731, 0.742] +A549_RGS_13head = [0.736, 0.733, 0.707, 0.718] +WTC11_13head = [0.791, 0.783, 0.777, 0.783] +GM23338_13head = [0.796, 0.786, 0.782, 0.799] +HG03432_13head = [0.74, 0.735, 0.738, 0.74] +MCF7_13head = [0.738, 0.734, 0.695, 0.742] +PC3_13head = [0.714, 0.721, 0.699, 0.714] +Panc1_13head = [0.766, 0.759, 0.746, 0.745] +RWPE2_13head = [0.717, 0.716, 0.692, 0.705] +GM12878_XSC_13head = [0.703, 0.699, 0.692, 0.69] +HEPG2_GJU_13head = [0.712, 0.711, 0.685, 0.708] +K562_FGK_13head = [0.714, 0.715, 0.722, 0.701] +IMR90_13head = [0.738, 0.726, 0.73, 0.709] +T_cell_f_21_lp = [0.681, 0.675, 0.666, 0.656] + +data_13heads_plus_lp = [ + HCT116_13head, + A549_RGS_13head, + WTC11_13head, + GM23338_13head, + HG03432_13head, + MCF7_13head, + PC3_13head, + Panc1_13head, + RWPE2_13head, + GM12878_XSC_13head, + HEPG2_GJU_13head, + K562_FGK_13head, + IMR90_13head, + T_cell_f_21_lp +] + +# 13 separate models +HCT116_single = [0.758] +A549_RGS_single = [0.737] +WTC11_single = [0.8] +GM23338_single = [0.821] +HG03432_single = [0.715] +MCF7_single = [0.738] +PC3_single = [0.7] +Panc1_single = [0.744] +RWPE2_single = [0.745] +GM12878_XSC_single = [0.701, 0.695, 0.692, 0.687] +HEPG2_GJU_single = [0.704, 0.704, 0.682, 0.701] +K562_FGK_single = [0.711, 0.713, 0.723, 0.698] +IMR90_single = [0.739, 0.734, 0.731, 0.714] +T_cell_f_21_separate = [0.7057, 0.701, 0.699, 0.692] + +data_single = [ + HCT116_single, + A549_RGS_single, + WTC11_single, + GM23338_single, + HG03432_single, + MCF7_single, + PC3_single, + Panc1_single, + RWPE2_single, + GM12878_XSC_single, + HEPG2_GJU_single, + K562_FGK_single, + IMR90_single, + T_cell_f_21_separate +] + + +data_14head_cl = [[0.77388, 0.77218, 0.76423, 0.76750], +[0.75115, 0.74464, 0.72093, 0.73558], +[0.80259, 0.79624, 0.79078, 0.79235], +[0.80860, 0.79816, 0.79233, 0.80874], +[0.75988, 0.75415, 0.75297, 0.75616], +[0.76864, 0.76478, 0.72526, 0.77316], +[0.73899, 0.74466, 0.72559, 0.74090], +[0.77968, 0.76998, 0.76067, 0.75639], +[0.74593, 0.74494, 0.71737, 0.73883], +[0.72286, 0.71904, 0.70844, 0.70955], +[0.72074, 0.71941, 0.69779, 0.71925], +[0.72675, 0.72961, 0.73622, 0.71564], +[0.75554, 0.74607, 0.74561, 0.72812], +[0.74056, 0.73343, 0.72759, 0.71919]] + + +# Compute mean and std across chroms +means_14head = [np.mean(x) for x in data_13heads_plus_lp] +stds_14head = [np.std(x) for x in data_13heads_plus_lp] +means_single = [np.mean(x) for x in data_single] +stds_single = [np.std(x) for x in data_single] +means_cl = [np.mean(x) for x in data_14head_cl] +stds_cl = [np.std(x) for x in data_14head_cl] + +# Plot +plt.figure(figsize=(12, 6)) + +width = 0.35 +x = np.arange(len(datasets)) + +plt.errorbar((x + width/4)[-5:], means_single[-5:], yerr=stds_single[-5:], + fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') + +plt.errorbar(x , means_14head, yerr=stds_14head, + fmt='o', capsize=5, label='1 model (13 heads, 1 head LP)', color='#1577b4') + +plt.errorbar(x - width/4, means_cl, yerr=stds_cl, + fmt='o', capsize=5, label='1 model (CL)', color='green') +for i in range(len(datasets) - 1): + plt.axvline( + x=i + 0.5, + color='gray', + linestyle=':', + linewidth=1, + alpha=0.7 + ) + +plt.xlabel("Dataset", fontsize=13) +plt.ylabel("Pearson's R", fontsize=13) +plt.title("Model Comparison Across Datasets (14 headmodel vs single models)") +plt.xticks(x, datasets, fontsize=8) +plt.legend(frameon=False) + +plt.tight_layout() +plt.savefig("14heads_eval.png", dpi=300, bbox_inches='tight') \ No newline at end of file diff --git a/plot_lp_13.py b/plot_lp_13.py index 46188f3..bf732eb 100644 --- a/plot_lp_13.py +++ b/plot_lp_13.py @@ -1,18 +1,111 @@ +# request rtx4090 +# change order and rerun 19heads +# lp on all new heads +# record runtimes +# keep finetuning 13 to 14 + +import numpy as np +import matplotlib.pyplot as plt +from scipy import stats + +datasets = ["T_cell_f_21", + "nk_cell_f_41", + "t_helper_17_m_50", + "td_CD8_ab_T_m_30", + "a_B_cell_m_22_treated", + "foreski_ker_m"] + T_cell_f_21_lp = [0.681, 0.675, 0.666, 0.656] #0.669 T_cell_f_21_separate = [0.7057, 0.701, 0.699, 0.692] #0.699 +nk_cell_f_41_lp = [0.702, 0.709, 0.691, 0.683] +nk_cell_f_41_separated = [0.721, 0.726, 0.716, 0.713] + +t_helper_17_m_50_lp = [0.698, 0.703, 0.684, 0.689] +t_helper_17_m_50_separated = [0.0,0,0,0] + +td_CD8_ab_T_m_30_lp = [0.695, 0.679, 0.671, 0.678] +td_CD8_ab_T_m_30_separated = [0.719, 0.703, 0.708, 0.696] + a_B_cell_m_22_treated_lp = [0.723, 0.712, 0.706, 0.730] #0.71775 a_B_cell_m_22_treated_separated = [0.702, 0.706, 0.703, 0.707] #0.704 -# request rtx4090 -# change order and rerun 19heads -# lp on all new heads -# record runtimes -# keep finetuning 13 to 14 +foreski_ker_m_lp = [ 0.706, 0.701, 0.685, 0.708] +foreski_ker_m_separated = [0.744, 0.742, 0.731, 0.738] + + +data_single = [T_cell_f_21_separate, + nk_cell_f_41_separated, + t_helper_17_m_50_separated, + td_CD8_ab_T_m_30_separated, + a_B_cell_m_22_treated_separated, + foreski_ker_m_separated] +data_lp = [T_cell_f_21_lp, + nk_cell_f_41_lp, + t_helper_17_m_50_lp, + td_CD8_ab_T_m_30_lp, + a_B_cell_m_22_treated_lp, + foreski_ker_m_lp] + +# Compute means and stds6 +means_single = [np.mean(d) for d in data_single] +stds_single = [np.std(d) for d in data_single] +means_lp = [np.mean(d) for d in data_lp] +stds_lp = [np.std(d) for d in data_lp] + +def get_sig_star(p): + if p < 0.001: return '***' + elif p < 0.01: return '**' + elif p < 0.05: return '*' + else: return 'ns' + +p_values = [] +for s, h in zip(data_single, data_lp): + _, p = stats.wilcoxon(s, h) # Wilcox + p_values.append(p) + +# Plotting +x = np.arange(len(datasets)) +width = 0.35 + +plt.figure(figsize=(12, 6)) + +# Plot the points with error bars +plt.errorbar(x - width/4, means_lp, yerr=stds_lp, + fmt='o', capsize=5, label='Linear Probing', color='#1f77b4') + +plt.errorbar(x + width/4, means_single, yerr=stds_single, + fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') + +# # 2. Add Significance Brackets +# for i in range(len(datasets)): +# # Determine height of the bracket +# y_max = max(means_lp[i] + stds_lp[i], means_single[i] + stds_single[i]) +# y_line = y_max + 0.005 # Bracket baseline +# h = 0.003 # Bracket tick height + +# # Draw bracket line +# # (x_start, x_end), (y_start, y_end) +# plt.plot([x[i] - width/4, x[i] - width/4, x[i] + width/4, x[i] + width/4], +# [y_line, y_line + h, y_line + h, y_line], lw=1, c='black') + +# # Add star/ns text +# star = get_sig_star(p_values[i]) +# plt.text(x[i], y_line + h, star, ha='center', va='bottom', fontsize=12) +# Labels and formatting +plt.xticks(x, datasets, fontsize=12) +plt.yticks(fontsize=12) +plt.xlabel("Dataset", fontsize=13) +plt.ylabel("Pearson's R", fontsize=13) +plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) +plt.legend(frameon=False) +plt.ylim(0.65, 0.78) # Adjusted to fit brackets +plt.tight_layout() -print(sum(T_cell_f_21_lp)/4) -print(sum(T_cell_f_21_separate)/4) +plt.savefig("13head-lp_on_primary.png") +plt.show() -print(sum(a_B_cell_m_22_treated_lp)/4) -print(sum(a_B_cell_m_22_treated_separated)/4) +# Print values for verification +for i, ds in enumerate(datasets): + print(f"{ds}: p-value = {p_values[i]:.4f} ({get_sig_star(p_values[i])})") diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index cfb2377..9cbbb35 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -388,7 +388,8 @@ def _fit( best_val_score = -1 for epoch in range(nr_epochs): - print(f'\nEpoch {epoch}: start training at {datetime.now()}') + before = datetime.now() + print(f'\nEpoch {epoch}: start training at {before}') if ddp_enabled: train_gen.sampler.set_epoch(epoch) train_log_payload = _train_epoch( @@ -401,7 +402,9 @@ def _fit( unmap_criterion, linear_probe, num_heads) - print(f'Epoch {epoch}: stop training at {datetime.now()}') + after = datetime.now() + print(f'Epoch {epoch}: stop training at {after}') + print(f'Epoch {epoch}: train duration {(after - before).total_seconds()}') if train_log_payload is not None and (not ddp_enabled or rank == 0): logger.log(train_log_payload, step=epoch) @@ -411,6 +414,8 @@ def _fit( # For synchronous loop breaking stop_early = torch.zeros(1).to(rank) + epoch_val_sum = 0.0 + if not ddp_enabled or rank == 0: logger.log({'lr': scheduler.get_last_lr()[0]}) predictions, true = val_res @@ -435,15 +440,16 @@ def _fit( print(f'\tTrain loss: {train_log_payload["train/loss"]}') print(f'\tVal pearson r: {val_log_payload["val/pearson_r"]}') print('-----------------------------------------') - if val_log_payload['val/pearson_r'] > best_val_score: - best_val_score = val_log_payload['val/pearson_r'] - logger.save_model(model, filename) - # handle early stopping - no_improvement_for = 0 - else: - no_improvement_for += 1 - if (epoch != nr_epochs -1) and early_stopping_after_no_improvement and no_improvement_for >= early_stopping_after_no_improvement: - stop_early += 1 + epoch_val_sum += val_log_payload['val/pearson_r'] + if(epoch_val_sum / (num_heads - start_head)) > best_val_score: + best_val_score = val_log_payload['val/pearson_r'] + logger.save_model(model, filename) + # handle early stopping + no_improvement_for = 0 + else: + no_improvement_for += 1 + if (epoch != nr_epochs -1) and early_stopping_after_no_improvement and no_improvement_for >= early_stopping_after_no_improvement: + stop_early += 1 if ddp_enabled: dist.all_reduce(stop_early) diff --git a/tutorials/train_t.py b/tutorials/train_t.py index 5d13cf0..0d1d40a 100644 --- a/tutorials/train_t.py +++ b/tutorials/train_t.py @@ -24,9 +24,9 @@ def main(): experiment_name = "a_B_cell_m_22_treated" # Training parameters - val_chroms = [1, 11, 20, 13] + test_chroms = [1, 11, 20, 13] train_chroms = [2, 10, 14, 19, 21] - test_chroms = [x for x in range(1, 23) if x not in val_chroms and x not in train_chroms] + val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] n_gpus = 1 From cc265ea13bc90745f5c647515d71abdd52ddd645 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Mon, 1 Jun 2026 15:50:19 +0200 Subject: [PATCH 18/37] Fine tuning --- src/asap/__init__.py | 2 +- src/asap/task.py | 61 +++++++++++++++++++++++++++++++++++++ src/asap/trainer/trainer.py | 35 +++++++++++++++------ 3 files changed, 87 insertions(+), 11 deletions(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index 9c1ca2c..7b667cf 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,5 +1,5 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head, train_multiheaded_model, eval_multihead_model +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head, train_new_head_ft, train_multiheaded_model, eval_multihead_model diff --git a/src/asap/task.py b/src/asap/task.py index 4b3c0c1..af072cb 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -75,6 +75,67 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) +def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): + ''' + Evaluate the model on the given dataset. + Args: + base_experiment_name (str): The name of the model whose weights will be loaded. + new_experiment_name (str): The name of the experiment for the new head. + model (str): The model to evaluate. + train_dataset: The training dataset. + val_dataset: The validation dataset for early stopping. + logs_dir (str): The directory to load model checkpoints from. + batch_size (int): The batch size for evaluation. + use_map (bool): If mappability information was used during training. + num_heads (int): The number of heads for the new model (base model number of heads+1) + ''' + if n_gpus > 0 and not torch.cuda.is_available(): + n_gpus = 0 + print("No GPU available, using CPU instead.") + + # Count the number of GPUs available + if n_gpus > torch.cuda.device_count(): + n_gpus = torch.cuda.device_count() + print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") + + # Initialize the model + model = _get_model(model, use_map=use_map, num_heads=num_heads) + + # freeze old heads + for param in model.core.heads[:-1].parameters(): + param.requires_grad = False + + # set modes + model.train() + for head in model.core.heads[:-1]: + head.eval() + + # Initialize the trainer with the model and datasets + trainer = Trainer( + filename=new_experiment_name, + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + fine_tune=True, + num_heads=num_heads, + ) + + print(f'Loading best model weights from {base_experiment_name}') + checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth' + + # Load the file manually to use strict=False + checkpoint = torch.load(checkpoint_path, map_location='cpu') # Load to CPU first to avoid OOM + trainer.model.load_state_dict(checkpoint, strict=False) + + + # Start training + trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) + print("trained a new model") + + def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): ''' Evaluate the model on the given dataset. diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 9cbbb35..f3df5c5 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -31,6 +31,7 @@ def __init__(self, nr_tracks: int = 1, num_heads: int = 1, linear_probe=False, + fine_tune=False, ): self.filename = filename self.model = model @@ -45,6 +46,7 @@ def __init__(self, self.batch_size = batch_size self.num_heads = num_heads self.linear_probe = linear_probe + self.fine_tune = fine_tune if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' @@ -77,6 +79,7 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.filename, self.nr_devices, self.linear_probe, + self.fine_tune, port, self.num_heads ), @@ -110,6 +113,7 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): ddp_enabled=False, num_heads=self.num_heads, linear_probe=self.linear_probe, + fine_tune=self.fine_tune ) def predict(self, gen): @@ -322,6 +326,7 @@ def _ddp_and_fit( filename, world_size, linear_probe, + fine_tune, port=12355, num_heads=1, ): @@ -353,6 +358,7 @@ def _ddp_and_fit( ddp_enabled=True, num_heads=num_heads, linear_probe=linear_probe, + fine_tune=fine_tune, ) dist.destroy_process_group() @@ -371,12 +377,12 @@ def _fit( ddp_enabled: bool, num_heads: int, linear_probe=False, + fine_tune=False ): - if linear_probe: - optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) - else: - optimizer = configure_adamw(model, lr=learning_rate) + + optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) + scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( optimizer=optimizer, warmup_steps=int(len(train_gen) * 0.25), @@ -401,6 +407,7 @@ def _fit( criterion, unmap_criterion, linear_probe, + fine_tune, num_heads) after = datetime.now() print(f'Epoch {epoch}: stop training at {after}') @@ -421,10 +428,10 @@ def _fit( predictions, true = val_res del val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() - start_head = num_heads-1 if linear_probe else 0 + start_head = num_heads-1 if (linear_probe or fine_tune) else 0 for head in range(start_head, num_heads): predictions_head = predictions[..., head].flatten().numpy() - if linear_probe: + if linear_probe or fine_tune: true_head = true[..., 0].flatten().numpy() else: true_head = true[..., head].flatten().numpy() @@ -465,10 +472,14 @@ def _fit( print('Completed training!') -def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, linear_probe, num_heads=1): +def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, linear_probe, fine_tune, num_heads=1): if linear_probe: model.eval() model.core.heads[-1].train() + elif fine_tune: + model.train() + for head in model.core.heads[:-1]: + head.eval() else: model.train() @@ -478,7 +489,7 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ # pbar if on rank 0 train_gen = tqdm(train_gen) - start_head = num_heads-1 if linear_probe else 0 + start_head = num_heads-1 if linear_probe or fine_tune else 0 for X_i, m_i, y_i in train_gen: X_i = X_i.to(rank) @@ -494,13 +505,17 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_ # only use the last head base_loss = 0 for head in range(start_head, num_heads): - base_loss += criterion(output[head], y_i[..., head:head+1]) + if linear_probe or fine_tune: + base_loss += criterion(output[head], y_i) + else: + base_loss += criterion(output[head], y_i[..., head:head+1]) loss = base_loss + unmap_loss else: output = model(X_i) loss = 0 for head in range(start_head, num_heads): - if linear_probe: + if linear_probe or fine_tune: + print("Training head ", head, "out of", len(output)) loss += criterion(output[head], y_i) else: loss += criterion(output[head], y_i[..., head:head+1]) From d994c5ecd3b9dd9dc59824c9261cc611fb5ec46c Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 30 Jun 2026 13:14:33 +0200 Subject: [PATCH 19/37] Remove local plotting and tutorial files from repository --- plot_13eval.py | 94 ------------------- plot_14heads.py | 128 -------------------------- plot_lp_13.py | 111 ---------------------- plot_validation.py | 213 ------------------------------------------- plots.py | 93 ------------------- tutorials/train20.py | 170 ---------------------------------- tutorials/train_t.py | 105 --------------------- 7 files changed, 914 deletions(-) delete mode 100644 plot_13eval.py delete mode 100644 plot_14heads.py delete mode 100644 plot_lp_13.py delete mode 100644 plot_validation.py delete mode 100644 plots.py delete mode 100644 tutorials/train20.py delete mode 100644 tutorials/train_t.py diff --git a/plot_13eval.py b/plot_13eval.py deleted file mode 100644 index 518fdc0..0000000 --- a/plot_13eval.py +++ /dev/null @@ -1,94 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt - -datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90'] -chroms = [1, 11, 20, 13] - -# combined, 13 epochs (time limit) -HCT116_13head = [0.748, 0.747, 0.731, 0.742] -A549_RGS_13head = [0.736, 0.733, 0.707, 0.718] -WTC11_13head = [0.791, 0.783, 0.777, 0.783] -GM23338_13head = [0.796, 0.786, 0.782, 0.799] -HG03432_13head = [0.74, 0.735, 0.738, 0.74] -MCF7_13head = [0.738, 0.734, 0.695, 0.742] -PC3_13head = [0.714, 0.721, 0.699, 0.714] -Panc1_13head = [0.766, 0.759, 0.746, 0.745] -RWPE2_13head = [0.717, 0.716, 0.692, 0.705] -GM12878_XSC_13head = [0.703, 0.699, 0.692, 0.69] -HEPG2_GJU_13head = [0.712, 0.711, 0.685, 0.708] -K562_FGK_13head = [0.714, 0.715, 0.722, 0.701] -IMR90_13head = [0.738, 0.726, 0.73, 0.709] - -data_13heads = [ - HCT116_13head, - A549_RGS_13head, - WTC11_13head, - GM23338_13head, - HG03432_13head, - MCF7_13head, - PC3_13head, - Panc1_13head, - RWPE2_13head, - GM12878_XSC_13head, - HEPG2_GJU_13head, - K562_FGK_13head, - IMR90_13head -] - -# 13 separate models -HCT116_single = [0.758] -A549_RGS_single = [0.737] -WTC11_single = [0.8] -GM23338_single = [0.821] -HG03432_single = [0.715] -MCF7_single = [0.738] -PC3_single = [0.7] -Panc1_single = [0.744] -RWPE2_single = [0.745] -GM12878_XSC_single = [0.701, 0.695, 0.692, 0.687] -HEPG2_GJU_single = [0.704, 0.704, 0.682, 0.701] -K562_FGK_single = [0.711, 0.713, 0.723, 0.698] -IMR90_single = [0.739, 0.734, 0.731, 0.714] - -data_single = [ - HCT116_single, - A549_RGS_single, - WTC11_single, - GM23338_single, - HG03432_single, - MCF7_single, - PC3_single, - Panc1_single, - RWPE2_single, - GM12878_XSC_single, - HEPG2_GJU_single, - K562_FGK_single, - IMR90_single -] - - -# Compute mean and std across chroms -means_13head = [np.mean(x) for x in data_13heads] -stds_13head = [np.std(x) for x in data_13heads] -means_single = [np.mean(x) for x in data_single] -stds_single = [np.std(x) for x in data_single] - -# Plot -plt.figure(figsize=(12, 6)) - -width = 0.35 -x = np.arange(len(datasets)) -plt.errorbar(x - width/4, means_13head, yerr=stds_13head, - fmt='o', capsize=5, label='1 model (13 heads)', color='#1f77b4') - -plt.errorbar((x + width/4)[-4:], means_single[-4:], yerr=stds_single[-4:], - fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') - -plt.xlabel("Dataset", fontsize=13) -plt.ylabel("Pearson's R", fontsize=13) -plt.title("Model Comparison Across Datasets (13 headmodel vs single models)") -plt.xticks(x, datasets, fontsize=8) -plt.legend(frameon=False) - -plt.tight_layout() -plt.savefig("13heads_eval.png", dpi=300, bbox_inches='tight') \ No newline at end of file diff --git a/plot_14heads.py b/plot_14heads.py deleted file mode 100644 index b1deedc..0000000 --- a/plot_14heads.py +++ /dev/null @@ -1,128 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt - -datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90', "T_cell_f_21_lp"] -chroms = [1, 11, 20, 13] - -# combined, 13 epochs (time limit) -HCT116_13head = [0.748, 0.747, 0.731, 0.742] -A549_RGS_13head = [0.736, 0.733, 0.707, 0.718] -WTC11_13head = [0.791, 0.783, 0.777, 0.783] -GM23338_13head = [0.796, 0.786, 0.782, 0.799] -HG03432_13head = [0.74, 0.735, 0.738, 0.74] -MCF7_13head = [0.738, 0.734, 0.695, 0.742] -PC3_13head = [0.714, 0.721, 0.699, 0.714] -Panc1_13head = [0.766, 0.759, 0.746, 0.745] -RWPE2_13head = [0.717, 0.716, 0.692, 0.705] -GM12878_XSC_13head = [0.703, 0.699, 0.692, 0.69] -HEPG2_GJU_13head = [0.712, 0.711, 0.685, 0.708] -K562_FGK_13head = [0.714, 0.715, 0.722, 0.701] -IMR90_13head = [0.738, 0.726, 0.73, 0.709] -T_cell_f_21_lp = [0.681, 0.675, 0.666, 0.656] - -data_13heads_plus_lp = [ - HCT116_13head, - A549_RGS_13head, - WTC11_13head, - GM23338_13head, - HG03432_13head, - MCF7_13head, - PC3_13head, - Panc1_13head, - RWPE2_13head, - GM12878_XSC_13head, - HEPG2_GJU_13head, - K562_FGK_13head, - IMR90_13head, - T_cell_f_21_lp -] - -# 13 separate models -HCT116_single = [0.758] -A549_RGS_single = [0.737] -WTC11_single = [0.8] -GM23338_single = [0.821] -HG03432_single = [0.715] -MCF7_single = [0.738] -PC3_single = [0.7] -Panc1_single = [0.744] -RWPE2_single = [0.745] -GM12878_XSC_single = [0.701, 0.695, 0.692, 0.687] -HEPG2_GJU_single = [0.704, 0.704, 0.682, 0.701] -K562_FGK_single = [0.711, 0.713, 0.723, 0.698] -IMR90_single = [0.739, 0.734, 0.731, 0.714] -T_cell_f_21_separate = [0.7057, 0.701, 0.699, 0.692] - -data_single = [ - HCT116_single, - A549_RGS_single, - WTC11_single, - GM23338_single, - HG03432_single, - MCF7_single, - PC3_single, - Panc1_single, - RWPE2_single, - GM12878_XSC_single, - HEPG2_GJU_single, - K562_FGK_single, - IMR90_single, - T_cell_f_21_separate -] - - -data_14head_cl = [[0.77388, 0.77218, 0.76423, 0.76750], -[0.75115, 0.74464, 0.72093, 0.73558], -[0.80259, 0.79624, 0.79078, 0.79235], -[0.80860, 0.79816, 0.79233, 0.80874], -[0.75988, 0.75415, 0.75297, 0.75616], -[0.76864, 0.76478, 0.72526, 0.77316], -[0.73899, 0.74466, 0.72559, 0.74090], -[0.77968, 0.76998, 0.76067, 0.75639], -[0.74593, 0.74494, 0.71737, 0.73883], -[0.72286, 0.71904, 0.70844, 0.70955], -[0.72074, 0.71941, 0.69779, 0.71925], -[0.72675, 0.72961, 0.73622, 0.71564], -[0.75554, 0.74607, 0.74561, 0.72812], -[0.74056, 0.73343, 0.72759, 0.71919]] - - -# Compute mean and std across chroms -means_14head = [np.mean(x) for x in data_13heads_plus_lp] -stds_14head = [np.std(x) for x in data_13heads_plus_lp] -means_single = [np.mean(x) for x in data_single] -stds_single = [np.std(x) for x in data_single] -means_cl = [np.mean(x) for x in data_14head_cl] -stds_cl = [np.std(x) for x in data_14head_cl] - -# Plot -plt.figure(figsize=(12, 6)) - -width = 0.35 -x = np.arange(len(datasets)) - -plt.errorbar((x + width/4)[-5:], means_single[-5:], yerr=stds_single[-5:], - fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') - -plt.errorbar(x , means_14head, yerr=stds_14head, - fmt='o', capsize=5, label='1 model (13 heads, 1 head LP)', color='#1577b4') - -plt.errorbar(x - width/4, means_cl, yerr=stds_cl, - fmt='o', capsize=5, label='1 model (CL)', color='green') -for i in range(len(datasets) - 1): - plt.axvline( - x=i + 0.5, - color='gray', - linestyle=':', - linewidth=1, - alpha=0.7 - ) - -plt.xlabel("Dataset", fontsize=13) -plt.ylabel("Pearson's R", fontsize=13) -plt.title("Model Comparison Across Datasets (14 headmodel vs single models)") -plt.xticks(x, datasets, fontsize=8) -plt.legend(frameon=False) - -plt.tight_layout() -plt.savefig("14heads_eval.png", dpi=300, bbox_inches='tight') \ No newline at end of file diff --git a/plot_lp_13.py b/plot_lp_13.py deleted file mode 100644 index bf732eb..0000000 --- a/plot_lp_13.py +++ /dev/null @@ -1,111 +0,0 @@ -# request rtx4090 -# change order and rerun 19heads -# lp on all new heads -# record runtimes -# keep finetuning 13 to 14 - -import numpy as np -import matplotlib.pyplot as plt -from scipy import stats - -datasets = ["T_cell_f_21", - "nk_cell_f_41", - "t_helper_17_m_50", - "td_CD8_ab_T_m_30", - "a_B_cell_m_22_treated", - "foreski_ker_m"] - -T_cell_f_21_lp = [0.681, 0.675, 0.666, 0.656] #0.669 -T_cell_f_21_separate = [0.7057, 0.701, 0.699, 0.692] #0.699 - -nk_cell_f_41_lp = [0.702, 0.709, 0.691, 0.683] -nk_cell_f_41_separated = [0.721, 0.726, 0.716, 0.713] - -t_helper_17_m_50_lp = [0.698, 0.703, 0.684, 0.689] -t_helper_17_m_50_separated = [0.0,0,0,0] - -td_CD8_ab_T_m_30_lp = [0.695, 0.679, 0.671, 0.678] -td_CD8_ab_T_m_30_separated = [0.719, 0.703, 0.708, 0.696] - -a_B_cell_m_22_treated_lp = [0.723, 0.712, 0.706, 0.730] #0.71775 -a_B_cell_m_22_treated_separated = [0.702, 0.706, 0.703, 0.707] #0.704 - -foreski_ker_m_lp = [ 0.706, 0.701, 0.685, 0.708] -foreski_ker_m_separated = [0.744, 0.742, 0.731, 0.738] - - -data_single = [T_cell_f_21_separate, - nk_cell_f_41_separated, - t_helper_17_m_50_separated, - td_CD8_ab_T_m_30_separated, - a_B_cell_m_22_treated_separated, - foreski_ker_m_separated] -data_lp = [T_cell_f_21_lp, - nk_cell_f_41_lp, - t_helper_17_m_50_lp, - td_CD8_ab_T_m_30_lp, - a_B_cell_m_22_treated_lp, - foreski_ker_m_lp] - -# Compute means and stds6 -means_single = [np.mean(d) for d in data_single] -stds_single = [np.std(d) for d in data_single] -means_lp = [np.mean(d) for d in data_lp] -stds_lp = [np.std(d) for d in data_lp] - -def get_sig_star(p): - if p < 0.001: return '***' - elif p < 0.01: return '**' - elif p < 0.05: return '*' - else: return 'ns' - -p_values = [] -for s, h in zip(data_single, data_lp): - _, p = stats.wilcoxon(s, h) # Wilcox - p_values.append(p) - -# Plotting -x = np.arange(len(datasets)) -width = 0.35 - -plt.figure(figsize=(12, 6)) - -# Plot the points with error bars -plt.errorbar(x - width/4, means_lp, yerr=stds_lp, - fmt='o', capsize=5, label='Linear Probing', color='#1f77b4') - -plt.errorbar(x + width/4, means_single, yerr=stds_single, - fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') - -# # 2. Add Significance Brackets -# for i in range(len(datasets)): -# # Determine height of the bracket -# y_max = max(means_lp[i] + stds_lp[i], means_single[i] + stds_single[i]) -# y_line = y_max + 0.005 # Bracket baseline -# h = 0.003 # Bracket tick height - -# # Draw bracket line -# # (x_start, x_end), (y_start, y_end) -# plt.plot([x[i] - width/4, x[i] - width/4, x[i] + width/4, x[i] + width/4], -# [y_line, y_line + h, y_line + h, y_line], lw=1, c='black') - -# # Add star/ns text -# star = get_sig_star(p_values[i]) -# plt.text(x[i], y_line + h, star, ha='center', va='bottom', fontsize=12) - -# Labels and formatting -plt.xticks(x, datasets, fontsize=12) -plt.yticks(fontsize=12) -plt.xlabel("Dataset", fontsize=13) -plt.ylabel("Pearson's R", fontsize=13) -plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) -plt.legend(frameon=False) -plt.ylim(0.65, 0.78) # Adjusted to fit brackets -plt.tight_layout() - -plt.savefig("13head-lp_on_primary.png") -plt.show() - -# Print values for verification -for i, ds in enumerate(datasets): - print(f"{ds}: p-value = {p_values[i]:.4f} ({get_sig_star(p_values[i])})") diff --git a/plot_validation.py b/plot_validation.py deleted file mode 100644 index 90cbc26..0000000 --- a/plot_validation.py +++ /dev/null @@ -1,213 +0,0 @@ -import re -import matplotlib.pyplot as plt - -raw_text = """ Val pearson r: 0.6964387753816834 - Val pearson r: 0.6870994218353818 - Val pearson r: 0.7291391529590228 - Val pearson r: 0.7444582223654782 - Val pearson r: 0.6241408496209283 - Val pearson r: 0.6560767789907609 - Val pearson r: 0.6151724151653086 - Val pearson r: 0.7034725605829238 - Val pearson r: 0.6497986103344944 - Val pearson r: 0.6155472763204469 - Val pearson r: 0.6552607604094044 - Val pearson r: 0.630054931594626 - Val pearson r: 0.6185279754623397 - Val pearson r: 0.7216549256354842 - Val pearson r: 0.7055036708491718 - Val pearson r: 0.7511119742668912 - Val pearson r: 0.7635181946177819 - Val pearson r: 0.6537835362569802 - Val pearson r: 0.6835558226295194 - Val pearson r: 0.6388907881539803 - Val pearson r: 0.7190023912652485 - Val pearson r: 0.6743860834015354 - Val pearson r: 0.6469024054589176 - Val pearson r: 0.6794421250988107 - Val pearson r: 0.6472221303262831 - Val pearson r: 0.640377184784484 - Val pearson r: 0.7295442255629667 - Val pearson r: 0.7136852801487562 - Val pearson r: 0.7589122433393578 - Val pearson r: 0.7707581572647464 - Val pearson r: 0.6622760109618123 - Val pearson r: 0.6936889889396246 - Val pearson r: 0.6511044945298297 - Val pearson r: 0.7263597983993362 - Val pearson r: 0.6833798111148207 - Val pearson r: 0.6560634099565105 - Val pearson r: 0.6966597691502155 - Val pearson r: 0.6540557195808048 - Val pearson r: 0.6567890556048219 - Val pearson r: 0.7355667790321698 - Val pearson r: 0.7181773762583913 - Val pearson r: 0.762306945735433 - Val pearson r: 0.7735634438674696 - Val pearson r: 0.6709062488003141 - Val pearson r: 0.7021438188342437 - Val pearson r: 0.6612441596265453 - Val pearson r: 0.7299128731387278 - Val pearson r: 0.6938203863166237 - Val pearson r: 0.6645958744627727 - Val pearson r: 0.7062905065624483 - Val pearson r: 0.6594368702854894 - Val pearson r: 0.6751533285462077 - Val pearson r: 0.7374728764598621 - Val pearson r: 0.7196044648882977 - Val pearson r: 0.7670996660386562 - Val pearson r: 0.7780080149901397 - Val pearson r: 0.6742898696492363 - Val pearson r: 0.7058934639201252 - Val pearson r: 0.6672949258734164 - Val pearson r: 0.7312364038403028 - Val pearson r: 0.6955165547211278 - Val pearson r: 0.6682926530595027 - Val pearson r: 0.7098759765598139 - Val pearson r: 0.6632259104204266 - Val pearson r: 0.6801936203807368 - Val pearson r: 0.7396180166522474 - Val pearson r: 0.7219697031367717 - Val pearson r: 0.7700792666119827 - Val pearson r: 0.7812709213412251 - Val pearson r: 0.6768472576134353 - Val pearson r: 0.7089772560114744 - Val pearson r: 0.6724248145994438 - Val pearson r: 0.7328345447485672 - Val pearson r: 0.70032733213575 - Val pearson r: 0.6708278062936569 - Val pearson r: 0.712862836202867 - Val pearson r: 0.6696105337131105 - Val pearson r: 0.6866884996048938 - Val pearson r: 0.7411891055219084 - Val pearson r: 0.7245902689326984 - Val pearson r: 0.7708178484437126 - Val pearson r: 0.7812205607572208 - Val pearson r: 0.6805039766603462 - Val pearson r: 0.7134462216944767 - Val pearson r: 0.6765413814193819 - Val pearson r: 0.7341053857499662 - Val pearson r: 0.7045731751043073 - Val pearson r: 0.6743634373727768 - Val pearson r: 0.716458337219881 - Val pearson r: 0.6729699794543645 - Val pearson r: 0.6892055798526996 - Val pearson r: 0.7405559298361586 - Val pearson r: 0.7237466444041802 - Val pearson r: 0.7730430125437816 - Val pearson r: 0.7841607376771775 - Val pearson r: 0.679896613150622 - Val pearson r: 0.7170980883605286 - Val pearson r: 0.679135088372655 - Val pearson r: 0.7333611224807183 - Val pearson r: 0.704946036120285 - Val pearson r: 0.6750229864727569 - Val pearson r: 0.7170640526514216 - Val pearson r: 0.6754801310366083 - Val pearson r: 0.6892097632370298 - Val pearson r: 0.7457581148412622 - Val pearson r: 0.7281806164610638 - Val pearson r: 0.775695809322687 - Val pearson r: 0.7861311517116034 - Val pearson r: 0.6837907068494478 - Val pearson r: 0.7254445193158302 - Val pearson r: 0.6826638561237102 - Val pearson r: 0.7371976392755624 - Val pearson r: 0.7129005819104313 - Val pearson r: 0.6785930010014146 - Val pearson r: 0.7203582832818095 - Val pearson r: 0.6786407513936481 - Val pearson r: 0.6946855095794354 - Val pearson r: 0.7439361343698927 - Val pearson r: 0.7271142408196352 - Val pearson r: 0.7758318372488294 - Val pearson r: 0.7857087072185104 - Val pearson r: 0.6817283255573829 - Val pearson r: 0.72494311855335 - Val pearson r: 0.6845077675512952 - Val pearson r: 0.7362547679616372 - Val pearson r: 0.7108034558528062 - Val pearson r: 0.6770417456291605 - Val pearson r: 0.7195329931669482 - Val pearson r: 0.6786599228741661 - Val pearson r: 0.6946550478293225 - Val pearson r: 0.7412733773670532 - Val pearson r: 0.7245576204965146 - Val pearson r: 0.7746532810950112 - Val pearson r: 0.7848534129186068 - Val pearson r: 0.6793816030317934 - Val pearson r: 0.7236412771844641 - Val pearson r: 0.6818868719419771 - Val pearson r: 0.7336746683379142 - Val pearson r: 0.7099762496923312 - Val pearson r: 0.6744342750812531 - Val pearson r: 0.7184549823385983 - Val pearson r: 0.6760434079584795 - Val pearson r: 0.6925774202111611 - Val pearson r: 0.7424062650662471 - Val pearson r: 0.7266368686876203 - Val pearson r: 0.7755794476134695 - Val pearson r: 0.7842987925516092 - Val pearson r: 0.68165292486155 - Val pearson r: 0.7258095444959102 - Val pearson r: 0.6834873311576279 - Val pearson r: 0.7354560756083379 - Val pearson r: 0.71199184196435 - Val pearson r: 0.6763740993440086 - Val pearson r: 0.7189250457504082 - Val pearson r: 0.6788830698904805 - Val pearson r: 0.6943106983132857 - Val pearson r: 0.7450365912906849 - Val pearson r: 0.7288023282895445 - Val pearson r: 0.7765684548470875 - Val pearson r: 0.7856408276313993 - Val pearson r: 0.6818362862528512 - Val pearson r: 0.7281351907168422 - Val pearson r: 0.6848795112165097 - Val pearson r: 0.736865135898467 - Val pearson r: 0.7149251550514664 - Val pearson r: 0.6773219620250895 - Val pearson r: 0.7227276417064363 - Val pearson r: 0.6809720803882103 - Val pearson r: 0.695572395715137 - Val pearson r: 0.7416012986516195 - Val pearson r: 0.7273640296381129 - Val pearson r: 0.7773683142944958 - Val pearson r: 0.7862939212395638 - Val pearson r: 0.6804398988843954 - Val pearson r: 0.7260217080921356 - Val pearson r: 0.6836087222999384 - Val pearson r: 0.7352362622176543 - Val pearson r: 0.7121195524652735 - Val pearson r: 0.674739671521881 - Val pearson r: 0.7216534767164122 - Val pearson r: 0.679829676185234 - Val pearson r: 0.6943985773305369""" - -datasets = ['HCT116', 'A549_RGS', 'WTC11', 'GM23338', 'HG03432', 'MCF7', 'PC3', 'Panc1', 'RWPE2', 'GM12878_XSC', 'HEPG2_GJU', 'K562_FGK', 'IMR90'] - -# Extract all floating point numbers -values = [float(x) for x in re.findall(r"\d+\.\d+", raw_text)] - -heads = 13 -epochs = len(values) // heads - -# Reshape into (epochs x heads) -data = [values[i*heads:(i+1)*heads] for i in range(epochs)] - -# Transpose to get per-head time series -per_head = list(zip(*data)) - -# Plot -plt.figure() - -for i, head_vals in enumerate(per_head): - plt.plot(range(1, epochs + 1), head_vals, label=f"{datasets[i]}") - -plt.xlabel("Epoch") -plt.ylabel("Validation Pearson r") -plt.title("Validation Scores per Head over Epochs") -plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left") -plt.tight_layout() - -plt.savefig("validation_scores_13head.png") \ No newline at end of file diff --git a/plots.py b/plots.py deleted file mode 100644 index 6b80c5f..0000000 --- a/plots.py +++ /dev/null @@ -1,93 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt -from scipy import stats - -test_chroms = [1, 11, 20, 13] -datasets = ["GM12878", "K562"] - -# 4 models -GM12878_single = [0.701, 0.695, 0.692, 0.687] # 11 epochs -K562_single = [0.711, 0.713, 0.723, 0.698] # 11 epochs -HepG2_single = [0.704, 0.704, 0.682, 0.701] # 12 epochs -IMR90_single = [0.739, 0.734, 0.731, 0.714] # 11 epochs - -# one model, 2 heads -# 8 epochs -GM12878_2head = [0.706, 0.699, 0.695, 0.688] -K562_2head = [0.715, 0.718, 0.722, 0.699] - -# one model, 4 heads -# 7 epochs -GM12878_4head = [0.703, 0.702, 0.692, 0.685] -K562_4head = [0.720, 0.723, 0.730, 0.705] -HepG2_4head = [0.716, 0.718, 0.69, 0.713] -IMR90_4head = [0.742, 0.733, 0.729, 0.708] - - - - -data_single = [GM12878_single, K562_single] -data_2head = [GM12878_2head, K562_2head] - -# Compute means and stds -means_single = [np.mean(d) for d in data_single] -stds_single = [np.std(d) for d in data_single] -means_2head = [np.mean(d) for d in data_2head] -stds_2head = [np.std(d) for d in data_2head] - -def get_sig_star(p): - if p < 0.001: return '***' - elif p < 0.01: return '**' - elif p < 0.05: return '*' - else: return 'ns' - -p_values = [] -for s, h in zip(data_single, data_2head): - _, p = stats.wilcoxon(s, h) # Wilcox - p_values.append(p) - -# Plotting -x = np.arange(len(datasets)) -width = 0.35 - -plt.figure(figsize=(8, 6)) - -# Plot the points with error bars -plt.errorbar(x - width/4, means_2head, yerr=stds_2head, - fmt='o', capsize=5, label='1 model (2 heads)', color='#1f77b4') - -plt.errorbar(x + width/4, means_single, yerr=stds_single, - fmt='o', capsize=5, label='Independent Models', color='#ff7f0e') - -# 2. Add Significance Brackets -for i in range(len(datasets)): - # Determine height of the bracket - y_max = max(means_2head[i] + stds_2head[i], means_single[i] + stds_single[i]) - y_line = y_max + 0.005 # Bracket baseline - h = 0.003 # Bracket tick height - - # Draw bracket line - # (x_start, x_end), (y_start, y_end) - plt.plot([x[i] - width/4, x[i] - width/4, x[i] + width/4, x[i] + width/4], - [y_line, y_line + h, y_line + h, y_line], lw=1, c='black') - - # Add star/ns text - star = get_sig_star(p_values[i]) - plt.text(x[i], y_line + h, star, ha='center', va='bottom', fontsize=12) - -# Labels and formatting -plt.xticks(x, datasets, fontsize=12) -plt.yticks(fontsize=12) -plt.xlabel("Dataset", fontsize=13) -plt.ylabel("Pearson's R", fontsize=13) -plt.title("Model Comparison Across Datasets", fontsize=14, pad=20) -plt.legend(frameon=False) -plt.ylim(0.67, 0.74) # Adjusted to fit brackets -plt.tight_layout() - -plt.savefig("2head-comparison-with-stats.png") -plt.show() - -# Print values for verification -for i, ds in enumerate(datasets): - print(f"{ds}: p-value = {p_values[i]:.4f} ({get_sig_star(p_values[i])})") \ No newline at end of file diff --git a/tutorials/train20.py b/tutorials/train20.py deleted file mode 100644 index 55b577e..0000000 --- a/tutorials/train20.py +++ /dev/null @@ -1,170 +0,0 @@ -import sys -import os -import numpy as np -import pyBigWig - -sys.path.append(os.path.abspath("src")) - -import asap - -# The below code is a complete script that sets up the training of a model using the ASAP library. -# It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. -# The script then trains the model using the specified parameters. - - -def main(): - - leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" - - cell_lines = [ - ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], - ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], - ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], - ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], - ["HG03432", "ENCFF993BIL.bigWig", "ENCFF831FGS.bed"], - ["MCF-7", "ENCFF976UNK.bigWig", "ENCFF821OEF.bed"], - ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], - ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], - ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], - ["GM12878_XSC", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], - ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], - ["K562_FGK", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], - ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] - ] - - signal_files_cell_lines = [f"{leomed_path}/{dataset[0]}.bigWig" for dataset in cell_lines] - signal_files_cell_lines[0] = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/HCT116.bigwig" - peak_files_cell_lines = [f"{leomed_path}/{dataset[0]}.bed" for dataset in cell_lines] - - signal_files_primary = ["/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bigWig", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/nk_cell_f_41.bigWig", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/t_helper_17_m_50.bigWig", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/td_CD8_ab_T_m_30.bigWig", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/foreski_ker_m.bigWig"] - peak_files_primary = ["/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/T_cell_f_21.bed", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/nk_cell_f_41.bed", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/t_helper_17_m_50.bed", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/td_CD8_ab_T_m_30.bed", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed", - "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/foreski_ker_m.bed"] - signal_files = signal_files_cell_lines + signal_files_primary - peak_files = peak_files_cell_lines + peak_files_primary - - genome = "data/hg38.fa" - blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] - unmap_file = "data/basenji_unmappable.bed" - generated = "tmp" - logs_dir = "tmp/logs" - - # Model parameters - model_name = "convnext_dcnn" - experiment_name = "all19" - - # Training parameters - val_chroms = [1, 11, 20, 13] - train_chroms = [2, 10, 14, 19, 21] - test_chroms = [x for x in range(1, 23) if x not in val_chroms and x not in train_chroms] - n_gpus = 1 - - - # Create the training and validation datasets - print("create the dataset") - train_comb, val_comb = asap.training_datasets( - signal_file=signal_files, - genome=genome, - train_chroms=train_chroms, - val_chroms=val_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - - - print("start to train!!!\n") - - # Train the model - asap.train_multiheaded_model( - experiment_name=experiment_name, - model=model_name, - num_heads=len(signal_files), - train_dataset=train_comb, - val_dataset=val_comb, - logs_dir=logs_dir, - n_gpus=n_gpus, - ) - - print("Training done.") - print("Create eval ds") - print() - - peak = [] - for i in range(len(signal_files)): - print("Evaluating ", peak_files[i]) - peak.append(asap.peak_dataset( - signal_file=signal_files[i], - peak_file=peak_files[i], - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - ) - - # Evaluate the model - peak_scores_head = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak[i], - logs_dir=logs_dir, - num_heads=len(signal_files), - target_head=i, - ) - print(f"Peak scores head {i}:", peak_scores_head) - print() - print() - - print("Peak scores bad") - peak_scores_bad = asap.eval_multihead_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak[0], - logs_dir=logs_dir, - num_heads=len(signal_files), - target_head=1, - ) - print("Peak scores bad:", peak_scores_bad) - print() - - print("Finished") - - - - -if __name__ == "__main__": - print("hi") - # sudo mount -a - main() - - # bigwig (signl) - - # import pyBigWig - # bw = pyBigWig.open(signal_file) - # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] - # print(bw.chroms()) [chr1:463772] {chrom:len} - - # fasta (genome): AGGGCAAAA... - - # bed (blacklist): chr1 10468 11447 - # (unmapabble region): if 65% overlap, remove 2046 window - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? diff --git a/tutorials/train_t.py b/tutorials/train_t.py deleted file mode 100644 index 0d1d40a..0000000 --- a/tutorials/train_t.py +++ /dev/null @@ -1,105 +0,0 @@ -import sys -import os -import numpy as np -import pyBigWig - -sys.path.append(os.path.abspath("src")) - -import asap - -# The below code is a complete script that sets up the training of a model using the ASAP library. -# It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. -# The script then trains the model using the specified parameters. - - -def main(): - genome = "data/hg38.fa" - blacklist_file = ["data/basenji_blacklist.bed"] - unmap_file = "data/basenji_unmappable.bed" - generated = "tmp" - logs_dir = "tmp/logs" - - # Model parameters - model_name = "convnext_dcnn" - experiment_name = "a_B_cell_m_22_treated" - - # Training parameters - test_chroms = [1, 11, 20, 13] - train_chroms = [2, 10, 14, 19, 21] - val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] - n_gpus = 1 - - - signal_file = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig" - peak_file = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed" - - train, val = asap.training_datasets( - signal_file=signal_file, - genome=genome, - train_chroms=train_chroms, - val_chroms=val_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - - print("training a new head") - asap.train_model( - experiment_name=experiment_name, - model=model_name, - train_dataset=train, - val_dataset=val, - logs_dir=logs_dir, - n_gpus=n_gpus, - max_epochs=20 - ) - print() - print("Create new eval dataset") - peak = asap.peak_dataset( - signal_file=signal_file, - peak_file=peak_file, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - print() - print("Eval") - peak_scores = asap.eval_model( - experiment_name=experiment_name, - model=model_name, - eval_dataset=peak, - logs_dir=logs_dir, - ) - print(f"Peak scores: ", peak_scores) - - - - -if __name__ == "__main__": - print("hi") - # sudo mount -a - main() - - # bigwig (signl) - - # import pyBigWig - # bw = pyBigWig.open(signal_file) - # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] - # print(bw.chroms()) [chr1:463772] {chrom:len} - - # fasta (genome): AGGGCAAAA... - - # bed (blacklist): chr1 10468 11447 - # (unmapabble region): if 65% overlap, remove 2046 window - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? From 2eaff9b87bdcc08a480b32085f33e02358137c32 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 30 Jun 2026 13:22:47 +0200 Subject: [PATCH 20/37] Update docstrings --- src/asap/__init__.py | 2 +- src/asap/task.py | 46 +++++++++++++++++++++++++++----------------- 2 files changed, 29 insertions(+), 19 deletions(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index 7b667cf..43c772f 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,5 +1,5 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head, train_new_head_ft, train_multiheaded_model, eval_multihead_model +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model diff --git a/src/asap/task.py b/src/asap/task.py index af072cb..ea2c099 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -77,17 +77,21 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): ''' - Evaluate the model on the given dataset. + Add a new head to a bse model and train it using finetuining + Args: - base_experiment_name (str): The name of the model whose weights will be loaded. - new_experiment_name (str): The name of the experiment for the new head. - model (str): The model to evaluate. + base_experiment_name(str): The name of the base model + new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved + model (str): The model to train. train_dataset: The training dataset. - val_dataset: The validation dataset for early stopping. - logs_dir (str): The directory to load model checkpoints from. - batch_size (int): The batch size for evaluation. - use_map (bool): If mappability information was used during training. - num_heads (int): The number of heads for the new model (base model number of heads+1) + val_dataset: The validation dataset. + logs_dir (str): The directory to save logs. + n_gpus (int): The number of GPUs to use for training. + max_epochs (int): The maximum number of epochs to train. + learning_rate (float): The learning rate for the optimizer. + batch_size (int): The batch size for training. + use_map (bool): Whether to use mappability for training. + num_heads (int): The number of heads (=num of prediction signals) of the new model (=num base heads + 1) ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 @@ -136,17 +140,23 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model print("trained a new model") -def train_new_head(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): +def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): ''' - Evaluate the model on the given dataset. + Add a new head to a bse model and train it using linear probing + Args: - base_experiment_name (str): The name of the model whose weights will be loaded. - new_experiment_name (str): The name of the experiment for the new head. - model (str): The model to evaluate. - eval_dataset: The evaluation dataset. - logs_dir (str): The directory to load model checkpoints from. - batch_size (int): The batch size for evaluation. - use_map (bool): If mappability information was used during training. + base_experiment_name(str): The name of the base model + new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved + model (str): The model to train. + train_dataset: The training dataset. + val_dataset: The validation dataset. + logs_dir (str): The directory to save logs. + n_gpus (int): The number of GPUs to use for training. + max_epochs (int): The maximum number of epochs to train. + learning_rate (float): The learning rate for the optimizer. + batch_size (int): The batch size for training. + use_map (bool): Whether to use mappability for training. + num_original_heads (int): The number of heads (=num of prediction signals) of the old model ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 From c8a6ca55d2186fc3fac0b89e0f0a701d092985ff Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Tue, 30 Jun 2026 13:26:23 +0200 Subject: [PATCH 21/37] Restore tutorials/train.py from main --- tutorials/train.py | 156 ++++++--------------------------------------- 1 file changed, 21 insertions(+), 135 deletions(-) diff --git a/tutorials/train.py b/tutorials/train.py index 1eeb6d7..3836023 100644 --- a/tutorials/train.py +++ b/tutorials/train.py @@ -1,73 +1,32 @@ -import sys -import os -import numpy as np -import pyBigWig - -sys.path.append(os.path.abspath("src")) - import asap # The below code is a complete script that sets up the training of a model using the ASAP library. # It includes data paths, model parameters, training parameters, and the creation of training and validation datasets. # The script then trains the model using the specified parameters. - def main(): - leomed_path = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw" - - datasets = [ - ["HCT116", "ENCFF624HRW.bigWig", "ENCFF296ZZB.bed"], - ["A549_RGS", "ENCFF399KCR.bigWig", "ENCFF899OMR.bed"], - ["WTC11", "ENCFF123YPY.bigWig", "ENCFF321VDH.bed"], - ["GM23338", "ENCFF234AYB.bigWig", "ENCFF567ZCX.bed"], - ["HG03432", "ENCFF993BIL.bigWig", "ENCFF831FGS.bed"], - ["MCF-7", "ENCFF976UNK.bigWig", "ENCFF821OEF.bed"], - ["PC-3", "ENCFF145UAD.bigWig", "ENCFF811MOZ.bed"], - ["Panc1", "ENCFF794CNJ.bigWig", "ENCFF182SSP.bed"], - ["RWPE2", "ENCFF881UWW.bigWig", "ENCFF729MMJ.bed"], - ["GM12878_XSC", "ENCFF667MDI.bigWig", "ENCFF748UZH.bed"], - ["HEPG2_GJU", "ENCFF262URW.bigWig", "ENCFF439EIO.bed"], - ["K562_FGK", "ENCFF357GNC.bigWig", "ENCFF333TAT.bed"], - ["IMR90", "ENCFF770EAV.bigWig", "ENCFF243NTP.bed"] - ] - - print(len(datasets)) - signal_files = [f"{leomed_path}/{dataset[0]}.bigWig" for dataset in datasets] - signal_files[0] = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/HCT116.bigwig" - peak_files = [f"{leomed_path}/{dataset[0]}.bed" for dataset in datasets] - - print(signal_files, peak_files) - - - genome = "data/hg38.fa" - blacklist_file = ["data/basenji_blacklist.bed", "data/example_snv.vcf"] - unmap_file = "data/basenji_unmappable.bed" - generated = "tmp" - logs_dir = "tmp/logs" + # Data paths + signal_file = "../data/TCGA-A6-A567/TCGA-A6-A567.nodup.no_chrM_MT.tn5.pval.signal.bigwig" + genome = "../data/hg38.fa" + blacklist_file = ["../data/basenji_blacklist.bed", "../data/example_snv.vcf"] + unmap_file = "../data/basenji_unmappable.bed" + generated = "../tmp" + logs_dir = "../tmp/logs" # Model parameters model_name = "convnext_dcnn" - experiment_name = "allCellLines13_1gpu" + experiment_name = "TCGA-A6-A567_convnext_dcnn" # Training parameters - test_chroms = [1, 11, 20, 13] - train_chroms = [2, 10, 14, 19, 21] - val_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in train_chroms] - n_gpus = 1 - - # change test for val for the next - - - # linear probing for a new head - print("create a new dataset") - experiment_name_new_head = f"{experiment_name}-linProbe-a_B_cell_m_22_treated" - - signal_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bigWig" - peak_file_new_head = "/cluster/work/boeva/mindilewitsc/UniversalEPI/data/atac/raw/a_B_cell_m_22_treated.bed" - - train_lp, val_lp = asap.training_datasets( - signal_file=signal_file_new_head, + test_chroms = [2, 10, 14, 19, 21] + val_chroms = [1, 11, 20, 13] + train_chroms = [x for x in range(1, 23) if x not in test_chroms and x not in val_chroms] + n_gpus = 4 + + # Create the training and validation datasets + train, val = asap.training_datasets( + signal_file=signal_file, genome=genome, train_chroms=train_chroms, val_chroms=val_chroms, @@ -76,88 +35,15 @@ def main(): unmap_file=unmap_file, ) - print("training a new head") - asap.train_new_head( - base_experiment_name=experiment_name, - new_experiment_name=experiment_name_new_head, + # Train the model + asap.train_model( + experiment_name=experiment_name, model=model_name, - train_dataset=train_lp, - val_dataset=val_lp, + train_dataset=train, + val_dataset=val, logs_dir=logs_dir, n_gpus=n_gpus, - num_original_heads=len(datasets), - max_epochs=20 - ) - - print("Create new eval dataset") - peak_new_head = asap.peak_dataset( - signal_file=signal_file_new_head, - peak_file=peak_file_new_head, - genome=genome, - chroms=test_chroms, - generated=generated, - blacklist_file=blacklist_file, - unmap_file=unmap_file, - ) - - print("Eval new head") - peak_scores_last_head = asap.eval_multihead_model( - experiment_name=experiment_name_new_head, - model=model_name, - eval_dataset=peak_new_head, - logs_dir=logs_dir, - num_heads=len(signal_files)+1, - target_head=len(signal_files), ) - print(f"Peak scores new head:", peak_scores_last_head) - - # # Re-evaluate to make sure the body was not modified - # for i in range(len(signal_files)): - # print("Reevaluating ", datasets[i][0]) - # peak_dataset = asap.peak_dataset( - # signal_file=signal_files[i], - # peak_file=peak_files[i], - # genome=genome, - # chroms=test_chroms, - # generated=generated, - # blacklist_file=blacklist_file, - # unmap_file=unmap_file, - # ) - # peak_scores_head = asap.eval_multihead_model( - # experiment_name=experiment_name_new_head, - # model=model_name, - # eval_dataset=peak_dataset, - # logs_dir=logs_dir, - # num_heads=len(signal_files)+1, - # target_head=i, - # ) - # print(f"Peak scores head {i}:", peak_scores_head) - - if __name__ == "__main__": - print("hi") - # sudo mount -a main() - - # bigwig (signl) - - # import pyBigWig - # bw = pyBigWig.open(signal_file) - # print(bw.values("chr1", 100000, 100100)) [0.1,0.9,...] - # print(bw.chroms()) [chr1:463772] {chrom:len} - - # fasta (genome): AGGGCAAAA... - - # bed (blacklist): chr1 10468 11447 - # (unmapabble region): if 65% overlap, remove 2046 window - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? - - -# send means for 2 and 4 -# Train 13 to validate -# Do Linear Probing On Primary cells for the 13 cells , compare to Alan? From 7dae871fde6f43b66a82b3a04bd13c4b30453132 Mon Sep 17 00:00:00 2001 From: Claudiu Craciun Date: Fri, 7 Aug 2026 20:49:33 +0200 Subject: [PATCH 22/37] Add target head to predict_snv_atac --- src/asap/snv/predict.py | 6 +++--- src/asap/task.py | 6 ++++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/asap/snv/predict.py b/src/asap/snv/predict.py index 619609d..8de81f4 100644 --- a/src/asap/snv/predict.py +++ b/src/asap/snv/predict.py @@ -20,7 +20,7 @@ def _idx_to_ohe(idx: np.ndarray) -> np.ndarray: one_hot_encoded = eye[idx] return one_hot_encoded -def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device) -> None: +def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device, target_head=0) -> None: # Ensure output columns exist output_columns = ['signal_true', 'signal_pred_ref', 'signal_pred_alt'] for col in output_columns: @@ -56,8 +56,8 @@ def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: n x = _idx_to_ohe(x).astype(np.float32) x_var = _idx_to_ohe(x_var).astype(np.float32) - ref_pred = model(torch.from_numpy(x).to(device)) - var_pred = model(torch.from_numpy(x_var).to(device)) + ref_pred = model(torch.from_numpy(x).to(device))[target_head] + var_pred = model(torch.from_numpy(x_var).to(device))[target_head] for (i, _), true, ref, var in zip(chr_df_i.iterrows(), y, ref_pred.cpu().detach().numpy(), var_pred.cpu().detach().numpy()): snv.loc[i, 'signal_true'] = json.dumps(true.flatten().tolist()) snv.loc[i, 'signal_pred_ref'] = json.dumps(ref.tolist()) diff --git a/src/asap/task.py b/src/asap/task.py index 5275769..9ecd287 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -400,7 +400,8 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, scores[chrom] = {'cov': float(np.nanmean(cov)), 'cov_per_bin': {f'bin_{i}': float(cov_per_bin[i]) for i in range(len(cov_per_bin))}} return scores -def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_file: str, logs_dir: str, out_dir: str, genome: str, chroms: List[int]=[*range(1,23)], use_map: bool=False, export_bigwig: str=None, scale: dict | float=1.0): +def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_file: str, logs_dir: str, out_dir: str, genome: str, chroms: List[int]=[*range(1,23)], use_map: bool=False, export_bigwig: str=None, scale: dict | float=1.0, num_heads=1, + target_head=0): """ Predict ATAC-seq for SNVs using the trained model. Args: @@ -427,7 +428,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil snv_file_name = pathlib.Path(snv_file).stem.split('.')[0] # Initialize the model - model = _get_model(model, use_map=use_map) + model = _get_model(model, use_map=use_map, num_heads=num_heads) checkpoint_path = pathlib.Path(logs_dir) / experiment_name / 'checkpoint.pth' device = 'cuda' if torch.cuda.is_available() else 'cpu' state_dict = torch.load(checkpoint_path, map_location=device) @@ -451,6 +452,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil window_size=window_size, bin_size=bin_size, device=device, + target_head=target_head ) # Save the results to a CSV file From ccbbf1df1dfa6a72a7ebd1f2088f4bf1f1aa9c26 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:11:45 +0200 Subject: [PATCH 23/37] restore correct optimizer --- src/asap/trainer/trainer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index f3df5c5..a622004 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -381,7 +381,10 @@ def _fit( ): - optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) + if linear_probe: + optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate) + else: + optimizer = configure_adamw(model, lr=learning_rate) scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR( optimizer=optimizer, From 46d2834ad5e4910cea2e800ba877c517d01c9c47 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:15:49 +0200 Subject: [PATCH 24/37] functions --- src/asap/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index 43c772f..7eb38e8 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,5 +1,5 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model, train_multiheaded_model_continually, extend_multiheaded_model_continually From 5de8c9ace83c4ea1bd6fb9c195f684b69f435fe3 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:23:30 +0200 Subject: [PATCH 25/37] update jt function names --- src/asap/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/asap/__init__.py b/src/asap/__init__.py index 7eb38e8..f2bdad5 100644 --- a/src/asap/__init__.py +++ b/src/asap/__init__.py @@ -1,5 +1,5 @@ from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset -from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model, train_multiheaded_model_continually, extend_multiheaded_model_continually +from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model, train_multiheaded_model_progressively, extend_multiheaded_model_progressively From 4c19ea63fdc1183a71819b42f45e5342fc625ca5 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:26:57 +0200 Subject: [PATCH 26/37] added progressive training from scratch --- src/asap/task.py | 122 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) diff --git a/src/asap/task.py b/src/asap/task.py index 9ecd287..924c106 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -347,6 +347,128 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs scores[chrom] = {key: result_metrics[key] for key in ['pearson_r', 'mse', 'poisson_nll', 'spearman_r', 'kendall_tau']} return scores +def train_multiheaded_model_progressively( + experiment_name: str, + model: str, + train_dataset: List[BaseDataset], + val_dataset: List[BaseDataset], + logs_dir: str, + n_gpus: int=0, + max_epochs: int=20, + learning_rate: float=1e-3, + batch_size: int=64, + use_map: bool=False, + num_heads: List[int] = [1, 2], +): + ''' + Train the model with the given datasets and parameters progressively. In Progressive Joint Training (PJT), + the number of heads should be strictly increasing at every step. In NJT*, the number of heads should be constant, + equal to the final number of heads. + + Args: + experiment_name (str): The base name of the experiments. Actual names are {experiment_name}_{i}_{num_heads[i]}, with i the current step. + model (str): The model to train. + train_dataset (List[BaseDataset]): The training dataset. + val_dataset (List[BaseDataset]): The validation dataset. + logs_dir (str): The directory to save logs. + n_gpus (int): The number of GPUs to use for training. + max_epochs (int): The maximum number of epochs to train (per step). + learning_rate (float): The learning rate for the optimizer. + batch_size (int): The batch size for training. + use_map (bool): Whether to use mappability for training. + num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. + ''' + + # Validate num_heads + if (len(num_heads) < 2): + raise ValueError("At least two values must be provided for num_heads when training progressively.") + for i in range(len(num_heads)): + if num_heads[i] <= 0: + raise ValueError("Cannot train on less than one head at any point.") + if i > 0: + if num_heads [i-1] > num_heads[i]: + raise ValueError("Number of heads must be nondecreasing.") + + + # Check if gpu is available + if n_gpus > 0 and not torch.cuda.is_available(): + n_gpus = 0 + print("No GPU available, using CPU instead.") + + # Count the number of GPUs available + if n_gpus > torch.cuda.device_count(): + n_gpus = torch.cuda.device_count() + print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") + + # Initialize the first model + model_type = model + model = _get_model(model_type, use_map=use_map, num_heads=num_heads[0]) + trainer = Trainer( + filename=experiment_name + f'_0_{num_heads[0]}', + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + num_heads=num_heads[0], + ) + + print(f'Starting training for the initial step with {num_heads[0]} heads.') + + # Train the first model + trainer.fit( + train_dset=train_dataset[0], + val_dset=val_dataset[0], + nr_epochs=max_epochs, + learning_rate=learning_rate, + ) + + print(f'Finished training for the initial step. Starting training for the next step with {num_heads[1]} heads.') + + # Progressively learn next heads + for i in range(1, len(num_heads)): + # print(f'Starting progressive training for step {i} with {num_heads[i]} heads.') + + # Initialize the next model with the previous model's weights + model_tmp = _get_model(model_type, use_map=use_map, num_heads=num_heads[i]) + + print(f'Loading best model weights from {trainer.filename}') + checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' + trainer.load_weights(checkpoint_path) + + print('Loading model weights from previous step.') + model_tmp.load_state_dict(trainer.model.state_dict(), strict=False) # strict=False should allow loading when number of heads changes + model = model_tmp + + print('Successfully loaded model weights.') + + # Initialize new trainer + print(f'Initializing trainer for step {i}.') + trainer = Trainer( + filename=experiment_name + f'_{i}_{num_heads[i]}', + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + num_heads=num_heads[i], + ) + + # Train the new model + print(f'Starting training for the step {i} with {num_heads[i]} heads.') + trainer.fit( + train_dset=train_dataset[0], + val_dset=val_dataset[0], + nr_epochs=max_epochs, + learning_rate=learning_rate, + ) + print(f'Finished training for the step {i} with {num_heads[i]} heads.') + + print('Finished progressively training the model.') + print(f'Final model is saved as {experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.') + def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17): ''' From 6a97a26892986701c0fe5e17b3e3a82663275893 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:31:28 +0200 Subject: [PATCH 27/37] added progressive extension of models --- src/asap/task.py | 107 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/src/asap/task.py b/src/asap/task.py index 924c106..e044bc4 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -469,6 +469,113 @@ def train_multiheaded_model_progressively( print('Finished progressively training the model.') print(f'Final model is saved as {experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.') +def extend_multiheaded_model_progressively( + base_experiment_name: str, + new_experiment_name: str, + model: str, + train_dataset: List[BaseDataset], + val_dataset: List[BaseDataset], + logs_dir: str, + n_gpus: int=0, + max_epochs: int=20, + learning_rate: float=1e-3, + batch_size: int=64, + use_map: bool=False, + num_heads: List[int] = [3, 4], # num_heads[0] is the number of heads in the already trained model, with num_heads[1] the first step of JT extension +): + ''' + Train the model with the given datasets and parameters progressively. Start with an already trained model + with some number of heads (num_heads[0]). The following stages can have an increasing number of heads (PJT), + or the same as the pretrained model (NJT*). + + Args: + base_experiment_name (str): The base name of the previous experiment. The model will be loaded from {base_experiment_name}_0_{num_heads[0]}. + new_experiment_name (str): The base name of the new experiment. Actual names are {new_experiment_name}_{i}_{num_heads[i]}, with i the current step (ignoring step 0). + model (str): The model to train. + train_dataset (List[BaseDataset]): The training dataset. + val_dataset (List[BaseDataset]): The validation dataset. + logs_dir (str): The directory to save logs. + n_gpus (int): The number of GPUs to use for training. + max_epochs (int): The maximum number of epochs to train (per step). + learning_rate (float): The learning rate for the optimizer. + batch_size (int): The batch size for training. + use_map (bool): Whether to use mappability for training. + num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. + ''' + # Validate num_heads + if (len(num_heads) < 2): + raise ValueError("At least two values must be provided for num_heads when training progressively.") + for i in range(len(num_heads)): + if num_heads[i] <= 0: + raise ValueError("Cannot train on less than one head at any point.") + if i > 0: + if num_heads [i-1] > num_heads[i]: + raise ValueError("Number of heads must be nondecreasing.") + + # Check if gpu is available + if n_gpus > 0 and not torch.cuda.is_available(): + n_gpus = 0 + print("No GPU available, using CPU instead.") + + # Count the number of GPUs available + if n_gpus > torch.cuda.device_count(): + n_gpus = torch.cuda.device_count() + print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.") + + # Create model with more heads + model_type = model + model = _get_model(model_type, use_map=use_map, num_heads=num_heads[0]) + + # Load the previous model + trainer = Trainer( + filename=base_experiment_name + f'_0_{num_heads[0]}', + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + num_heads=num_heads[0], + ) + + for i in range(1, len(num_heads)): + model_tmp = _get_model(model_type, use_map=use_map, num_heads=num_heads[i]) + + print(f'Loading best model weights from {trainer.filename}') + checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' + trainer.load_weights(checkpoint_path) + + print('Loading model weights from previous step.') + model_tmp.load_state_dict(trainer.model.state_dict(), strict=False) # strict=False should allow loading when number of heads changes + model = model_tmp + + print('Successfully loaded model weights.') + + # Initialize new trainer + print(f'Initializing trainer for step {i}.') + trainer = Trainer( + filename=new_experiment_name + f'_{i}_{num_heads[i]}', + model=model, + criterion=nn.PoissonNLLLoss(log_input=False), + unmap_criterion=use_map, + batch_size=batch_size, + logger=TextLogger(logs_dir=logs_dir), + n_gpus=n_gpus, + num_heads=num_heads[i], + ) + + # Train the new model + print(f'Starting training for the step {i} with {num_heads[i]} heads.') + trainer.fit( + train_dset=train_dataset[0], + val_dset=val_dataset[0], + nr_epochs=max_epochs, + learning_rate=learning_rate, + ) + print(f'Finished training for the step {i} with {num_heads[i]} heads.') + print('Finished progressively extending the model.') + print(f'Final model is saved as {new_experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.') + def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17): ''' From b1bbad6f653d4e026b88528aa1ccf25e78bf3b99 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 14:54:46 +0200 Subject: [PATCH 28/37] add validation on new heads in task.py --- src/asap/task.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/asap/task.py b/src/asap/task.py index e044bc4..cf4213e 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -359,6 +359,7 @@ def train_multiheaded_model_progressively( batch_size: int=64, use_map: bool=False, num_heads: List[int] = [1, 2], + checkpoint_on_new_heads_only: bool=False, ): ''' Train the model with the given datasets and parameters progressively. In Progressive Joint Training (PJT), @@ -377,6 +378,7 @@ def train_multiheaded_model_progressively( batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. + checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. ''' # Validate num_heads @@ -422,6 +424,7 @@ def train_multiheaded_model_progressively( val_dset=val_dataset[0], nr_epochs=max_epochs, learning_rate=learning_rate, + # validation always occurs on all heads in the very first step ) print(f'Finished training for the initial step. Starting training for the next step with {num_heads[1]} heads.') @@ -463,6 +466,7 @@ def train_multiheaded_model_progressively( val_dset=val_dataset[0], nr_epochs=max_epochs, learning_rate=learning_rate, + val_on_heads=list(range(num_heads[i-1], num_heads[i])) if checkpoint_on_new_heads_only else None, ) print(f'Finished training for the step {i} with {num_heads[i]} heads.') @@ -482,6 +486,7 @@ def extend_multiheaded_model_progressively( batch_size: int=64, use_map: bool=False, num_heads: List[int] = [3, 4], # num_heads[0] is the number of heads in the already trained model, with num_heads[1] the first step of JT extension + checkpoint_on_new_heads_only: bool=False, ): ''' Train the model with the given datasets and parameters progressively. Start with an already trained model @@ -501,6 +506,7 @@ def extend_multiheaded_model_progressively( batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. + checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. ''' # Validate num_heads if (len(num_heads) < 2): @@ -571,6 +577,7 @@ def extend_multiheaded_model_progressively( val_dset=val_dataset[0], nr_epochs=max_epochs, learning_rate=learning_rate, + val_on_heads=list(range(num_heads[i-1], num_heads[i])) if checkpoint_on_new_heads_only else None, ) print(f'Finished training for the step {i} with {num_heads[i]} heads.') print('Finished progressively extending the model.') From 762c18b8880993c571a50ca287406e41b4e15536 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Thu, 6 Aug 2026 15:02:46 +0200 Subject: [PATCH 29/37] val on new heads in trainer.py --- src/asap/trainer/trainer.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index a622004..4f4308f 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -60,7 +60,7 @@ def __init__(self, self.device = 'cpu' self.model.to(self.device) - def fit(self, train_dset, val_dset, nr_epochs, learning_rate): + def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None): print(f'Training {self.filename}...') if self.nr_devices > 1: port = 10000 + randint(0,2355) @@ -81,7 +81,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): self.linear_probe, self.fine_tune, port, - self.num_heads + self.num_heads, + val_on_heads ), nprocs=self.nr_devices ) @@ -113,7 +114,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate): ddp_enabled=False, num_heads=self.num_heads, linear_probe=self.linear_probe, - fine_tune=self.fine_tune + fine_tune=self.fine_tune, + val_on_heads=val_on_heads ) def predict(self, gen): @@ -329,6 +331,7 @@ def _ddp_and_fit( fine_tune, port=12355, num_heads=1, + val_on_heads=None ): #model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = setup_ddp(rank, world_size, model, port) @@ -359,6 +362,7 @@ def _ddp_and_fit( num_heads=num_heads, linear_probe=linear_probe, fine_tune=fine_tune, + val_on_heads=val_on_heads ) dist.destroy_process_group() @@ -377,7 +381,8 @@ def _fit( ddp_enabled: bool, num_heads: int, linear_probe=False, - fine_tune=False + fine_tune=False, + val_on_heads: Union[None, list]=None ): @@ -432,6 +437,9 @@ def _fit( del val_res predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu() start_head = num_heads-1 if (linear_probe or fine_tune) else 0 + if start_head == 0 and val_on_heads is not None: + start_head = val_on_heads[0] + # TODO verify for head in range(start_head, num_heads): predictions_head = predictions[..., head].flatten().numpy() if linear_probe or fine_tune: From d3e65987f6237f7a767c0893b0e95cb68dfac539 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Fri, 7 Aug 2026 14:40:19 +0200 Subject: [PATCH 30/37] clauded num_workers to nonzero --- src/asap/task.py | 98 +++++++++++++++++++++++-------------- src/asap/trainer/trainer.py | 29 +++++++---- 2 files changed, 81 insertions(+), 46 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index cf4213e..33d4001 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -32,10 +32,10 @@ def _get_model(model_name: str, use_map: bool = False, num_heads=1): raise ValueError(f'Unknown model name: {model_name}') -def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False): +def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_workers: int=4): """ Train the model with the given datasets and parameters. - + Args: experiment_name (str): The name of the experiment. model (str): The model to train. @@ -47,12 +47,13 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v learning_rate (float): The learning rate for the optimizer. batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. + num_workers (int): The number of DataLoader worker processes per GPU. """ # Check if gpu is available if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 print("No GPU available, using CPU instead.") - + # Count the number of GPUs available if n_gpus > torch.cuda.device_count(): n_gpus = torch.cuda.device_count() @@ -63,24 +64,25 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Initialize the trainer with the model and datasets trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, + num_workers=num_workers, ) # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): +def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): ''' Add a new head to a bse model and train it using finetuining - + Args: - base_experiment_name(str): The name of the base model + base_experiment_name(str): The name of the base model new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved model (str): The model to train. train_dataset: The training dataset. @@ -92,6 +94,7 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (int): The number of heads (=num of prediction signals) of the new model (=num base heads + 1) + num_workers (int): The number of DataLoader worker processes per GPU. ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 @@ -125,6 +128,7 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model n_gpus=n_gpus, fine_tune=True, num_heads=num_heads, + num_workers=num_workers, ) print(f'Loading best model weights from {base_experiment_name}') @@ -140,12 +144,12 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model print("trained a new model") -def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): +def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1, num_workers: int=4): ''' Add a new head to a bse model and train it using linear probing - + Args: - base_experiment_name(str): The name of the base model + base_experiment_name(str): The name of the base model new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved model (str): The model to train. train_dataset: The training dataset. @@ -157,6 +161,7 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_original_heads (int): The number of heads (=num of prediction signals) of the old model + num_workers (int): The number of DataLoader worker processes per GPU. ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 @@ -193,6 +198,7 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model n_gpus=n_gpus, linear_probe=True, num_heads=num_original_heads+1, + num_workers=num_workers, ) print(f'Loading best model weights from {base_experiment_name}') @@ -207,10 +213,10 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) print("trained a new model") -def train_multiheaded_model(experiment_name : str, model: str, train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): +def train_multiheaded_model(experiment_name : str, model: str, train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): """ Train the model with the given datasets and parameters. - + Args: experiment_name (str): The name of the experiment. model (str): The model to train. @@ -223,13 +229,14 @@ def train_multiheaded_model(experiment_name : str, model: str, train_dataset: L batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (int): The number of heads (=num of prediction signals) + num_workers (int): The number of DataLoader worker processes per GPU. """ # Check if gpu is available if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 print("No GPU available, using CPU instead.") - + # Count the number of GPUs available if n_gpus > torch.cuda.device_count(): n_gpus = torch.cuda.device_count() @@ -240,20 +247,21 @@ def train_multiheaded_model(experiment_name : str, model: str, train_dataset: L # Initialize the trainer with the model and datasets trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, - num_heads=num_heads + num_heads=num_heads, + num_workers=num_workers, ) # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, target_head:int = 0): +def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, target_head:int = 0, num_workers: int=4): ''' Evaluate the model on the given dataset. Args: @@ -263,6 +271,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat logs_dir (str): The directory to load model checkpoints from. batch_size (int): The batch size for evaluation. use_map (bool): If mappability information was used during training. + num_workers (int): The number of DataLoader worker processes. ''' n_gpus = 1 if torch.cuda.is_available() else 0 @@ -270,12 +279,12 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat model = _get_model(model, use_map=use_map, num_heads=num_heads) trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads ) @@ -284,7 +293,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat print(f'Loading best model weights from {trainer.filename}') checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' trainer.load_weights(checkpoint_path) - + test_chroms = eval_dataset.chroms scores = {} for chrom in test_chroms: @@ -292,8 +301,9 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False + batch_size=batch_size, + is_train=False, + num_workers=num_workers ) _, _, result_metrics = trainer.predict_and_evaluate_multihead(test_gen, target_head=target_head) @@ -301,7 +311,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat return scores -def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1): +def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): ''' Evaluate the model on the given dataset. Args: @@ -311,6 +321,7 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs logs_dir (str): The directory to load model checkpoints from. batch_size (int): The batch size for evaluation. use_map (bool): If mappability information was used during training. + num_workers (int): The number of DataLoader worker processes. ''' n_gpus = 1 if torch.cuda.is_available() else 0 @@ -318,12 +329,12 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs model = _get_model(model, use_map=use_map, num_heads=num_heads) trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, ) @@ -331,7 +342,7 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs print(f'Loading best model weights from {trainer.filename}') checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' trainer.load_weights(checkpoint_path) - + test_chroms = eval_dataset.chroms scores = {} for chrom in test_chroms: @@ -339,8 +350,9 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False + batch_size=batch_size, + is_train=False, + num_workers=num_workers ) _, _, result_metrics = trainer.predict_and_evaluate(test_gen) @@ -360,11 +372,12 @@ def train_multiheaded_model_progressively( use_map: bool=False, num_heads: List[int] = [1, 2], checkpoint_on_new_heads_only: bool=False, + num_workers: int=4, ): ''' Train the model with the given datasets and parameters progressively. In Progressive Joint Training (PJT), the number of heads should be strictly increasing at every step. In NJT*, the number of heads should be constant, - equal to the final number of heads. + equal to the final number of heads. Args: experiment_name (str): The base name of the experiments. Actual names are {experiment_name}_{i}_{num_heads[i]}, with i the current step. @@ -379,6 +392,7 @@ def train_multiheaded_model_progressively( use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. + num_workers (int): The number of DataLoader worker processes per GPU. ''' # Validate num_heads @@ -414,6 +428,7 @@ def train_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[0], + num_workers=num_workers, ) print(f'Starting training for the initial step with {num_heads[0]} heads.') @@ -457,6 +472,7 @@ def train_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[i], + num_workers=num_workers, ) # Train the new model @@ -487,9 +503,10 @@ def extend_multiheaded_model_progressively( use_map: bool=False, num_heads: List[int] = [3, 4], # num_heads[0] is the number of heads in the already trained model, with num_heads[1] the first step of JT extension checkpoint_on_new_heads_only: bool=False, + num_workers: int=4, ): ''' - Train the model with the given datasets and parameters progressively. Start with an already trained model + Train the model with the given datasets and parameters progressively. Start with an already trained model with some number of heads (num_heads[0]). The following stages can have an increasing number of heads (PJT), or the same as the pretrained model (NJT*). @@ -507,6 +524,7 @@ def extend_multiheaded_model_progressively( use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. + num_workers (int): The number of DataLoader worker processes per GPU. ''' # Validate num_heads if (len(num_heads) < 2): @@ -542,6 +560,7 @@ def extend_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[0], + num_workers=num_workers, ) for i in range(1, len(num_heads)): @@ -568,6 +587,7 @@ def extend_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[i], + num_workers=num_workers, ) # Train the new model @@ -584,7 +604,7 @@ def extend_multiheaded_model_progressively( print(f'Final model is saved as {new_experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.') -def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17): +def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17, num_workers: int=4): ''' Evaluate the robustness of the model on the given dataset. Args: @@ -595,6 +615,7 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, batch_size (int): The batch size for evaluation. use_map (bool): Whether to use mappability for evaluation. nr_samples_for_var (int): The number of samples for variance calculation. + num_workers (int): The number of DataLoader worker processes. ''' # Fixed margin size for robustness evaluation margin = 768 @@ -627,8 +648,9 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size // (nr_samples_for_var - 1), - is_train=False + batch_size=batch_size // (nr_samples_for_var - 1), + is_train=False, + num_workers=num_workers ) _, _, cov, cov_per_bin = trainer.predict_robust_batch(test_gen, nr_samples_for_var=nr_samples_for_var, window=eval_dataset.window_size, margin=margin) @@ -837,7 +859,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil -def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, out_dir: str, batch_size: int=64, use_map: bool=False): +def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, out_dir: str, batch_size: int=64, use_map: bool=False, num_workers: int=4): """ Export the predictions to a file. Args: @@ -848,6 +870,7 @@ def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDatas out_dir (str): The output directory for predictions. batch_size (int): The batch size for evaluation. use_map (bool): Whether to use mappability for evaluation. + num_workers (int): The number of DataLoader worker processes. """ n_gpus = 1 if torch.cuda.is_available() else 0 @@ -879,8 +902,9 @@ def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDatas test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False + batch_size=batch_size, + is_train=False, + num_workers=num_workers ) _, predictions, _ = trainer.predict_and_evaluate(test_gen, no_eval=True) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 4f4308f..ea2273f 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -32,21 +32,23 @@ def __init__(self, num_heads: int = 1, linear_probe=False, fine_tune=False, + num_workers: int = 4, ): self.filename = filename self.model = model - self.criterion = criterion + self.criterion = criterion self.unmap_criterion = nn.MSELoss() if unmap_criterion is True else None self.train_unmap = not self.unmap_criterion is False self.logger: Logger = logger - self.logspace = True + self.logspace = True self.nr_tracks = 1 self.nr_devices = n_gpus self.batch_size = batch_size self.num_heads = num_heads self.linear_probe = linear_probe self.fine_tune = fine_tune + self.num_workers = num_workers if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' @@ -82,7 +84,8 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None) self.fine_tune, port, self.num_heads, - val_on_heads + val_on_heads, + self.num_workers ), nprocs=self.nr_devices ) @@ -92,13 +95,15 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None) ddp_enabled=False, dataset=train_dset, batch_size=self.batch_size, - is_train=True + is_train=True, + num_workers=self.num_workers ) val_gen = make_dataloader( ddp_enabled=False, dataset=val_dset, batch_size=self.batch_size, - is_train=False + is_train=False, + num_workers=self.num_workers ) _fit( self.device, @@ -276,6 +281,7 @@ def load_weights(self, path): self.model.load_state_dict(state_dict) def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_workers: int = 0, pin_memory: bool = True): + persistent_workers = num_workers > 0 # if using DDP, use DistributedSampler if ddp_enabled: sampler = torch.utils.data.distributed.DistributedSampler( @@ -288,6 +294,7 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w pin_memory=pin_memory, shuffle=False, # shuffling handled by sampler num_workers=num_workers, + persistent_workers=persistent_workers, sampler=sampler ) else: @@ -297,7 +304,8 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w batch_size=batch_size, pin_memory=pin_memory, shuffle=is_train, - num_workers=num_workers + num_workers=num_workers, + persistent_workers=persistent_workers ) @@ -331,7 +339,8 @@ def _ddp_and_fit( fine_tune, port=12355, num_heads=1, - val_on_heads=None + val_on_heads=None, + num_workers=4, ): #model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = setup_ddp(rank, world_size, model, port) @@ -339,13 +348,15 @@ def _ddp_and_fit( ddp_enabled=True, dataset=train_dset, batch_size=batch_size, - is_train=True + is_train=True, + num_workers=num_workers ) val_gen = make_dataloader( ddp_enabled=True, dataset=val_dset, batch_size=batch_size, - is_train=False + is_train=False, + num_workers=num_workers ) _fit( rank=rank, From ed96b450229375c072d018952c68a48a5a224c21 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Fri, 7 Aug 2026 15:50:03 +0200 Subject: [PATCH 31/37] memmap --- src/asap/dataset.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/src/asap/dataset.py b/src/asap/dataset.py index 162b971..8bb2faa 100644 --- a/src/asap/dataset.py +++ b/src/asap/dataset.py @@ -1,7 +1,7 @@ from .dataloader import WGDataset, PeakDataset from typing import List -def training_datasets(signal_file: str, genome: str, train_chroms: List[int], val_chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): +def training_datasets(signal_file: str, genome: str, train_chroms: List[int], val_chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): ''' Create training and validation datasets for the model. Args: @@ -12,8 +12,9 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. + memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' - + train_dataset = WGDataset( genome=genome, signal_files=signal_file, @@ -29,7 +30,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va unmap_threshold=0.35, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_train=True, is_robustness=False, @@ -51,7 +52,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va unmap_threshold=0, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_train=False, is_robustness=False, @@ -60,7 +61,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va return train_dataset, val_dataset -def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): +def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): ''' Create a peak dataset for evaluation. Args: @@ -71,6 +72,7 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. + memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = PeakDataset( genome=genome, @@ -88,7 +90,7 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int unmap_threshold=0, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_robustness=False, ) @@ -96,7 +98,7 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int return dataset -def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): +def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): ''' Create a peak dataset for robustness evaluation. Args: @@ -107,6 +109,7 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. + memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = PeakDataset( genome=genome, @@ -124,7 +127,7 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom unmap_threshold=0, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_robustness=True, ) @@ -132,7 +135,7 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom return dataset -def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): +def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): ''' Create a whole genome dataset for evaluation. Args: @@ -142,6 +145,7 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. + memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = WGDataset( genome=genome, @@ -158,7 +162,7 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, unmap_threshold=0, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_train=False, is_robustness=False, @@ -167,7 +171,7 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, return dataset -def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): +def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): ''' Create a whole genome dataset for robustness evaluation. Args: @@ -177,6 +181,7 @@ def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], gene generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. + memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = WGDataset( genome=genome, @@ -193,7 +198,7 @@ def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], gene unmap_threshold=0, logspace=True, output_format="ohe", - memmap=False, + memmap=memmap, generated=generated, is_train=False, is_robustness=True, From a305c589148e1d0ac3c96c374f69cb09d792909e Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Fri, 7 Aug 2026 17:01:45 +0200 Subject: [PATCH 32/37] de-duplicated chromosome windows --- src/asap/dataloader/base.py | 56 +++--- src/asap/dataloader/bounded.py | 6 +- src/asap/dataloader/bw_to_data.py | 277 ++++++++++++++++++++++-------- src/asap/dataloader/peak.py | 6 +- 4 files changed, 239 insertions(+), 106 deletions(-) diff --git a/src/asap/dataloader/base.py b/src/asap/dataloader/base.py index e6545c6..e151e1e 100644 --- a/src/asap/dataloader/base.py +++ b/src/asap/dataloader/base.py @@ -87,15 +87,16 @@ def setup(self): pathlib.Path(self.generated).mkdir(parents=True, exist_ok=True) # chroms have been grouped by amount of bounded data per chromosome to balance folds self.chrom_lengths = [] - self.X, self.y, self.seq_starts = [], [], [] + self.chrom_seq, self.mappability, self.chrom_y, self.seq_starts = [], [], [], [] if self.signal_files is None: - self.y = None + self.chrom_y = None for chrom in self.chroms: - X_, y_, seq_starts_ = self._generate_chrom_data(chrom) - self.chrom_lengths.append(len(X_)) - self.X.append(X_) - if y_ is not None: - self.y.append(y_) + chrom_seq_, mappability_, chrom_y_, seq_starts_ = self._generate_chrom_data(chrom) + self.chrom_lengths.append(len(seq_starts_)) + self.chrom_seq.append(chrom_seq_) + self.mappability.append(mappability_) + if chrom_y_ is not None: + self.chrom_y.append(chrom_y_) self.seq_starts.append(seq_starts_) self.cum_chrom_lengths = np.cumsum(self.chrom_lengths) @@ -128,30 +129,35 @@ def __getitem__(self, index) -> Tuple[torch.Tensor, torch.Tensor]: idx = index - self.cum_chrom_lengths[chrom_idx - 1] else: idx = index - y = None + + start = self.seq_starts[chrom_idx][idx] if self.random_shift: - shift = random.randint( - 0, self.window_size // self.bin_size - ) # take a random window within step_size - x_shift = shift * self.bin_size - X = self.X[chrom_idx][ - idx, x_shift : x_shift + (self.window_size + 2 * self.margin_size) - ] - if self.y is not None: - y = self.y[chrom_idx][ - idx, shift : shift + (self.window_size // self.bin_size) - ] + # take a random window within step_size, in bp + start = start + random.randint(0, self.window_size // self.bin_size) * self.bin_size + + chrom_seq = self.chrom_seq[chrom_idx] + X = chrom_seq[start - self.margin_size : start + self.window_size + self.margin_size] + + # mappability channel is only meaningful (not all-ones filler) when + # unmap_threshold is a soft threshold -- matches the original behaviour, + # where strict filtering (unmap_threshold == 0) never produced a real + # per-window mappability channel either. + track = self.mappability[chrom_idx] + if track is not None and self.unmap_threshold != 0: + m = track[start - self.margin_size : start + self.window_size + self.margin_size].astype(np.float32) else: - X = self.X[chrom_idx][idx] - if self.y is not None: - y = self.y[chrom_idx][idx] + m = np.ones_like(X, dtype=np.float32) + m = m[..., np.newaxis] + + y = None + if self.chrom_y is not None: + y_raw = self.chrom_y[chrom_idx][start : start + self.window_size] + nr_bins = self.window_size // self.bin_size + y = y_raw.reshape(nr_bins, self.bin_size, y_raw.shape[-1]).max(axis=1) if self.logspace and y is not None: y = np.log(y + 1) - m = X[..., [1]].astype(np.float32) - X = X[..., 0] - for aug in self.augmentations: X, y = aug(X, y) diff --git a/src/asap/dataloader/bounded.py b/src/asap/dataloader/bounded.py index 413c414..5af27a6 100644 --- a/src/asap/dataloader/bounded.py +++ b/src/asap/dataloader/bounded.py @@ -51,8 +51,8 @@ def __init__( if self.chroms: self.setup() - def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: - X, y, seq_starts = bw_to_data.get_wg_filtered_data( + def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + chrom_seq, mappability, chrom_y, seq_starts = bw_to_data.get_wg_filtered_data( genome=self.genome, signal_files=self.signal_files, chrom=chrom, @@ -67,4 +67,4 @@ def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: memmap=self.memmap, generated=self.generated, ) - return X, y, seq_starts \ No newline at end of file + return chrom_seq, mappability, chrom_y, seq_starts \ No newline at end of file diff --git a/src/asap/dataloader/bw_to_data.py b/src/asap/dataloader/bw_to_data.py index 1ce5684..2e4aa52 100644 --- a/src/asap/dataloader/bw_to_data.py +++ b/src/asap/dataloader/bw_to_data.py @@ -1,9 +1,10 @@ -from typing import Tuple, List +from typing import Tuple, List, Optional import numpy as np +import pandas as pd import pyBigWig -from asap.dataloader.utils.data_bed import filter_idx_by_bed, filter_idx_by_unmap_threshold, whole_genome_idx +from asap.dataloader.utils.data_bed import filter_idx_by_bed, whole_genome_idx from asap.dataloader.utils.data_bw import get_binned_signal from asap.dataloader.utils.seq import get_chr_seq from asap.dataloader.utils.io import get_bw_from_file @@ -11,10 +12,190 @@ from tqdm import tqdm +# --------------------------------------------------------------------------- +# Whole-chromosome caches (sequence, mappability, signal). +# +# Unlike the per-window caches these used to be, these depend only on +# (genome, chrom[, signal_files/unmap_file]) -- not on window_size, margin_size, +# step_size or bin_size. They're generated once per chromosome and reused +# across every window/step/margin configuration and every epoch. +# --------------------------------------------------------------------------- + +def get_cached_chrom_seq(genome: str, chrom: int, generated: str, memmap: bool = True) -> np.ndarray: + mmap_mode = 'r' if memmap else None + id_ = hash_dn(f'chromseq: genome={genome}, chrom={chrom}', salt='0') + path = f'{generated}/{id_}_chromseq.npy' + try: + return np.load(path, mmap_mode=mmap_mode) + except FileNotFoundError: + seq = get_chr_seq(genome, chrom).astype('int8') + np.save(path, seq) + return np.load(path, mmap_mode=mmap_mode) + + +def get_cached_chrom_mappability(genome: str, chrom: int, unmappable_bed_file: Optional[str], + chrom_len: int, generated: str, memmap: bool = True) -> Optional[np.ndarray]: + if unmappable_bed_file is None: + return None + mmap_mode = 'r' if memmap else None + id_ = hash_dn(f'chrommap: genome={genome}, chrom={chrom}, unmap={unmappable_bed_file}', salt='0') + path = f'{generated}/{id_}_chrommap.npy' + try: + return np.load(path, mmap_mode=mmap_mode) + except FileNotFoundError: + track = np.ones(chrom_len, dtype='int8') + unmap = pd.read_csv(unmappable_bed_file, delimiter='\t', header=None, names=['chr', 'start', 'end']) + unmap = unmap[(unmap.chr == f'chr{chrom}') | (unmap.chr.astype(str) == f'{chrom}')] + for gap_start, gap_end in zip(unmap.start, unmap.end): + s = max(0, int(gap_start)) + e = min(chrom_len, int(gap_end)) + if e > s: + track[s:e] = 0 + np.save(path, track) + return np.load(path, mmap_mode=mmap_mode) + + +def get_cached_chrom_signal(genome: str, chrom: int, signal_files: Optional[List[str]], chrom_len: int, + generated: str, memmap: bool = True) -> Optional[np.ndarray]: + if signal_files is None: + return None + mmap_mode = 'r' if memmap else None + id_ = hash_dn(f'chromsignal: genome={genome}, chrom={chrom}, files={signal_files}', salt='0') + path = f'{generated}/{id_}_chromsignal.npy' + try: + return np.load(path, mmap_mode=mmap_mode) + except FileNotFoundError: + whole_signals = _get_binned_whole_signals(signal_files, chrom, bin_size=1, start=0, end=chrom_len) + y = whole_signals.astype('float32') + np.save(path, y) + return np.load(path, mmap_mode=mmap_mode) + + +def get_cached_chrom_data(genome: str, chrom: int, signal_files: Optional[List[str]], + unmappable_bed_file: Optional[str], generated: str, + memmap: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]: + chrom_seq = get_cached_chrom_seq(genome, chrom, generated, memmap=memmap) + chrom_len = len(chrom_seq) + mappability = get_cached_chrom_mappability(genome, chrom, unmappable_bed_file, chrom_len, generated, memmap=memmap) + chrom_y = get_cached_chrom_signal(genome, chrom, signal_files, chrom_len, generated, memmap=memmap) + return chrom_seq, mappability, chrom_y + + +# --------------------------------------------------------------------------- +# Index building: which positions are valid samples, given filtering rules. +# This is the only piece that still depends on window_size/margin_size/ +# step_size/bin_size/blacklist/unmap_threshold -- and it's cheap (an array of +# positions), so it's cached separately from the (much larger) chromosome data. +# --------------------------------------------------------------------------- + +def _unmappable_fraction(mappability_track: np.ndarray, starts: np.ndarray, extent: int) -> np.ndarray: + # Fraction of unmappable bases in [start, start + extent) for each start. + # NOTE: mirrors the original filter_idx_by_unmap_threshold behaviour exactly, + # which checks [start, start + extent) -- i.e. anchored at `start`, not + # symmetric around it -- even though the data actually read is centered + # (start - margin, start + window + margin). Preserved intentionally rather + # than "fixed", since changing it would silently change which windows are + # kept. + unmap = (1 - mappability_track).astype(np.int64) + prefix = np.concatenate(([0], np.cumsum(unmap))) + counts = prefix[starts + extent] - prefix[starts] + return counts / extent + + +def build_seq_starts_index(genome: str, chrom: int, seq_starts: np.ndarray, window_size: int, margin_size: int, + chrom_len: int, mappability_track: Optional[np.ndarray], chrom_y: Optional[np.ndarray], + bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, + unmap_threshold: float = None, lower_bound: int = None, generated: str = None) -> np.ndarray: + fingerprint = f'{len(seq_starts)}:{int(seq_starts[0]) if len(seq_starts) else 0}:' \ + f'{int(seq_starts[-1]) if len(seq_starts) else 0}:{int(seq_starts.sum())}' + str_ = (f'idx: chr={chrom}, window={window_size}, margin={margin_size}, bin={bin_size}, ' + f'candidates={fingerprint}, blacklist={blacklist_bed_files}, unmappable={unmappable_bed_file}, ' + f'unmap_th={unmap_threshold}, lower_bound={lower_bound}, genome={genome}') + id_ = hash_dn(str_, salt='0') + path = f'{generated}/{id_}_idx.npy' + try: + print(f'Attempting to load index from file with {str_}') + seq_starts = np.load(path) + print('\t...done!') + return seq_starts + except FileNotFoundError: + print('\t...not found.') + print('Generating index...') + + # Filter by blacklists + if blacklist_bed_files is not None: + for bl_bed in blacklist_bed_files: + seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size, + blacklist_bed_file=bl_bed) + + # Filter out sequences extending beyond the chrom + seq_starts = seq_starts[seq_starts + window_size + margin_size < chrom_len] + seq_starts = seq_starts[seq_starts - margin_size > 0] + + # Filter by unmappable, using the whole-chromosome mappability track. + # Extent matches the original: [seq_starts, seq_starts + window_size + 2*margin_size). + if mappability_track is not None and unmap_threshold is not None: + extent = window_size + 2 * margin_size + if unmap_threshold == 0: + # Strict mode: reuse filter_idx_by_bed directly against the unmap bed + # file (exactly as the original did), rather than deriving "zero + # overlap" from the mappability track -- the two have different + # (inclusive vs. half-open) boundary handling at exact edges, and + # this keeps strict-mode behavior bit-identical to before. + seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=extent, + blacklist_bed_file=unmappable_bed_file) + else: + # --- Correct behavior (current) --- + # Fraction of bases in [start, start+extent) that are actually + # unmappable, computed once for the whole chromosome via a prefix + # sum over the mappability track, then sliced per candidate window. + frac = _unmappable_fraction(mappability_track, seq_starts, extent) + seq_starts = seq_starts[frac <= unmap_threshold] + + # --- Old (buggy) behavior, kept here for reference --- + # The original filter_idx_by_unmap_threshold (data_bed.py) computed, + # per window, `overlap_length = sum(overlaps.end - overlaps.start)` + # using the RAW, unclamped length of every unmap interval that + # touched the window at all -- even when only a sliver of that + # interval actually fell inside the window. So a single long + # unmappable region (e.g. 60bp) could cause a window that only + # overlapped it by a handful of bp to be dropped entirely, because + # the full 60bp counted against the threshold budget instead of the + # true (much smaller) overlap. This was a bug, not an intentional + # design choice -- verified against the original by differential + # testing (see conversation/PR notes). If bit-identical window + # selection against an old cached index or trained checkpoint is + # ever needed, this reproduces the dominant effect of that bug + # (raw interval length, not clipped overlap) without reintroducing + # the original's slow per-window pandas loop -- it loops over + # unmap intervals instead, which are typically far fewer than + # candidate windows. Uncomment in place of the two lines above: + # + # unmap_df = pd.read_csv(unmappable_bed_file, delimiter='\t', header=None, + # names=['chr', 'start', 'end']) + # unmap_df = unmap_df[(unmap_df.chr == f'chr{chrom}') | (unmap_df.chr.astype(str) == f'{chrom}')] + # window_end = seq_starts + extent + # raw_overlap_length = np.zeros(len(seq_starts), dtype=np.int64) + # for gs, ge in zip(unmap_df.start.to_numpy(), unmap_df.end.to_numpy()): + # touches = (seq_starts < ge) & (gs < window_end) + # raw_overlap_length[touches] += (ge - gs) # bug: full interval length, not clipped + # seq_starts = seq_starts[raw_overlap_length <= unmap_threshold * extent] + + # Filter by lower bound + if lower_bound is not None and chrom_y is not None: + idx = seq_starts[:, np.newaxis] + np.arange(window_size) + max_signal = chrom_y[idx].max(axis=(1, 2)) + seq_starts = seq_starts[max_signal >= lower_bound] + + print(f'Generated index: {seq_starts.shape}') + np.save(path, seq_starts) + return seq_starts + + def get_wg_filtered_data(genome: str, signal_files: List[str], chrom: int, window_size: int, margin_size: int, step_size: int, bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, unmap_threshold: float = None, - lower_bound: int = None, memmap=True, generated=None) -> Tuple[np.ndarray, np.ndarray]: + lower_bound: int = None, memmap=True, generated=None): idx = whole_genome_idx(genome=genome, chrom=chrom, step_window=step_size) return idx_to_filtered_data(genome=genome, signal_files=signal_files, seq_starts=idx, chrom=chrom, window_size=window_size, @@ -25,79 +206,25 @@ def get_wg_filtered_data(genome: str, signal_files: List[str], chrom: int, windo def idx_to_filtered_data(genome: str, signal_files: List[str], seq_starts: np.ndarray, chrom: int, window_size: int, margin_size: int, bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, - unmap_threshold: float = None, lower_bound: int = None, memmap=True, generated=None) -> Tuple[np.ndarray, np.ndarray]: - mmap_mode = 'r' if memmap else None - str_ = (f'chr={chrom}, window={window_size}, margin={margin_size}, bin={bin_size}, idx={seq_starts.shape}, ' - f'blacklist={blacklist_bed_files}, unmappable={unmappable_bed_file}, th={unmap_threshold}, ' - f'lower_bound={lower_bound}, files={signal_files}, include_map={True}...') - id_ = hash_dn(str_, salt='0') - y = None - try: - print(f'Attempting to load data from file with {str_}') - x = np.load(f'{generated}/{id_}_x.npy', mmap_mode=mmap_mode) - if signal_files is not None: - y = np.load(f'{generated}/{id_}_y.npy', mmap_mode=mmap_mode) - seq_starts = np.load(f'{generated}/{id_}_seq.npy', mmap_mode=mmap_mode) - print('\t...done!') - - except FileNotFoundError: - print('\t...not found.') - print('Generating data...') - # Filter by blacklists - if blacklist_bed_files is not None: - for bl_bed in blacklist_bed_files: - seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size, - blacklist_bed_file=bl_bed) - - - # Filter out sequences extending beyond the chrom - seq = get_chr_seq(genome, chrom) - seq_starts = seq_starts[seq_starts + window_size + margin_size < len(seq)] - seq_starts = seq_starts[seq_starts - margin_size > 0] - - # Filter by unmappable - mappability = None - if unmappable_bed_file is not None: - if unmap_threshold == 0: - seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size + 2* margin_size, - blacklist_bed_file=unmappable_bed_file) - else: - seq_starts, mappability = filter_idx_by_unmap_threshold(chrom=chrom, seq_starts=seq_starts, window_size=window_size + 2* margin_size, - unmappable_bed_file=unmappable_bed_file, - threshold=unmap_threshold, return_unmap=True) - x, y = get_data_by_idx(genome, signal_files, chrom, seq_starts, window_size, margin_size, bin_size) - - # Filter by lower bound - if lower_bound is not None: - indices_over_bound = y.max(axis=(1, 2), initial=0) >= lower_bound - x = x[indices_over_bound] - if y is not None: - y = y[indices_over_bound] - seq_starts = seq_starts[indices_over_bound] - if mappability is not None: - mappability = mappability[indices_over_bound] - print(f'Data after filtering bound {lower_bound}: x={x.shape} y={y.shape}') - - x = x.astype('int8') - if mappability is None: - mappability = np.ones_like(x) - else: - mappability = mappability.astype('int8') - x = np.stack([x, mappability], axis=-1) - if y is not None: - y = y.astype('float32') - np.save(f'{generated}/{id_}_y.npy', y) - y = np.load(f'{generated}/{id_}_y.npy', mmap_mode=mmap_mode) - - # save and reload in mmap mode - np.save(f'{generated}/{id_}_x.npy', x) - np.save(f'{generated}/{id_}_seq.npy', seq_starts) - - x = np.load(f'{generated}/{id_}_x.npy', mmap_mode=mmap_mode) - seq_starts = np.load(f'{generated}/{id_}_seq.npy', mmap_mode=mmap_mode) - print('\t...done!') - return x, y, seq_starts + unmap_threshold: float = None, lower_bound: int = None, memmap=True, generated=None): + chrom_seq, mappability, chrom_y = get_cached_chrom_data( + genome=genome, chrom=chrom, signal_files=signal_files, unmappable_bed_file=unmappable_bed_file, + generated=generated, memmap=memmap, + ) + seq_starts = build_seq_starts_index( + genome=genome, chrom=chrom, seq_starts=seq_starts, window_size=window_size, margin_size=margin_size, + chrom_len=len(chrom_seq), mappability_track=mappability, chrom_y=chrom_y, bin_size=bin_size, + blacklist_bed_files=blacklist_bed_files, unmappable_bed_file=unmappable_bed_file, + unmap_threshold=unmap_threshold, lower_bound=lower_bound, generated=generated, + ) + return chrom_seq, mappability, chrom_y, seq_starts + +# --------------------------------------------------------------------------- +# Unchanged below: ad hoc per-position window extraction (used by snv/predict.py +# for one-off SNV lookups, not by the cached bulk dataset path above) and +# bigwig export. +# --------------------------------------------------------------------------- def get_data_by_idx(genome: str, signal_files: List[str], chrom: int, seq_starts: np.ndarray, window: int, margin:int, bin_size: int) -> Tuple[ np.ndarray, np.ndarray]: diff --git a/src/asap/dataloader/peak.py b/src/asap/dataloader/peak.py index 8d45c1b..5770752 100644 --- a/src/asap/dataloader/peak.py +++ b/src/asap/dataloader/peak.py @@ -51,10 +51,10 @@ def __init__( if self.chroms: self.setup() - def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: + def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: peak_centers = _get_merged_peaks(self.bed_files, chrom) seq_starts = peak_centers - self.pre_process_window_size // 2 - X, y, seq_starts = bw_to_data.idx_to_filtered_data( + chrom_seq, mappability, chrom_y, seq_starts = bw_to_data.idx_to_filtered_data( genome=self.genome, signal_files=self.signal_files, chrom=chrom, @@ -69,7 +69,7 @@ def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: memmap=self.memmap, generated=self.generated, ) - return X, y, seq_starts + return chrom_seq, mappability, chrom_y, seq_starts def _get_merged_peaks( From 3c04238bbf102d2bfa9c670d17eb90e9665db2d7 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Fri, 7 Aug 2026 17:26:41 +0200 Subject: [PATCH 33/37] added comments on potential future issues --- src/asap/trainer/trainer.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index ea2273f..9c70e21 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -52,6 +52,8 @@ def __init__(self, if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' + + # According to Claude, this does not really do anything. TODO reconsider torch.backends.cudnn.enabled = False elif self.nr_devices == 1: self.ddp_enabled = False @@ -284,6 +286,20 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w persistent_workers = num_workers > 0 # if using DDP, use DistributedSampler if ddp_enabled: + # NOTE: DistributedSampler defaults to drop_last=False, which pads the + # index list by repeating the first few indices so every rank gets an + # equal sample count when len(dataset) isn't evenly divisible by + # world_size. Harmless for training (a few samples just get seen + # twice in some epoch). For validation/eval it's not: _predict() + # gathers every rank's predictions via dist.all_gather before rank 0 + # computes metrics, so those padded/repeated samples get counted + # TWICE in the reported val metric each epoch - a small + # double-counting bias that grows (relatively) with more GPUs against + # a fixed val set size. + # TODO fix: add a `drop_last` param to this function (default False) + # and pass drop_last=True for the eval/val DistributedSampler specifically + # - that drops up to world_size-1 trailing samples instead of duplicating + # them, so a few samples go unevaluated rather than being double-weighted. sampler = torch.utils.data.distributed.DistributedSampler( dataset, shuffle=is_train From bb3f5f9ecd564621913a81e64eb294c2558800f8 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Tue, 11 Aug 2026 18:03:01 +0200 Subject: [PATCH 34/37] Revert "added comments on potential future issues" This reverts commit 3c04238bbf102d2bfa9c670d17eb90e9665db2d7. --- src/asap/trainer/trainer.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index 9c70e21..ea2273f 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -52,8 +52,6 @@ def __init__(self, if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' - - # According to Claude, this does not really do anything. TODO reconsider torch.backends.cudnn.enabled = False elif self.nr_devices == 1: self.ddp_enabled = False @@ -286,20 +284,6 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w persistent_workers = num_workers > 0 # if using DDP, use DistributedSampler if ddp_enabled: - # NOTE: DistributedSampler defaults to drop_last=False, which pads the - # index list by repeating the first few indices so every rank gets an - # equal sample count when len(dataset) isn't evenly divisible by - # world_size. Harmless for training (a few samples just get seen - # twice in some epoch). For validation/eval it's not: _predict() - # gathers every rank's predictions via dist.all_gather before rank 0 - # computes metrics, so those padded/repeated samples get counted - # TWICE in the reported val metric each epoch - a small - # double-counting bias that grows (relatively) with more GPUs against - # a fixed val set size. - # TODO fix: add a `drop_last` param to this function (default False) - # and pass drop_last=True for the eval/val DistributedSampler specifically - # - that drops up to world_size-1 trailing samples instead of duplicating - # them, so a few samples go unevaluated rather than being double-weighted. sampler = torch.utils.data.distributed.DistributedSampler( dataset, shuffle=is_train From 370d64f8d05d0711abac38933f7aef4aebb71086 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Tue, 11 Aug 2026 18:03:01 +0200 Subject: [PATCH 35/37] Revert "de-duplicated chromosome windows" This reverts commit a305c589148e1d0ac3c96c374f69cb09d792909e. --- src/asap/dataloader/base.py | 56 +++--- src/asap/dataloader/bounded.py | 6 +- src/asap/dataloader/bw_to_data.py | 277 ++++++++---------------------- src/asap/dataloader/peak.py | 6 +- 4 files changed, 106 insertions(+), 239 deletions(-) diff --git a/src/asap/dataloader/base.py b/src/asap/dataloader/base.py index e151e1e..e6545c6 100644 --- a/src/asap/dataloader/base.py +++ b/src/asap/dataloader/base.py @@ -87,16 +87,15 @@ def setup(self): pathlib.Path(self.generated).mkdir(parents=True, exist_ok=True) # chroms have been grouped by amount of bounded data per chromosome to balance folds self.chrom_lengths = [] - self.chrom_seq, self.mappability, self.chrom_y, self.seq_starts = [], [], [], [] + self.X, self.y, self.seq_starts = [], [], [] if self.signal_files is None: - self.chrom_y = None + self.y = None for chrom in self.chroms: - chrom_seq_, mappability_, chrom_y_, seq_starts_ = self._generate_chrom_data(chrom) - self.chrom_lengths.append(len(seq_starts_)) - self.chrom_seq.append(chrom_seq_) - self.mappability.append(mappability_) - if chrom_y_ is not None: - self.chrom_y.append(chrom_y_) + X_, y_, seq_starts_ = self._generate_chrom_data(chrom) + self.chrom_lengths.append(len(X_)) + self.X.append(X_) + if y_ is not None: + self.y.append(y_) self.seq_starts.append(seq_starts_) self.cum_chrom_lengths = np.cumsum(self.chrom_lengths) @@ -129,35 +128,30 @@ def __getitem__(self, index) -> Tuple[torch.Tensor, torch.Tensor]: idx = index - self.cum_chrom_lengths[chrom_idx - 1] else: idx = index - - start = self.seq_starts[chrom_idx][idx] + y = None if self.random_shift: - # take a random window within step_size, in bp - start = start + random.randint(0, self.window_size // self.bin_size) * self.bin_size - - chrom_seq = self.chrom_seq[chrom_idx] - X = chrom_seq[start - self.margin_size : start + self.window_size + self.margin_size] - - # mappability channel is only meaningful (not all-ones filler) when - # unmap_threshold is a soft threshold -- matches the original behaviour, - # where strict filtering (unmap_threshold == 0) never produced a real - # per-window mappability channel either. - track = self.mappability[chrom_idx] - if track is not None and self.unmap_threshold != 0: - m = track[start - self.margin_size : start + self.window_size + self.margin_size].astype(np.float32) + shift = random.randint( + 0, self.window_size // self.bin_size + ) # take a random window within step_size + x_shift = shift * self.bin_size + X = self.X[chrom_idx][ + idx, x_shift : x_shift + (self.window_size + 2 * self.margin_size) + ] + if self.y is not None: + y = self.y[chrom_idx][ + idx, shift : shift + (self.window_size // self.bin_size) + ] else: - m = np.ones_like(X, dtype=np.float32) - m = m[..., np.newaxis] - - y = None - if self.chrom_y is not None: - y_raw = self.chrom_y[chrom_idx][start : start + self.window_size] - nr_bins = self.window_size // self.bin_size - y = y_raw.reshape(nr_bins, self.bin_size, y_raw.shape[-1]).max(axis=1) + X = self.X[chrom_idx][idx] + if self.y is not None: + y = self.y[chrom_idx][idx] if self.logspace and y is not None: y = np.log(y + 1) + m = X[..., [1]].astype(np.float32) + X = X[..., 0] + for aug in self.augmentations: X, y = aug(X, y) diff --git a/src/asap/dataloader/bounded.py b/src/asap/dataloader/bounded.py index 5af27a6..413c414 100644 --- a/src/asap/dataloader/bounded.py +++ b/src/asap/dataloader/bounded.py @@ -51,8 +51,8 @@ def __init__( if self.chroms: self.setup() - def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - chrom_seq, mappability, chrom_y, seq_starts = bw_to_data.get_wg_filtered_data( + def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: + X, y, seq_starts = bw_to_data.get_wg_filtered_data( genome=self.genome, signal_files=self.signal_files, chrom=chrom, @@ -67,4 +67,4 @@ def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.n memmap=self.memmap, generated=self.generated, ) - return chrom_seq, mappability, chrom_y, seq_starts \ No newline at end of file + return X, y, seq_starts \ No newline at end of file diff --git a/src/asap/dataloader/bw_to_data.py b/src/asap/dataloader/bw_to_data.py index 2e4aa52..1ce5684 100644 --- a/src/asap/dataloader/bw_to_data.py +++ b/src/asap/dataloader/bw_to_data.py @@ -1,10 +1,9 @@ -from typing import Tuple, List, Optional +from typing import Tuple, List import numpy as np -import pandas as pd import pyBigWig -from asap.dataloader.utils.data_bed import filter_idx_by_bed, whole_genome_idx +from asap.dataloader.utils.data_bed import filter_idx_by_bed, filter_idx_by_unmap_threshold, whole_genome_idx from asap.dataloader.utils.data_bw import get_binned_signal from asap.dataloader.utils.seq import get_chr_seq from asap.dataloader.utils.io import get_bw_from_file @@ -12,190 +11,10 @@ from tqdm import tqdm -# --------------------------------------------------------------------------- -# Whole-chromosome caches (sequence, mappability, signal). -# -# Unlike the per-window caches these used to be, these depend only on -# (genome, chrom[, signal_files/unmap_file]) -- not on window_size, margin_size, -# step_size or bin_size. They're generated once per chromosome and reused -# across every window/step/margin configuration and every epoch. -# --------------------------------------------------------------------------- - -def get_cached_chrom_seq(genome: str, chrom: int, generated: str, memmap: bool = True) -> np.ndarray: - mmap_mode = 'r' if memmap else None - id_ = hash_dn(f'chromseq: genome={genome}, chrom={chrom}', salt='0') - path = f'{generated}/{id_}_chromseq.npy' - try: - return np.load(path, mmap_mode=mmap_mode) - except FileNotFoundError: - seq = get_chr_seq(genome, chrom).astype('int8') - np.save(path, seq) - return np.load(path, mmap_mode=mmap_mode) - - -def get_cached_chrom_mappability(genome: str, chrom: int, unmappable_bed_file: Optional[str], - chrom_len: int, generated: str, memmap: bool = True) -> Optional[np.ndarray]: - if unmappable_bed_file is None: - return None - mmap_mode = 'r' if memmap else None - id_ = hash_dn(f'chrommap: genome={genome}, chrom={chrom}, unmap={unmappable_bed_file}', salt='0') - path = f'{generated}/{id_}_chrommap.npy' - try: - return np.load(path, mmap_mode=mmap_mode) - except FileNotFoundError: - track = np.ones(chrom_len, dtype='int8') - unmap = pd.read_csv(unmappable_bed_file, delimiter='\t', header=None, names=['chr', 'start', 'end']) - unmap = unmap[(unmap.chr == f'chr{chrom}') | (unmap.chr.astype(str) == f'{chrom}')] - for gap_start, gap_end in zip(unmap.start, unmap.end): - s = max(0, int(gap_start)) - e = min(chrom_len, int(gap_end)) - if e > s: - track[s:e] = 0 - np.save(path, track) - return np.load(path, mmap_mode=mmap_mode) - - -def get_cached_chrom_signal(genome: str, chrom: int, signal_files: Optional[List[str]], chrom_len: int, - generated: str, memmap: bool = True) -> Optional[np.ndarray]: - if signal_files is None: - return None - mmap_mode = 'r' if memmap else None - id_ = hash_dn(f'chromsignal: genome={genome}, chrom={chrom}, files={signal_files}', salt='0') - path = f'{generated}/{id_}_chromsignal.npy' - try: - return np.load(path, mmap_mode=mmap_mode) - except FileNotFoundError: - whole_signals = _get_binned_whole_signals(signal_files, chrom, bin_size=1, start=0, end=chrom_len) - y = whole_signals.astype('float32') - np.save(path, y) - return np.load(path, mmap_mode=mmap_mode) - - -def get_cached_chrom_data(genome: str, chrom: int, signal_files: Optional[List[str]], - unmappable_bed_file: Optional[str], generated: str, - memmap: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]: - chrom_seq = get_cached_chrom_seq(genome, chrom, generated, memmap=memmap) - chrom_len = len(chrom_seq) - mappability = get_cached_chrom_mappability(genome, chrom, unmappable_bed_file, chrom_len, generated, memmap=memmap) - chrom_y = get_cached_chrom_signal(genome, chrom, signal_files, chrom_len, generated, memmap=memmap) - return chrom_seq, mappability, chrom_y - - -# --------------------------------------------------------------------------- -# Index building: which positions are valid samples, given filtering rules. -# This is the only piece that still depends on window_size/margin_size/ -# step_size/bin_size/blacklist/unmap_threshold -- and it's cheap (an array of -# positions), so it's cached separately from the (much larger) chromosome data. -# --------------------------------------------------------------------------- - -def _unmappable_fraction(mappability_track: np.ndarray, starts: np.ndarray, extent: int) -> np.ndarray: - # Fraction of unmappable bases in [start, start + extent) for each start. - # NOTE: mirrors the original filter_idx_by_unmap_threshold behaviour exactly, - # which checks [start, start + extent) -- i.e. anchored at `start`, not - # symmetric around it -- even though the data actually read is centered - # (start - margin, start + window + margin). Preserved intentionally rather - # than "fixed", since changing it would silently change which windows are - # kept. - unmap = (1 - mappability_track).astype(np.int64) - prefix = np.concatenate(([0], np.cumsum(unmap))) - counts = prefix[starts + extent] - prefix[starts] - return counts / extent - - -def build_seq_starts_index(genome: str, chrom: int, seq_starts: np.ndarray, window_size: int, margin_size: int, - chrom_len: int, mappability_track: Optional[np.ndarray], chrom_y: Optional[np.ndarray], - bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, - unmap_threshold: float = None, lower_bound: int = None, generated: str = None) -> np.ndarray: - fingerprint = f'{len(seq_starts)}:{int(seq_starts[0]) if len(seq_starts) else 0}:' \ - f'{int(seq_starts[-1]) if len(seq_starts) else 0}:{int(seq_starts.sum())}' - str_ = (f'idx: chr={chrom}, window={window_size}, margin={margin_size}, bin={bin_size}, ' - f'candidates={fingerprint}, blacklist={blacklist_bed_files}, unmappable={unmappable_bed_file}, ' - f'unmap_th={unmap_threshold}, lower_bound={lower_bound}, genome={genome}') - id_ = hash_dn(str_, salt='0') - path = f'{generated}/{id_}_idx.npy' - try: - print(f'Attempting to load index from file with {str_}') - seq_starts = np.load(path) - print('\t...done!') - return seq_starts - except FileNotFoundError: - print('\t...not found.') - print('Generating index...') - - # Filter by blacklists - if blacklist_bed_files is not None: - for bl_bed in blacklist_bed_files: - seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size, - blacklist_bed_file=bl_bed) - - # Filter out sequences extending beyond the chrom - seq_starts = seq_starts[seq_starts + window_size + margin_size < chrom_len] - seq_starts = seq_starts[seq_starts - margin_size > 0] - - # Filter by unmappable, using the whole-chromosome mappability track. - # Extent matches the original: [seq_starts, seq_starts + window_size + 2*margin_size). - if mappability_track is not None and unmap_threshold is not None: - extent = window_size + 2 * margin_size - if unmap_threshold == 0: - # Strict mode: reuse filter_idx_by_bed directly against the unmap bed - # file (exactly as the original did), rather than deriving "zero - # overlap" from the mappability track -- the two have different - # (inclusive vs. half-open) boundary handling at exact edges, and - # this keeps strict-mode behavior bit-identical to before. - seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=extent, - blacklist_bed_file=unmappable_bed_file) - else: - # --- Correct behavior (current) --- - # Fraction of bases in [start, start+extent) that are actually - # unmappable, computed once for the whole chromosome via a prefix - # sum over the mappability track, then sliced per candidate window. - frac = _unmappable_fraction(mappability_track, seq_starts, extent) - seq_starts = seq_starts[frac <= unmap_threshold] - - # --- Old (buggy) behavior, kept here for reference --- - # The original filter_idx_by_unmap_threshold (data_bed.py) computed, - # per window, `overlap_length = sum(overlaps.end - overlaps.start)` - # using the RAW, unclamped length of every unmap interval that - # touched the window at all -- even when only a sliver of that - # interval actually fell inside the window. So a single long - # unmappable region (e.g. 60bp) could cause a window that only - # overlapped it by a handful of bp to be dropped entirely, because - # the full 60bp counted against the threshold budget instead of the - # true (much smaller) overlap. This was a bug, not an intentional - # design choice -- verified against the original by differential - # testing (see conversation/PR notes). If bit-identical window - # selection against an old cached index or trained checkpoint is - # ever needed, this reproduces the dominant effect of that bug - # (raw interval length, not clipped overlap) without reintroducing - # the original's slow per-window pandas loop -- it loops over - # unmap intervals instead, which are typically far fewer than - # candidate windows. Uncomment in place of the two lines above: - # - # unmap_df = pd.read_csv(unmappable_bed_file, delimiter='\t', header=None, - # names=['chr', 'start', 'end']) - # unmap_df = unmap_df[(unmap_df.chr == f'chr{chrom}') | (unmap_df.chr.astype(str) == f'{chrom}')] - # window_end = seq_starts + extent - # raw_overlap_length = np.zeros(len(seq_starts), dtype=np.int64) - # for gs, ge in zip(unmap_df.start.to_numpy(), unmap_df.end.to_numpy()): - # touches = (seq_starts < ge) & (gs < window_end) - # raw_overlap_length[touches] += (ge - gs) # bug: full interval length, not clipped - # seq_starts = seq_starts[raw_overlap_length <= unmap_threshold * extent] - - # Filter by lower bound - if lower_bound is not None and chrom_y is not None: - idx = seq_starts[:, np.newaxis] + np.arange(window_size) - max_signal = chrom_y[idx].max(axis=(1, 2)) - seq_starts = seq_starts[max_signal >= lower_bound] - - print(f'Generated index: {seq_starts.shape}') - np.save(path, seq_starts) - return seq_starts - - def get_wg_filtered_data(genome: str, signal_files: List[str], chrom: int, window_size: int, margin_size: int, step_size: int, bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, unmap_threshold: float = None, - lower_bound: int = None, memmap=True, generated=None): + lower_bound: int = None, memmap=True, generated=None) -> Tuple[np.ndarray, np.ndarray]: idx = whole_genome_idx(genome=genome, chrom=chrom, step_window=step_size) return idx_to_filtered_data(genome=genome, signal_files=signal_files, seq_starts=idx, chrom=chrom, window_size=window_size, @@ -206,25 +25,79 @@ def get_wg_filtered_data(genome: str, signal_files: List[str], chrom: int, windo def idx_to_filtered_data(genome: str, signal_files: List[str], seq_starts: np.ndarray, chrom: int, window_size: int, margin_size: int, bin_size: int, blacklist_bed_files: List[str] = None, unmappable_bed_file: str = None, - unmap_threshold: float = None, lower_bound: int = None, memmap=True, generated=None): - chrom_seq, mappability, chrom_y = get_cached_chrom_data( - genome=genome, chrom=chrom, signal_files=signal_files, unmappable_bed_file=unmappable_bed_file, - generated=generated, memmap=memmap, - ) - seq_starts = build_seq_starts_index( - genome=genome, chrom=chrom, seq_starts=seq_starts, window_size=window_size, margin_size=margin_size, - chrom_len=len(chrom_seq), mappability_track=mappability, chrom_y=chrom_y, bin_size=bin_size, - blacklist_bed_files=blacklist_bed_files, unmappable_bed_file=unmappable_bed_file, - unmap_threshold=unmap_threshold, lower_bound=lower_bound, generated=generated, - ) - return chrom_seq, mappability, chrom_y, seq_starts - + unmap_threshold: float = None, lower_bound: int = None, memmap=True, generated=None) -> Tuple[np.ndarray, np.ndarray]: + mmap_mode = 'r' if memmap else None + str_ = (f'chr={chrom}, window={window_size}, margin={margin_size}, bin={bin_size}, idx={seq_starts.shape}, ' + f'blacklist={blacklist_bed_files}, unmappable={unmappable_bed_file}, th={unmap_threshold}, ' + f'lower_bound={lower_bound}, files={signal_files}, include_map={True}...') + id_ = hash_dn(str_, salt='0') + y = None + try: + print(f'Attempting to load data from file with {str_}') + x = np.load(f'{generated}/{id_}_x.npy', mmap_mode=mmap_mode) + if signal_files is not None: + y = np.load(f'{generated}/{id_}_y.npy', mmap_mode=mmap_mode) + seq_starts = np.load(f'{generated}/{id_}_seq.npy', mmap_mode=mmap_mode) + print('\t...done!') + + except FileNotFoundError: + print('\t...not found.') + print('Generating data...') + # Filter by blacklists + if blacklist_bed_files is not None: + for bl_bed in blacklist_bed_files: + seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size, + blacklist_bed_file=bl_bed) + + + # Filter out sequences extending beyond the chrom + seq = get_chr_seq(genome, chrom) + seq_starts = seq_starts[seq_starts + window_size + margin_size < len(seq)] + seq_starts = seq_starts[seq_starts - margin_size > 0] + + # Filter by unmappable + mappability = None + if unmappable_bed_file is not None: + if unmap_threshold == 0: + seq_starts = filter_idx_by_bed(chrom=chrom, seq_starts=seq_starts, window_size=window_size + 2* margin_size, + blacklist_bed_file=unmappable_bed_file) + else: + seq_starts, mappability = filter_idx_by_unmap_threshold(chrom=chrom, seq_starts=seq_starts, window_size=window_size + 2* margin_size, + unmappable_bed_file=unmappable_bed_file, + threshold=unmap_threshold, return_unmap=True) + x, y = get_data_by_idx(genome, signal_files, chrom, seq_starts, window_size, margin_size, bin_size) + + # Filter by lower bound + if lower_bound is not None: + indices_over_bound = y.max(axis=(1, 2), initial=0) >= lower_bound + x = x[indices_over_bound] + if y is not None: + y = y[indices_over_bound] + seq_starts = seq_starts[indices_over_bound] + if mappability is not None: + mappability = mappability[indices_over_bound] + print(f'Data after filtering bound {lower_bound}: x={x.shape} y={y.shape}') + + x = x.astype('int8') + if mappability is None: + mappability = np.ones_like(x) + else: + mappability = mappability.astype('int8') + x = np.stack([x, mappability], axis=-1) + if y is not None: + y = y.astype('float32') + np.save(f'{generated}/{id_}_y.npy', y) + y = np.load(f'{generated}/{id_}_y.npy', mmap_mode=mmap_mode) + + # save and reload in mmap mode + np.save(f'{generated}/{id_}_x.npy', x) + np.save(f'{generated}/{id_}_seq.npy', seq_starts) + + x = np.load(f'{generated}/{id_}_x.npy', mmap_mode=mmap_mode) + seq_starts = np.load(f'{generated}/{id_}_seq.npy', mmap_mode=mmap_mode) + print('\t...done!') + return x, y, seq_starts -# --------------------------------------------------------------------------- -# Unchanged below: ad hoc per-position window extraction (used by snv/predict.py -# for one-off SNV lookups, not by the cached bulk dataset path above) and -# bigwig export. -# --------------------------------------------------------------------------- def get_data_by_idx(genome: str, signal_files: List[str], chrom: int, seq_starts: np.ndarray, window: int, margin:int, bin_size: int) -> Tuple[ np.ndarray, np.ndarray]: diff --git a/src/asap/dataloader/peak.py b/src/asap/dataloader/peak.py index 5770752..8d45c1b 100644 --- a/src/asap/dataloader/peak.py +++ b/src/asap/dataloader/peak.py @@ -51,10 +51,10 @@ def __init__( if self.chroms: self.setup() - def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray]: peak_centers = _get_merged_peaks(self.bed_files, chrom) seq_starts = peak_centers - self.pre_process_window_size // 2 - chrom_seq, mappability, chrom_y, seq_starts = bw_to_data.idx_to_filtered_data( + X, y, seq_starts = bw_to_data.idx_to_filtered_data( genome=self.genome, signal_files=self.signal_files, chrom=chrom, @@ -69,7 +69,7 @@ def _generate_chrom_data(self, chrom: int) -> Tuple[np.ndarray, np.ndarray, np.n memmap=self.memmap, generated=self.generated, ) - return chrom_seq, mappability, chrom_y, seq_starts + return X, y, seq_starts def _get_merged_peaks( From ab91befdde3a354527763544ddd486ab1b73f627 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Tue, 11 Aug 2026 18:03:01 +0200 Subject: [PATCH 36/37] Revert "memmap" This reverts commit ed96b450229375c072d018952c68a48a5a224c21. --- src/asap/dataset.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/src/asap/dataset.py b/src/asap/dataset.py index 8bb2faa..162b971 100644 --- a/src/asap/dataset.py +++ b/src/asap/dataset.py @@ -1,7 +1,7 @@ from .dataloader import WGDataset, PeakDataset from typing import List -def training_datasets(signal_file: str, genome: str, train_chroms: List[int], val_chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): +def training_datasets(signal_file: str, genome: str, train_chroms: List[int], val_chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): ''' Create training and validation datasets for the model. Args: @@ -12,9 +12,8 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. - memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' - + train_dataset = WGDataset( genome=genome, signal_files=signal_file, @@ -30,7 +29,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va unmap_threshold=0.35, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_train=True, is_robustness=False, @@ -52,7 +51,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va unmap_threshold=0, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_train=False, is_robustness=False, @@ -61,7 +60,7 @@ def training_datasets(signal_file: str, genome: str, train_chroms: List[int], va return train_dataset, val_dataset -def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): +def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): ''' Create a peak dataset for evaluation. Args: @@ -72,7 +71,6 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. - memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = PeakDataset( genome=genome, @@ -90,7 +88,7 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int unmap_threshold=0, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_robustness=False, ) @@ -98,7 +96,7 @@ def peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int return dataset -def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): +def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): ''' Create a peak dataset for robustness evaluation. Args: @@ -109,7 +107,6 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. - memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = PeakDataset( genome=genome, @@ -127,7 +124,7 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom unmap_threshold=0, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_robustness=True, ) @@ -135,7 +132,7 @@ def robustness_peak_dataset(signal_file: str, peak_file: str, genome: str, chrom return dataset -def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): +def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): ''' Create a whole genome dataset for evaluation. Args: @@ -145,7 +142,6 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. - memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = WGDataset( genome=genome, @@ -162,7 +158,7 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, unmap_threshold=0, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_train=False, is_robustness=False, @@ -171,7 +167,7 @@ def wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, return dataset -def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None, memmap: bool = True): +def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], generated: str, blacklist_file: List[str] = None, unmap_file: str = None): ''' Create a whole genome dataset for robustness evaluation. Args: @@ -181,7 +177,6 @@ def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], gene generated (str): Path to the generated data. blacklist_file (List[str]): List of paths to blacklist files (including SNVs). unmap_file (str): Path to the unmapped regions file. - memmap (bool): Whether to memory-map the cached data arrays instead of loading them fully into RAM. ''' dataset = WGDataset( genome=genome, @@ -198,7 +193,7 @@ def robustness_wg_dataset(signal_file: str, genome: str, chroms: List[int], gene unmap_threshold=0, logspace=True, output_format="ohe", - memmap=memmap, + memmap=False, generated=generated, is_train=False, is_robustness=True, From 3ccb7060a6cc832e5b8a2aabf648a6eb24e21433 Mon Sep 17 00:00:00 2001 From: alankuznicki Date: Tue, 11 Aug 2026 18:03:01 +0200 Subject: [PATCH 37/37] Revert "clauded num_workers to nonzero" This reverts commit d3e65987f6237f7a767c0893b0e95cb68dfac539. --- src/asap/task.py | 98 ++++++++++++++----------------------- src/asap/trainer/trainer.py | 29 ++++------- 2 files changed, 46 insertions(+), 81 deletions(-) diff --git a/src/asap/task.py b/src/asap/task.py index 33d4001..cf4213e 100644 --- a/src/asap/task.py +++ b/src/asap/task.py @@ -32,10 +32,10 @@ def _get_model(model_name: str, use_map: bool = False, num_heads=1): raise ValueError(f'Unknown model name: {model_name}') -def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_workers: int=4): +def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False): """ Train the model with the given datasets and parameters. - + Args: experiment_name (str): The name of the experiment. model (str): The model to train. @@ -47,13 +47,12 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v learning_rate (float): The learning rate for the optimizer. batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. - num_workers (int): The number of DataLoader worker processes per GPU. """ # Check if gpu is available if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 print("No GPU available, using CPU instead.") - + # Count the number of GPUs available if n_gpus > torch.cuda.device_count(): n_gpus = torch.cuda.device_count() @@ -64,25 +63,24 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v # Initialize the trainer with the model and datasets trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, - num_workers=num_workers, ) # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): +def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): ''' Add a new head to a bse model and train it using finetuining - + Args: - base_experiment_name(str): The name of the base model + base_experiment_name(str): The name of the base model new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved model (str): The model to train. train_dataset: The training dataset. @@ -94,7 +92,6 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (int): The number of heads (=num of prediction signals) of the new model (=num base heads + 1) - num_workers (int): The number of DataLoader worker processes per GPU. ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 @@ -128,7 +125,6 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model n_gpus=n_gpus, fine_tune=True, num_heads=num_heads, - num_workers=num_workers, ) print(f'Loading best model weights from {base_experiment_name}') @@ -144,12 +140,12 @@ def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model print("trained a new model") -def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1, num_workers: int=4): +def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1): ''' Add a new head to a bse model and train it using linear probing - + Args: - base_experiment_name(str): The name of the base model + base_experiment_name(str): The name of the base model new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved model (str): The model to train. train_dataset: The training dataset. @@ -161,7 +157,6 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_original_heads (int): The number of heads (=num of prediction signals) of the old model - num_workers (int): The number of DataLoader worker processes per GPU. ''' if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 @@ -198,7 +193,6 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model n_gpus=n_gpus, linear_probe=True, num_heads=num_original_heads+1, - num_workers=num_workers, ) print(f'Loading best model weights from {base_experiment_name}') @@ -213,10 +207,10 @@ def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) print("trained a new model") -def train_multiheaded_model(experiment_name : str, model: str, train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): +def train_multiheaded_model(experiment_name : str, model: str, train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1): """ Train the model with the given datasets and parameters. - + Args: experiment_name (str): The name of the experiment. model (str): The model to train. @@ -229,14 +223,13 @@ def train_multiheaded_model(experiment_name : str, model: str, train_dataset: L batch_size (int): The batch size for training. use_map (bool): Whether to use mappability for training. num_heads (int): The number of heads (=num of prediction signals) - num_workers (int): The number of DataLoader worker processes per GPU. """ # Check if gpu is available if n_gpus > 0 and not torch.cuda.is_available(): n_gpus = 0 print("No GPU available, using CPU instead.") - + # Count the number of GPUs available if n_gpus > torch.cuda.device_count(): n_gpus = torch.cuda.device_count() @@ -247,21 +240,20 @@ def train_multiheaded_model(experiment_name : str, model: str, train_dataset: L # Initialize the trainer with the model and datasets trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, - num_heads=num_heads, - num_workers=num_workers, + num_heads=num_heads ) # Start training trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate) -def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, target_head:int = 0, num_workers: int=4): +def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, target_head:int = 0): ''' Evaluate the model on the given dataset. Args: @@ -271,7 +263,6 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat logs_dir (str): The directory to load model checkpoints from. batch_size (int): The batch size for evaluation. use_map (bool): If mappability information was used during training. - num_workers (int): The number of DataLoader worker processes. ''' n_gpus = 1 if torch.cuda.is_available() else 0 @@ -279,12 +270,12 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat model = _get_model(model, use_map=use_map, num_heads=num_heads) trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads ) @@ -293,7 +284,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat print(f'Loading best model weights from {trainer.filename}') checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' trainer.load_weights(checkpoint_path) - + test_chroms = eval_dataset.chroms scores = {} for chrom in test_chroms: @@ -301,9 +292,8 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False, - num_workers=num_workers + batch_size=batch_size, + is_train=False ) _, _, result_metrics = trainer.predict_and_evaluate_multihead(test_gen, target_head=target_head) @@ -311,7 +301,7 @@ def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDat return scores -def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1, num_workers: int=4): +def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, num_heads: int=1): ''' Evaluate the model on the given dataset. Args: @@ -321,7 +311,6 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs logs_dir (str): The directory to load model checkpoints from. batch_size (int): The batch size for evaluation. use_map (bool): If mappability information was used during training. - num_workers (int): The number of DataLoader worker processes. ''' n_gpus = 1 if torch.cuda.is_available() else 0 @@ -329,12 +318,12 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs model = _get_model(model, use_map=use_map, num_heads=num_heads) trainer = Trainer( - filename=experiment_name, + filename=experiment_name, model=model, criterion=nn.PoissonNLLLoss(log_input=False), unmap_criterion=use_map, batch_size=batch_size, - logger=TextLogger(logs_dir=logs_dir), + logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, ) @@ -342,7 +331,7 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs print(f'Loading best model weights from {trainer.filename}') checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth' trainer.load_weights(checkpoint_path) - + test_chroms = eval_dataset.chroms scores = {} for chrom in test_chroms: @@ -350,9 +339,8 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False, - num_workers=num_workers + batch_size=batch_size, + is_train=False ) _, _, result_metrics = trainer.predict_and_evaluate(test_gen) @@ -372,12 +360,11 @@ def train_multiheaded_model_progressively( use_map: bool=False, num_heads: List[int] = [1, 2], checkpoint_on_new_heads_only: bool=False, - num_workers: int=4, ): ''' Train the model with the given datasets and parameters progressively. In Progressive Joint Training (PJT), the number of heads should be strictly increasing at every step. In NJT*, the number of heads should be constant, - equal to the final number of heads. + equal to the final number of heads. Args: experiment_name (str): The base name of the experiments. Actual names are {experiment_name}_{i}_{num_heads[i]}, with i the current step. @@ -392,7 +379,6 @@ def train_multiheaded_model_progressively( use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. - num_workers (int): The number of DataLoader worker processes per GPU. ''' # Validate num_heads @@ -428,7 +414,6 @@ def train_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[0], - num_workers=num_workers, ) print(f'Starting training for the initial step with {num_heads[0]} heads.') @@ -472,7 +457,6 @@ def train_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[i], - num_workers=num_workers, ) # Train the new model @@ -503,10 +487,9 @@ def extend_multiheaded_model_progressively( use_map: bool=False, num_heads: List[int] = [3, 4], # num_heads[0] is the number of heads in the already trained model, with num_heads[1] the first step of JT extension checkpoint_on_new_heads_only: bool=False, - num_workers: int=4, ): ''' - Train the model with the given datasets and parameters progressively. Start with an already trained model + Train the model with the given datasets and parameters progressively. Start with an already trained model with some number of heads (num_heads[0]). The following stages can have an increasing number of heads (PJT), or the same as the pretrained model (NJT*). @@ -524,7 +507,6 @@ def extend_multiheaded_model_progressively( use_map (bool): Whether to use mappability for training. num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on. checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads. - num_workers (int): The number of DataLoader worker processes per GPU. ''' # Validate num_heads if (len(num_heads) < 2): @@ -560,7 +542,6 @@ def extend_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[0], - num_workers=num_workers, ) for i in range(1, len(num_heads)): @@ -587,7 +568,6 @@ def extend_multiheaded_model_progressively( logger=TextLogger(logs_dir=logs_dir), n_gpus=n_gpus, num_heads=num_heads[i], - num_workers=num_workers, ) # Train the new model @@ -604,7 +584,7 @@ def extend_multiheaded_model_progressively( print(f'Final model is saved as {new_experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.') -def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17, num_workers: int=4): +def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17): ''' Evaluate the robustness of the model on the given dataset. Args: @@ -615,7 +595,6 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, batch_size (int): The batch size for evaluation. use_map (bool): Whether to use mappability for evaluation. nr_samples_for_var (int): The number of samples for variance calculation. - num_workers (int): The number of DataLoader worker processes. ''' # Fixed margin size for robustness evaluation margin = 768 @@ -648,9 +627,8 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size // (nr_samples_for_var - 1), - is_train=False, - num_workers=num_workers + batch_size=batch_size // (nr_samples_for_var - 1), + is_train=False ) _, _, cov, cov_per_bin = trainer.predict_robust_batch(test_gen, nr_samples_for_var=nr_samples_for_var, window=eval_dataset.window_size, margin=margin) @@ -859,7 +837,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil -def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, out_dir: str, batch_size: int=64, use_map: bool=False, num_workers: int=4): +def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, out_dir: str, batch_size: int=64, use_map: bool=False): """ Export the predictions to a file. Args: @@ -870,7 +848,6 @@ def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDatas out_dir (str): The output directory for predictions. batch_size (int): The batch size for evaluation. use_map (bool): Whether to use mappability for evaluation. - num_workers (int): The number of DataLoader worker processes. """ n_gpus = 1 if torch.cuda.is_available() else 0 @@ -902,9 +879,8 @@ def export_predictions(experiment_name: str, model: str, eval_dataset: BaseDatas test_gen = make_dataloader( ddp_enabled=False, dataset=eval_dataset, - batch_size=batch_size, - is_train=False, - num_workers=num_workers + batch_size=batch_size, + is_train=False ) _, predictions, _ = trainer.predict_and_evaluate(test_gen, no_eval=True) diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py index ea2273f..4f4308f 100644 --- a/src/asap/trainer/trainer.py +++ b/src/asap/trainer/trainer.py @@ -32,23 +32,21 @@ def __init__(self, num_heads: int = 1, linear_probe=False, fine_tune=False, - num_workers: int = 4, ): self.filename = filename self.model = model - self.criterion = criterion + self.criterion = criterion self.unmap_criterion = nn.MSELoss() if unmap_criterion is True else None self.train_unmap = not self.unmap_criterion is False self.logger: Logger = logger - self.logspace = True + self.logspace = True self.nr_tracks = 1 self.nr_devices = n_gpus self.batch_size = batch_size self.num_heads = num_heads self.linear_probe = linear_probe self.fine_tune = fine_tune - self.num_workers = num_workers if self.nr_devices > 1: self.ddp_enabled = True self.device = 'cuda' @@ -84,8 +82,7 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None) self.fine_tune, port, self.num_heads, - val_on_heads, - self.num_workers + val_on_heads ), nprocs=self.nr_devices ) @@ -95,15 +92,13 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None) ddp_enabled=False, dataset=train_dset, batch_size=self.batch_size, - is_train=True, - num_workers=self.num_workers + is_train=True ) val_gen = make_dataloader( ddp_enabled=False, dataset=val_dset, batch_size=self.batch_size, - is_train=False, - num_workers=self.num_workers + is_train=False ) _fit( self.device, @@ -281,7 +276,6 @@ def load_weights(self, path): self.model.load_state_dict(state_dict) def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_workers: int = 0, pin_memory: bool = True): - persistent_workers = num_workers > 0 # if using DDP, use DistributedSampler if ddp_enabled: sampler = torch.utils.data.distributed.DistributedSampler( @@ -294,7 +288,6 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w pin_memory=pin_memory, shuffle=False, # shuffling handled by sampler num_workers=num_workers, - persistent_workers=persistent_workers, sampler=sampler ) else: @@ -304,8 +297,7 @@ def make_dataloader(ddp_enabled, dataset, batch_size: int, is_train: bool, num_w batch_size=batch_size, pin_memory=pin_memory, shuffle=is_train, - num_workers=num_workers, - persistent_workers=persistent_workers + num_workers=num_workers ) @@ -339,8 +331,7 @@ def _ddp_and_fit( fine_tune, port=12355, num_heads=1, - val_on_heads=None, - num_workers=4, + val_on_heads=None ): #model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = setup_ddp(rank, world_size, model, port) @@ -348,15 +339,13 @@ def _ddp_and_fit( ddp_enabled=True, dataset=train_dset, batch_size=batch_size, - is_train=True, - num_workers=num_workers + is_train=True ) val_gen = make_dataloader( ddp_enabled=True, dataset=val_dset, batch_size=batch_size, - is_train=False, - num_workers=num_workers + is_train=False ) _fit( rank=rank,