scripts/train.py

# train.py
"""
Main training script for TahoeFormer.

This script handles:
- Loading configuration from a YAML file.
- Setting up logging (Weights & Biases).
- Initializing the model (LitEnformerSMILES, using Morgan Fingerprints).
- Initializing dataloaders (TahoeSMILESDataset).
- Setting up PyTorch Lightning Callbacks (ModelCheckpoint, EarlyStopping, MetricLogger).
- Running the training and testing loops using PyTorch Lightning Trainer.
"""

import argparse
import yaml
import os
import torch
import pandas as pd
import lightning.pytorch as pl
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor
from lightning.pytorch import Trainer
from lightning.pytorch.loggers import WandbLogger
from torch.utils.data import DataLoader, random_split
import wandb

from pl_models import LitEnformerSMILES, MetricLogger
from datasets import TahoeSMILESDataset, ENFORMER_INPUT_SEQ_LENGTH

import warnings
warnings.filterwarnings('ignore', '.*does not have many workers.*')
warnings.filterwarnings('ignore', '.*Detecting val_dataloader.*')

# --- Default Configs --- (can be overridden by config YAML)
DEFAULT_CONFIG = {
    'data': {
        'regions_csv_path': 'data/Enformer_genomic_regions_TSSCenteredGenes_FixedOverlapRemoval.csv',
        'pbulk_parquet_path': 'data/pseudoBulk_celllineXdrug_top3k_for_testing.parquet',
        'drug_meta_csv_path': 'data/drug_metadata.csv',
        'fasta_file_path': 'data/hg38.fa',
        'enformer_input_seq_length': 196_608,
        'morgan_fp_radius': 2, # For TahoeSMILESDataset
        'morgan_fp_nbits': 2048, # For TahoeSMILESDataset
        'filter_drugs_by_ids': None,
        # Column name defaults for TahoeSMILESDataset
        'regions_gene_col': 'gene_id',
        'regions_chr_col': 'seqnames',
        'regions_start_col': 'start',
        'regions_end_col': 'end',
        'regions_strand_col': None,
        'regions_set_col': 'set', # column name for train/val/test split in regions_csv
        'pbulk_gene_col': 'gene_id',
        'pbulk_dose_col': 'drug_dose',
        'pbulk_expr_col': 'expression',
        'pbulk_cell_line_col': 'cell_line_id',
        'drug_meta_id_col': 'drug_id',
        'drug_meta_smiles_col': 'canonical_smiles'
    },
    'model': {
        'enformer_model_name': 'EleutherAI/enformer-official-rough',
        'enformer_target_length': -1,
        'morgan_fingerprint_dim': 2048, 
        'dose_input_dim': 1,
        'fusion_hidden_dim': 256,
        'final_output_tracks': 1,
        'learning_rate': 5e-6, 
        'loss_alpha': 1.0,
        'weight_decay': 0.01,
        'eval_gene_sets': None 
    },
    'training': {
        'batch_size': 2,
        'num_workers': 0,
        'pin_memory': False,
        'max_epochs': 50,
        'gpus': -1, # -1 for all available GPUs, or specify count e.g., 1, 2
        'accelerator': 'auto',
        'strategy': 'ddp_find_unused_parameters_true', 
        'precision': '16-mixed', # '32' or '16-mixed' or 'bf16-mixed'
        'val_check_interval': 1.0, 
        'limit_train_batches': 1.0, 
        'limit_val_batches': 1.0,   
        'limit_test_batches': 1.0,  
        'deterministic': True, 
        'seed': 42
    },
    'logging': {
        'wandb_project': 'TahoeformerDebug',
        'wandb_entity': None, # W&B info (username or team)
        'save_dir': 'outputs/model_checkpoints',
        'checkpoint_monitor_metric': 'validation_pearson_epoch',
        'checkpoint_monitor_mode': 'max',
        'save_top_k': 1,
        'early_stopping_metric': 'validation_pearson_epoch',
        'early_stopping_mode': 'max',
        'early_stopping_patience': 10
    }
}

def delete_checkpoint_at_end(trainer):
    """ Delete checkpoint after training and testing if desired """
    checkpoint_callbacks = [cb for cb in trainer.callbacks if isinstance(cb, ModelCheckpoint)]
    if checkpoint_callbacks:
        checkpoint_callback = checkpoint_callbacks[0]
        if hasattr(checkpoint_callback, 'best_model_path') and checkpoint_callback.best_model_path and os.path.exists(checkpoint_callback.best_model_path):
            print(f"Deleting best checkpoint: {checkpoint_callback.best_model_path}")
            os.remove(checkpoint_callback.best_model_path)
        else:
            print("No best model path found to delete or path does not exist.")
    else:
        print("No ModelCheckpoint callback found.")

