subsampling_config.py

#!/usr/bin/env python3
"""
Configuration management for subsampling datasets.
Provides easy way to run multiple datasets with different configurations.
"""

import argparse
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional


# Dataset configurations
# Add new datasets here with their specific parameters
DATASET_CONFIGS = {
    # Default configurations - subsample cells only (no read subsampling)
    "aging": {
        "input_path": "/fh/fast/setty_m/user/ethieme/repositories/mds-aging-private/02_CITE_aging/anndatas/processed_filtered_HSPCandMature_withcorrection_withregression_postcelltype.h5ad",
        "fraction_reads": 0.99,  # Ignored due to --no-subsample-reads
        "fraction_cells": 0.99,
        "min_reads_per_cell": 1000,
        "min_total_cells": 100,
        "raw_count_layer": "raw_counts",
        "embedding_key": "X_umap",
        "pca_components": 50,
        "max_write_workers": 4,
        "slurm_time": "20:00:00",
        "slurm_memory": "32G",
        "slurm_cpus": 4,
        "additional_args": ["--no-subsample-reads"],
        "description": "Aging dataset - subsample cells only"
    },
    
    "covid": {
        "input_path": "/fh/fast/setty_m/user/dotto/benchmarkDA/data/real/covid19-pbmc/meyer_nikolic_covid_pbmc.cellxgene.20210813.h5ad",
        "fraction_reads": 0.99,  # Ignored due to --no-subsample-reads
        "fraction_cells": 0.99,
        "min_reads_per_cell": 1000,
        "min_total_cells": 100,
        "raw_count_layer": "X",
        "embedding_key": "X_umap (after harmony ADT)",
        "pca_components": 50,
        "max_write_workers": 4,
        "slurm_time": "20:00:00",
        "slurm_memory": "100G",
        "slurm_cpus": 4,
        "additional_args": ["--no-subsample-reads", "--skip-processing"],
        "description": "COVID-19 PBMC dataset - subsample cells only, data already processed"
    },
    
    "bcr-xl": {
        "input_path": "/fh/fast/setty_m/user/dotto/benchmarkDA/data/real/bcr-xl/bcr_xl_preprocessed.h5ad",
        "fraction_reads": 0.99,  # Ignored due to --no-subsample-reads
        "fraction_cells": 0.99,
        "min_reads_per_cell": 1000,
        "min_total_cells": 100,
        "raw_count_layer": "X",
        "embedding_key": "X_umap",
        "pca_components": 50,
        "max_write_workers": 4,
        "slurm_time": "4:00:00",
        "slurm_memory": "16G",
        "slurm_cpus": 4,
        "additional_args": ["--no-subsample-reads", "--skip-processing"],
        "description": "BCR-XL stimulation dataset - subsample cells only, data already processed"
    },
    
    # Special variant: aging with both reads AND cells subsampling
    "aging_with_reads": {
        "input_path": "/fh/fast/setty_m/user/ethieme/repositories/mds-aging-private/02_CITE_aging/anndatas/processed_filtered_HSPCandMature_withcorrection_withregression_postcelltype.h5ad",
        "fraction_reads": 0.99,
        "fraction_cells": 0.99,
        "min_reads_per_cell": 1000,
        "min_total_cells": 100,
        "raw_count_layer": "raw_counts",
        "embedding_key": "X_umap",
        "pca_components": 50,
        "max_write_workers": 4,
        "slurm_time": "10:00:00",  # Longer time for read subsampling
        "slurm_memory": "40G",     # More memory for read subsampling
        "slurm_cpus": 4,
        "description": "Aging dataset - subsample BOTH reads and cells (original method)"
    },
    
    # Special variant: aging with precomputed embeddings (PCA + diffusion maps from full dataset)
    "aging_precomputed_embeddings": {
        "input_path": "/fh/fast/setty_m/user/ethieme/repositories/mds-aging-private/02_CITE_aging/anndatas/processed_filtered_HSPCandMature_withcorrection_withregression_postcelltype.h5ad",
        "fraction_reads": 0.99,  # Ignored due to --no-subsample-reads
        "fraction_cells": 0.99,
        "min_reads_per_cell": 1000,
        "min_total_cells": 100,
        "raw_count_layer": "raw_counts",
        "embedding_key": "X_umap",
        "pca_components": 50,
        "max_write_workers": 4,
        "slurm_time": "30:00:00",  # Longer time for embedding computation on full dataset
        "slurm_memory": "64G",     # More memory for full dataset PCA + diffusion maps
        "slurm_cpus": 4,
        "additional_args": ["--no-subsample-reads", "--precompute-embeddings"],
        "description": "Aging dataset - subsample cells only, keeping original PCA and diffusion maps from full dataset"
    },
}


