nextflow.config

/*
========================================================================================
    SORAT - Nextflow Configuration
========================================================================================
    Configuration file for the Segmentation Orchestration and Reproducible Analysis Toolkit.
    
    Profiles:
      - local: Run on local machine with Apptainer/Singularity
      - slurm: Run on SLURM cluster (e.g., UBC Sockeye) with Apptainer
      - test:  Quick test with single patient
----------------------------------------------------------------------------------------
*/

// Global default parameters
params {
    // Input/Output
    // input is optional; when omitted, model-dependent defaults are used.
    input                   = null
    outdir                  = './results'

    // Frame mode: 'auto' uses ED/ES when info_cfg present, else all frames.
    // 'ed_es' forces 2-frame mode, 'all' always segments all temporal frames.
    frames_mode             = System.getenv('SORAT_FRAMES_MODE') ?: 'auto'
    max_frames              = System.getenv('SORAT_MAX_FRAMES') ? (System.getenv('SORAT_MAX_FRAMES') as int) : null

    // Default samplesheets used when --input is omitted.
    // `sax`/`atrial` are preferred generic keys; `acdc`/`mbas` are legacy aliases.
    default_inputs {
        sax                 = System.getenv('SORAT_SAX_SAMPLESHEET') ?: null
        atrial              = System.getenv('SORAT_ATRIAL_SAMPLESHEET') ?: null
        acdc                = System.getenv('SORAT_ACDC_SAMPLESHEET') ?: null
        mbas                = System.getenv('SORAT_MBAS_SAMPLESHEET') ?: null
    }
    
    // Generic SAX dataset location (ACDC-style layout) for auto-generating defaults.
    sax_data_root           = System.getenv('SORAT_SAX_DATA_ROOT') ?: System.getenv('SORAT_ACDC_DIR') ?: null
    sax_data_split          = System.getenv('SORAT_SAX_DATA_SPLIT') ?: 'testing'

    // Legacy aliases retained for backward compatibility.
    acdc_dir                = System.getenv('SORAT_ACDC_DIR') ?: null
    acdc_dataset            = System.getenv('SORAT_ACDC_DATASET') ?: 'testing'
    
    // Model selection: comma-separated list or 'all'
    // Options: 'cinema', 'nnformer', 'vsa3l', 'atrial_nnunet', 'all'
    // Note: 'all' intentionally runs only cinema, nnformer, and vsa3l.
    models                  = 'all'
    
    // Enable model comparison and report generation
    compare                 = true

    // Inference-only mode: skip metrics and report generation
    inference_only          = false

    // Ground-truth/evaluation contract for the current run.
    evaluation {
        label_schema        = System.getenv('SORAT_LABEL_SCHEMA') ?: 'architecture_default'
        // Raw ventricular GT label values, e.g. 'rv=1,myo=2,lv=3' (ACDC) or
        // 'lv=1,myo=2,rv=3' (M&Ms, M&Ms-2). 'auto' infers them from anatomy.
        gt_label_map        = System.getenv('SORAT_GT_LABEL_MAP') ?: 'auto'
    }

    // Single-patient mode: process only one patient from the samplesheet
    // --single_patient runs the first patient; --patient_name overrides which patient to run
    single_patient          = false
    patient_name            = null

    // Debug mode: generate execution/resource/quality analytics report
    debug                   = false

    // Results-only debug mode support. These parameters are primarily used with
    // `-entry DEBUG_ONLY` to regenerate debug outputs from an existing outdir.
    debug_run_name          = null
    debug_workflow_duration = null
    debug_workflow_start    = null
    debug_workflow_success  = null
    debug_source_outdir     = null
    debug_source_summary    = null
    debug_section3_outdir   = null
    debug_refresh_outdir    = null
    debug_reference_summaries = null

    // Auto-discover all locally available models inside containers
    auto_discover_models    = true
    
    // Model registry file for custom models
    model_registry          = "${projectDir}/conf/model_registry.json"
    
    // SLURM account (for slurm profile)
    slurm_account           = System.getenv('SORAT_SLURM_ACCOUNT') ?: null

    // Runtime mount/cache configuration (override per-user in .sorat/user.config)
    singularity_cache_dir   = System.getenv('SORAT_SINGULARITY_CACHEDIR') ?: "${System.getProperty('user.home')}/.singularity_cache"
    singularity_local_binds = "${projectDir}/bin:/app/bin,${projectDir}/nnformer:/app/nnformer"
    singularity_slurm_binds = "/arc:/arc,/scratch:/scratch,${projectDir}/bin:/app/bin,${projectDir}/nnformer:/app/nnformer"

    // SLURM tuning knobs (to reduce scheduler/orchestration overhead)
    slurm_max_forks         = 30
    slurm_queue_size        = 64
    slurm_submit_rate       = '50/1min'
    slurm_poll_interval     = '30 sec'
    slurm_queue_stat_interval = '60 sec'

    // Preprocess cache controls (Phase 3)
    preprocess_cache_enabled = true
    preprocess_cache_dir     = "${projectDir}/.cache/preprocess"
    
    // CineMA-specific parameters
    cinema {
        trained_dataset     = 'acdc'
        seeds               = [0, 1, 2]
        ensemble            = true
    }
    
    // nnFormer-specific parameters
    nnformer {
        fold                = 0
        tta                 = true
        mixed_precision     = true
    }
    
    // VSA-3L-specific parameters
    vsa3l {
        input_size          = [256, 256]
    }

    // Atrial nnUNet-specific parameters
    atrial_nnunet {
        dataset_id          = 'Dataset001_LGE'
        configuration       = '2d'
        folds               = [0, 1, 2, 3, 4]
        save_probabilities  = false
        // Preferred generic atrial dataset root. Accepts either:
        // 1) nnUNet layout root containing imagesTr/ and labelsTr/
        // 2) MBAS training-style root containing MBAS_### subfolders
        dataset_root        = System.getenv('SORAT_ATRIAL_DATA_ROOT') ?: null

        // Legacy alias retained for backward compatibility.
        mbas_root           = System.getenv('SORAT_MBAS_ROOT') ?: null
        // TEMP workaround currently uses wrapper-based checkpoint compatibility.
        // Remove after rebuilding atrial container with aligned dependencies.
    }

    // Optional postprocessing: dark LV -> MYO relabel
    postprocess {
        enabled             = false
        use_for_metrics     = false          // also score post-processed masks (as '<model>_pp'); raw output is always scored
        strength            = 1.0            // [0,1], conservative to aggressive

        // Standalone postprocess-only workflow controls
        samplesheet         = null
        results_dir         = null
    }

    // Segmentation preview generation (quick-look PNG overlays)
    visualization {
        enabled             = true
        // When true, generate previews for all frames in all-frames mode.
        // When false (default), only the first frame (frame00) gets a preview.
        // ED/ES frames are always previewed regardless of this flag.
        all_frames          = false
    }

    // Optional feature extraction from predicted or post-processed masks
    feature_extraction {
        enabled             = false
        // Options: predictions, postprocess, prompt
        mask_source         = 'prompt'

        // Optional explicit output directory for feature CSVs.
        // When unset: the main workflow writes to <results_dir or outdir>/features/{raw,postprocessed};
        // -entry FEATURES_ONLY writes to this run's <outdir>/features/{raw,postprocessed},
        // never into the --feature_extraction.results_dir it reads segmentations from.
        output_dir          = null

        // Optional exact model tag filter (comma-separated), e.g.
        // cinema__acdc_ensemble or cinema__acdc_seed1.
        model_tag           = null

        // Seed preference when model_tag is not provided.
        // Options:
        //   ensemble (default): choose ensemble per family, fallback to single available model
        //   all: use all discovered variants for each selected architecture
        //   seedN or N: choose that seed (e.g., seed0 or 0)
        seed                = 'ensemble'

        // Optional pre-built virtualenv containing pyradiomics.
        // Recommended on offline HPC nodes (e.g., Sockeye) where runtime pip installs are blocked.
        virtualenv_path     = System.getenv('SORAT_FEATURE_VENV') ?: null
        require_virtualenv  = false

        // PyRadiomics settings. Defaults reproduce earlier runs (no normalization,
        // binWidth 25, 3-D texture, native spacing). A harmonized cross-scanner setup:
        //   normalize=true normalize_scale=100 bin_count=32
        //   resample_spacing='1.25,1.25,0' force2d=true
        radiomics {
            normalize           = false
            normalize_scale     = null    // e.g. 100
            remove_outliers     = null    // e.g. 3 (standard deviations)
            bin_count           = null    // mutually exclusive with bin_width
            bin_width           = null    // PyRadiomics default is 25
            resample_spacing    = null    // 'x,y,z' in mm; 0 keeps an axis
            force2d             = false
            force2d_dimension   = 0
            // Per-image intensity reference instead of PyRadiomics normalize:
            // 'lv_bloodpool' divides by the LV blood-pool mean (x scale, default 100).
            intensity_reference       = null
            intensity_reference_scale = null
        }

        // Standalone FEATURES_ONLY workflow controls
        samplesheet         = null
        results_dir         = null
    }

    // Heuristic ETA configuration for mode-level runtime estimates.
    eta {
        enabled                     = false
        minutes_per_case_main       = 6.0
        minutes_per_case_postprocess_only = 1.0
        minutes_per_case_features_only = 0.6
        minutes_debug_only          = 2.0
    }
    
    // Resource defaults
    max_memory              = '128.GB'
    max_cpus                = 16
    max_time                = '24.h'
    
    // Help
    help                    = false
}