def parse_optional_gene_list(filepath):
    """ Parses a file containing one gene name per row, returns a list. Returns empty list if path is None or invalid. """
    if filepath is None or not os.path.exists(filepath):
        return []
    gene_list = []
    with open(filepath, 'r') as file:
        for gene in file:
            gene_list.append(gene.strip())
    return gene_list

def load_config(config_path=None):
    """Loads configuration from YAML, merging with defaults. Ensures deep copy of defaults."""
    # Manual deep copy for 2 levels, as DEFAULT_CONFIG is structured
    config = {}
    for k, v in DEFAULT_CONFIG.items():
        if isinstance(v, dict):
            config[k] = v.copy() # Copies the inner dictionary
        else:
            config[k] = v

    if config_path:
        with open(config_path, 'r') as f:
            user_config = yaml.safe_load(f)
        if user_config: # Ensure user_config is not None (e.g. if YAML is empty)
            for key, value in user_config.items():
                if isinstance(value, dict) and key in config and isinstance(config[key], dict):
                    config[key].update(value) # Merge level 2 dicts
                else:
                    config[key] = value # Overwrite or add new keys/values
    return config

def build_model(config):
    """
    Builds the LitEnformerSMILES model using Morgan Fingerprints.
    Model parameters are sourced from the 'model' section of the config.
    """
    model_params = config['model']
    # Ensure morgan_fingerprint_dim from data config (for dataset) matches model config
    # Model will use its own `morgan_fingerprint_dim` parameter.
    # The dataset's `morgan_fp_nbits` should align with this.
    print(f"Building LitEnformerSMILES model with morgan_fingerprint_dim: {model_params.get('morgan_fingerprint_dim')}")

    return LitEnformerSMILES(
        enformer_model_name=model_params.get('enformer_model_name'),
        enformer_target_length=model_params.get('enformer_target_length'),
        num_output_tracks_enformer_head=model_params.get('num_output_tracks_enformer_head'),
        morgan_fingerprint_dim=model_params.get('morgan_fingerprint_dim', 2048), # Default from model if not in config
        dose_input_dim=model_params.get('dose_input_dim'),
        fusion_hidden_dim=model_params.get('fusion_hidden_dim'),
        final_output_tracks=model_params.get('final_output_tracks'),
        learning_rate=model_params.get('learning_rate'),
        loss_alpha=model_params.get('loss_alpha'),
        weight_decay=model_params.get('weight_decay'),
        eval_gene_sets=model_params.get('eval_gene_sets')
    )

def load_tahoe_smiles_dataloaders(config):
    """
    Initializes TahoeSMILESDataset (now using Morgan Fingerprints) and creates DataLoaders.
    Dataset parameters are sourced from the 'data' section of the config.
    Training parameters (batch_size, num_workers) from 'training' section.
    """
    data_config = config['data']
    train_config = config['training']

    # Pass Morgan fingerprint params to TahoeSMILESDataset
    dataset_args = {
        'regions_csv_path': data_config['regions_csv_path'],
        'pbulk_parquet_path': data_config['pbulk_parquet_path'],
        'drug_meta_csv_path': data_config['drug_meta_csv_path'],
        'fasta_file_path': data_config['fasta_file_path'],
        'enformer_input_seq_length': data_config.get('enformer_input_seq_length'),
        'morgan_fp_radius': data_config.get('morgan_fp_radius', 2),
        'morgan_fp_nbits': data_config.get('morgan_fp_nbits', 2048),
        'filter_drugs_by_ids': data_config.get('filter_drugs_by_ids'),
        # Pass column name configurations
        'regions_gene_col': data_config.get('regions_gene_col', 'gene_name'),
        'regions_chr_col': data_config.get('regions_chr_col', 'seqnames'),
        'regions_start_col': data_config.get('regions_start_col', 'starts'),
        'regions_end_col': data_config.get('regions_end_col', 'ends'),
        'regions_strand_col': data_config.get('regions_strand_col', None),
        'regions_set_col': data_config.get('regions_set_col', 'set'), # Added for set-based splitting
        'pbulk_gene_col': data_config.get('pbulk_gene_col', 'gene_id'),
        'pbulk_dose_col': data_config.get('pbulk_dose_col', 'dose_nM'),
        'pbulk_expr_col': data_config.get('pbulk_expr_col', 'value'),
        'pbulk_cell_line_col': data_config.get('pbulk_cell_line_col', 'cell_line_id'),
        'drug_meta_id_col': data_config.get('drug_meta_id_col', 'drug_id'),
        'drug_meta_smiles_col': data_config.get('drug_meta_smiles_col', 'canonical_smiles')
    }
    
    # print(f"Initializing TahoeSMILESDataset with morgan_fp_nbits: {dataset_args['morgan_fp_nbits']}")

    # Instantiate dataset for each split using the 'target_set' parameter
    print("Initializing train dataset...")
    train_dataset = TahoeSMILESDataset(**dataset_args, target_set='train')
    print("Initializing validation dataset...")
    val_dataset = TahoeSMILESDataset(**dataset_args, target_set='valid')
    # In the original Enformer_genomic_regions_TSSCenteredGenes_FixedOverlapRemoval.csv, 
    # the validation set is often named 'valid'. If it's 'validation' in your file, adjust accordingly.
    print("Initializing test dataset...")
    test_dataset = TahoeSMILESDataset(**dataset_args, target_set='test')

    
    if len(train_dataset) == 0:
        print("WARNING: Train dataset is empty. This could be due to filtering by set='train' or other data issues. Training might fail or be skipped.")
    if len(val_dataset) == 0:
        print("WARNING: Validation dataset is empty (set='valid'). Validation loop will likely be skipped.")
    if len(test_dataset) == 0:
        print("WARNING: Test dataset is empty (set='test'). Testing loop will likely be skipped.")

    train_loader = DataLoader(
        train_dataset,
        batch_size=train_config.get('batch_size', 2),
        shuffle=True,
        num_workers=train_config.get('num_workers', 0),
        pin_memory=train_config.get('pin_memory', False),
        drop_last=True # Important for DDP and BatchNorm if batches can be size 1 per GPU
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=train_config.get('batch_size', 2) * 2, # Often use larger batch for val
        shuffle=False,
        num_workers=train_config.get('num_workers', 0),
        pin_memory=train_config.get('pin_memory', False)
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=train_config.get('batch_size', 2) * 2,
        shuffle=False,
        num_workers=train_config.get('num_workers', 0),
        pin_memory=train_config.get('pin_memory', False)
    )
    return train_loader, val_loader, test_loader