def validate_config(config: Dict) -> List[str]:
    """Validate a dataset configuration and return any errors."""
    errors = []
    
    required_fields = [
        "input_path", "fraction_reads", "fraction_cells", 
        "min_reads_per_cell", "min_total_cells", "raw_count_layer",
        "embedding_key", "pca_components", "max_write_workers"
    ]
    
    for field in required_fields:
        if field not in config:
            errors.append(f"Missing required field: {field}")
    
    # Validate file existence
    if "input_path" in config:
        input_path = Path(config["input_path"])
        if not input_path.exists():
            errors.append(f"Input file does not exist: {config['input_path']}")
    
    # Validate numeric ranges
    if "fraction_reads" in config:
        if not 0 < config["fraction_reads"] <= 1:
            errors.append("fraction_reads must be between 0 and 1")
    
    if "fraction_cells" in config:
        if not 0 < config["fraction_cells"] <= 1:
            errors.append("fraction_cells must be between 0 and 1")
    
    if "min_reads_per_cell" in config:
        if config["min_reads_per_cell"] < 0:
            errors.append("min_reads_per_cell must be non-negative")
    
    if "min_total_cells" in config:
        if config["min_total_cells"] < 1:
            errors.append("min_total_cells must be at least 1")
    
    if "pca_components" in config:
        if config["pca_components"] < 1:
            errors.append("pca_components must be at least 1")
    
    if "max_write_workers" in config:
        if config["max_write_workers"] < 1:
            errors.append("max_write_workers must be at least 1")
    
    return errors


def build_command_args(dataset_name: str, config: Dict, use_slurm: bool = True) -> List[str]:
    """Build command line arguments from configuration."""
    
    if use_slurm:
        cmd = ["bash", "submit_subsampling.sh"]
    else:
        cmd = ["python", "subsample_dataset.py"]
    
    cmd.extend([dataset_name, config["input_path"]])
    
    # Add standard arguments
    standard_args = [
        "fraction_reads", "fraction_cells", "min_reads_per_cell", 
        "min_total_cells", "raw_count_layer", "embedding_key",
        "pca_components", "max_write_workers"
    ]
    
    for arg in standard_args:
        if arg in config:
            arg_name = arg.replace("_", "-")
            cmd.extend([f"--{arg_name}", str(config[arg])])
    
    # Add any additional arguments
    if "additional_args" in config:
        cmd.extend(config["additional_args"])
    
    return cmd


def submit_dataset(dataset_name: str, config: Dict, dry_run: bool = False, use_slurm: bool = True) -> bool:
    """Submit a dataset for subsampling."""
    
    # Validate configuration
    errors = validate_config(config)
    if errors:
        print(f"Configuration errors for {dataset_name}:")
        for error in errors:
            print(f"  - {error}")
        return False
    
    # Build command
    cmd = build_command_args(dataset_name, config, use_slurm)
    
    print(f"Dataset: {dataset_name}")
    print(f"Description: {config.get('description', 'No description')}")
    print(f"Command: {' '.join(cmd)}")
    
    if dry_run:
        print("  [DRY RUN - not executed]")
        return True
    
    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        print("  ✓ Submitted successfully")
        if result.stdout:
            print(f"  Output: {result.stdout.strip()}")
        return True
    
    except subprocess.CalledProcessError as e:
        print(f"  ✗ Failed to submit: {e}")
        if e.stderr:
            print(f"  Error: {e.stderr.strip()}")
        return False
    
    except FileNotFoundError as e:
        print(f"  ✗ Command not found: {e}")
        return False