// Process-specific resource configuration
process {
    // Default resources
    cpus   = { check_max( 2 * task.attempt, 'cpus' ) }
    memory = { check_max( 8.GB * task.attempt, 'memory' ) }
    time   = { check_max( 4.h * task.attempt, 'time' ) }
    
    errorStrategy = { task.exitStatus in [143,137,104,134,135,139,141] ? 'retry' : 'finish' }
    maxRetries    = 2
    maxErrors     = '-1'
    
    // Process-specific resource requirements
    withLabel: process_low {
        cpus   = { check_max( 2 * task.attempt, 'cpus' ) }
        memory = { check_max( 4.GB * task.attempt, 'memory' ) }
        time   = { check_max( 2.h * task.attempt, 'time' ) }
    }
    
    withLabel: process_medium {
        cpus   = { check_max( 4 * task.attempt, 'cpus' ) }
        memory = { check_max( 16.GB * task.attempt, 'memory' ) }
        time   = { check_max( 8.h * task.attempt, 'time' ) }
    }
    
    withLabel: process_high {
        cpus   = { check_max( 8 * task.attempt, 'cpus' ) }
        memory = { check_max( 32.GB * task.attempt, 'memory' ) }
        time   = { check_max( 16.h * task.attempt, 'time' ) }
    }

    withLabel: process_gpu {
        cpus   = { check_max( 8 * task.attempt, 'cpus' ) }
        memory = { check_max( 32.GB * task.attempt, 'memory' ) }
        time   = { check_max( 8.h * task.attempt, 'time' ) }
    }

    // Stage-specific labels for tighter resource shaping
    withLabel: process_preprocess {
        cpus   = { check_max( 6 * task.attempt, 'cpus' ) }
        memory = { check_max( 8.GB * task.attempt, 'memory' ) }
        time   = { check_max( 4.h * task.attempt, 'time' ) }
    }

    withLabel: process_gpu_light {
        cpus   = { check_max( 8 * task.attempt, 'cpus' ) }
        memory = { check_max( 32.GB * task.attempt, 'memory' ) }
        time   = { check_max( 8.h * task.attempt, 'time' ) }
    }

    withLabel: process_gpu_heavy {
        cpus   = { check_max( 12 * task.attempt, 'cpus' ) }
        memory = { check_max( 48.GB * task.attempt, 'memory' ) }
        time   = { check_max( 10.h * task.attempt, 'time' ) }
    }
    
    // Model-specific container assignments (local .sif files)
    withName: 'CINEMA_.*' {
        container = "${projectDir}/containers/sorat-cinema.sif"
    }
    
    withName: 'NNFORMER_.*' {
        container = "${projectDir}/containers/sorat-nnformer.sif"
    }
    
    withName: 'VSA3L_.*' {
        container = "${projectDir}/containers/sorat-vsa3l.sif"
    }

    withName: 'ATRIAL_NNUNET_.*' {
        container = "${projectDir}/containers/sorat-atrial-nnunet.sif"
    }
    
    withName: 'PREPARE_INPUT|COMPUTE_METRICS|AGGREGATE_RESULTS|GENERATE_REPORT|GENERATE_DEBUG_REPORT|POSTPROCESS_.*|VISUALIZE_POSTPROCESS_.*|DISCOVER_POSTPROCESS_.*|GENERATE_SEGMENTATION_PREVIEWS|EXTRACT_FEATURES' {
        container = "${projectDir}/containers/sorat-cinema.sif"
    }

    // Architecture-specific runtime shaping
    withName: 'CINEMA_PREPROCESS|NNFORMER_PREPROCESS|VSA3L_PREPROCESS|ATRIAL_NNUNET_PREPROCESS' {
        label = 'process_preprocess'
    }

    withName: 'CINEMA_SEGMENT' {
        label = 'process_gpu_heavy'
    }

    withName: 'NNFORMER_SEGMENT|VSA3L_SEGMENT|ATRIAL_NNUNET_SEGMENT' {
        label = 'process_gpu_light'
    }
}