def load_trainer_and_callbacks(config, experiment_name_for_wandb, run_name_for_wandb):
    """ Loads PyTorch Lightning Trainer and associated callbacks. """
    

    metric_logger = MetricLogger(save_dir_prefix=os.path.join(config['logging']['save_dir'], "metrics"))


    checkpoint_dir = os.path.join(config['logging']['save_dir'], 'checkpoints')
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    monitor_metric = config['logging']['checkpoint_monitor_metric'] # MetricLogger logs with _epoch suffix
    monitor_mode = config['logging']['checkpoint_monitor_mode']
    save_top_k = config['logging'].get('save_top_k', 1)

    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_dir,
        filename=f"{{epoch}}-{{{monitor_metric}:.4f}}",
        save_top_k=save_top_k,
        monitor=monitor_metric,
        mode=monitor_mode
    )


    early_stop_monitor_metric = config['logging']['early_stopping_metric']
    early_stop_mode = config['logging']['early_stopping_mode']
    min_delta = 0.001
    patience = config['logging']['early_stopping_patience']

    early_stopping_callback = EarlyStopping(
        monitor=early_stop_monitor_metric,
        min_delta=min_delta,
        patience=patience,
        verbose=True,
        mode=early_stop_mode
    )

    lr_monitor = LearningRateMonitor(logging_interval='step')

    callbacks = [checkpoint_callback, metric_logger, early_stopping_callback, lr_monitor]


    wandb_logger = None
    if config['logging'].get('wandb_project'):
        wandb_logger = WandbLogger(
            name=run_name_for_wandb,
            project=config['logging']['wandb_project'],
            group=experiment_name_for_wandb,
            config=config, # Log the entire config dictionary
            save_dir=config['logging']['save_dir'], # Optional: ensure logger saves to the same base dir
            id=run_name_for_wandb # Use the unique run_name as the W&B run ID
        )

    trainer = Trainer(
        max_epochs=config['training']['max_epochs'],
        precision=config['training']['precision'],
        accumulate_grad_batches=config['training'].get('accumulate_grad_batches', 1),
        gradient_clip_val=config['training'].get('gradient_clip_val', 0.5),
        callbacks=callbacks,
        logger=wandb_logger, # Use the configured logger
        num_sanity_val_steps=config['training'].get('num_sanity_val_steps', 0), # Often 0 if val metrics are complex
        log_every_n_steps=config['training'].get('log_every_n_steps', 50),
        check_val_every_n_epoch=config['training'].get('check_val_every_n_epoch', 1),
        deterministic=config['training']['deterministic'], # For reproducibility
        strategy=config['training']['strategy'],
        accelerator=config['training']['accelerator'],
        devices=config['training'].get('gpus', 'auto')
    )
    

    if config['training'].get('accumulate_grad_batches', 1) > 1:
        effective_batch_size = config['training']['batch_size'] * config['training'].get('accumulate_grad_batches', 1)
        print(f"Gradient Accumulation: Effective batch size will be {effective_batch_size}")
        # Log to wandb config if logger is active and wandb.run exists
        if wandb_logger and wandb.run:
             wandb.config.update({'effective_train_batch_size': effective_batch_size}, allow_val_change=True)

    return trainer