def list_datasets():
    """List all available dataset configurations."""
    print("Available dataset configurations:")
    print("=" * 50)
    
    for name, config in DATASET_CONFIGS.items():
        print(f"\nDataset: {name}")
        print(f"  Description: {config.get('description', 'No description')}")
        print(f"  Input: {config['input_path']}")
        print(f"  Fraction reads: {config['fraction_reads']}")
        print(f"  Fraction cells: {config['fraction_cells']}")
        print(f"  Min reads/cell: {config['min_reads_per_cell']}")
        print(f"  Min total cells: {config['min_total_cells']}")
        print(f"  PCA components: {config['pca_components']}")
        
        if "additional_args" in config:
            print(f"  Additional args: {' '.join(config['additional_args'])}")
        
        if "slurm_time" in config:
            print(f"  SLURM time: {config['slurm_time']}")
            print(f"  SLURM memory: {config['slurm_memory']}")
            print(f"  SLURM CPUs: {config['slurm_cpus']}")


def main():
    """Main entry point for configuration management."""
    parser = argparse.ArgumentParser(
        description="Subsampling configuration management",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List all available configurations
  python subsampling_config.py --list
  
  # Submit a specific dataset
  python subsampling_config.py --submit aging
  
  # Submit multiple datasets
  python subsampling_config.py --submit aging covid bcr-xl
  
  # Submit all datasets
  python subsampling_config.py --submit-all
  
  # Dry run to see commands without executing
  python subsampling_config.py --submit aging --dry-run
  
  # Run locally instead of SLURM
  python subsampling_config.py --submit aging --no-slurm
        """
    )
    
    parser.add_argument("--list", action="store_true",
                       help="List all available dataset configurations")
    parser.add_argument("--submit", nargs="+", metavar="DATASET",
                       help="Submit specific datasets for subsampling")
    parser.add_argument("--submit-all", action="store_true",
                       help="Submit all datasets for subsampling")
    parser.add_argument("--dry-run", action="store_true",
                       help="Show commands without executing them")
    parser.add_argument("--no-slurm", action="store_true",
                       help="Run locally instead of submitting to SLURM")
    parser.add_argument("--validate", action="store_true",
                       help="Validate all configurations without submitting")
    
    args = parser.parse_args()
    
    if args.list:
        list_datasets()
        return
    
    if args.validate:
        print("Validating all configurations...")
        all_valid = True
        for name, config in DATASET_CONFIGS.items():
            errors = validate_config(config)
            if errors:
                print(f"\n❌ {name}: Configuration errors:")
                for error in errors:
                    print(f"    - {error}")
                all_valid = False
            else:
                print(f"✅ {name}: Valid")
        
        if all_valid:
            print(f"\nAll {len(DATASET_CONFIGS)} configurations are valid!")
        else:
            print(f"\nSome configurations have errors. Please fix them before submitting.")
            sys.exit(1)
        return
    
    # Determine which datasets to submit
    datasets_to_submit = []
    
    if args.submit_all:
        datasets_to_submit = list(DATASET_CONFIGS.keys())
    elif args.submit:
        datasets_to_submit = args.submit
        # Validate dataset names
        invalid_datasets = [d for d in datasets_to_submit if d not in DATASET_CONFIGS]
        if invalid_datasets:
            print(f"Error: Unknown datasets: {', '.join(invalid_datasets)}")
            print(f"Available datasets: {', '.join(DATASET_CONFIGS.keys())}")
            sys.exit(1)
    else:
        parser.print_help()
        return
    
    # Submit datasets
    use_slurm = not args.no_slurm
    print(f"Submitting {len(datasets_to_submit)} dataset(s) {'via SLURM' if use_slurm else 'locally'}...")
    if args.dry_run:
        print("(DRY RUN MODE - commands will be shown but not executed)")
    
    print("=" * 60)
    
    success_count = 0
    for dataset_name in datasets_to_submit:
        config = DATASET_CONFIGS[dataset_name]
        success = submit_dataset(dataset_name, config, args.dry_run, use_slurm)
        if success:
            success_count += 1
        print()
    
    print("=" * 60)
    print(f"Summary: {success_count}/{len(datasets_to_submit)} datasets submitted successfully")
    
    if not args.dry_run and use_slurm and success_count > 0:
        print("\nMonitor jobs with: squeue -u $USER")
        print("Check logs in: logs/subsample_*")


if __name__ == "__main__":
    main()