diff --git a/CHANGELOG b/CHANGELOG index 8f2584b..38442ac 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,5 +1,33 @@ # Changelog +## Unreleased + +- Taxonomic validation: `tv_local_blast.py` now batches queries into chunks of 50 + sequences per `blastn` invocation instead of running one process per sequence, + so a large reference database is loaded once per chunk rather than once per + barcode. Per-sequence TSV outputs, the summary CSV layout and resume behaviour + are unchanged. A failed chunk retries its sequences individually, and + sequences that still fail now cause a non-zero exit (with no summary CSV) + rather than being reported as no-match; `--allow-partial-failures` restores the + previous leniency. +- Fixed the summary CSV dropping all hits for sequences whose FASTA header + carries a description, where BLAST's `qseqid` (first token only) did not match + the sanitized full header used as the CSV key. +- Removed the unused `blast_options` entry from the Snakefile's + `taxonomic_validation` fallback, and corrected the documented BLAST hit limit + (100, not 500). +- Structural validation: fixed a temp-file leak in `structural_validation.py`, + where the per-sequence nhmmer query and tabular files were created with + `delete=False` and never removed, leaving two files per sequence in `TMPDIR` + (48 per sample, since each sample is validated across 6 parameter + combinations and 4 input FASTAs). +- Structural validation: added `--threads`, passed from the rule, setting + nhmmer's `--cpu` per invocation. Defaults to 1, so standalone behaviour is + unchanged. +- Right-sized the `structural_validation` rule resources from 32 GB / 1 thread / + himem to 4 GB / 8 threads / medium. Peak memory scales with sequence count + rather than read depth, measured at roughly 110 MB per 1000 samples. + ## v3.0.4 (2026-08-13) Initial PyPI release. diff --git a/README.md b/README.md index bb28297..a978738 100644 --- a/README.md +++ b/README.md @@ -137,7 +137,7 @@ bash run_local.sh - Cleaned consensus generation and metrics aggregation ([`05_consensus_generator.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/05_consensus_generator.py)) 6. **Barcode validation and selection** (see [Validation process](#validation-process)): - **Structural validation** - HMM-based barcode extraction, reading frame analysis, stop codon detection and quality ranking ([`structural_validation.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/structural_validation.py)) - - **Local BLASTn search** - parallel searches of structurally validated barcodes against a local reference database ([`tv_local_blast.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/tv_local_blast.py)) + - **Local BLASTn search** - chunked, parallel searches of structurally validated barcodes against a local reference database ([`tv_local_blast.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/tv_local_blast.py)) - **Taxonomic validation** - hierarchical matching of BLAST results against expected taxonomy, selecting the best sequence per sample ([`tv_blast2taxonomy.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/tv_blast2taxonomy.py)) 7. **Statistics compilation** - QC, recovery, cleaning, filtering and validation metrics aggregated into CSV reports ([`compile_barcoding_stats.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/compile_barcoding_stats.py)). 8. **Final integration** - all pipeline metrics merged into a unified output CSV ([`val_csv_merger.py`](https://github.com/bge-barcoding/BeeGees/blob/main/workflow/scripts/val_csv_merger.py)). @@ -504,7 +504,7 @@ Structural validation (via `structural_validation.py`) assesses every barcode co ## Taxonomic validation ## Taxonomic validation runs in two steps, via `tv_local_blast.py` and `tv_blast2taxonomy.py`. -**1. Local BLASTn search.** Parallel BLASTn searches against a local database, either built from a multi-FASTA with `makeblastdb` or supplied pre-built. The e-value threshold is hardcoded to 1e-5. Per-sequence TSV outputs (outfmt 6) hold the top 500 hits ordered by descending percent identity; the top 100 are carried into the summary CSV. +**1. Local BLASTn search.** BLASTn searches against a local database, either built from a multi-FASTA with `makeblastdb` or supplied pre-built. The e-value threshold is hardcoded to 1e-5 and `max_target_seqs` to 100. Queries are batched into chunks of 50 sequences per `blastn` invocation, with up to `threads` chunks running concurrently, so the database is loaded once per chunk rather than once per sequence. Results are demultiplexed back into one TSV per sequence (outfmt 6), each holding up to 100 hits ordered by descending percent identity, and all 100 are carried into the summary CSV. Sequences whose TSV already exists are skipped, so an interrupted run resumes without repeating finished work. **2. Taxonomic assignment validation.** BLASTn results are checked against expected taxonomy using hierarchical matching and quality-based filtering: 1. Parse the local BLASTn summary CSV, per-sample expected lineages, database taxonomy mappings, and structurally validated sequences. diff --git a/beegees/config/config.yaml b/beegees/config/config.yaml index 439634d..0d96cab 100644 --- a/beegees/config/config.yaml +++ b/beegees/config/config.yaml @@ -97,72 +97,81 @@ taxonomic_validation: # taxdump_dir: "/path/to/writable/ncbi_taxdump" -## Resource allocation for each rule. Memory shown in Mb (mem_mb), e.g. 2048 = 2G memory. Rules have dynamic memory scaling upon retry (mem_mb * retry #). +## Resource allocation for each rule +# +# mem_mb is in MB (2048 = 2 GB) and is multiplied by the attempt number on retry. +# A rule that dies at 4096 is resubmitted with 8192, etc. Therefore, size the base +# value for a typical run and let retries absorb any outliers. +# +# PER-SAMPLE rules run once per sample (MitoGeneExtractor runs once per sample per +# r/s combination). Their cost is fixed no matter how many samples are run, so +# size them for your deepest single sample - a 1-sample and a 500-sample run need +# the same values. +# +# AGGREGATE rules run once for the whole run, so their memory grows with sample +# count. The values below carry headroom to ~500 samples; the per-rule comments +# highlight what drives each one. +# +# partition is only read for gene_fetch, MitoGeneExtractor, structural_validation +# and taxonomic_validation. Every other rule takes the profile default (medium). +# Do not delete those four keys: the Snakefile indexes them directly and will fail +# with a KeyError. rules: - gene_fetch: # Gene Fetch: reference retrieval from NCBI GenBank - mem_mb: 8192 - threads: 4 - partition: long + # ===== PER-SAMPLE: cost per job is independent of how many samples are in the run ===== + gene_fetch: # gene-fetch: reference retrieval from NCBI GenBank + mem_mb: 4096 + partition: PLACEHOLDER # Runtime scales on sample number. fastp_qc: # fastp: adapter trimming, QC, merging/deduplication (PE merge and concat modes) - mem_mb: 16384 - threads: 4 - clean_headers_merge: # Clean fastq headers for MitoGeneExtractor compatibility (merge mode) mem_mb: 8192 - threads: 1 + threads: 4 # fastp --thread. Note gains flatten past ~8 even on deep samples + clean_headers_merge: # Clean fastq headers for MitoGeneExtractor compatibility (merge mode) + mem_mb: 4096 fastq_concat: # Concatenation of fastp QC'd R1+R2 fastq files (concat mode) - mem_mb: 8192 - threads: 1 + mem_mb: 2048 quality_trim: # Trim Galore: secondary quality trimming of concatenated reads (concat mode) mem_mb: 8192 - threads: 4 + threads: 4 # trim_galore --cores. Note it spawns ~3-4x this many processes (cutadapt + pigz), so 4 occupies ~15 CPUs downsample: # reformat.sh: optional read downsampling (concat/merge/SE modes) mem_mb: 8192 - threads: 4 - MitoGeneExtractor: # MitoGeneExtractor: protein-guided barcode extraction (concat/merge/SE modes); single-threaded (Exonerate) - mem_mb: 32768 - threads: 1 - partition: himem - rename_and_combine_cons: # Standardise and combine MGE consensus FASTA headers across parameter combinations - mem_mb: 8192 - threads: 4 + MitoGeneExtractor: # MitoGeneExtractor: protein-guided barcode extraction (concat/merge/SE modes) + mem_mb: 32768 # Per sample per r/s combination; scales with read depth. + partition: PLACEHOLDER # Submit to high memory nodes (iv available) to allow memory scaling upon retry. + # ===== AGGREGATE: one job for the whole run, so memory grows with sample count ===== + rename_and_combine_cons: # rename_headers.py: standardise and combine MGE consensus headers across parameter combinations + mem_mb: 4096 + threads: 4 # rename_headers.py --threads. >4 is likely to have dimishing speed returns gzip_merged_clean: # Gzip *_merged_clean.fastq files to reclaim disk space - mem_mb: 8192 - threads: 4 + mem_mb: 2048 # Flat memory requirement regardless of sample number human_cox1_filter: # 01_human_cox1_filter.py: remove human COX1 contamination - mem_mb: 16384 + mem_mb: 8192 # Process pool over per-sample alignments: peak is [threads] x [largest alignment], so it tracks read depth, NOT sample count threads: 4 at_content_filter: # 02_at_content_filter.py: remove reads with aberrant AT content - mem_mb: 16384 + mem_mb: 8192 # As above: [threads] x [largest alignment] threads: 4 statistical_outlier_filter: # 03_statistical_outlier_filter.py: remove reads statistically dissimilar to the consensus - mem_mb: 16384 + mem_mb: 8192 # As above: [threads] x [largest alignment], plus per-alignment numpy distance arrays - the heaviest of the three filters threads: 4 reference_filter: # 04_reference_filter.py: optional reference-based sequence retention or contaminant removal mem_mb: 4096 - threads: 1 + threads: 4 consensus_generation: # 05_consensus_generator.py: generate cleaned consensus sequences and aggregate metrics - mem_mb: 8192 + mem_mb: 8192 # As above: [threads] x [largest alignment]. threads: 4 - extract_stats_to_csv: # mge_stats.py: extract per-sample MGE alignment statistics to CSV - mem_mb: 8192 - threads: 1 - structural_validation: # structural_validation.py: HMM-based extraction, reading frame and stop-codon checks; nhmmer is single-threaded - mem_mb: 32768 - threads: 1 - partition: himem - taxonomic_validation: # tv_local_blast.py: parallel BLASTn searches against local reference database - mem_mb: 8192 - threads: 8 - partition: medium + extract_stats_to_csv: # compile_barcoding_stats.py: extract per-sample MGE alignment statistics to CSV + mem_mb: 4096 # Builds one table across all samples, so grows with sample count + structural_validation: # structural_validation.py: HMM-based extraction, reading frame and stop-codon checks + mem_mb: 4096 + threads: 4 # CPUs given to each nhmmer call (--cpu). Processing is sub-linear, so >4 threads has diminishing returns + partition: PLACEHOLDER + taxonomic_validation: # tv_local_blast.py: chunked BLASTn searches against local reference database + mem_mb: 8192 # ~250 MB per concurrent (50 consensus seq) chunk, plus the summary table over all barcodes + threads: 8 # Max concurrent single-threaded BLASTn chunks (50 sequences per chunk). + partition: PLACEHOLDER blast2taxonomy: # tv_blast2taxonomy.py: hierarchical taxonomic matching and best-barcode selection - mem_mb: 16384 - threads: 4 - download_taxdump: # Download NCBI taxdump for blast2taxonomy mem_mb: 4096 - threads: 1 - multiqc_plots: # Generate per-sample summary plots from pipeline metrics - mem_mb: 8192 - threads: 2 + download_taxdump: # Download and extract the NCBI taxdump for blast2taxonomy + mem_mb: 2048 + multiqc_plots: # multiqc_plots.R: per-sample summary plots from pipeline metrics + mem_mb: 4096 # R data frames spanning all samples; grows with sample count multiqc: # MultiQC: aggregate all QC reports into a single interactive HTML report - mem_mb: 8192 - threads: 1 + mem_mb: 8192 # Parses every report in the run - the largest aggregate consumer at high sample counts \ No newline at end of file diff --git a/beegees/config/multiqc_config.yaml b/beegees/config/multiqc_config.yaml index f63b534..15126c9 100644 --- a/beegees/config/multiqc_config.yaml +++ b/beegees/config/multiqc_config.yaml @@ -6,8 +6,8 @@ title: "BeeGees Pipeline Report" report_comment: > This MultiQC report was generated by the - BeeGees (Barcode gene Extraction and Evaluation from Genome Skims) Snakemake workflow - developed at NHMUK. The report outlines overall barcoding success per sample, as well as read Quality Control (QC) + BeeGees (Barcode Gene Extraction and Evaluation from Genome Skims) pipeline + developed at NHMUK by Dr. Dan Parsons. The report outlines overall barcoding success per sample, as well as read Quality Control (QC) statistics, the outcome of reference retrieval (Gene Fetch), and barcode validation results. This report is complemented by the *_final_metrics.csv and plots within 05_barcoding_outcome/plots/. diff --git a/beegees/workflow/Snakefile b/beegees/workflow/Snakefile index 7c1fb9c..655b41e 100644 --- a/beegees/workflow/Snakefile +++ b/beegees/workflow/Snakefile @@ -759,8 +759,9 @@ if downsampling_enabled: echo "Placeholder output created: $(date)" >> {log.out} exit 0 fi - - reformat.sh in={input.merged_reads} \\ + + reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\ + in={input.merged_reads} \\ out={output.downsampled} \\ samplereadstarget={params.max_reads} \\ sampleseed=12345 \\ @@ -1060,7 +1061,8 @@ if downsampling_enabled: exit 0 fi - reformat.sh in={input.concat_trimmed} \\ + reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\ + in={input.concat_trimmed} \\ out={output.downsampled} \\ samplereadstarget={params.max_reads} \\ sampleseed=12345 \\ @@ -1186,7 +1188,6 @@ if downsampling_enabled: log: out=os.path.join(preprocessing_dir_se, "logs/downsample/{sample}.out"), err=os.path.join(preprocessing_dir_se, "logs/downsample/{sample}.err") - threads: rule_resources["downsample"]["threads"] resources: mem_mb=lambda wildcards, attempt: rule_resources["downsample"]["mem_mb"] * attempt retries: 3 @@ -1206,7 +1207,8 @@ if downsampling_enabled: exit 0 fi - reformat.sh in={input.se_trimmed} \\ + reformat.sh -Xmx$(( {resources.mem_mb} * 80 / 100 ))m -eoom \\ + in={input.se_trimmed} \\ out={output.downsampled} \\ samplereadstarget={params.max_reads} \\ sampleseed=12345 \\ @@ -3064,11 +3066,11 @@ rule multiqc_plots: {params.taxdump_dir} \ {params.run_mode} \ >> {log} 2>&1 - + # Move static PNGs from mqc_dir to plots_dir find {params.mqc_dir} -maxdepth 1 -name "*.png" \ -exec mv {{}} {params.plots_dir}/ \; - + echo "Completed: $(date)" >> {log} """ @@ -3091,29 +3093,29 @@ rule multiqc: shell: """ set -euo pipefail - + mkdir -p {params.output_dir} mkdir -p $(dirname {log}) - + echo "Running MultiQC: $(date)" > {log} echo "Data dir: {params.mqc_data_dir}" >> {log} echo "Config: {input.mqc_config}" >> {log} echo "Output dir: {params.output_dir}" >> {log} - + multiqc {params.mqc_data_dir} \ --config {input.mqc_config} \ --outdir {params.output_dir} \ --filename multiqc_report.html \ --force \ >> {log} 2>&1 - + # Copy report back to main output directory cp {params.output_dir}/multiqc_report.html \ {main_output_dir}/multiqc_report.html - + echo "MultiQC complete: $(date)" >> {log} """ - + # Final clean up superfluous files rule cleanup_files: input: @@ -3234,9 +3236,6 @@ rule cleanup_files: f.write("\nPreprocessing mode: concat") - - - # ----- SE MitoGeneExtractor (mirrors MitoGeneExtractor_concat) ----- rule MitoGeneExtractor_se: input: @@ -3255,7 +3254,6 @@ rule MitoGeneExtractor_se: t=t, output_dir=barcode_recovery_dir_se, vulgar_dir=lambda wildcards: os.path.join(barcode_recovery_dir_se, f"logs/mge/{wildcards.sample}_r_{wildcards.r}_s_{wildcards.s}/") - threads: rule_resources["MitoGeneExtractor"]["threads"] resources: mem_mb=lambda wildcards, attempt: rule_resources["MitoGeneExtractor"]["mem_mb"] * attempt, slurm_partition=rule_resources["MitoGeneExtractor"]["partition"] diff --git a/beegees/workflow/scripts/structural_validation.py b/beegees/workflow/scripts/structural_validation.py index 2ee7164..0e6b42c 100644 --- a/beegees/workflow/scripts/structural_validation.py +++ b/beegees/workflow/scripts/structural_validation.py @@ -111,6 +111,8 @@ Optional: --code/-c: Genetic code table for translation (default: 1 for standard genetic code) +--threads/-t: CPUs passed to each nhmmer call via --cpu (default: 1). nhmmer is invoked once per + sequence, so this sets threads per invocation rather than the number of concurrent invocations. --log-file LOG_FILE: Specify a custom path for the log file (default: creates timestamped log) --verbose, -v: Enable detailed debug logging --disable-selection: Skip best-per-process selection and output all passing sequences instead @@ -320,28 +322,33 @@ def get_complete_codons(seq, offset): complete_codons += codon return complete_codons -def run_nhmmer_on_sequence(sequence, seq_id, hmm_file): +def run_nhmmer_on_sequence(sequence, seq_id, hmm_file, threads=1): try: - # Create FASTA file with the sequence - with tempfile.NamedTemporaryFile(mode='w+', suffix='.fasta', delete=False) as temp_input, \ - tempfile.NamedTemporaryFile(mode='w+', suffix='.tbl', delete=False) as temp_tblout: - - temp_input.write(f">{seq_id}\n{sequence}\n") - temp_input.flush() - + # One temporary directory per call, removed on exit (including on error). + # The scratch query and tabular files previously used + # NamedTemporaryFile(delete=False) and were never unlinked, leaking two + # files per sequence into TMPDIR for the lifetime of the job. + with tempfile.TemporaryDirectory(prefix='structval_nhmmer_') as temp_dir: + input_path = os.path.join(temp_dir, 'query.fasta') + tblout_path = os.path.join(temp_dir, 'hits.tbl') + + # Create FASTA file with the sequence + with open(input_path, 'w') as temp_input: + temp_input.write(f">{seq_id}\n{sequence}\n") + # Run nhmmer with separate tabular output file nhmmer_cmd = [ 'nhmmer', - '--tblout', temp_tblout.name, + '--tblout', tblout_path, '--incE', '1e-3', - '--cpu', '1', + '--cpu', str(threads), str(hmm_file), - temp_input.name + input_path ] - + logging.debug(f"Running nhmmer on sequence {seq_id}: {' '.join(nhmmer_cmd)}") result = subprocess.run(nhmmer_cmd, capture_output=True, text=True) - + if result.returncode != 0: logging.error(f"nhmmer failed for {seq_id}: {result.stderr}") return None @@ -349,17 +356,24 @@ def run_nhmmer_on_sequence(sequence, seq_id, hmm_file): # Log complete nhmmer output for each sample logging.debug(f"=== COMPLETE nhmmer OUTPUT for {seq_id} ===") logging.debug(result.stdout) - - # Parse the clean tabular output file - with open(temp_tblout.name, 'r') as f: - tabular_content = f.read() - + + # Parse the clean tabular output file. nhmmer writes this whenever + # --tblout is given, but treat an absent file as "no hits" rather + # than letting it surface as the FileNotFoundError handler below, + # which would misreport it as a missing nhmmer binary. + if os.path.exists(tblout_path): + with open(tblout_path, 'r') as f: + tabular_content = f.read() + else: + logging.warning(f"nhmmer produced no tabular output for {seq_id}") + tabular_content = '' + logging.debug(f"=== TABULAR OUTPUT for {seq_id} ===") logging.debug(tabular_content) - + # Parse tabular output to get best alignment alignment_result = parse_nhmmer_result(tabular_content, seq_id) - + return alignment_result except FileNotFoundError: @@ -506,7 +520,7 @@ def trim_sequence_ends(sequence): return trimmed -def align_sequence_with_nhmmer(record, hmm_file, hmm_length): +def align_sequence_with_nhmmer(record, hmm_file, hmm_length, threads=1): """ Process sequence by: 1. Removing tilde characters (preserve gaps) @@ -529,7 +543,7 @@ def align_sequence_with_nhmmer(record, hmm_file, hmm_length): logging.debug(f"After gap-to-N replacement: {n_padded_seq[:100]}...") # Step 3: Run nhmmer on the single N-padded sequence - nhmmer_result = run_nhmmer_on_sequence(n_padded_seq, record.id, hmm_file) + nhmmer_result = run_nhmmer_on_sequence(n_padded_seq, record.id, hmm_file, threads) if not nhmmer_result: logging.warning(f"No significant nhmmer alignment found for sequence {record.id}") return None, original_seq_before_nhmmer @@ -773,7 +787,7 @@ def format_barcode_gaps(sequence): def format_sequence_id(process_id, parameters): return f"{process_id}_{parameters}" if parameters else process_id -def analyse_fasta(file_path, hmm_file, hmm_length, trans_table): +def analyse_fasta(file_path, hmm_file, hmm_length, trans_table, threads=1): try: # Initialise dictionaries results = {} @@ -856,7 +870,7 @@ def analyse_fasta(file_path, hmm_file, hmm_length, trans_table): longest_stretch = calculate_longest_stretch_full_seq(seq) # Use nhmmer to extract and align barcode region - aligned_barcode, original_seq_before_nhmmer = align_sequence_with_nhmmer(record, hmm_file, hmm_length) + aligned_barcode, original_seq_before_nhmmer = align_sequence_with_nhmmer(record, hmm_file, hmm_length, threads) # Initialise translation variables reading_frame = -1 @@ -1171,6 +1185,8 @@ def main(): # Optional arguments parser.add_argument('--code', '-c', type=int, default=1, help='Genetic code table for translation (default: 1 for standard code)') + parser.add_argument('--threads', '-t', type=int, default=1, + help='CPUs passed to each nhmmer call via --cpu (default: 1). nhmmer is invoked once per sequence, so this sets threads per invocation, not concurrent invocations.') parser.add_argument('--log-file', help='Path to log file (optional)') parser.add_argument('--verbose', '-v', action='store_true', help='Enable verbose output') parser.add_argument('--disable-selection', action='store_true', help='Skip best-per-process selection and output all passing sequences instead') @@ -1234,7 +1250,7 @@ def main(): # Analyse each FASTA file for file in args.input: logging.info(f"Processing file: {file}") - results = analyse_fasta(file, args.hmm, hmm_length, args.code) + results = analyse_fasta(file, args.hmm, hmm_length, args.code, args.threads) for seq_id, result in results.items(): result['seq_id'] = seq_id all_results.append(result) diff --git a/beegees/workflow/scripts/tv_local_blast.py b/beegees/workflow/scripts/tv_local_blast.py index c30d151..6a2d9f3 100644 --- a/beegees/workflow/scripts/tv_local_blast.py +++ b/beegees/workflow/scripts/tv_local_blast.py @@ -6,6 +6,19 @@ It can handle single sequences, multi-FASTA files, or entire directories of FASTA files, and automatically creates BLAST databases from FASTA files when needed. +Chunked execution: +Query sequences are batched into chunks of CHUNK_SIZE (50) sequences and one blastn process is +launched per chunk, with up to --processes chunks running concurrently. The final chunk holds +whatever remainder is left over. Batching matters because every blastn invocation re-opens and +re-loads the BLAST database; against a large database that fixed startup cost dominates the +search itself, so a 500-sequence run drops from 500 invocations to 10. + +Each chunk is written with synthetic query IDs (>q1, >q2, ...) so that results can be +demultiplexed back to their source sequence by exact lookup rather than by re-deriving the +sequence name from BLAST's qseqid (which is only the first whitespace token of the header). +The real query ID is written back into the output before the per-sequence TSV is saved, so +per-sequence output is unchanged. + Database Handling: The script accepts three types of database input via the -d/--database argument: 1. Directory containing BLAST databases (auto-detects if only one database present) @@ -21,39 +34,48 @@ Input Processing: - Single FASTA files: Processes as single sequence or splits multi-FASTA automatically - Directories: Processes all .fasta and .fa files found in the directory -- Multi-FASTA files: Automatically splits into individual sequences for parallel processing +- Multi-FASTA files: Automatically chunked for parallel processing - Output files: Named based on sequence headers, organized in subdirectories for multi-FASTA input +- Sequences whose output TSV already exists are skipped, so an interrupted run resumes cheaply Output Format: - Tab-separated values (TSV) format with standard BLAST fields and headers, ordered by pident (highest first) - Default output format: qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore stitle -- One output file per input sequence +- One output file per input sequence (sequences with no hits get a header-only file) - Files named using sanitized sequence headers - Optional summary CSV file with top hits per query in a flattened format (includes sequences with no hits) +Failure Handling: +If a chunk's blastn call fails, its sequences are retried individually so that one bad record +cannot take down its 49 neighbours. Sequences that still fail are reported and no output file is +written for them (an empty file would be indistinguishable from a genuine no-hit result). The +script then exits non-zero without writing the summary CSV, so a caller such as Snakemake can +retry; successful per-sequence TSVs remain on disk, so each retry only redoes the failures. +Pass --allow-partial-failures to downgrade this to a warning. + Requirements: - Python 3.6+ - BLAST+ suite (blastn and makeblastdb commands must be in PATH) Usage Examples: # Use existing database directory (auto-detect) - python blast_parallel.py -i sequences.fasta -d /path/to/databases/ -o results/ - + python tv_local_blast.py -i sequences.fasta -d /path/to/databases/ -o results/ + # Use specific existing database - python blast_parallel.py -i sequences.fasta -d /path/to/databases/nt -o results/ - + python tv_local_blast.py -i sequences.fasta -d /path/to/databases/nt -o results/ + # Create database from FASTA file - python blast_parallel.py -i queries.fasta -d reference_genome.fasta -o results/ - + python tv_local_blast.py -i queries.fasta -d reference_genome.fasta -o results/ + # Process directory with custom settings and CSV summary - python blast_parallel.py -i /fasta_dir/ -d database.fasta -o results/ -p 16 --output-csv summary.csv - + python tv_local_blast.py -i /fasta_dir/ -d database.fasta -o results/ -p 16 --output-csv summary.csv + # Custom BLAST parameters - python blast_parallel.py -i input.fasta -d db.fasta -o out/ --blast-opts "-evalue 1e-10 -max_target_seqs 5" + python tv_local_blast.py -i input.fasta -d db.fasta -o out/ --blast-opts "-evalue 1e-10 -max_target_seqs 5" Author: D. Parsons @NHMUK License: MIT -Version: 2.3 +Version: 2.4 """ import os import sys @@ -62,11 +84,12 @@ import tempfile import shutil import csv +import time from pathlib import Path from multiprocessing import cpu_count -from concurrent.futures import ProcessPoolExecutor, as_completed +from concurrent.futures import ThreadPoolExecutor, as_completed import logging -from typing import List, Tuple, Optional, Dict, Set +from typing import List, Tuple, Optional, Dict, Set, NamedTuple import re # Configure logging @@ -76,40 +99,101 @@ ) logger = logging.getLogger(__name__) +# Number of query sequences per blastn invocation. The final chunk of a run holds +# the remainder (< CHUNK_SIZE). Every blastn call reloads the BLAST database, so +# batching queries is what buys the speed-up; retune here if a measurement on a +# particular database calls for it. +CHUNK_SIZE = 50 + +# The outfmt 6 column specification, and the matching header line written at the +# top of every per-sequence TSV. Kept together so the two cannot drift apart. +BLAST_OUTFMT_COLS = ("qseqid sseqid pident length mismatch gapopen qstart qend " + "sstart send evalue bitscore stitle") +BLAST_TSV_HEADER = "\t".join(BLAST_OUTFMT_COLS.split()) + "\n" + + +class DemuxError(RuntimeError): + """A chunk's BLAST output referenced a query ID that is not a chunk member.""" + + +class QueryRec(NamedTuple): + """One query sequence and where its result should be written.""" + stem: str # sanitized full header; the summary CSV keys on this + full_header: str # original '>...' line, verbatim + qseqid: str # first whitespace token, i.e. what BLAST would report + out_path: Path # destination TSV + + +class ChunkSpec(NamedTuple): + """One blastn invocation: a query FASTA, its output tabular, and its members.""" + index: object # int for a normal chunk, 'N.M' for a retry singleton + query_fa: Path + tabular: Path + members: Dict[str, QueryRec] # synthetic BLAST id ('q7') -> QueryRec + + +class ScanResult(NamedTuple): + """Outcome of scanning an input FASTA for work to do.""" + records: int # '>' lines seen + pending: List[QueryRec] # sequences that still need BLASTing, in file order + headers: Dict[str, str] # stem -> full header, for every record seen + qseqid_to_stem: Dict[str, str] # qseqid -> stem, for every record seen + resumed: int # skipped because their output already exists + skipped: int # skipped as duplicate or unusable + + +class ChunkResult(NamedTuple): + """Outcome of one chunk (including any individual retries it triggered).""" + processed: int + with_hits: int + without_hits: int + failed: List[str] # qseqids with no output written + seconds: float + + class BLASTRunner: - def __init__(self, db_path: str, output_dir: str, blast_options: str, num_processes: int = None, output_csv: str = None): + def __init__(self, db_path: str, output_dir: str, blast_options: str, num_processes: int = None, + output_csv: str = None, allow_partial_failures: bool = False): """ Initialize BLAST runner with configuration parameters. - + Args: db_path: Path to BLAST database directory, specific database, or FASTA file to create database from output_dir: Directory to save BLAST results blast_options: Additional BLAST options string - num_processes: Number of parallel processes (default: CPU count) + num_processes: Number of concurrent blastn chunks (default: CPU count) output_csv: Filename for summary CSV file (optional) + allow_partial_failures: Continue (with a warning) when some sequences fail BLAST """ self.db_input = Path(db_path) self.output_dir = Path(output_dir) self.blast_options = blast_options self.num_processes = num_processes or cpu_count() self.output_csv = output_csv - + self.allow_partial_failures = allow_partial_failures + # Track all processed sequences for CSV summary self.processed_sequences: Dict[str, str] = {} # seq_id -> original_header - + # Map BLAST's qseqid back to the sanitized seq_id used above. BLAST reports + # only the first whitespace token, so for a header with a description the + # two differ and the summary CSV cannot be keyed off qseqid alone. + self.qseqid_to_stem: Dict[str, str] = {} + # qseqids that failed BLAST and therefore have no output file + self.failed_sequences: List[str] = [] + # Create output directory self.output_dir.mkdir(parents=True, exist_ok=True) - + # Find/create and validate database self.db_path, self.db_name = self._handle_database() - + # Validate setup self._validate_setup() - + def _handle_database(self) -> Tuple[Path, str]: """ Handle database input - either find existing database or create from FASTA. - + Returns: Tuple of (full_db_path, db_name) """ @@ -120,18 +204,18 @@ def _handle_database(self) -> Tuple[Path, str]: else: # Original behavior - directory or database path return self._find_database() - + def _handle_fasta_input(self) -> Tuple[Path, str]: """ Handle FASTA file input - check for existing database or create new one. - + Returns: Tuple of (full_db_path, db_name) """ fasta_path = self.db_input db_name = fasta_path.stem # filename without extension db_path = fasta_path.parent / db_name - + # Check if database already exists if self._is_valid_database(db_path): logger.info(f"Database already exists for {fasta_path.name}, using existing database: {db_name}") @@ -140,11 +224,11 @@ def _handle_fasta_input(self) -> Tuple[Path, str]: logger.info(f"Creating BLAST database from {fasta_path.name}") self._create_blast_database(fasta_path, db_path) return db_path, db_name - + def _create_blast_database(self, fasta_path: Path, db_path: Path) -> None: """ Create BLAST database from FASTA file. - + Args: fasta_path: Path to input FASTA file db_path: Path for output database (without extension) @@ -157,34 +241,34 @@ def _create_blast_database(self, fasta_path: Path, db_path: Path) -> None: '-out', str(db_path), '-title', f"Database created from {fasta_path.name}" ] - + logger.info("Running makeblastdb...") result = subprocess.run(cmd, capture_output=True, text=True, check=True) - + if result.stdout: logger.debug(f"makeblastdb output: {result.stdout}") - + logger.info(f"✓ Database created successfully: {db_path}") - + except subprocess.CalledProcessError as e: error_msg = f"Failed to create BLAST database: {e.stderr}" logger.error(error_msg) raise RuntimeError(error_msg) except FileNotFoundError: raise RuntimeError("makeblastdb command not found. Please make sure BLAST+ is installed and in your PATH") - + def _find_database(self) -> Tuple[Path, str]: """ Find BLAST database from the given path. Can handle both directory (auto-detect) and specific database path. - + Returns: Tuple of (full_db_path, db_name) """ if self.db_input.is_dir(): # Directory provided - find databases automatically available_dbs = self._find_available_databases_in_dir(self.db_input) - + if not available_dbs: raise RuntimeError(f"No BLAST databases found in directory: {self.db_input}") elif len(available_dbs) == 1: @@ -206,25 +290,25 @@ def _find_database(self) -> Tuple[Path, str]: # Check if it's a database name in current directory or if parent dir exists parent_dir = self.db_input.parent db_name = self.db_input.name - + if parent_dir.exists(): # Check if database exists in the parent directory potential_db = parent_dir / db_name if self._is_valid_database(potential_db): return potential_db, db_name - + raise FileNotFoundError(f"Database not found: {self.db_input}") - + def _find_available_databases_in_dir(self, db_dir: Path) -> List[str]: """Find available BLAST databases in a directory.""" if not db_dir.exists(): return [] - + db_files = [] # Look for nucleotide database files for pattern in ['*.nhr', '*.00.nhr', '*.nal']: db_files.extend(db_dir.glob(pattern)) - + # Extract database names (remove extensions) db_names = set() for db_file in db_files: @@ -233,37 +317,37 @@ def _find_available_databases_in_dir(self, db_dir: Path) -> List[str]: name = re.sub(r'\.(nhr|nal)$', '', name) name = re.sub(r'\.00\.nhr$', '', name) db_names.add(name) - + return sorted(list(db_names)) - + def _is_valid_database(self, db_path: Path) -> bool: """Check if a database path points to a valid BLAST database.""" db_extensions = ['.nhr', '.00.nhr', '.nal'] return any((db_path.parent / f"{db_path.name}{ext}").exists() for ext in db_extensions) - + def _validate_setup(self): """Validate BLAST installation and database availability.""" # Check if blastn is available try: - subprocess.run(['blastn', '-version'], + subprocess.run(['blastn', '-version'], capture_output=True, check=True) except (subprocess.CalledProcessError, FileNotFoundError): raise RuntimeError("Error: blastn command not found. Please make sure BLAST+ is installed and in your PATH") - + # Database validation is now handled in _handle_database() if not self._is_valid_database(self.db_path): available_dbs = self._find_available_databases_in_dir(self.db_path.parent) raise RuntimeError(f"Error: BLAST database '{self.db_name}' not found\n" f"Available databases in {self.db_path.parent}: {available_dbs}") - + @staticmethod def sanitize_header(header: str) -> str: """ Sanitize FASTA header for use as filename. - + Args: header: FASTA header string - + Returns: Sanitized string safe for use as filename """ @@ -272,294 +356,462 @@ def sanitize_header(header: str) -> str: sanitized = re.sub(r'[/|:*"<>? ]', '_', sanitized) # Limit length to avoid filename issues return sanitized[:100] - - def _extract_sequences_from_fasta(self, fasta_file: Path) -> Dict[str, str]: + + # ------------------------------------------------------------------ + # Step 1 - scan the input for work to do + # ------------------------------------------------------------------ + @staticmethod + def collect_pending(fasta_file: Path, output_dir: Path, prefix: Optional[str] = None) -> ScanResult: """ - Extract sequence IDs and headers from a FASTA file. - + Scan a (multi-)FASTA once and decide which sequences still need BLASTing. + + Every record seen is recorded in `headers` and `qseqid_to_stem` - including + records skipped below - because the summary CSV emits a row for every input + sequence, whether or not it produced hits. + + A record is excluded from `pending` when: + - its output TSV already exists (resume; counted in `resumed`) + - its sanitized name duplicates an earlier record (counted in `skipped`) + - its header yields no usable query ID (counted in `skipped`) + Args: - fasta_file: Path to FASTA file - + fasta_file: Path to the input FASTA + output_dir: Directory the per-sequence TSVs are written to + prefix: Optional filename prefix (used for single-sequence inputs) + Returns: - Dictionary mapping sanitized sequence IDs to original headers + A ScanResult. """ - sequences = {} - try: - with open(fasta_file, 'r') as f: - for line in f: - line = line.strip() - if line.startswith('>'): - header = line - seq_id = self.sanitize_header(header) - sequences[seq_id] = header - except Exception as e: - logger.warning(f"Error reading FASTA file {fasta_file}: {str(e)}") - - return sequences - - def _sort_blast_output(self, output_file: Path) -> None: - """ - Sort BLAST output file by pident in descending order. - - Args: - output_file: Path to BLAST output TSV file - """ - try: - # Read the file - with open(output_file, 'r') as f: - lines = f.readlines() - - if not lines: - return - - # Separate header from data - header_line = None - data_lines = [] - - for line in lines: - if line.strip(): - parts = line.strip().split('\t') - # Check if this is the header line (contains 'qseqid' as first column) - if parts[0] == 'qseqid': - header_line = line - else: - # Parse data lines and sort by pident (column index 2, descending) - if len(parts) >= 3: - try: - pident = float(parts[2]) - data_lines.append((pident, line)) - except (ValueError, IndexError): - # Keep malformed lines at the end - data_lines.append((0.0, line)) - - # Sort by pident (descending) - data_lines.sort(key=lambda x: x[0], reverse=True) - - # Write back sorted data with header - with open(output_file, 'w') as f: - if header_line: - f.write(header_line) - for _, line in data_lines: - f.write(line) - - except Exception as e: - logger.warning(f"Failed to sort BLAST output {output_file}: {str(e)}") - - def _run_blast_single(self, query_file: Path, output_file: Path, header: str) -> Tuple[bool, str]: + fasta_file = Path(fasta_file) + output_dir = Path(output_dir) + + pending: List[QueryRec] = [] + headers: Dict[str, str] = {} + qseqid_to_stem: Dict[str, str] = {} + seen: Set[str] = set() + records = 0 + resumed = 0 + skipped = 0 + + with open(fasta_file, 'r') as fh: + for line in fh: + if not line.startswith('>'): + continue + records += 1 + full_header = line.strip() + stem = BLASTRunner.sanitize_header(full_header) + body = full_header.lstrip('>').strip() + qseqid = body.split(None, 1)[0] if body else '' + + if not qseqid or not stem: + skipped += 1 + logger.warning(f"Skipping record with unusable header: {full_header!r}") + continue + + # Track every record for the summary CSV, even if skipped below. + headers[stem] = full_header + qseqid_to_stem.setdefault(qseqid, stem) + + if stem in seen: + skipped += 1 + logger.warning(f"Duplicate sequence name after sanitisation, " + f"skipping repeat: {stem}") + continue + seen.add(stem) + + # NOTE: the filename carries the prefix but the summary CSV key + # (stem) deliberately does not. + out_name = f"{prefix}_{stem}.tsv" if prefix else f"{stem}.tsv" + out_path = output_dir / out_name + if out_path.exists(): + resumed += 1 + logger.debug(f"○ Skipping {stem} - output file already exists") + continue + + pending.append(QueryRec(stem=stem, full_header=full_header, + qseqid=qseqid, out_path=out_path)) + + return ScanResult(records=records, pending=pending, headers=headers, + qseqid_to_stem=qseqid_to_stem, resumed=resumed, skipped=skipped) + + # ------------------------------------------------------------------ + # Step 2 - write the chunk FASTAs + # ------------------------------------------------------------------ + @staticmethod + def write_chunks(fasta_file: Path, pending: List[QueryRec], chunk_size: int, + temp_dir: Path) -> List[ChunkSpec]: """ - Run BLAST on a single sequence file. - + Stream the input FASTA and write the pending sequences out in contiguous + chunks of `chunk_size`, in file order. The final chunk holds the remainder. + + Records are written with synthetic query IDs ('>q1', '>q2', ...) so that + results can be demultiplexed by exact lookup. Sequence lines are copied + byte-for-byte, so the search itself is unaffected - blastn scores sequence, + not the defline. + Args: - query_file: Path to query FASTA file - output_file: Path to output file - header: Sequence header for logging - + fasta_file: Path to the input FASTA + pending: Sequences to include (from collect_pending) + chunk_size: Maximum sequences per chunk + temp_dir: Directory for the chunk FASTAs and their tabular outputs + Returns: - Tuple of (success, message) + List of ChunkSpec, in file order. """ - try: - # Construct BLAST command - cmd = [ - 'blastn', - '-query', str(query_file), - '-db', str(self.db_path), - '-out', str(output_file), - '-outfmt', '6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore stitle' - ] + self.blast_options.split() - - # Run BLAST - result = subprocess.run(cmd, capture_output=True, text=True, check=True) - - # Add header to the output file - self._add_header_to_tsv(output_file) - - # Sort output by pident - self._sort_blast_output(output_file) - - return True, f"✓ BLAST completed for {header}" - - except subprocess.CalledProcessError as e: - error_msg = f"✗ BLAST failed for {header}: {e.stderr}" - logger.error(error_msg) - return False, error_msg - except Exception as e: - error_msg = f"✗ Unexpected error for {header}: {str(e)}" - logger.error(error_msg) - return False, error_msg - - def _add_header_to_tsv(self, output_file: Path) -> None: + temp_dir = Path(temp_dir) + pending_by_stem = {rec.stem: rec for rec in pending} + chunks: List[ChunkSpec] = [] + if not pending_by_stem: + return chunks + + handle = None + query_fa = None + chunk_index = 0 + members: Dict[str, QueryRec] = {} + global_index = 0 + emitted: Set[str] = set() + include = False + + with open(fasta_file, 'r') as fh: + for line in fh: + if line.startswith('>'): + stem = BLASTRunner.sanitize_header(line.strip()) + rec = pending_by_stem.get(stem) + include = rec is not None and stem not in emitted + if not include: + continue + emitted.add(stem) + + # Start a new chunk on the first record and whenever the + # current one is full. + if handle is None or len(members) >= chunk_size: + if handle is not None: + handle.close() + chunks.append(ChunkSpec( + index=chunk_index, + query_fa=query_fa, + tabular=temp_dir / f"chunk_{chunk_index}.tab", + members=members, + )) + chunk_index += 1 + query_fa = temp_dir / f"chunk_{chunk_index}.fa" + handle = open(query_fa, 'w') + members = {} + + global_index += 1 + blast_id = f"q{global_index}" + members[blast_id] = rec + handle.write(f">{blast_id}\n") + elif include and handle is not None: + handle.write(line) + + if handle is not None: + handle.close() + chunks.append(ChunkSpec( + index=chunk_index, + query_fa=query_fa, + tabular=temp_dir / f"chunk_{chunk_index}.tab", + members=members, + )) + + return chunks + + @staticmethod + def split_chunk_fasta(spec: ChunkSpec, temp_dir: Path) -> List[ChunkSpec]: """ - Add header line to BLAST TSV output file. - - Args: - output_file: Path to BLAST output TSV file + Split an already-written chunk FASTA into one single-record ChunkSpec per + member, reusing the synthetic IDs. Used to retry a failed chunk one + sequence at a time so a single bad record cannot take down its neighbours. """ + temp_dir = Path(temp_dir) + singles: List[ChunkSpec] = [] + handle = None + + with open(spec.query_fa, 'r') as fh: + for line in fh: + if line.startswith('>'): + if handle is not None: + handle.close() + handle = None + blast_id = line[1:].strip() + rec = spec.members.get(blast_id) + if rec is None: + continue + n = len(singles) + 1 + fa = temp_dir / f"chunk_{spec.index}_{n}.fa" + handle = open(fa, 'w') + handle.write(line) + singles.append(ChunkSpec( + index=f"{spec.index}.{n}", + query_fa=fa, + tabular=temp_dir / f"chunk_{spec.index}_{n}.tab", + members={blast_id: rec}, + )) + elif handle is not None: + handle.write(line) + + if handle is not None: + handle.close() + + return singles + + # ------------------------------------------------------------------ + # Step 3 - run blastn on a chunk + # ------------------------------------------------------------------ + def _run_blast_chunk(self, spec: ChunkSpec) -> None: + """ + Run blastn over one chunk FASTA. Raises subprocess.CalledProcessError on + a non-zero exit. + """ + cmd = [ + 'blastn', + '-query', str(spec.query_fa), + '-db', str(self.db_path), + '-out', str(spec.tabular), + '-outfmt', f'6 {BLAST_OUTFMT_COLS}' + ] + self.blast_options.split() + + logger.debug(f"Running: {' '.join(cmd)}") + subprocess.run(cmd, capture_output=True, text=True, check=True) + + # ------------------------------------------------------------------ + # Step 4 - demultiplex a chunk's results back to per-sequence TSVs + # ------------------------------------------------------------------ + @staticmethod + def _pident_key(row: str) -> float: + """ + Sort key placing higher percent identity first. Rows whose pident cannot + be parsed sort last, matching the previous per-sequence sort behaviour. + Used with a stable sort, so BLAST's own ordering is preserved among hits + of equal percent identity - downstream taxonomy picks the first matching + hit, so tie order is significant. + """ + parts = row.rstrip('\n').split('\t') try: - # Read existing content - with open(output_file, 'r') as f: - content = f.read() - - # Write header + content - with open(output_file, 'w') as f: - f.write("qseqid\tsseqid\tpident\tlength\tmismatch\tgapopen\tqstart\tqend\tsstart\tsend\tevalue\tbitscore\tstitle\n") - f.write(content) - - except Exception as e: - logger.warning(f"Failed to add header to {output_file}: {str(e)}") - - def process_single_sequence(self, args: Tuple[Path, str, Path]) -> Tuple[bool, str, str]: + return -float(parts[2]) + except (ValueError, IndexError): + return 0.0 + + @staticmethod + def demux_chunk_tabular(tabular: Path, members: Dict[str, QueryRec]) -> Tuple[int, int]: """ - Process a single sequence (wrapper for multiprocessing). - + Split one chunk's BLAST tabular into the per-sequence TSVs. + + Rows are grouped by their synthetic query ID, column 0 is rewritten to the + real query ID, and each group is written pident-descending under the standard + header. Members that returned no rows get a header-only file, which is what + a no-hit query produced previously and what the summary CSV expects to find. + + Files are written to a '.part' path and renamed into place, so an interrupted + run cannot leave a truncated TSV that a later run mistakes for complete. + Args: - args: Tuple of (fasta_file, base_output_prefix, output_dir) - + tabular: Chunk output written by blastn + members: Synthetic query ID -> QueryRec for this chunk + Returns: - Tuple of (success, message, header) + Tuple of (queries_with_hits, queries_without_hits) + + Raises: + DemuxError: if BLAST reported a query ID that is not a chunk member. """ - fasta_file, base_output_prefix, output_dir = args - - try: - # Read the first line to get header - with open(fasta_file, 'r') as f: - header = f.readline().strip() - - if not header.startswith('>'): - return False, f"✗ Error: File doesn't appear to be in FASTA format", str(fasta_file) - - # Create output filename - safe_header = self.sanitize_header(header) - if base_output_prefix: - output_file = output_dir / f"{base_output_prefix}_{safe_header}.tsv" + groups: Dict[str, List[str]] = {} + + with open(tabular, 'r') as fh: + for line in fh: + line = line.rstrip('\n') + if not line.strip(): + continue + parts = line.split('\t') + # Rows with fewer than 3 columns carry no usable pident and were + # discarded by the previous sort step too. + if len(parts) < 3: + continue + rec = members.get(parts[0]) + if rec is None: + raise DemuxError( + f"BLAST reported query id {parts[0]!r} which is not a member of " + f"{tabular}; refusing to discard its hits" + ) + parts[0] = rec.qseqid + groups.setdefault(rec.stem, []).append('\t'.join(parts) + '\n') + + with_hits = 0 + without_hits = 0 + for rec in members.values(): + rows = groups.get(rec.stem, []) + rows.sort(key=BLASTRunner._pident_key) + part_path = Path(str(rec.out_path) + '.part') + with open(part_path, 'w') as out: + out.write(BLAST_TSV_HEADER) + out.writelines(rows) + os.replace(part_path, rec.out_path) + if rows: + with_hits += 1 else: - output_file = output_dir / f"{safe_header}.tsv" - - # Track this sequence for CSV summary - self.processed_sequences[safe_header] = header - - # Check if output already exists - if output_file.exists(): - return True, f"○ Skipping {header} - output file already exists", header - - # Run BLAST - success, message = self._run_blast_single(fasta_file, output_file, header) - return success, message, header - - except Exception as e: - return False, f"✗ Error processing {fasta_file}: {str(e)}", str(fasta_file) - - def split_multifasta(self, fasta_file: Path) -> Tuple[List[Path], Path]: + without_hits += 1 + + return with_hits, without_hits + + # ------------------------------------------------------------------ + # Step 5 - chunk orchestration + # ------------------------------------------------------------------ + def _process_chunk(self, spec: ChunkSpec, temp_dir: Path) -> ChunkResult: """ - Split multi-FASTA file into individual sequence files. - - Args: - fasta_file: Path to multi-FASTA file - - Returns: - Tuple of (list of paths to individual sequence files, temp directory path) + Run one chunk and demultiplex its output. On a blastn failure, retry the + chunk's sequences individually so one bad record does not lose the rest. """ - temp_dir = Path(tempfile.mkdtemp()) - split_files = [] - - current_file = None - current_handle = None - seq_count = 0 - + start = time.monotonic() try: - with open(fasta_file, 'r') as f: - for line in f: - if line.startswith('>'): - # Close previous file if open - if current_handle: - current_handle.close() - - # Track sequence for CSV summary - header = line.strip() - safe_header = self.sanitize_header(header) - self.processed_sequences[safe_header] = header - - # Start new file - seq_count += 1 - current_file = temp_dir / f"seq_{seq_count}.fa" - current_handle = open(current_file, 'w') - split_files.append(current_file) - - if current_handle: - current_handle.write(line) - - # Close last file - if current_handle: - current_handle.close() - - except Exception as e: - # Clean up on error - if current_handle: - current_handle.close() - shutil.rmtree(temp_dir, ignore_errors=True) - raise e - - return split_files, temp_dir - - def process_multifasta_parallel(self, fasta_file: Path, output_dir: Path) -> Tuple[int, int]: - """ - Process multi-FASTA file in parallel. - + self._run_blast_chunk(spec) + except subprocess.CalledProcessError as exc: + stderr = (exc.stderr or '')[:2000] + logger.error(f"✗ chunk {spec.index} ({len(spec.members)} queries) " + f"FAILED rc={exc.returncode}: {stderr}") + logger.error(f" members: " + f"{', '.join(r.qseqid for r in spec.members.values())}") + return self._retry_chunk_individually(spec, temp_dir, start) + except Exception as exc: + logger.error(f"✗ chunk {spec.index} ({len(spec.members)} queries) " + f"raised {type(exc).__name__}: {exc}") + return self._retry_chunk_individually(spec, temp_dir, start) + + try: + with_hits, without_hits = self.demux_chunk_tabular(spec.tabular, spec.members) + except DemuxError as exc: + # Retrying will not help - the output does not match the input we sent. + logger.error(f"✗ chunk {spec.index}: {exc}") + return ChunkResult(processed=0, with_hits=0, without_hits=0, + failed=[r.qseqid for r in spec.members.values()], + seconds=time.monotonic() - start) + finally: + # Keep peak temp usage to one live tabular per worker. + try: + os.remove(spec.tabular) + except OSError: + pass + + return ChunkResult(processed=len(spec.members), with_hits=with_hits, + without_hits=without_hits, failed=[], + seconds=time.monotonic() - start) + + def _retry_chunk_individually(self, spec: ChunkSpec, temp_dir: Path, + start: float) -> ChunkResult: + """Re-run a failed chunk's sequences one at a time to isolate the culprit.""" + singles = self.split_chunk_fasta(spec, temp_dir) + logger.info(f"[chunk {spec.index}] Retrying {len(singles)} sequences individually...") + + processed = 0 + with_hits = 0 + without_hits = 0 + failed: List[str] = [] + + for single in singles: + rec = next(iter(single.members.values())) + try: + self._run_blast_chunk(single) + hits, no_hits = self.demux_chunk_tabular(single.tabular, single.members) + except subprocess.CalledProcessError as exc: + stderr = (exc.stderr or '')[:1000] + logger.error(f"[chunk {spec.index}] ✗ {rec.qseqid} failed on individual " + f"retry rc={exc.returncode}: {stderr}") + failed.append(rec.qseqid) + continue + except Exception as exc: + logger.error(f"[chunk {spec.index}] ✗ {rec.qseqid} failed on individual " + f"retry: {type(exc).__name__}: {exc}") + failed.append(rec.qseqid) + continue + finally: + for path in (single.query_fa, single.tabular): + try: + os.remove(path) + except OSError: + pass + processed += 1 + with_hits += hits + without_hits += no_hits + + logger.info(f"[chunk {spec.index}] Recovered {processed}/{len(singles)} sequences") + return ChunkResult(processed=processed, with_hits=with_hits, + without_hits=without_hits, failed=failed, + seconds=time.monotonic() - start) + + def process_fasta_chunked(self, fasta_file: Path, output_dir: Path, + prefix: Optional[str] = None) -> Tuple[int, int, int]: + """ + BLAST one (multi-)FASTA in chunks of CHUNK_SIZE sequences. + Args: - fasta_file: Path to multi-FASTA file - output_dir: Directory to save results - + fasta_file: Path to the input FASTA + output_dir: Directory to write the per-sequence TSVs into + prefix: Optional filename prefix (used for single-sequence inputs) + Returns: - Tuple of (processed_count, skipped_count) - """ - # Split multi-FASTA file - split_files, temp_dir = self.split_multifasta(fasta_file) - total_sequences = len(split_files) - - logger.info(f"Found {total_sequences} sequences in the multi-FASTA file") - + Tuple of (processed_count, skipped_count, failed_count) + """ + fasta_file = Path(fasta_file) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + scan = self.collect_pending(fasta_file, output_dir, prefix) + # Populate shared state here, in the main thread, before any worker starts. + self.processed_sequences.update(scan.headers) + self.qseqid_to_stem.update(scan.qseqid_to_stem) + + logger.info(f"Scanned {scan.records} records in {fasta_file.name}: " + f"{len(scan.pending)} pending, {scan.resumed} already complete (resume), " + f"{scan.skipped} skipped (duplicate/unusable)") + + if not scan.pending: + logger.info("Nothing to BLAST - all outputs already present or no usable records") + return 0, scan.resumed, 0 + + temp_dir = Path(tempfile.mkdtemp()) try: - # Prepare arguments for parallel processing - process_args = [(f, None, output_dir) for f in split_files] - - processed_count = 0 - skipped_count = 0 - - # Process in parallel - with ProcessPoolExecutor(max_workers=self.num_processes) as executor: - # Submit all jobs - future_to_file = { - executor.submit(self.process_single_sequence, args): args[0] - for args in process_args - } - - # Collect results - for i, future in enumerate(as_completed(future_to_file), 1): - file_path = future_to_file[future] + chunks = self.write_chunks(fasta_file, scan.pending, CHUNK_SIZE, temp_dir) + final_size = len(chunks[-1].members) if chunks else 0 + logger.info(f"Chunking {len(scan.pending)} pending sequences into " + f"{len(chunks)} chunk(s) of {CHUNK_SIZE} (final chunk: {final_size})") + + processed = 0 + failed: List[str] = [] + workers = max(1, min(self.num_processes, len(chunks))) + logger.info(f"Running up to {workers} concurrent blastn chunk(s)") + + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = {executor.submit(self._process_chunk, spec, temp_dir): spec + for spec in chunks} + for i, future in enumerate(as_completed(futures), 1): + spec = futures[future] try: - success, message, header = future.result() - - logger.info(f"[{i}/{total_sequences}] {message}") - - if success: - if "Skipping" in message: - skipped_count += 1 - else: - processed_count += 1 - - except Exception as e: - logger.error(f"[{i}/{total_sequences}] ✗ Error processing {file_path}: {str(e)}") - - logger.info(f"Completed multi-FASTA processing: {processed_count} processed, {skipped_count} skipped") - return processed_count, skipped_count - + result = future.result() + except Exception as exc: + logger.error(f"[{i}/{len(chunks)}] ✗ chunk {spec.index} raised " + f"{type(exc).__name__}: {exc}") + failed.extend(r.qseqid for r in spec.members.values()) + continue + processed += result.processed + failed.extend(result.failed) + glyph = '✗' if result.failed else '✓' + suffix = f", {len(result.failed)} failed" if result.failed else '' + logger.info(f"[{i}/{len(chunks)}] {glyph} chunk {spec.index} " + f"({len(spec.members)} queries) - {result.with_hits} with hits, " + f"{result.without_hits} no hits{suffix} - {result.seconds:.1f}s") + + self.failed_sequences.extend(failed) + logger.info(f"Completed {fasta_file.name}: {processed} processed, " + f"{scan.resumed} skipped, {len(failed)} failed - " + f"{len(chunks)} blastn invocation(s)") + return processed, scan.resumed, len(failed) finally: # Clean up temporary files shutil.rmtree(temp_dir, ignore_errors=True) - + def process_directory_parallel(self, input_dir: Path) -> None: """ - Process all FASTA files in a directory in parallel. - + Process all FASTA files in a directory. + Args: input_dir: Directory containing FASTA files """ @@ -568,65 +820,48 @@ def process_directory_parallel(self, input_dir: Path) -> None: fasta_files = [] for pattern in fasta_patterns: fasta_files.extend(input_dir.glob(pattern)) - + if not fasta_files: raise ValueError(f"No FASTA files found in {input_dir}") - - # Extract all sequences from all FASTA files for tracking - for fasta_file in fasta_files: - sequences = self._extract_sequences_from_fasta(fasta_file) - self.processed_sequences.update(sequences) - + total_files = len(fasta_files) logger.info(f"Found {total_files} FASTA files to process") logger.info("Starting BLAST searches...") - + dir_processed_count = 0 dir_skipped_count = 0 - + for i, fasta_file in enumerate(fasta_files, 1): filename = fasta_file.name base_name = fasta_file.stem - + logger.info(f"[{i}/{total_files}] Processing file: {filename}") - + # Count sequences in file with open(fasta_file, 'r') as f: seq_count = sum(1 for line in f if line.startswith('>')) - + if seq_count == 0: logger.error(f" ✗ Error: No FASTA sequences found in {fasta_file}") continue elif seq_count == 1: logger.info(" Single sequence FASTA file") - - # Check if output already exists - with open(fasta_file, 'r') as f: - header = f.readline().strip() - safe_header = self.sanitize_header(header) - output_file = self.output_dir / f"{base_name}_{safe_header}.tsv" - - if output_file.exists(): - logger.info(f" ○ Skipping {filename} - output file already exists") - dir_skipped_count += 1 - else: - success, message, _ = self.process_single_sequence((fasta_file, base_name, self.output_dir)) - logger.info(f" {message}") - if success: - dir_processed_count += 1 + processed, skipped, _ = self.process_fasta_chunked( + fasta_file, self.output_dir, prefix=base_name) else: logger.info(f" Multi-FASTA file with {seq_count} sequences") - + # Create subdirectory for this file's results file_output_dir = self.output_dir / base_name file_output_dir.mkdir(exist_ok=True) - - processed, skipped = self.process_multifasta_parallel(fasta_file, file_output_dir) - dir_processed_count += processed - dir_skipped_count += skipped - + + processed, skipped, _ = self.process_fasta_chunked( + fasta_file, file_output_dir) + + dir_processed_count += processed + dir_skipped_count += skipped + logger.info(f"Directory processing completed: {dir_processed_count} sequences processed, {dir_skipped_count} sequences skipped") - def _generate_summary_csv(self) -> None: """ @@ -636,48 +871,48 @@ def _generate_summary_csv(self) -> None: """ if not self.output_csv: return - + logger.info("Generating summary CSV file...") - + # Find all TSV files recursively in output directory tsv_files = list(self.output_dir.rglob("*.tsv")) - + if not tsv_files: logger.warning("No TSV files found for CSV summary generation") return - + # Limit max hits in CSV to 100 max_csv_hits = 100 - + # Extract max_target_seqs from blast_options (for processing TSV files) max_hits = 1000 # default if 'max_target_seqs' in self.blast_options: match = re.search(r'max_target_seqs\s+(\d+)', self.blast_options) if match: max_hits = int(match.group(1)) - + # Create column headers - now includes hitN_gaps, limited to 100 hits headers = ['seq_id', 'original_header'] for i in range(1, max_csv_hits + 1): headers.extend([f'hit{i}', f'hit{i}_pident', f'hit{i}_length', f'hit{i}_mismatch', f'hit{i}_gaps', f'hit{i}_evalue']) - + # Dictionary to store results for each sequence sequence_results = {} - + # Initialize all sequences with empty results for seq_id, original_header in self.processed_sequences.items(): sequence_results[seq_id] = { 'original_header': original_header, 'hits': [] } - + # Process TSV files to populate hits for tsv_file in tsv_files: try: # Read TSV file with open(tsv_file, 'r') as f: lines = f.readlines() - + if not lines: # Empty file - find corresponding sequence ID from filename tsv_filename = tsv_file.stem @@ -687,24 +922,24 @@ def _generate_summary_csv(self) -> None: logger.debug(f"No hits found for sequence: {seq_id}") break continue - + # Parse TSV data (skip header if present) current_qseqid = None current_hits = [] - + for line in lines: line = line.strip() if not line: continue - + # Skip header line parts = line.split('\t') if parts[0] == 'qseqid': continue - + if len(parts) < 13: # Ensure we have all required columns continue - + qseqid = parts[0] sseqid = parts[1] pident = parts[2] @@ -712,54 +947,50 @@ def _generate_summary_csv(self) -> None: mismatch = parts[4] gapopen = parts[5] # Extract gapopen evalue = parts[10] - + # If this is a new query sequence, process the previous one if current_qseqid is not None and qseqid != current_qseqid: - # Find matching sequence ID in processed_sequences - sanitized_qseqid = self.sanitize_header(current_qseqid) - matching_seq_id = sanitized_qseqid if sanitized_qseqid in sequence_results else None - + matching_seq_id = self._resolve_seq_id(current_qseqid, sequence_results) + if matching_seq_id and matching_seq_id in sequence_results: # Limit to max_hits when saving from the file, but will be further limited in CSV output sequence_results[matching_seq_id]['hits'] = current_hits[:max_hits] - - current_hits = [] - + + current_hits = [] + current_qseqid = qseqid current_hits.append([sseqid, pident, length, mismatch, gapopen, evalue]) - + # Process the last query sequence if current_qseqid is not None: - # Find matching sequence ID in processed_sequences - sanitized_qseqid = self.sanitize_header(current_qseqid) - matching_seq_id = sanitized_qseqid if sanitized_qseqid in sequence_results else None - + matching_seq_id = self._resolve_seq_id(current_qseqid, sequence_results) + if matching_seq_id and matching_seq_id in sequence_results: # Limit to max_hits when saving from the file, but will be further limited in CSV output sequence_results[matching_seq_id]['hits'] = current_hits[:max_hits] - + except Exception as e: logger.warning(f"Error processing TSV file {tsv_file}: {str(e)}") continue - + # Count sequences with and without hits correctly sequences_with_hits = len([seq for seq in sequence_results.values() if seq['hits']]) sequences_without_hits = len([seq for seq in sequence_results.values() if not seq['hits']]) - + # Prepare CSV data - limiting to max_csv_hits csv_data = [] for seq_id, result in sequence_results.items(): row_data = [seq_id, result['original_header']] - + # Add hit data (up to max_csv_hits) - includes gapopen for i in range(max_csv_hits): if i < len(result['hits']): row_data.extend(result['hits'][i]) else: row_data.extend(['', '', '', '', '', '']) # Empty values for missing hits (6 fields) - + csv_data.append(row_data) - + # Write CSV file try: csv_path = self.output_dir / self.output_csv @@ -767,107 +998,126 @@ def _generate_summary_csv(self) -> None: writer = csv.writer(csvfile) writer.writerow(headers) writer.writerows(csv_data) - + logger.info(f"✓ Summary CSV file created: {csv_path}") logger.info(f" - Total sequences: {len(sequence_results)}") logger.info(f" - Sequences with hits: {sequences_with_hits}") logger.info(f" - Sequences without hits: {sequences_without_hits}") logger.info(f" - Included up to {max_csv_hits} hits per query (out of maximum {max_hits})") - + except Exception as e: logger.error(f"Failed to create summary CSV file: {str(e)}") - + + def _resolve_seq_id(self, qseqid: str, sequence_results: Dict) -> Optional[str]: + """ + Map a BLAST qseqid back to the key used in the summary CSV. + + BLAST reports only the first whitespace token of a header, whereas the CSV + is keyed on the sanitized *full* header, so for '>SEQ1 some description' + the two differ ('SEQ1' vs 'SEQ1_some_description'). The index built while + scanning the input resolves this; the sanitize fallback covers TSVs found + on disk from an earlier run whose input was not scanned this time. + """ + seq_id = self.qseqid_to_stem.get(qseqid) + if seq_id is not None and seq_id in sequence_results: + return seq_id + fallback = self.sanitize_header(qseqid) + return fallback if fallback in sequence_results else None + def process_input(self, input_path: Path) -> None: """ Process input (either directory or single file). - + Args: input_path: Path to input directory or file """ if not input_path.exists(): raise FileNotFoundError(f"Input path {input_path} does not exist") - + # Display configuration logger.info("===== BLAST Configuration =====") logger.info(f"Database: {self.db_path}") logger.info(f"Input: {input_path}") logger.info(f"Output directory: {self.output_dir}") logger.info(f"BLAST options: {self.blast_options}") - logger.info(f"Parallel processes: {self.num_processes}") + logger.info(f"Sequences per chunk: {CHUNK_SIZE}") + logger.info(f"Concurrent chunks: {self.num_processes}") if self.output_csv: logger.info(f"Summary CSV: {self.output_csv}") logger.info("===============================") - + if input_path.is_dir(): logger.info("Input is a directory. Processing all FASTA files...") self.process_directory_parallel(input_path) else: logger.info("Input is a single file...") - + # Validate file extension (with warning) if input_path.suffix not in ['.fa', '.fasta']: logger.warning("Input file does not have .fa or .fasta extension. Continuing anyway...") - + filename = input_path.name base_name = input_path.stem - - # Extract sequences for tracking - sequences = self._extract_sequences_from_fasta(input_path) - self.processed_sequences.update(sequences) - + # Count sequences with open(input_path, 'r') as f: seq_count = sum(1 for line in f if line.startswith('>')) - + if seq_count == 0: raise ValueError(f"No FASTA sequences found in {input_path}") elif seq_count == 1: logger.info(f"Processing single sequence FASTA file: {filename}") - - with open(input_path, 'r') as f: - header = f.readline().strip() - safe_header = self.sanitize_header(header) - output_file = self.output_dir / f"{base_name}_{safe_header}.tsv" - - if output_file.exists(): - logger.info(f" ○ Skipping {filename} - output file already exists") - else: - success, message, _ = self.process_single_sequence((input_path, base_name, self.output_dir)) - logger.info(f" {message}") + self.process_fasta_chunked(input_path, self.output_dir, prefix=base_name) else: logger.info(f"Processing multi-FASTA file: {filename} with {seq_count} sequences") - processed, skipped = self.process_multifasta_parallel(input_path, self.output_dir) - + self.process_fasta_chunked(input_path, self.output_dir) + + # A sequence with no output file is not the same as a sequence with no + # hits: writing the summary CSV anyway would present a BLAST failure as a + # genuine no-match to downstream taxonomy assignment. Fail loudly instead + # so the caller can retry - the TSVs that did succeed stay on disk, so a + # retry only redoes the failures. + if self.failed_sequences: + logger.error(f"{len(self.failed_sequences)} sequence(s) failed BLAST and have no output:") + for qseqid in self.failed_sequences: + logger.error(f" - {qseqid}") + if not self.allow_partial_failures: + logger.error("Not writing summary CSV. Re-run to retry only the failed sequences.") + sys.exit(1) + logger.warning("--allow-partial-failures set: continuing with an incomplete result set") + # Generate summary CSV if requested if self.output_csv: self._generate_summary_csv() - + logger.info(f"BLAST searches completed. Results saved to {self.output_dir}") - + + def main(): """Main function to run the BLAST script.""" parser = argparse.ArgumentParser( - description="Run BLASTn on FASTA files in parallel", + description="Run BLASTn on FASTA files in parallel, in chunks of " + f"{CHUNK_SIZE} sequences per blastn invocation", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Auto-detect database in directory (works if only one database present) - python blast_parallel.py -i sequences.fasta -d /path/to/databases/ -o results/ - + python tv_local_blast.py -i sequences.fasta -d /path/to/databases/ -o results/ + # Specify exact database path - python blast_parallel.py -i sequences.fasta -d /path/to/databases/nt -o results/ - + python tv_local_blast.py -i sequences.fasta -d /path/to/databases/nt -o results/ + # Create database from FASTA file (or use existing if present) - python blast_parallel.py -i queries.fasta -d reference.fasta -o results/ - + python tv_local_blast.py -i queries.fasta -d reference.fasta -o results/ + # Process directory with custom options and generate summary CSV - python blast_parallel.py -i /fasta_dir/ -d /db_path/nr -o results/ -p 16 --output-csv summary.csv - + python tv_local_blast.py -i /fasta_dir/ -d /db_path/nr -o results/ -p 16 --output-csv summary.csv + # Custom BLAST parameters - python blast_parallel.py -i input.fasta -d db/mydb -o out/ --blast-opts "-evalue 1e-10 -max_target_seqs 10" + python tv_local_blast.py -i input.fasta -d db/mydb -o out/ --blast-opts "-evalue 1e-10 -max_target_seqs 10" """ ) - + # Required arguments parser.add_argument('-i', '--input', required=True, type=Path, help='Input FASTA file or directory containing FASTA files') @@ -875,23 +1125,25 @@ def main(): help='Path to BLAST database directory (auto-detects databases), specific database path, or FASTA file to create database from') parser.add_argument('-o', '--output', required=True, type=Path, help='Output directory for BLAST results') - + # Optional arguments parser.add_argument('-p', '--processes', type=int, default=None, - help=f'Number of parallel processes (default: {cpu_count()})') + help=f'Number of concurrent blastn chunks (default: {cpu_count()})') parser.add_argument('--blast-opts', default='-evalue 1e-5 -max_target_seqs 100 -num_threads 1', help='Additional BLAST options (default: "-evalue 1e-5 -max_target_seqs 100 -num_threads 1")') parser.add_argument('--output-csv', type=str, default=None, help='Generate summary CSV file with specified filename (e.g. "summary.csv")') + parser.add_argument('--allow-partial-failures', action='store_true', + help='Warn instead of failing when some sequences could not be BLASTed') parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose logging') - + args = parser.parse_args() - + # Set logging level if args.verbose: logging.getLogger().setLevel(logging.DEBUG) - + try: # Create BLAST runner blast_runner = BLASTRunner( @@ -899,12 +1151,15 @@ def main(): output_dir=str(args.output), blast_options=args.blast_opts, num_processes=args.processes, - output_csv=args.output_csv + output_csv=args.output_csv, + allow_partial_failures=args.allow_partial_failures ) - + # Process input blast_runner.process_input(args.input) - + + except SystemExit: + raise except Exception as e: logger.error(f"Error: {str(e)}") sys.exit(1) diff --git a/tests/conftest.py b/tests/conftest.py index a23a730..da67b57 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,11 @@ """Shared fixtures and path setup for BeeGees tests.""" +import functools +import importlib.util import sys from pathlib import Path + +import pytest + from beegees.utils.configs import get_package_dir # Resolve scripts from the installed beegees package (single source of truth) @@ -8,3 +13,42 @@ SCRIPTS_DIR = get_package_dir() / "workflow" / "scripts" if str(SCRIPTS_DIR) not in sys.path: sys.path.insert(0, str(SCRIPTS_DIR)) + + +@functools.lru_cache(maxsize=None) +def load_script_module(name): + """Load a bundled workflow script as a module, once per session. + + The scripts are shipped as package data and run via `python