// Execution profiles
profiles {
    
    // Local execution with Apptainer/Singularity
    local {
        process.executor       = 'local'
        singularity.enabled    = true
        singularity.autoMounts = true
        docker.enabled         = false
        
        // Mount local override code into containers for patchability.
        singularity.runOptions = "--bind ${params.singularity_local_binds}"
    }
    
    // SLURM cluster profile (for HPC like UBC Sockeye)
    slurm {
        process.executor       = 'slurm'
        singularity.enabled    = true
        singularity.autoMounts = true
        docker.enabled         = false

        executor {
            queueSize         = params.slurm_queue_size
            submitRateLimit   = params.slurm_submit_rate
            pollInterval      = params.slurm_poll_interval
            queueStatInterval = params.slurm_queue_stat_interval
        }
        
        process {
            maxForks = params.slurm_max_forks

            // Load required modules in each job
            beforeScript = 'module purge && module load CVMFS_CC && module load apptainer/1.3.4'
            
            // SLURM requires --nodes=1 explicitly on some clusters
            clusterOptions = { "--account=${params.slurm_account} --nodes=1 --ntasks=1" }
            
            // GPU jobs use -gpu suffix on account
            withLabel: process_gpu {
                clusterOptions = { "--account=${params.slurm_account}-gpu --nodes=1 --ntasks=1 --gpus-per-node=1 --gpu-bind=per_task:1" }
                containerOptions = '--nv'
            }

            withLabel: process_gpu_light {
                clusterOptions = { "--account=${params.slurm_account}-gpu --nodes=1 --ntasks=1 --gpus-per-node=1 --gpu-bind=per_task:1" }
                containerOptions = '--nv'
            }

            withLabel: process_gpu_heavy {
                clusterOptions = { "--account=${params.slurm_account}-gpu --nodes=1 --ntasks=1 --gpus-per-node=1 --gpu-bind=per_task:1" }
                containerOptions = '--nv'
            }
        }
        
        // Bind paths and mount local bin/ and nnformer/ over container's versions.
        // GPU support is added only to GPU-labelled jobs via containerOptions.
        // --writable-tmpfs allows pip install inside container at runtime
        singularity.runOptions = "--writable-tmpfs --bind ${params.singularity_slurm_binds}"
        singularity.cacheDir   = params.singularity_cache_dir
    }
    
    // Test profile with minimal resources
    test {
        params {
            input   = "${projectDir}/test/samplesheet.csv"
            outdir  = "${projectDir}/test/results"
            models  = 'cinema'
            compare = false
        }
        
        process {
            cpus   = 2
            memory = 4.GB
            time   = 1.h
        }
    }
}