def run_experiment(config: wandb.config): 
    """ Main training and evaluation loop. """
    print("Starting experiment with configuration:")
    for key, value in config.items():
        print(f"  {key}: {value}")


    train_loader, val_loader, test_loader = load_tahoe_smiles_dataloaders(config)


    eval_gene_sets = {
        'train_eval_set': parse_optional_gene_list(config.get('eval_train_gene_path')),
        'valid_eval_set': parse_optional_gene_list(config.get('eval_valid_gene_path')),
        'test_eval_set': parse_optional_gene_list(config.get('eval_test_gene_path'))
    }
    eval_gene_sets = {k: v for k, v in eval_gene_sets.items() if v} # Keep only non-empty lists


    model = build_model(config)


    experiment_name = config.get('experiment_name', 'DefaultExperiment')
    run_name = config.get('run_name', f"{experiment_name}_default_run_id") # Fallback run_name

    trainer = load_trainer_and_callbacks(config, experiment_name, run_name)


    if config.get('validate_before_train', False) and val_loader.dataset:
        print("Running pre-training validation loop...")
        trainer.validate(model, dataloaders=val_loader)


    print("Starting training...")
    if train_loader.dataset:
        trainer.fit(
            model=model,
            train_dataloaders=train_loader,
            val_dataloaders=val_loader if val_loader.dataset else None
        )
    else:
        print("Skipping training as train_loader is empty.")


    print("Starting testing...")
    if test_loader.dataset:
         
        best_model_path = trainer.checkpoint_callback.best_model_path if hasattr(trainer.checkpoint_callback, 'best_model_path') else None
        if best_model_path and os.path.exists(best_model_path):
            print(f"Loading best model for testing from: {best_model_path}")
            trainer.test(model, dataloaders=test_loader, ckpt_path=best_model_path)
        elif not best_model_path:
             print("No best_model_path found from checkpoint callback. Testing with current model state (if any training happened).")
             trainer.test(model, dataloaders=test_loader) # Test with current model if no checkpoint or if training was skipped
        else: # path exists but is false for some reason or doesnt exist
            print(f"Best model path {best_model_path} not found. Testing with current model state.")
            trainer.test(model, dataloaders=test_loader)
    else:
        print("Skipping testing as test_loader is empty.")

    if config.get('delete_checkpoint_after_run', False):
        delete_checkpoint_at_end(trainer)


def main():
    parser = argparse.ArgumentParser(description='Run PyTorch Lightning Enformer-SMILES experiment.')
    parser.add_argument('--config_path', type=str, required=True, help='Path to the YAML configuration file.')
    # allow seed override from command line, though config is primary source
    parser.add_argument("--seed", type=int, help="Override seed from config file.") 
    args = parser.parse_args()


    effective_config = load_config(args.config_path)
    

    if args.seed is not None:
        # Ensure 'training' key exists if seed is to be put there
        if 'training' not in effective_config:
            effective_config['training'] = {}
        effective_config['training']['seed'] = args.seed
    

    seed = effective_config.get('training', {}).get('seed', 42)
    pl.seed_everything(seed, workers=True)
    

    if effective_config.get('training', {}).get('deterministic', False):
         torch.use_deterministic_algorithms(True, warn_only=True) # Ensure deterministic ops if requested

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

    experiment_name = effective_config.get('experiment_name', 'EnformerSMILESExperiment')
    

    run_name = f"{experiment_name}_Seed-{seed}_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"
    effective_config['run_name'] = run_name
    effective_config['experiment_name'] = experiment_name # Ensure experiment_name is also in config


    default_save_dir = os.path.join(current_script_dir, f"../results/{experiment_name}/{run_name}")
    
    if 'logging' not in effective_config:
        effective_config['logging'] = {}
    effective_config['logging']['save_dir'] = effective_config.get('logging', {}).get('save_dir', default_save_dir)
    os.makedirs(effective_config['logging']['save_dir'], exist_ok=True)
    print(f"Results and checkpoints will be saved in: {effective_config['logging']['save_dir']}")

 
    run_experiment(effective_config) 

    if effective_config.get('logging', {}).get('wandb_project') and wandb.run is not None:
        wandb.finish()

if __name__ == '__main__':
    main()