opencode-knowledge-work 0.0.0-stage → 0.1.0

package/content/bio-research/skills/bio-research-nextflow-development/scripts/generate_samplesheet.py

@@ -0,0 +1,455 @@
+#!/usr/bin/env python3
+"""
+Enhanced nf-core samplesheet generator.
+
+Features:
+- FASTQ, BAM, and CRAM support
+- Tumor/normal status inference for sarek
+- Robust R1/R2 matching with scoring
+- Pre-write validation with clear error messages
+- Pipeline config-driven column generation
+
+Usage:
+    python generate_samplesheet.py /path/to/data rnaseq -o samplesheet.csv
+    python generate_samplesheet.py /path/to/bams sarek --input-type bam
+    python generate_samplesheet.py --validate samplesheet.csv rnaseq
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import yaml
+
+# Add parent directory to path for utils import
+sys.path.insert(0, str(Path(__file__).parent))
+
+from utils.file_discovery import discover_files, detect_input_type, find_index_file
+from utils.sample_inference import (
+    extract_sample_info,
+    infer_tumor_normal_status,
+    match_read_pairs,
+    extract_replicate_number
+)
+from utils.validators import validate_samplesheet, ValidationResult
+
+
+def load_pipeline_config(pipeline: str) -> Dict:
+    """Load pipeline configuration from YAML."""
+    config_dir = Path(__file__).parent / "config" / "pipelines"
+    config_file = config_dir / f"{pipeline}.yaml"
+
+    if not config_file.exists():
+        available = [f.stem for f in config_dir.glob("*.yaml") if not f.stem.startswith("_")]
+        raise ValueError(f"Unknown pipeline '{pipeline}'. Available: {', '.join(available)}")
+
+    with open(config_file) as f:
+        return yaml.safe_load(f)
+
+
+def generate_samplesheet(
+    input_dir: str,
+    pipeline: str,
+    output_file: Optional[str] = None,
+    input_type: str = "auto",
+    single_end: bool = False,
+    interactive: bool = True
+) -> Tuple[Optional[str], ValidationResult]:
+    """
+    Generate samplesheet for specified pipeline.
+
+    Args:
+        input_dir: Directory containing sequencing files
+        pipeline: Pipeline name (rnaseq, sarek, atacseq)
+        output_file: Output CSV path (default: samplesheet_{pipeline}.csv)
+        input_type: File type (auto, fastq, bam, cram)
+        single_end: Suppress pairing warnings for single-end data
+        interactive: Prompt for missing info
+
+    Returns:
+        Tuple of (output_path, validation_result)
+    """
+    config = load_pipeline_config(pipeline)
+    samplesheet_config = config.get("samplesheet", {})
+    supported_types = samplesheet_config.get("input_types", ["fastq"])
+
+    # Determine input type
+    if input_type == "auto":
+        input_type = detect_input_type(input_dir)
+        print(f"Auto-detected input type: {input_type.upper()}")
+
+    if input_type not in supported_types:
+        return None, ValidationResult(
+            valid=False,
+            errors=[f"Pipeline '{pipeline}' does not support {input_type.upper()} input. "
+                    f"Supported: {supported_types}"]
+        )
+
+    # Discover files
+    try:
+        files = discover_files(input_dir, input_type)
+    except ValueError as e:
+        return None, ValidationResult(valid=False, errors=[str(e)])
+
+    if not files:
+        return None, ValidationResult(
+            valid=False,
+            errors=[f"No {input_type.upper()} files found in {input_dir}"],
+            suggestions=[
+                "Check directory path is correct",
+                "Verify file extensions (.fastq.gz, .fq.gz, .bam, .cram)",
+                f"Run: ls {input_dir}"
+            ]
+        )
+
+    print(f"Found {len(files)} {input_type.upper()} files")
+
+    # Process based on input type
+    if input_type == "fastq":
+        rows = _process_fastq_files(files, config, single_end)
+    else:
+        rows = _process_alignment_files(files, config, input_type)
+
+    if not rows:
+        return None, ValidationResult(
+            valid=False,
+            errors=["Could not generate any samplesheet rows from files"]
+        )
+
+    print(f"Generated {len(rows)} samplesheet rows")
+
+    # Pipeline-specific processing
+    if pipeline == "sarek":
+        rows = _process_sarek_samples(rows, interactive)
+    elif pipeline == "atacseq":
+        rows = _process_atacseq_samples(rows)
+
+    # Validate before writing
+    validation = validate_samplesheet(rows, pipeline, config)
+
+    if not validation.valid:
+        print("\nValidation errors:")
+        for error in validation.errors:
+            print(f"  - {error}")
+
+        if interactive:
+            response = input("\nProceed anyway? [y/N]: ").strip().lower()
+            if response != 'y':
+                return None, validation
+    elif validation.warnings:
+        print("\nWarnings:")
+        for warning in validation.warnings:
+            print(f"  - {warning}")
+
+    # Determine output path
+    output_path = output_file or f"samplesheet_{pipeline}.csv"
+
+    # Write samplesheet
+    _write_samplesheet(rows, config, output_path)
+
+    print(f"\nGenerated: {output_path}")
+    print(f"  Pipeline: {pipeline} v{config.get('version', 'unknown')}")
+    print(f"  Samples: {len(set(r.get('sample', r.get('patient', '')) for r in rows))}")
+    print(f"  Rows: {len(rows)}")
+
+    # Preview
+    _print_preview(rows, config)
+
+    return output_path, validation
+
+
+def _process_fastq_files(files, config: Dict, single_end: bool) -> List[Dict]:
+    """Process FASTQ files into samplesheet rows."""
+    pairs = match_read_pairs(files)
+
+    if not pairs:
+        return []
+
+    # Check for unpaired files
+    unpaired = [k for k, v in pairs.items() if v.get('r1') and not v.get('r2')]
+    if unpaired and not single_end:
+        print(f"\nNote: {len(unpaired)} samples appear to be single-end (no R2)")
+
+    rows = []
+    columns = config.get("samplesheet", {}).get("columns", [])
+
+    for sample_key, pair_info in sorted(pairs.items()):
+        if not pair_info.get('r1'):
+            continue  # Skip entries with only R2
+
+        info = pair_info.get('info', {})
+
+        row = {
+            'sample': info.get('sample', sample_key),
+            'fastq_1': str(Path(pair_info['r1']).absolute()),
+            'fastq_2': str(Path(pair_info['r2']).absolute()) if pair_info.get('r2') else '',
+        }
+
+        # Add additional info from filename
+        if 'patient' in [c['name'] for c in columns]:
+            row['patient'] = info.get('patient', info.get('sample', sample_key))
+
+        if 'lane' in [c['name'] for c in columns]:
+            row['lane'] = info.get('lane', 'L001')
+
+        # Apply defaults from config
+        for col in columns:
+            if col['name'] not in row and 'default' in col:
+                row[col['name']] = col['default']
+
+        rows.append(row)
+
+    return rows
+
+
+def _process_alignment_files(files, config: Dict, input_type: str) -> List[Dict]:
+    """Process BAM/CRAM files into samplesheet rows."""
+    rows = []
+    columns = config.get("samplesheet", {}).get("columns", [])
+
+    for file_info in files:
+        # Find index file
+        index_path = find_index_file(file_info.path)
+
+        info = extract_sample_info(file_info.path)
+
+        row = {
+            'sample': info.get('sample', file_info.stem),
+            'bam': str(Path(file_info.path).absolute()),
+            'bai': str(Path(index_path).absolute()) if index_path else '',
+        }
+
+        # Add patient for sarek
+        if 'patient' in [c['name'] for c in columns]:
+            row['patient'] = info.get('patient', info.get('sample', file_info.stem))
+
+        # Apply defaults
+        for col in columns:
+            if col['name'] not in row and 'default' in col:
+                row[col['name']] = col['default']
+
+        # Warn if no index found
+        if not index_path:
+            print(f"  Warning: No index found for {file_info.name}")
+
+        rows.append(row)
+
+    return rows
+
+
+def _process_sarek_samples(rows: List[Dict], interactive: bool) -> List[Dict]:
+    """Process sarek samples: infer and confirm tumor/normal status."""
+    # Auto-infer status from sample names
+    for row in rows:
+        sample_name = row.get('sample', '')
+        inferred = infer_tumor_normal_status(sample_name)
+        if inferred is not None:
+            row['status'] = inferred
+
+    # Report inference results
+    inferred_tumor = [r for r in rows if r.get('status') == 1]
+    inferred_normal = [r for r in rows if r.get('status') == 0]
+    unknown = [r for r in rows if 'status' not in r]
+
+    if inferred_tumor or inferred_normal:
+        print(f"\nTumor/normal inference:")
+        print(f"  Tumor samples: {len(inferred_tumor)}")
+        print(f"  Normal samples: {len(inferred_normal)}")
+
+    # Handle unknown samples
+    if unknown and interactive:
+        print(f"\n{len(unknown)} sample(s) with unknown status:")
+        for r in unknown:
+            print(f"  - {r.get('sample')}")
+
+        print("\nSpecify status for each (0=normal, 1=tumor, Enter=skip):")
+        for r in unknown:
+            response = input(f"  {r.get('sample')} [0/1/Enter]: ").strip()
+            if response in ['0', '1']:
+                r['status'] = int(response)
+            else:
+                r['status'] = 0  # Default to normal
+                print(f"    Defaulting to normal (0)")
+    elif unknown:
+        # Non-interactive: default to normal
+        for r in unknown:
+            r['status'] = 0
+
+    return rows
+
+
+def _process_atacseq_samples(rows: List[Dict]) -> List[Dict]:
+    """Process ATAC-seq samples: ensure replicate numbers."""
+    # Group by sample name
+    sample_counts = {}
+    for row in rows:
+        sample = row.get('sample', '')
+        if sample not in sample_counts:
+            sample_counts[sample] = 0
+        sample_counts[sample] += 1
+
+    # Assign replicate numbers if not present
+    sample_rep = {}
+    for row in rows:
+        sample = row.get('sample', '')
+
+        if 'replicate' not in row or not row['replicate']:
+            # Try to extract from filename
+            extracted = extract_replicate_number(row.get('fastq_1', ''))
+            if extracted:
+                row['replicate'] = extracted
+            else:
+                # Auto-assign sequential
+                if sample not in sample_rep:
+                    sample_rep[sample] = 0
+                sample_rep[sample] += 1
+                row['replicate'] = sample_rep[sample]
+
+    return rows
+
+
+def _write_samplesheet(rows: List[Dict], config: Dict, output_path: str):
+    """Write samplesheet to CSV file."""
+    columns = config.get("samplesheet", {}).get("columns", [])
+    column_names = [c['name'] for c in columns]
+
+    # Filter to columns that have data
+    active_columns = [c for c in column_names if any(c in row and row[c] for row in rows)]
+
+    # Ensure fastq_1/fastq_2 or bam/bai are included
+    for required in ['fastq_1', 'bam']:
+        if required in column_names and required not in active_columns:
+            if any(required in row for row in rows):
+                active_columns.append(required)
+
+    # Maintain original column order
+    active_columns = [c for c in column_names if c in active_columns]
+
+    with open(output_path, 'w') as f:
+        f.write(','.join(active_columns) + '\n')
+        for row in rows:
+            values = [str(row.get(col, '')) for col in active_columns]
+            f.write(','.join(values) + '\n')
+
+
+def _print_preview(rows: List[Dict], config: Dict):
+    """Print preview of generated samplesheet."""
+    columns = config.get("samplesheet", {}).get("columns", [])
+    column_names = [c['name'] for c in columns]
+    active_columns = [c for c in column_names if any(c in row for row in rows)]
+
+    print(f"\nPreview (first 3 rows):")
+    print(','.join(active_columns))
+    for row in rows[:3]:
+        values = [str(row.get(col, ''))[:40] for col in active_columns]  # Truncate long paths
+        print(','.join(values))
+    if len(rows) > 3:
+        print(f"... ({len(rows) - 3} more rows)")
+
+
+def validate_existing_samplesheet(csv_path: str, pipeline: str) -> ValidationResult:
+    """Validate an existing samplesheet file."""
+    import csv
+
+    if not os.path.exists(csv_path):
+        return ValidationResult(valid=False, errors=[f"File not found: {csv_path}"])
+
+    try:
+        with open(csv_path, 'r') as f:
+            reader = csv.DictReader(f)
+            rows = list(reader)
+    except Exception as e:
+        return ValidationResult(valid=False, errors=[f"Failed to read CSV: {e}"])
+
+    if not rows:
+        return ValidationResult(valid=False, errors=["Samplesheet is empty"])
+
+    config = load_pipeline_config(pipeline)
+    return validate_samplesheet(rows, pipeline, config)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Generate nf-core samplesheet from data directory',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Generate samplesheet for RNA-seq
+    %(prog)s ./fastqs rnaseq -o samples.csv
+
+    # Generate samplesheet for sarek from BAM files
+    %(prog)s ./bams sarek --input-type bam
+
+    # Validate existing samplesheet
+    %(prog)s --validate samplesheet.csv rnaseq
+
+Supported pipelines: rnaseq, sarek, atacseq
+        """
+    )
+
+    parser.add_argument('input', help='Directory with data files, or CSV path for --validate')
+    parser.add_argument('pipeline', help='Pipeline name (rnaseq, sarek, atacseq)')
+    parser.add_argument('-o', '--output', help='Output CSV filename')
+    parser.add_argument('--input-type', choices=['auto', 'fastq', 'bam', 'cram'],
+                        default='auto', help='Input file type (default: auto-detect)')
+    parser.add_argument('--single-end', action='store_true',
+                        help='Treat as single-end data (suppress pairing warnings)')
+    parser.add_argument('--validate', action='store_true',
+                        help='Validate existing samplesheet instead of generating')
+    parser.add_argument('--no-interactive', action='store_true',
+                        help='Non-interactive mode (use defaults)')
+
+    args = parser.parse_args()
+
+    try:
+        if args.validate:
+            # Validate existing samplesheet
+            result = validate_existing_samplesheet(args.input, args.pipeline)
+            if result.valid:
+                print(f"✓ Samplesheet is valid for {args.pipeline}")
+                if result.warnings:
+                    print("\nWarnings:")
+                    for w in result.warnings:
+                        print(f"  - {w}")
+                sys.exit(0)
+            else:
+                print(f"✗ Samplesheet validation failed")
+                print(result.summary())
+                sys.exit(1)
+        else:
+            # Generate new samplesheet
+            if not os.path.isdir(args.input):
+                print(f"Error: Not a directory: {args.input}")
+                sys.exit(1)
+
+            output_path, result = generate_samplesheet(
+                args.input,
+                args.pipeline,
+                args.output,
+                args.input_type,
+                args.single_end,
+                interactive=not args.no_interactive
+            )
+
+            if output_path is None:
+                print("\nFailed to generate samplesheet.")
+                if result.suggestions:
+                    print("\nSuggestions:")
+                    for s in result.suggestions:
+                        print(f"  - {s}")
+                sys.exit(1)
+
+            sys.exit(0)
+
+    except ValueError as e:
+        print(f"Error: {e}")
+        sys.exit(1)
+    except KeyboardInterrupt:
+        print("\nAborted.")
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()