// Manifest
manifest {
    name            = 'SORAT'
    author          = 'Your Name'
    homePage        = 'https://github.com/your-org/SORAT'
    description     = 'Segmentation Orchestration and Reproducible Analysis Toolkit'
    mainScript      = 'main.nf'
    nextflowVersion = '>=23.04.0'
    version         = '1.0.0'
}

// Load additional config files
includeConfig 'conf/base.config'

// Optionally load per-user local overrides from a gitignored file.
def soratUserConfig = new File("${projectDir}/.sorat/user.config")
if (soratUserConfig.exists()) {
    includeConfig soratUserConfig.toString()
}

// Function to check max resources
def check_max(obj, type) {
    if (type == 'memory') {
        try {
            if (obj.compareTo(params.max_memory as nextflow.util.MemoryUnit) == 1)
                return params.max_memory as nextflow.util.MemoryUnit
            else
                return obj
        } catch (all) {
            println "WARNING: Max memory '${params.max_memory}' is not valid. Using default value: $obj"
            return obj
        }
    } else if (type == 'time') {
        try {
            if (obj.compareTo(params.max_time as nextflow.util.Duration) == 1)
                return params.max_time as nextflow.util.Duration
            else
                return obj
        } catch (all) {
            println "WARNING: Max time '${params.max_time}' is not valid. Using default value: $obj"
            return obj
        }
    } else if (type == 'cpus') {
        try {
            return Math.min( obj, params.max_cpus as int )
        } catch (all) {
            println "WARNING: Max cpus '${params.max_cpus}' is not valid. Using default value: $obj"
            return obj
        }
    }
}