From bf9273f878e54ffae450f92e54161e090ce2af7b Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Wed, 25 Mar 2026 14:12:58 +1100 Subject: [PATCH 01/20] Added in logging for snakemake steps --- tests/test_snakemake_workflow.py | 25 ++++++++++ vartracker/Snakefile | 81 +++++++++++++++++++++++--------- 2 files changed, 84 insertions(+), 22 deletions(-) create mode 100644 tests/test_snakemake_workflow.py diff --git a/tests/test_snakemake_workflow.py b/tests/test_snakemake_workflow.py new file mode 100644 index 0000000..fb8778b --- /dev/null +++ b/tests/test_snakemake_workflow.py @@ -0,0 +1,25 @@ +from pathlib import Path + + +def test_snakemake_rules_write_logs_under_outdir(): + snakefile = ( + Path(__file__).resolve().parents[1] / "vartracker" / "Snakefile" + ).read_text(encoding="utf-8") + + assert "/dev/null" not in snakefile + assert 'return f"{OUTDIR}/logs/{rule_name}.log"' in snakefile + assert 'return f"{OUTDIR}/{sample}/logs/{rule_name}.log"' in snakefile + + expected_logs = [ + 'log:\n _workflow_log("bwa_index")', + 'lambda w: _sample_log(w.sample, "fastp")', + 'lambda w: _sample_log(w.sample, "bwa_mem")', + 'lambda w: _sample_log(w.sample, "ampliconclip")', + 'lambda w: _sample_log(w.sample, "lofreq_indelqual")', + 'lambda w: _sample_log(w.sample, "samtools_depth")', + 'lambda w: _sample_log(w.sample, "lofreq_call")', + 'log:\n _workflow_log("update_csv")', + ] + + for expected in expected_logs: + assert expected in snakefile diff --git a/vartracker/Snakefile b/vartracker/Snakefile index 843c041..64ff841 100644 --- a/vartracker/Snakefile +++ b/vartracker/Snakefile @@ -12,6 +12,14 @@ PRIMER_BED = config.get("primer_bed", None) MODE = config.get("mode", "reads") +def _workflow_log(rule_name): + return f"{OUTDIR}/logs/{rule_name}.log" + + +def _sample_log(sample, rule_name): + return f"{OUTDIR}/{sample}/logs/{rule_name}.log" + + def _value_or_blank(sample, key, allow_blank=False): value = SAMPLES[sample].get(key) if value is None or (isinstance(value, float) and pd.isna(value)): @@ -56,9 +64,12 @@ if MODE == "reads": bwt = REF + ".bwt", pac = REF + ".pac", sa = REF + ".sa" + log: + _workflow_log("bwa_index") shell: """ - bwa index {input.ref} 2> /dev/null + mkdir -p $(dirname {log}) + bwa index {input.ref} 2> {log} """ rule fastp: @@ -71,10 +82,13 @@ if MODE == "reads": r2 = temp(f"{OUTDIR}/{{sample}}/trimmed_R2.fastq.gz"), html = temp(f"{OUTDIR}/{{sample}}/fastp.html"), json = temp(f"{OUTDIR}/{{sample}}/fastp.json") + log: + lambda w: _sample_log(w.sample, "fastp") threads: max(1, int(workflow.cores * 0.5)) shell: """ mkdir -p {OUTDIR}/{wildcards.sample} + mkdir -p $(dirname {log}) if [ -n "{params.r2}" ]; then fastp -i {input.r1} -I {params.r2} \ -o {output.r1} -O {output.r2} \ @@ -86,7 +100,7 @@ if MODE == "reads": --correction \ --length_required 50 \ -h {output.html} \ - -j {output.json} 2> /dev/null + -j {output.json} 2> {log} else fastp -i {input.r1} \ -o {output.r1} \ @@ -96,7 +110,7 @@ if MODE == "reads": --cut_mean_quality 20 \ --length_required 50 \ -h {output.html} \ - -j {output.json} 2> /dev/null + -j {output.json} 2> {log} touch {output.r2} fi """ @@ -112,20 +126,23 @@ if MODE == "reads": bai = temp(f"{OUTDIR}/{{sample}}/aligned.raw.bam.bai") params: rg = lambda w: f"@RG\\tID:{w.sample}\\tSM:{w.sample}\\tPL:ILLUMINA" + log: + lambda w: _sample_log(w.sample, "bwa_mem") threads: max(1, workflow.cores) shell: """ mkdir -p {OUTDIR}/{wildcards.sample} + mkdir -p $(dirname {log}) if [ -s {input.r2} ]; then - bwa mem -t {threads} -R '{params.rg}' {input.ref} {input.r1} {input.r2} 2> /dev/null | \ + bwa mem -t {threads} -R '{params.rg}' {input.ref} {input.r1} {input.r2} 2> {log} | \ samtools view -b - | \ - samtools sort -@ {threads} -o {output.bam} 2> /dev/null + samtools sort -@ {threads} -o {output.bam} 2>> {log} else - bwa mem -t {threads} -R '{params.rg}' {input.ref} {input.r1} 2> /dev/null | \ + bwa mem -t {threads} -R '{params.rg}' {input.ref} {input.r1} 2> {log} | \ samtools view -b - | \ - samtools sort -@ {threads} -o {output.bam} 2> /dev/null + samtools sort -@ {threads} -o {output.bam} 2>> {log} fi - samtools index {output.bam} + samtools index {output.bam} 2>> {log} """ rule ampliconclip: @@ -133,24 +150,27 @@ if MODE == "reads": bam = f"{OUTDIR}/{{sample}}/aligned.raw.bam" output: bam = temp(f"{OUTDIR}/{{sample}}/aligned.clipped.bam"), - bai = temp(f"{OUTDIR}/{{sample}}/aligned.clipped.bam.bai"), - log = temp(f"{OUTDIR}/{{sample}}/ampliconclip.log") + bai = temp(f"{OUTDIR}/{{sample}}/aligned.clipped.bam.bai") params: bed = PRIMER_BED if PRIMER_BED else "" + log: + lambda w: _sample_log(w.sample, "ampliconclip") threads: max(1, int(workflow.cores * 0.5)) run: if PRIMER_BED: shell(""" + mkdir -p $(dirname {log}) samtools ampliconclip -b {params.bed} -@ {threads} \ - --strand --both-ends -o - {input.bam} 2> {output.log} \ - | samtools sort -@ {threads} -o {output.bam} - 2> /dev/null - samtools index {output.bam} + --strand --both-ends -o - {input.bam} 2> {log} \ + | samtools sort -@ {threads} -o {output.bam} - 2>> {log} + samtools index {output.bam} 2>> {log} """) else: shell(""" - cp {input.bam} {output.bam} - samtools index {output.bam} - touch {output.log} + mkdir -p $(dirname {log}) + : > {log} + cp {input.bam} {output.bam} 2>> {log} + samtools index {output.bam} 2>> {log} """) rule lofreq_indelqual: @@ -159,11 +179,14 @@ rule lofreq_indelqual: ref = REF output: bam = f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam" + log: + lambda w: _sample_log(w.sample, "lofreq_indelqual") shell: """ mkdir -p {OUTDIR}/{wildcards.sample} - lofreq indelqual --dindel -f {input.ref} -o {output.bam} {input.bam} 2> /dev/null - samtools index {output.bam} + mkdir -p $(dirname {log}) + lofreq indelqual --dindel -f {input.ref} -o {output.bam} {input.bam} 2> {log} + samtools index {output.bam} 2>> {log} """ rule samtools_depth: @@ -171,9 +194,12 @@ rule samtools_depth: bam = f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam" output: depth = f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt" + log: + lambda w: _sample_log(w.sample, "samtools_depth") shell: """ - samtools depth -aa {input.bam} > {output.depth} 2> /dev/null + mkdir -p $(dirname {log}) + samtools depth -aa {input.bam} > {output.depth} 2> {log} """ rule lofreq_call: @@ -184,14 +210,17 @@ rule lofreq_call: vcf_raw = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf"), vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", csi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi" + log: + lambda w: _sample_log(w.sample, "lofreq_call") threads: max(1, workflow.cores) shell: """ + mkdir -p $(dirname {log}) lofreq call-parallel --no-baq --call-indels --pp-threads {threads} \ - -f {input.ref} -o {output.vcf_raw} {input.bam} 2> /dev/null + -f {input.ref} -o {output.vcf_raw} {input.bam} 2> {log} - bgzip -c {output.vcf_raw} > {output.vcf} - bcftools index {output.vcf} + bgzip -c {output.vcf_raw} > {output.vcf} 2>> {log} + bcftools index {output.vcf} 2>> {log} """ rule update_csv: @@ -204,10 +233,14 @@ rule update_csv: params: original_csv = config["samples_csv"], outdir = OUTDIR + log: + _workflow_log("update_csv") run: import pandas as pd import os + Path(log[0]).parent.mkdir(parents=True, exist_ok=True) + # Read original CSV df = pd.read_csv(params.original_csv) @@ -224,3 +257,7 @@ rule update_csv: # Write updated CSV df.to_csv(output.csv, index=False) + Path(log[0]).write_text( + f"Updated execution spreadsheet written to {output.csv}\n", + encoding="utf-8", + ) From e2c9ac67ee96160e291693f866f5eadb85bdd376 Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Thu, 26 Mar 2026 09:29:14 +1100 Subject: [PATCH 02/20] Plotting and lofreq adjustments --- README.md | 14 +++- tests/test_analysis.py | 135 +++++++++++++++++++++++++++++++ tests/test_main.py | 83 +++++++++++++++++++ tests/test_snakemake_workflow.py | 19 +++-- vartracker/Snakefile | 32 +++----- vartracker/analysis.py | 35 +++++++- vartracker/main.py | 17 ++++ 7 files changed, 301 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index dda2869..6c8334e 100755 --- a/README.md +++ b/README.md @@ -217,6 +217,11 @@ vartracker end-to-end --test All modes understand `--test`, which copies the example dataset from `vartracker/test_data` into a temporary directory, resolves relative paths, and runs the appropriate workflow. +Temporary LoFreq note: +- In `bam` and `end-to-end` mode, `vartracker` currently caps `lofreq call-parallel` at 8 threads even if `--cores` is higher. +- This is a temporary workaround for an older Bioconda LoFreq build that can fail during `call-parallel` final filtering when many shards produce an excessively long merged VCF header. +- The cap will be revisited once an updated LoFreq build is available through Bioconda. + ### Input Spreadsheets Every CLI mode reads the same canonical columns: @@ -244,11 +249,18 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test ### Mode-specific options - `vartracker vcf` – accepts plotting and filtering options such as `--min-snv-freq`, `--min-indel-freq`, - `--allele-frequency-tag`, `--name`, `--outdir`, `--passage-cap`, `--manifest-level`, and literature controls + `--allele-frequency-tag`, `--heatmap-exclude`, `--name`, `--outdir`, `--passage-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. - `vartracker bam` – everything from `vcf`, plus Snakemake options: `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. - `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. + +Heatmap filtering: +- By default, all consequence classes are included in the heatmaps. +- Use `--heatmap-exclude` with a comma-separated list of `type_of_change` values to omit those consequence classes from the heatmaps. +- Any listed value is excluded; this is not limited to a fixed set of consequence classes. +- Shell-style wildcard matching is supported, so patterns such as `*frameshift*` will exclude any matching consequence class. +- Example: `--heatmap-exclude "synonymous,frameshift,stop_gained"` - `vartracker prepare spreadsheet` – specify `--mode` (`vcf`, `bam`, or `e2e`), `--dir` to scan, `--out` for the CSV, and `--dry-run` to preview without writing a file. - `vartracker prepare reference` – build a merged FASTA/GFF3 bundle from GenBank nucleotide accessions. diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 8d5bba1..2c1c6d3 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -7,6 +7,7 @@ from vartracker.analysis import ( search_literature, _prepare_variant_heatmap_matrix, + process_joint_variants, generate_variant_heatmap, ) @@ -142,6 +143,140 @@ def test_prepare_variant_heatmap_matrix_orders_variants_by_genome(): assert matrix.loc["S:D215G\n(A22206G)", "P1"] == 1.0 +def test_prepare_variant_heatmap_matrix_excludes_selected_consequence_types(): + table = pd.DataFrame( + [ + { + "gene": "S", + "amino_acid_consequence": "D215G", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "alt_freq": "0.0 / 0.7", + "samples": "P0 / P1", + "variant": "A22206G", + "start": 22206, + }, + { + "gene": "S", + "amino_acid_consequence": "T716T", + "nsp_aa_change": "", + "type_of_change": "synonymous", + "type_of_variant": "snp", + "alt_freq": "0.0 / 0.8", + "samples": "P0 / P1", + "variant": "C23403T", + "start": 23403, + }, + { + "gene": "S", + "amino_acid_consequence": "145del", + "nsp_aa_change": "", + "type_of_change": "frameshift", + "type_of_variant": "indel", + "alt_freq": "0.0 / 0.9", + "samples": "P0 / P1", + "variant": "A22029-", + "start": 22029, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix( + table, + ["P0", "P1"], + 0.2, + 0.2, + excluded_consequence_types=["synonymous", "frameshift"], + ) + + assert list(matrix.index) == ["S:D215G\n(A22206G)"] + + +def test_prepare_variant_heatmap_matrix_excludes_wildcard_consequence_types(): + table = pd.DataFrame( + [ + { + "gene": "S", + "amino_acid_consequence": "D215G", + "nsp_aa_change": "", + "type_of_change": "joint_frameshift", + "type_of_variant": "indel", + "alt_freq": "0.0 / 0.7", + "samples": "P0 / P1", + "variant": "A22206G", + "start": 22206, + }, + { + "gene": "S", + "amino_acid_consequence": "N501Y", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "alt_freq": "0.0 / 0.8", + "samples": "P0 / P1", + "variant": "A23063T", + "start": 23063, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix( + table, + ["P0", "P1"], + 0.2, + 0.2, + excluded_consequence_types=["*frameshift*"], + ) + + assert list(matrix.index) == ["S:N501Y\n(A23063T)"] + + +def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): + csv_path = tmp_path / "results.csv" + pd.DataFrame( + [ + { + "start": 100, + "gene": "S", + "amino_acid_consequence": "N501Y", + "nsp_aa_change": "", + "bcsq_nt_notation": "c.1A>T", + "bcsq_aa_notation": "p.N501Y", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "joint_joint_frameshift", + }, + { + "start": 101, + "gene": "", + "amino_acid_consequence": "", + "nsp_aa_change": "", + "bcsq_nt_notation": "", + "bcsq_aa_notation": "@100", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "frameshift", + }, + ] + ).to_csv(csv_path, index=False) + + result = process_joint_variants(str(csv_path)) + + assert result.loc[0, "type_of_change"] == "joint_frameshift" + assert result.loc[1, "type_of_change"] == "joint_frameshift" + + def test_generate_variant_heatmap_creates_interactive_html(tmp_path, monkeypatch): mpl_dir = tmp_path / "mpl" mpl_dir.mkdir() diff --git a/tests/test_main.py b/tests/test_main.py index 2850d61..243168b 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -274,6 +274,89 @@ def fake_parse_pokay(argv): assert (outdir / "literature_database.csv").exists() +def test_vcf_heatmap_exclude_option_is_forwarded_to_heatmap( + monkeypatch, tmp_path, minimal_vcf +): + coverage_path = tmp_path / "sample.cov.txt" + coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") + + csv_path = tmp_path / "inputs.csv" + csv_path.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + "Sample1,0,,,,sample.vcf,sample.cov.txt\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + + def fake_setup(args): + args.reference = "/tmp/mock_reference.fasta" + args.gff3 = "/tmp/mock_annotation.gff3" + return args + + monkeypatch.setattr(main_module, "setup_default_paths", fake_setup) + monkeypatch.setattr( + main_module, "validate_reference_and_annotation", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "generate_cumulative_lineplot", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "process_joint_variants", lambda path: pd.read_csv(path) + ) + monkeypatch.setattr( + main_module, "generate_gene_table", lambda table, *_a, **_k: table + ) + monkeypatch.setattr(main_module, "plot_gene_table", lambda *a, **k: None) + monkeypatch.setattr(main_module, "search_literature", lambda *a, **k: None) + + recorded = {} + + def fake_heatmap(*args, **kwargs): + recorded["excluded"] = kwargs.get("excluded_consequence_types") + + monkeypatch.setattr(main_module, "generate_variant_heatmap", fake_heatmap) + + formatted_csq = tmp_path / "formatted.csq.vcf.gz" + monkeypatch.setattr( + main_module, + "format_vcf", + lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), + ) + monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr( + main_module, + "process_vcf", + lambda *a, **k: pd.DataFrame( + { + "gene": ["S"], + "variant": ["A266C"], + "amino_acid_consequence": ["S:A1C"], + "nsp_aa_change": [""], + "presence_absence": ["Y"], + "variant_status": ["new"], + "persistence_status": ["new_persistent"], + "samples": ["Sample1"], + "alt_freq": ["0.5"], + } + ), + ) + + exit_code = main_module.main( + [ + "vcf", + str(csv_path), + "--outdir", + str(tmp_path / "results"), + "--heatmap-exclude", + "synonymous,frameshift,stop_gained", + ] + ) + + assert exit_code == 0 + assert recorded["excluded"] == ["synonymous", "frameshift", "stop_gained"] + + def test_e2e_runs_snakemake_then_vcf(monkeypatch, tmp_path): updated_csv = tmp_path / "samples_updated.csv" vcf_out = tmp_path / "vcf.gz" diff --git a/tests/test_snakemake_workflow.py b/tests/test_snakemake_workflow.py index fb8778b..184d7a4 100644 --- a/tests/test_snakemake_workflow.py +++ b/tests/test_snakemake_workflow.py @@ -7,18 +7,17 @@ def test_snakemake_rules_write_logs_under_outdir(): ).read_text(encoding="utf-8") assert "/dev/null" not in snakefile - assert 'return f"{OUTDIR}/logs/{rule_name}.log"' in snakefile - assert 'return f"{OUTDIR}/{sample}/logs/{rule_name}.log"' in snakefile + assert "threads: min(8, max(1, workflow.cores))" in snakefile expected_logs = [ - 'log:\n _workflow_log("bwa_index")', - 'lambda w: _sample_log(w.sample, "fastp")', - 'lambda w: _sample_log(w.sample, "bwa_mem")', - 'lambda w: _sample_log(w.sample, "ampliconclip")', - 'lambda w: _sample_log(w.sample, "lofreq_indelqual")', - 'lambda w: _sample_log(w.sample, "samtools_depth")', - 'lambda w: _sample_log(w.sample, "lofreq_call")', - 'log:\n _workflow_log("update_csv")', + 'f"{OUTDIR}/logs/bwa_index.log"', + 'f"{OUTDIR}/{{sample}}/logs/fastp.log"', + 'f"{OUTDIR}/{{sample}}/logs/bwa_mem.log"', + 'f"{OUTDIR}/{{sample}}/logs/ampliconclip.log"', + 'f"{OUTDIR}/{{sample}}/logs/lofreq_indelqual.log"', + 'f"{OUTDIR}/{{sample}}/logs/samtools_depth.log"', + 'f"{OUTDIR}/{{sample}}/logs/lofreq_call.log"', + 'f"{OUTDIR}/logs/update_csv.log"', ] for expected in expected_logs: diff --git a/vartracker/Snakefile b/vartracker/Snakefile index 64ff841..8aae8ac 100644 --- a/vartracker/Snakefile +++ b/vartracker/Snakefile @@ -11,15 +11,6 @@ OUTDIR = config.get("outdir", "results") PRIMER_BED = config.get("primer_bed", None) MODE = config.get("mode", "reads") - -def _workflow_log(rule_name): - return f"{OUTDIR}/logs/{rule_name}.log" - - -def _sample_log(sample, rule_name): - return f"{OUTDIR}/{sample}/logs/{rule_name}.log" - - def _value_or_blank(sample, key, allow_blank=False): value = SAMPLES[sample].get(key) if value is None or (isinstance(value, float) and pd.isna(value)): @@ -65,7 +56,7 @@ if MODE == "reads": pac = REF + ".pac", sa = REF + ".sa" log: - _workflow_log("bwa_index") + f"{OUTDIR}/logs/bwa_index.log" shell: """ mkdir -p $(dirname {log}) @@ -83,7 +74,7 @@ if MODE == "reads": html = temp(f"{OUTDIR}/{{sample}}/fastp.html"), json = temp(f"{OUTDIR}/{{sample}}/fastp.json") log: - lambda w: _sample_log(w.sample, "fastp") + f"{OUTDIR}/{{sample}}/logs/fastp.log" threads: max(1, int(workflow.cores * 0.5)) shell: """ @@ -127,7 +118,7 @@ if MODE == "reads": params: rg = lambda w: f"@RG\\tID:{w.sample}\\tSM:{w.sample}\\tPL:ILLUMINA" log: - lambda w: _sample_log(w.sample, "bwa_mem") + f"{OUTDIR}/{{sample}}/logs/bwa_mem.log" threads: max(1, workflow.cores) shell: """ @@ -154,7 +145,7 @@ if MODE == "reads": params: bed = PRIMER_BED if PRIMER_BED else "" log: - lambda w: _sample_log(w.sample, "ampliconclip") + f"{OUTDIR}/{{sample}}/logs/ampliconclip.log" threads: max(1, int(workflow.cores * 0.5)) run: if PRIMER_BED: @@ -180,7 +171,7 @@ rule lofreq_indelqual: output: bam = f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam" log: - lambda w: _sample_log(w.sample, "lofreq_indelqual") + f"{OUTDIR}/{{sample}}/logs/lofreq_indelqual.log" shell: """ mkdir -p {OUTDIR}/{wildcards.sample} @@ -195,7 +186,7 @@ rule samtools_depth: output: depth = f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt" log: - lambda w: _sample_log(w.sample, "samtools_depth") + f"{OUTDIR}/{{sample}}/logs/samtools_depth.log" shell: """ mkdir -p $(dirname {log}) @@ -211,8 +202,8 @@ rule lofreq_call: vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", csi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi" log: - lambda w: _sample_log(w.sample, "lofreq_call") - threads: max(1, workflow.cores) + f"{OUTDIR}/{{sample}}/logs/lofreq_call.log" + threads: min(8, max(1, workflow.cores)) shell: """ mkdir -p $(dirname {log}) @@ -234,12 +225,13 @@ rule update_csv: original_csv = config["samples_csv"], outdir = OUTDIR log: - _workflow_log("update_csv") + f"{OUTDIR}/logs/update_csv.log" run: import pandas as pd import os - Path(log[0]).parent.mkdir(parents=True, exist_ok=True) + log_path = Path(str(log[0])) + log_path.parent.mkdir(parents=True, exist_ok=True) # Read original CSV df = pd.read_csv(params.original_csv) @@ -257,7 +249,7 @@ rule update_csv: # Write updated CSV df.to_csv(output.csv, index=False) - Path(log[0]).write_text( + log_path.write_text( f"Updated execution spreadsheet written to {output.csv}\n", encoding="utf-8", ) diff --git a/vartracker/analysis.py b/vartracker/analysis.py index e13f3ec..2abe96e 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -5,6 +5,7 @@ """ import html +import fnmatch import os import re import string @@ -23,6 +24,14 @@ mpl.rcParams["pdf.fonttype"] = 42 +def _ensure_joint_prefix(change_type: object) -> str: + text = str(change_type or "").strip() + if not text: + return "joint" + normalised = re.sub(r"^(joint_)+", "", text) + return f"joint_{normalised}" if normalised else "joint" + + def process_joint_variants(path): """ Process joint variants from bcftools csq output. @@ -75,8 +84,12 @@ def process_joint_variants(path): tab.at[i, col] = tab.at[j, col] # Update type of change for joint variants - tab.at[i, "type_of_change"] = "joint_" + tab.at[j, "type_of_change"] - tab.at[j, "type_of_change"] = "joint_" + tab.at[j, "type_of_change"] + tab.at[i, "type_of_change"] = _ensure_joint_prefix( + tab.at[j, "type_of_change"] + ) + tab.at[j, "type_of_change"] = _ensure_joint_prefix( + tab.at[j, "type_of_change"] + ) except (ValueError, IndexError, KeyError) as e: print(f"Warning: Could not process joint variant at index {i}: {str(e)}") @@ -490,6 +503,7 @@ def _prepare_variant_heatmap_matrix( sample_names: Sequence[str], min_snv_freq: float, min_indel_freq: float, + excluded_consequence_types: Sequence[str] | None = None, gene_lengths: Dict[str, int] | None = None, ) -> pd.DataFrame: """Prepare a matrix of allele frequencies for heatmap plotting.""" @@ -519,6 +533,11 @@ def _prepare_variant_heatmap_matrix( use_nsps = table["nsp_aa_change"].astype(str).str.contains(":").any() gene_order_map, _ = _build_gene_order_map(gene_lengths, include_nsps=use_nsps) + excluded_patterns = { + str(value).strip().lower() + for value in (excluded_consequence_types or []) + if str(value).strip() + } records: List[Dict[str, Union[str, float, int]]] = [] seen_labels = set() @@ -533,6 +552,10 @@ def _prepare_variant_heatmap_matrix( if base_label in seen_labels: continue + change_type = str(getattr(row, "type_of_change", "")).strip().lower() + if any(fnmatch.fnmatch(change_type, pattern) for pattern in excluded_patterns): + continue + variant_type = str(getattr(row, "type_of_variant", "")).lower() sample_tokens = [ @@ -1068,6 +1091,7 @@ def generate_variant_heatmap( project_name: str, min_snv_freq: float, min_indel_freq: float, + excluded_consequence_types: Sequence[str] | None = None, gene_lengths: Dict[str, int] | None = None, literature_hits: Optional[pd.DataFrame] = None, literature_table_path: Optional[str] = None, @@ -1077,7 +1101,12 @@ def generate_variant_heatmap( try: heatmap_data = _prepare_variant_heatmap_matrix( - table, sample_names, min_snv_freq, min_indel_freq, gene_lengths + table, + sample_names, + min_snv_freq, + min_indel_freq, + excluded_consequence_types, + gene_lengths, ) if heatmap_data.empty: print("No variant data available for heatmap; skipping plot.") diff --git a/vartracker/main.py b/vartracker/main.py index 394f72e..fa32c82 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -103,6 +103,12 @@ ] +def _parse_csv_option_list(value: str | None) -> list[str]: + if value is None: + return [] + return [item.strip() for item in str(value).split(",") if item.strip()] + + def _print_dependency_error(error: DependencyError) -> None: """Render a dependency error with optional remediation tips.""" message = str(error) @@ -417,6 +423,16 @@ def _configure_vcf_parser( default="AF", help="INFO tag name for allele frequency (default: AF)", ) + vt_group.add_argument( + "--heatmap-exclude", + action="store", + required=False, + default="", + help=( + "Comma-separated amino-acid consequence types to exclude from heatmaps " + "(for example: synonymous,frameshift,stop_gained)" + ), + ) def _add_vcf_subparser(subparsers): @@ -1592,6 +1608,7 @@ def _process_files( pname, args.min_snv_freq, args.min_indel_freq, + excluded_consequence_types=_parse_csv_option_list(args.heatmap_exclude), gene_lengths=gene_lengths, literature_hits=literature_hits_df, literature_table_path=literature_full_csv_path, From 565f4304f0649835407d6f0d935b902fcf873aa6 Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Thu, 26 Mar 2026 11:42:15 +1100 Subject: [PATCH 03/20] Arguments refactor and plotting refinements --- README.md | 28 ++- tests/test_analysis.py | 73 +++++++ tests/test_main.py | 92 ++++++++- vartracker/__init__.py | 7 +- vartracker/analysis.py | 185 ++++++++++++++++- vartracker/main.py | 453 ++++++++++++++++++++++++++++++++++------- 6 files changed, 746 insertions(+), 92 deletions(-) diff --git a/README.md b/README.md index 6c8334e..0bfa6fb 100755 --- a/README.md +++ b/README.md @@ -202,6 +202,11 @@ vartracker end-to-end path/to/read_inputs.csv \ --cores 12 \ --outdir results/e2e_summary +# Re-plot a heatmap from an existing vartracker results file +vartracker plot heatmap results/results.csv \ + --heatmap-aa-exclude "*frameshift*" \ + --outdir results/replots + # Generate a template spreadsheet for a directory of files vartracker prepare spreadsheet --mode e2e --dir data/passaging --out inputs.csv @@ -249,18 +254,29 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test ### Mode-specific options - `vartracker vcf` – accepts plotting and filtering options such as `--min-snv-freq`, `--min-indel-freq`, - `--allele-frequency-tag`, `--heatmap-exclude`, `--name`, `--outdir`, `--passage-cap`, `--manifest-level`, and literature controls + `--allele-frequency-tag`, `--heatmap-aa-exclude`, `--heatmap-aa-include`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. - `vartracker bam` – everything from `vcf`, plus Snakemake options: `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. - `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. +- `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV. Heatmap filtering: -- By default, all consequence classes are included in the heatmaps. -- Use `--heatmap-exclude` with a comma-separated list of `type_of_change` values to omit those consequence classes from the heatmaps. -- Any listed value is excluded; this is not limited to a fixed set of consequence classes. -- Shell-style wildcard matching is supported, so patterns such as `*frameshift*` will exclude any matching consequence class. -- Example: `--heatmap-exclude "synonymous,frameshift,stop_gained"` +- By default, all consequence classes are included except joint variants. Use `--heatmap-include-joint` to show joint variants. +- `--heatmap-aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. +- `--heatmap-aa-include`: comma-separated `type_of_change` patterns to include. +- `--heatmap-only-persistent`: only include `new_persistent` variants. +- `--heatmap-only-new`: only include variants with `variant_status == new`. +- `--heatmap-gene-include` and `--heatmap-gene-exclude`: comma-separated gene patterns. +- `--heatmap-variant-type`: comma-separated variant-type patterns such as `snp` or `indel`. +- `--heatmap-qc`: comma-separated `overall_variant_qc` patterns to include. +- `--heatmap-min-persistence`: minimum number of included samples in which the variant must be present. +- `--heatmap-min-max-af`: minimum maximum allele frequency across included samples. +- `--heatmap-min-sample-af`: minimum allele frequency that must be reached in at least one included sample. +- `--heatmap-sample-subset`: comma-separated sample-name patterns to plot. +- `--heatmap-hide-singletons`: hide variants present in only one included sample. +- `--heatmap-min-depth`: minimum site depth a variant must reach in at least one included sample. +- Example: `--heatmap-aa-exclude "synonymous,*frameshift*,stop_gained"` - `vartracker prepare spreadsheet` – specify `--mode` (`vcf`, `bam`, or `e2e`), `--dir` to scan, `--out` for the CSV, and `--dry-run` to preview without writing a file. - `vartracker prepare reference` – build a merged FASTA/GFF3 bundle from GenBank nucleotide accessions. diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 2c1c6d3..e636338 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -232,6 +232,79 @@ def test_prepare_variant_heatmap_matrix_excludes_wildcard_consequence_types(): assert list(matrix.index) == ["S:N501Y\n(A23063T)"] +def test_prepare_variant_heatmap_matrix_applies_extended_filters(): + table = pd.DataFrame( + [ + { + "gene": "S", + "amino_acid_consequence": "D215G", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "variant_status": "new", + "persistence_status": "new_persistent", + "overall_variant_qc": "PASS", + "presence_absence": "N / Y / Y", + "alt_freq": "0.0 / 0.6 / 0.7", + "variant_site_depth": "0 / 120 / 125", + "samples": "P0 / P1 / P2", + "variant": "A22206G", + "start": 22206, + }, + { + "gene": "S", + "amino_acid_consequence": "145fs", + "nsp_aa_change": "", + "type_of_change": "joint_frameshift", + "type_of_variant": "indel", + "variant_status": "new", + "persistence_status": "new_persistent", + "overall_variant_qc": "PASS", + "presence_absence": "N / N / Y", + "alt_freq": "0.0 / 0.0 / 0.8", + "variant_site_depth": "0 / 0 / 130", + "samples": "P0 / P1 / P2", + "variant": "A22029-", + "start": 22029, + }, + { + "gene": "N", + "amino_acid_consequence": "R203K", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "variant_status": "original", + "persistence_status": "original_retained", + "overall_variant_qc": "FAIL", + "presence_absence": "Y / Y / Y", + "alt_freq": "0.5 / 0.5 / 0.5", + "variant_site_depth": "110 / 115 / 120", + "samples": "P0 / P1 / P2", + "variant": "G28881A", + "start": 28881, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix( + table, + ["P0", "P1", "P2"], + 0.2, + 0.2, + only_new=True, + gene_include=["S"], + variant_type_include=["snp"], + qc_include=["pass"], + min_persistence=2, + min_max_af=0.6, + sample_subset=["P1", "P2"], + hide_singletons=True, + min_depth=100, + ) + + assert list(matrix.index) == ["S:D215G\n(A22206G)"] + + def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): csv_path = tmp_path / "results.csv" pd.DataFrame( diff --git a/tests/test_main.py b/tests/test_main.py index 243168b..2df85c5 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -274,7 +274,7 @@ def fake_parse_pokay(argv): assert (outdir / "literature_database.csv").exists() -def test_vcf_heatmap_exclude_option_is_forwarded_to_heatmap( +def test_vcf_heatmap_options_are_forwarded_to_heatmap( monkeypatch, tmp_path, minimal_vcf ): coverage_path = tmp_path / "sample.cov.txt" @@ -313,7 +313,7 @@ def fake_setup(args): recorded = {} def fake_heatmap(*args, **kwargs): - recorded["excluded"] = kwargs.get("excluded_consequence_types") + recorded.update(kwargs) monkeypatch.setattr(main_module, "generate_variant_heatmap", fake_heatmap) @@ -348,13 +348,95 @@ def fake_heatmap(*args, **kwargs): str(csv_path), "--outdir", str(tmp_path / "results"), - "--heatmap-exclude", - "synonymous,frameshift,stop_gained", + "--heatmap-aa-exclude", + "synonymous,*frameshift*,stop_gained", + "--heatmap-aa-include", + "missense", + "--heatmap-only-persistent", + "--heatmap-only-new", + "--heatmap-gene-include", + "S", + "--heatmap-gene-exclude", + "N", + "--heatmap-variant-type", + "snp", + "--heatmap-qc", + "PASS", + "--heatmap-min-persistence", + "2", + "--heatmap-min-max-af", + "0.4", + "--heatmap-min-sample-af", + "0.3", + "--heatmap-sample-subset", + "Sample1", + "--heatmap-hide-singletons", + "--heatmap-min-depth", + "20", + ] + ) + + assert exit_code == 0 + assert recorded["excluded_consequence_types"] == [ + "synonymous", + "*frameshift*", + "stop_gained", + ] + assert recorded["included_consequence_types"] == ["missense"] + assert recorded["include_joint"] is False + assert recorded["only_persistent"] is True + assert recorded["only_new"] is True + assert recorded["gene_include"] == ["S"] + assert recorded["gene_exclude"] == ["N"] + assert recorded["variant_type_include"] == ["snp"] + assert recorded["qc_include"] == ["PASS"] + assert recorded["min_persistence"] == 2 + assert recorded["min_max_af"] == 0.4 + assert recorded["min_sample_af"] == 0.3 + assert recorded["sample_subset"] == ["Sample1"] + assert recorded["hide_singletons"] is True + assert recorded["min_depth"] == 20 + + +def test_plot_heatmap_replots_results_csv(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + results_csv.write_text( + "samples,name,alt_freq,variant_site_depth,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,start\n" + "P0 / P1,Example,0.0 / 0.5,0 / 100,N / Y,new,new_persistent,snp,missense,S,A266C,266\n", + encoding="utf-8", + ) + + recorded = {} + + def fake_heatmap(*args, **kwargs): + recorded["table"] = args[0] + recorded["sample_names"] = args[1] + recorded["outdir"] = args[3] + recorded["project_name"] = args[4] + recorded["kwargs"] = kwargs + + monkeypatch.setattr(main_module, "generate_variant_heatmap", fake_heatmap) + + outdir = tmp_path / "plots" + exit_code = main_module.main( + [ + "plot", + "heatmap", + str(results_csv), + "--outdir", + str(outdir), + "--heatmap-aa-exclude", + "*frameshift*", + "--heatmap-include-joint", ] ) assert exit_code == 0 - assert recorded["excluded"] == ["synonymous", "frameshift", "stop_gained"] + assert list(recorded["sample_names"]) == ["P0", "P1"] + assert recorded["outdir"] == str(outdir) + assert recorded["project_name"] == "Example" + assert recorded["kwargs"]["excluded_consequence_types"] == ["*frameshift*"] + assert recorded["kwargs"]["include_joint"] is True def test_e2e_runs_snakemake_then_vcf(monkeypatch, tmp_path): diff --git a/vartracker/__init__.py b/vartracker/__init__.py index 38451d0..5e2c631 100644 --- a/vartracker/__init__.py +++ b/vartracker/__init__.py @@ -12,6 +12,11 @@ __author__ = "Dr Charles Foster" __email__ = "github.com/charlesfoster" -from .main import main # noqa: F401 + +def main(): + from .main import main as _main + + return _main() + __all__ = ["__version__", "main"] diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 2abe96e..a65272d 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -32,6 +32,17 @@ def _ensure_joint_prefix(change_type: object) -> str: return f"joint_{normalised}" if normalised else "joint" +def _parse_slash_separated_tokens(value: object) -> list[str]: + return [token.strip() for token in str(value or "").split(" / ")] + + +def _match_any_pattern(value: object, patterns: Sequence[str]) -> bool: + text = str(value or "").strip().lower() + if not text: + return False + return any(fnmatch.fnmatch(text, pattern) for pattern in patterns) + + def process_joint_variants(path): """ Process joint variants from bcftools csq output. @@ -504,6 +515,20 @@ def _prepare_variant_heatmap_matrix( min_snv_freq: float, min_indel_freq: float, excluded_consequence_types: Sequence[str] | None = None, + included_consequence_types: Sequence[str] | None = None, + include_joint: bool = False, + only_persistent: bool = False, + only_new: bool = False, + gene_include: Sequence[str] | None = None, + gene_exclude: Sequence[str] | None = None, + variant_type_include: Sequence[str] | None = None, + qc_include: Sequence[str] | None = None, + min_persistence: int | None = None, + min_max_af: float | None = None, + min_sample_af: float | None = None, + sample_subset: Sequence[str] | None = None, + hide_singletons: bool = False, + min_depth: int | None = None, gene_lengths: Dict[str, int] | None = None, ) -> pd.DataFrame: """Prepare a matrix of allele frequencies for heatmap plotting.""" @@ -518,6 +543,19 @@ def _prepare_variant_heatmap_matrix( ordered_samples = [ s.strip() for s in str(table.iloc[0]["samples"]).split(" / ") ] + sample_subset_patterns = [ + str(value).strip().lower() + for value in (sample_subset or []) + if str(value).strip() + ] + if sample_subset_patterns: + ordered_samples = [ + sample + for sample in ordered_samples + if _match_any_pattern(sample, sample_subset_patterns) + ] + if not ordered_samples: + return pd.DataFrame(columns=[]) use_nsps = True if gene_lengths is not None: @@ -538,6 +576,37 @@ def _prepare_variant_heatmap_matrix( for value in (excluded_consequence_types or []) if str(value).strip() } + included_patterns = { + str(value).strip().lower() + for value in (included_consequence_types or []) + if str(value).strip() + } + gene_include_patterns = { + str(value).strip().lower() + for value in (gene_include or []) + if str(value).strip() + } + gene_exclude_patterns = { + str(value).strip().lower() + for value in (gene_exclude or []) + if str(value).strip() + } + variant_type_patterns = { + str(value).strip().lower() + for value in (variant_type_include or []) + if str(value).strip() + } + qc_patterns = { + str(value).strip().lower() for value in (qc_include or []) if str(value).strip() + } + effective_min_af = max( + value + for value in ( + 0.0, + min_max_af if min_max_af is not None else 0.0, + min_sample_af if min_sample_af is not None else 0.0, + ) + ) records: List[Dict[str, Union[str, float, int]]] = [] seen_labels = set() @@ -547,29 +616,91 @@ def _prepare_variant_heatmap_matrix( if str(gene_value) in {"5' UTR", "3' UTR", "INTERGENIC"}: continue + if gene_include_patterns and not _match_any_pattern( + gene_value, list(gene_include_patterns) + ): + continue + if gene_exclude_patterns and _match_any_pattern( + gene_value, list(gene_exclude_patterns) + ): + continue + gene_label, display_label, base_label = _resolve_variant_labels(row) if base_label in seen_labels: continue change_type = str(getattr(row, "type_of_change", "")).strip().lower() + if not include_joint and change_type.startswith("joint"): + continue + if included_patterns and not any( + fnmatch.fnmatch(change_type, pattern) for pattern in included_patterns + ): + continue if any(fnmatch.fnmatch(change_type, pattern) for pattern in excluded_patterns): continue + if ( + only_persistent + and str(getattr(row, "persistence_status", "")).strip().lower() + != "new_persistent" + ): + continue + if ( + only_new + and str(getattr(row, "variant_status", "")).strip().lower() != "new" + ): + continue + variant_type = str(getattr(row, "type_of_variant", "")).lower() + if variant_type_patterns and not any( + fnmatch.fnmatch(variant_type, pattern) for pattern in variant_type_patterns + ): + continue + if qc_patterns and not _match_any_pattern( + getattr(row, "overall_variant_qc", ""), list(qc_patterns) + ): + continue - sample_tokens = [ - s.strip() for s in str(getattr(row, "samples", "")).split(" / ") - ] + sample_tokens = _parse_slash_separated_tokens(getattr(row, "samples", "")) freq_tokens = [ _coerce_frequency(token) - for token in str(getattr(row, "alt_freq", "")).split(" / ") + for token in _parse_slash_separated_tokens(getattr(row, "alt_freq", "")) ] + presence_tokens = _parse_slash_separated_tokens( + getattr(row, "presence_absence", "") + ) + depth_tokens = _parse_slash_separated_tokens( + getattr(row, "variant_site_depth", "") + ) if not freq_tokens: continue - max_freq = max(freq_tokens) + sample_freq_map = { + sample: freq for sample, freq in zip(sample_tokens, freq_tokens) + } + sample_presence_map = { + sample: token for sample, token in zip(sample_tokens, presence_tokens) + } + sample_depth_map: dict[str, float] = {} + for sample, token in zip(sample_tokens, depth_tokens): + text = str(token).strip() + if text in {"", ".", "M"}: + sample_depth_map[sample] = 0.0 + continue + try: + sample_depth_map[sample] = float(text) + except ValueError: + sample_depth_map[sample] = 0.0 + + row_values = [sample_freq_map.get(sample, 0.0) for sample in ordered_samples] + row_presence = [ + sample_presence_map.get(sample, "N") for sample in ordered_samples + ] + row_depths = [sample_depth_map.get(sample, 0.0) for sample in ordered_samples] + + max_freq = max(row_values) if row_values else 0.0 threshold = min_snv_freq if "indel" in variant_type: @@ -577,11 +708,17 @@ def _prepare_variant_heatmap_matrix( if max_freq < threshold: continue - - sample_freq_map = { - sample: freq for sample, freq in zip(sample_tokens, freq_tokens) - } - row_values = [sample_freq_map.get(sample, 0.0) for sample in ordered_samples] + if effective_min_af and max_freq < effective_min_af: + continue + if ( + min_persistence is not None + and sum(token == "Y" for token in row_presence) < min_persistence + ): + continue + if hide_singletons and sum(token == "Y" for token in row_presence) <= 1: + continue + if min_depth is not None and max(row_depths, default=0.0) < min_depth: + continue record: Dict[str, Union[str, float, int]] = { "label": display_label, @@ -1092,6 +1229,20 @@ def generate_variant_heatmap( min_snv_freq: float, min_indel_freq: float, excluded_consequence_types: Sequence[str] | None = None, + included_consequence_types: Sequence[str] | None = None, + include_joint: bool = False, + only_persistent: bool = False, + only_new: bool = False, + gene_include: Sequence[str] | None = None, + gene_exclude: Sequence[str] | None = None, + variant_type_include: Sequence[str] | None = None, + qc_include: Sequence[str] | None = None, + min_persistence: int | None = None, + min_max_af: float | None = None, + min_sample_af: float | None = None, + sample_subset: Sequence[str] | None = None, + hide_singletons: bool = False, + min_depth: int | None = None, gene_lengths: Dict[str, int] | None = None, literature_hits: Optional[pd.DataFrame] = None, literature_table_path: Optional[str] = None, @@ -1106,6 +1257,20 @@ def generate_variant_heatmap( min_snv_freq, min_indel_freq, excluded_consequence_types, + included_consequence_types, + include_joint, + only_persistent, + only_new, + gene_include, + gene_exclude, + variant_type_include, + qc_include, + min_persistence, + min_max_af, + min_sample_af, + sample_subset, + hide_singletons, + min_depth, gene_lengths, ) if heatmap_data.empty: diff --git a/vartracker/main.py b/vartracker/main.py index fa32c82..3dbf9d7 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -109,6 +109,148 @@ def _parse_csv_option_list(value: str | None) -> list[str]: return [item.strip() for item in str(value).split(",") if item.strip()] +def _add_heatmap_option_arguments(group: argparse._ArgumentGroup) -> None: + group.add_argument( + "--heatmap-aa-exclude", + action="store", + required=False, + default="", + help=( + "Comma-separated `type_of_change` patterns to exclude from heatmaps " + "(wildcards supported, e.g. synonymous,*frameshift*)" + ), + ) + group.add_argument( + "--heatmap-aa-include", + action="store", + required=False, + default="", + help=( + "Comma-separated `type_of_change` patterns to include in heatmaps " + "(wildcards supported)" + ), + ) + group.add_argument( + "--heatmap-include-joint", + action="store_true", + default=False, + help="Include joint variants in heatmaps (default: exclude them)", + ) + group.add_argument( + "--heatmap-only-persistent", + action="store_true", + default=False, + help="Only include variants with persistence_status == new_persistent", + ) + group.add_argument( + "--heatmap-only-new", + action="store_true", + default=False, + help="Only include variants with variant_status == new", + ) + group.add_argument( + "--heatmap-gene-include", + action="store", + required=False, + default="", + help="Comma-separated gene patterns to include in heatmaps", + ) + group.add_argument( + "--heatmap-gene-exclude", + action="store", + required=False, + default="", + help="Comma-separated gene patterns to exclude from heatmaps", + ) + group.add_argument( + "--heatmap-variant-type", + action="store", + required=False, + default="", + help="Comma-separated variant type patterns to include (e.g. snp,indel)", + ) + group.add_argument( + "--heatmap-qc", + action="store", + required=False, + default="", + help="Comma-separated overall QC patterns to include (e.g. PASS)", + ) + group.add_argument( + "--heatmap-min-persistence", + action="store", + type=int, + default=None, + help="Minimum number of samples in which a variant must be present", + ) + group.add_argument( + "--heatmap-min-max-af", + action="store", + type=float, + default=None, + help="Minimum maximum allele frequency across included samples", + ) + group.add_argument( + "--heatmap-min-sample-af", + action="store", + type=float, + default=None, + help="Minimum allele frequency that must be reached in at least one included sample", + ) + group.add_argument( + "--heatmap-sample-subset", + action="store", + required=False, + default="", + help="Comma-separated sample name patterns to plot", + ) + group.add_argument( + "--heatmap-hide-singletons", + action="store_true", + default=False, + help="Hide variants present in only one included sample", + ) + group.add_argument( + "--heatmap-min-depth", + action="store", + type=int, + default=None, + help="Minimum site depth a variant must reach in at least one included sample", + ) + + +def _collect_heatmap_kwargs(args) -> dict[str, object]: + return { + "excluded_consequence_types": _parse_csv_option_list( + getattr(args, "heatmap_aa_exclude", "") + ), + "included_consequence_types": _parse_csv_option_list( + getattr(args, "heatmap_aa_include", "") + ), + "include_joint": getattr(args, "heatmap_include_joint", False), + "only_persistent": getattr(args, "heatmap_only_persistent", False), + "only_new": getattr(args, "heatmap_only_new", False), + "gene_include": _parse_csv_option_list( + getattr(args, "heatmap_gene_include", "") + ), + "gene_exclude": _parse_csv_option_list( + getattr(args, "heatmap_gene_exclude", "") + ), + "variant_type_include": _parse_csv_option_list( + getattr(args, "heatmap_variant_type", "") + ), + "qc_include": _parse_csv_option_list(getattr(args, "heatmap_qc", "")), + "min_persistence": getattr(args, "heatmap_min_persistence", None), + "min_max_af": getattr(args, "heatmap_min_max_af", None), + "min_sample_af": getattr(args, "heatmap_min_sample_af", None), + "sample_subset": _parse_csv_option_list( + getattr(args, "heatmap_sample_subset", "") + ), + "hide_singletons": getattr(args, "heatmap_hide_singletons", False), + "min_depth": getattr(args, "heatmap_min_depth", None), + } + + def _print_dependency_error(error: DependencyError) -> None: """Render a dependency error with optional remediation tips.""" message = str(error) @@ -303,9 +445,19 @@ def _configure_vcf_parser( else: parser.add_argument("input_csv", nargs="?", help="Input CSV file") - vt_group = parser.add_argument_group("Vartracker options") + analysis_group = parser.add_argument_group("Vartracker Analysis Options") + output_group = parser.add_argument_group("Vartracker Output Options") + heatmap_group = parser.add_argument_group("Heatmap") - vt_group.add_argument( + analysis_group.add_argument( + "-r", + "--reference", + action="store", + required=False, + default=None, + help="Reference genome (default: uses packaged SARS-CoV-2 reference)", + ) + analysis_group.add_argument( "-g", "--gff3", action="store", @@ -313,7 +465,14 @@ def _configure_vcf_parser( default=None, help="GFF3 annotations to use (default: packaged SARS-CoV-2 annotations)", ) - vt_group.add_argument( + analysis_group.add_argument( + "--allele-frequency-tag", + action="store", + required=False, + default="AF", + help="INFO tag name for allele frequency (default: AF)", + ) + analysis_group.add_argument( "-m", "--min-snv-freq", action="store", @@ -322,7 +481,7 @@ def _configure_vcf_parser( default=0.03, help="Minimum allele frequency of SNV variants to keep (default: 0.03)", ) - vt_group.add_argument( + analysis_group.add_argument( "-M", "--min-indel-freq", action="store", @@ -331,7 +490,7 @@ def _configure_vcf_parser( default=0.1, help="Minimum allele frequency of indel variants to keep (default: 0.1)", ) - vt_group.add_argument( + analysis_group.add_argument( "-d", "--min-depth", action="store", @@ -340,7 +499,34 @@ def _configure_vcf_parser( default=10, help="Minimum depth threshold for variant QC (default: 10)", ) - vt_group.add_argument( + analysis_group.add_argument( + "--sample-cap", + action="store", + type=int, + help="Only analyse samples with sample_number less than or equal to this value", + default=None, + ) + analysis_group.add_argument( + "--literature-csv", + action="store", + required=False, + default=None, + help="Path to a literature CSV file (see README for file structure)", + ) + analysis_group.add_argument( + "--search-pokay", + action="store_true", + help='Automatically download and search against the "pokay" SARS-CoV-2 literature database.', + default=False, + ) + analysis_group.add_argument( + "--test", + action="store_true", + help="Run vartracker against the bundled demonstration dataset", + default=False, + ) + + output_group.add_argument( "-n", "--name", action="store", @@ -348,7 +534,7 @@ def _configure_vcf_parser( default=None, help="Optional: add a column to results with the name specified here", ) - vt_group.add_argument( + output_group.add_argument( "-o", "--outdir", action="store", @@ -356,13 +542,13 @@ def _configure_vcf_parser( default=".", help="Output directory for vartracker results (default: current directory)", ) - vt_group.add_argument( + output_group.add_argument( "--manifest-level", choices=["light", "deep"], default="light", help="Manifest detail level for run metadata (default: light)", ) - vt_group.add_argument( + output_group.add_argument( "-f", "--filename", action="store", @@ -370,69 +556,41 @@ def _configure_vcf_parser( default="results.csv", help="Output file name (default: results.csv)", ) - vt_group.add_argument( - "-r", - "--reference", - action="store", - required=False, - default=None, - help="Reference genome (default: uses packaged SARS-CoV-2 reference)", - ) - vt_group.add_argument( - "--passage-cap", - action="store", - type=int, - help="Cap the number of passages at this number", - default=None, - ) - vt_group.add_argument( + + output_group.add_argument( "--debug", action="store_true", help="Print commands being run for debugging", default=False, ) - vt_group.add_argument( + output_group.add_argument( "--keep-temp", action="store_true", help="Keep temporary files for debugging", default=False, ) - vt_group.add_argument( - "--literature-csv", - action="store", - required=False, - default=None, - help="Path to a literature CSV file (see README for file structure)", - ) - vt_group.add_argument( - "--search-pokay", - action="store_true", - help='Automatically download and search against the "pokay" SARS-CoV-2 literature database.', - default=False, - ) - vt_group.add_argument( - "--test", - action="store_true", - help="Run vartracker against the bundled demonstration dataset", - default=False, - ) - vt_group.add_argument( - "--allele-frequency-tag", - action="store", - required=False, - default="AF", - help="INFO tag name for allele frequency (default: AF)", - ) - vt_group.add_argument( - "--heatmap-exclude", - action="store", - required=False, - default="", - help=( - "Comma-separated amino-acid consequence types to exclude from heatmaps " - "(for example: synonymous,frameshift,stop_gained)" - ), - ) + + _add_heatmap_option_arguments(heatmap_group) + + +def _move_action_group_after( + parser: argparse.ArgumentParser, group_title: str, anchor_title: str +) -> None: + groups = parser._action_groups + try: + group_index = next( + i for i, grp in enumerate(groups) if grp.title == group_title + ) + anchor_index = next( + i for i, grp in enumerate(groups) if grp.title == anchor_title + ) + except StopIteration: + return + + group = groups.pop(group_index) + if group_index < anchor_index: + anchor_index -= 1 + groups.insert(anchor_index + 1, group) def _add_vcf_subparser(subparsers): @@ -498,8 +656,6 @@ def _add_bam_subparser(subparsers): """, ) - _configure_vcf_parser(bam_parser, include_input_csv=True, input_csv_required=False) - snk_group = bam_parser.add_argument_group("Snakemake options") snk_group.add_argument( "--snakemake-outdir", @@ -536,6 +692,11 @@ def _add_bam_subparser(subparsers): default=False, ) + _configure_vcf_parser(bam_parser, include_input_csv=True, input_csv_required=False) + _move_action_group_after( + bam_parser, "Snakemake options", "Vartracker Analysis Options" + ) + bam_parser.set_defaults(handler=_run_bam_command, _subparser=bam_parser) @@ -704,6 +865,72 @@ def _add_schema_subparser(subparsers): schema_parser.set_defaults(handler=_run_describe_output_command) +def _run_plot_command(args): + if getattr(args, "handler", None) is None or args.command == "plot": + args._subparser.print_help() + return 1 + return args.handler(args) + + +def _add_plot_heatmap_subparser(subparsers): + parser = subparsers.add_parser( + "heatmap", + aliases=["hm"], + help="Regenerate the variant heatmap from an existing results CSV", + description="Read a vartracker results CSV and regenerate the heatmap outputs.", + formatter_class=HelpFormatter, + ) + parser.add_argument("results_csv", help="Path to a vartracker results CSV") + parser.add_argument( + "--outdir", + default=None, + help="Output directory for regenerated heatmap files (default: results CSV directory)", + ) + parser.add_argument( + "--name", + default=None, + help="Optional plot title prefix (default: use the `name` column if present)", + ) + parser.add_argument( + "--literature-csv", + default=None, + help="Optional literature hits CSV to link from the interactive heatmap", + ) + parser.add_argument( + "-m", + "--min-snv-freq", + action="store", + required=False, + type=float, + default=0.03, + help="Minimum allele frequency of SNV variants to keep (default: 0.03)", + ) + parser.add_argument( + "-M", + "--min-indel-freq", + action="store", + required=False, + type=float, + default=0.1, + help="Minimum allele frequency of indel variants to keep (default: 0.1)", + ) + heatmap_group = parser.add_argument_group("Heatmap options") + _add_heatmap_option_arguments(heatmap_group) + parser.set_defaults(handler=_run_plot_heatmap_command) + + +def _add_plot_subparser(subparsers): + plot_parser = subparsers.add_parser( + "plot", + help="Regenerate plots from existing vartracker outputs", + description="Regenerate selected plots from an existing vartracker results file.", + formatter_class=HelpFormatter, + ) + plot_subparsers = plot_parser.add_subparsers(dest="plot_command") + _add_plot_heatmap_subparser(plot_subparsers) + plot_parser.set_defaults(handler=_run_plot_command, _subparser=plot_parser) + + def create_parser(): """Create and return the top-level argument parser with subcommands.""" @@ -719,6 +946,7 @@ def create_parser(): _add_vcf_subparser(subparsers) _add_bam_subparser(subparsers) _add_e2e_subparser(subparsers) + _add_plot_subparser(subparsers) _add_prep_subparser(subparsers) _add_schema_subparser(subparsers) @@ -903,9 +1131,9 @@ def resolve_path(value: str) -> str: optional_empty={"reads1", "reads2", "bam"}, ) - # Apply passage cap if specified - if args.passage_cap is not None: - input_file = input_file[input_file["sample_number"] <= args.passage_cap] + # Apply sample cap if specified + if args.sample_cap is not None: + input_file = input_file[input_file["sample_number"] <= args.sample_cap] literature = None if args.search_pokay and args.literature_csv is not None: @@ -1056,8 +1284,6 @@ def _add_e2e_subparser(subparsers): aliases=["e2e"], ) - _configure_vcf_parser(e2e_parser, include_input_csv=False) - e2e_parser.add_argument( "samples_csv", nargs="?", @@ -1104,6 +1330,11 @@ def _add_e2e_subparser(subparsers): default=False, ) + _configure_vcf_parser(e2e_parser, include_input_csv=False) + _move_action_group_after( + e2e_parser, "Snakemake options", "Vartracker Analysis Options" + ) + e2e_parser.set_defaults(handler=_run_e2e_command, _subparser=e2e_parser) @@ -1391,6 +1622,88 @@ def _run_bam_command(args): test_context.cleanup() +def _run_plot_heatmap_command(args): + try: + results_csv = Path(args.results_csv).expanduser().resolve() + if not results_csv.exists(): + raise InputValidationError(f"Results CSV not found: {results_csv}") + + outdir = ( + Path(args.outdir).expanduser().resolve() + if args.outdir + else results_csv.parent + ) + outdir.mkdir(parents=True, exist_ok=True) + + table = pd.read_csv(results_csv, keep_default_na=False) + if table.empty: + raise InputValidationError("Results CSV is empty") + if "samples" not in table.columns: + raise InputValidationError( + "Results CSV must contain a 'samples' column to regenerate the heatmap" + ) + + sample_names = [ + token.strip() + for token in str(table.iloc[0]["samples"]).split(" / ") + if token.strip() + ] + if not sample_names: + raise InputValidationError( + "Could not determine sample names from the results CSV" + ) + + project_name = args.name + if project_name is None and "name" in table.columns: + names = [ + str(value).strip() + for value in table["name"].unique() + if str(value).strip() + ] + if len(names) == 1: + project_name = names[0] + if project_name is None: + project_name = "" + + literature_df = None + literature_path = None + if args.literature_csv: + literature_path = str(Path(args.literature_csv).expanduser().resolve()) + try: + literature_df = pd.read_csv(literature_path) + except Exception as exc: + raise InputValidationError( + f"Could not read literature CSV: {exc}" + ) from exc + + cli_command = getattr(args, "_invocation", None) + generate_variant_heatmap( + table, + sample_names, + sample_names, + str(outdir), + project_name, + args.min_snv_freq, + args.min_indel_freq, + literature_hits=literature_df, + literature_table_path=literature_path, + cli_command=cli_command, + **_collect_heatmap_kwargs(args), + ) + print(f"\nFinished: find results in {outdir}\n") + return 0 + except (InputValidationError, ProcessingError) as exc: + print(f"\nERROR: {exc}\n") + return 1 + except Exception as exc: + print(f"\nUnexpected error: {exc}\n") + if getattr(args, "debug", False): + import traceback + + traceback.print_exc() + return 1 + + def _run_generate_command(args): directory = Path(args.dir).expanduser() output_path = Path(args.out).expanduser().resolve() @@ -1608,11 +1921,11 @@ def _process_files( pname, args.min_snv_freq, args.min_indel_freq, - excluded_consequence_types=_parse_csv_option_list(args.heatmap_exclude), gene_lengths=gene_lengths, literature_hits=literature_hits_df, literature_table_path=literature_full_csv_path, cli_command=cli_command, + **_collect_heatmap_kwargs(args), ) # Write specialized tables From f4a2b84d50c3668fa585ec9c2c2f9b5810dbb8ac Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Thu, 26 Mar 2026 12:49:40 +1100 Subject: [PATCH 04/20] Fixing potential results duplication with some samples; made variant_qc more interpretable --- README.md | 3 +- tests/test_analysis.py | 30 ++++- tests/test_main.py | 22 ++++ vartracker/analysis.py | 113 +++++++++++++++++- vartracker/main.py | 26 +++- vartracker/schemas.py | 17 ++- .../test_data/precomputed/test_results.csv | 6 +- vartracker/vcf_processing.py | 29 +++-- 8 files changed, 219 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 0bfa6fb..e50b859 100755 --- a/README.md +++ b/README.md @@ -269,7 +269,8 @@ Heatmap filtering: - `--heatmap-only-new`: only include variants with `variant_status == new`. - `--heatmap-gene-include` and `--heatmap-gene-exclude`: comma-separated gene patterns. - `--heatmap-variant-type`: comma-separated variant-type patterns such as `snp` or `indel`. -- `--heatmap-qc`: comma-separated `overall_variant_qc` patterns to include. +- `--heatmap-qc`: comma-separated `all_samples_pass_qc` patterns to include. Accepted values include `true`, `false`, `pass`, and `fail`. +- `--min-prop-passing-qc`: minimum fraction of samples that must pass per-sample QC. - `--heatmap-min-persistence`: minimum number of included samples in which the variant must be present. - `--heatmap-min-max-af`: minimum maximum allele frequency across included samples. - `--heatmap-min-sample-af`: minimum allele frequency that must be reached in at least one included sample. diff --git a/tests/test_analysis.py b/tests/test_analysis.py index e636338..b52a732 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -3,8 +3,10 @@ from __future__ import annotations import pandas as pd +import pytest from vartracker.analysis import ( + _heatmap_figure_size, search_literature, _prepare_variant_heatmap_matrix, process_joint_variants, @@ -243,7 +245,8 @@ def test_prepare_variant_heatmap_matrix_applies_extended_filters(): "type_of_variant": "snp", "variant_status": "new", "persistence_status": "new_persistent", - "overall_variant_qc": "PASS", + "all_samples_pass_qc": True, + "proportion_samples_passing_qc": 1.0, "presence_absence": "N / Y / Y", "alt_freq": "0.0 / 0.6 / 0.7", "variant_site_depth": "0 / 120 / 125", @@ -259,7 +262,8 @@ def test_prepare_variant_heatmap_matrix_applies_extended_filters(): "type_of_variant": "indel", "variant_status": "new", "persistence_status": "new_persistent", - "overall_variant_qc": "PASS", + "all_samples_pass_qc": True, + "proportion_samples_passing_qc": 0.33, "presence_absence": "N / N / Y", "alt_freq": "0.0 / 0.0 / 0.8", "variant_site_depth": "0 / 0 / 130", @@ -275,7 +279,8 @@ def test_prepare_variant_heatmap_matrix_applies_extended_filters(): "type_of_variant": "snp", "variant_status": "original", "persistence_status": "original_retained", - "overall_variant_qc": "FAIL", + "all_samples_pass_qc": False, + "proportion_samples_passing_qc": 0.67, "presence_absence": "Y / Y / Y", "alt_freq": "0.5 / 0.5 / 0.5", "variant_site_depth": "110 / 115 / 120", @@ -294,7 +299,8 @@ def test_prepare_variant_heatmap_matrix_applies_extended_filters(): only_new=True, gene_include=["S"], variant_type_include=["snp"], - qc_include=["pass"], + qc_include=["true"], + min_prop_passing_qc=0.9, min_persistence=2, min_max_af=0.6, sample_subset=["P1", "P2"], @@ -364,6 +370,7 @@ def test_generate_variant_heatmap_creates_interactive_html(tmp_path, monkeypatch "type_of_change": "missense", "type_of_variant": "snp", "alt_freq": "0.0 / 0.5", + "per_sample_variant_qc": "P / F", "samples": "P0 / P1", "variant": "A22206G", "start": 22206, @@ -377,6 +384,7 @@ def test_generate_variant_heatmap_creates_interactive_html(tmp_path, monkeypatch "type_of_change": "missense", "type_of_variant": "snp", "alt_freq": "0.2 / 0.6", + "per_sample_variant_qc": "P / P", "samples": "P0 / P1", "variant": "C21575T", "start": 21575, @@ -430,8 +438,22 @@ def test_generate_variant_heatmap_creates_interactive_html(tmp_path, monkeypatch assert "Literature results" in content assert "table-scroll" in content assert "heatmap-anchor" in content + assert "cell-qc-fail" in content + assert "AF=0.50, QC=FAIL" in content assert 'data-anchor="s:d215g' in content assert 'data-anchor="nsp6:l37f' in content assert ">10.1000/xyz123" in content assert ".cell:hover .cell-value" in content assert "clearActive" in content + + +def test_heatmap_figure_size_enforces_minimum_row_height(): + width, height = _heatmap_figure_size(6, 2) + + assert width == 4.0 + assert height == 4.0 + + width, height = _heatmap_figure_size(8, 2) + + assert width == 4.0 + assert height == pytest.approx(5.2) diff --git a/tests/test_main.py b/tests/test_main.py index 2df85c5..f05c9a4 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -31,6 +31,25 @@ def test_setup_default_paths_preserves_explicit_values(): assert updated.gff3 == "/tmp/custom.gff3" +def test_drop_exact_duplicate_result_rows_removes_only_exact_duplicates(capsys): + table = pd.DataFrame( + [ + {"variant": "A1C", "gene": "S", "amino_acid_consequence": "S:A1C"}, + {"variant": "A1C", "gene": "S", "amino_acid_consequence": "S:A1C"}, + {"variant": "A1C", "gene": "N", "amino_acid_consequence": "N:A1C"}, + ] + ) + + deduped = main_module._drop_exact_duplicate_result_rows(table) + + assert len(deduped) == 2 + assert deduped.to_dict(orient="records") == [ + {"variant": "A1C", "gene": "S", "amino_acid_consequence": "S:A1C"}, + {"variant": "A1C", "gene": "N", "amino_acid_consequence": "N:A1C"}, + ] + assert "Removed 1 exact duplicate result rows." in capsys.readouterr().out + + def test_prepare_reference_command_invokes_bundle(monkeypatch, tmp_path): recorded = {} @@ -362,6 +381,8 @@ def fake_heatmap(*args, **kwargs): "snp", "--heatmap-qc", "PASS", + "--min-prop-passing-qc", + "0.75", "--heatmap-min-persistence", "2", "--heatmap-min-max-af", @@ -390,6 +411,7 @@ def fake_heatmap(*args, **kwargs): assert recorded["gene_exclude"] == ["N"] assert recorded["variant_type_include"] == ["snp"] assert recorded["qc_include"] == ["PASS"] + assert recorded["min_prop_passing_qc"] == 0.75 assert recorded["min_persistence"] == 2 assert recorded["min_max_af"] == 0.4 assert recorded["min_sample_af"] == 0.3 diff --git a/vartracker/analysis.py b/vartracker/analysis.py index a65272d..5286625 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -15,6 +15,7 @@ import numpy as np import matplotlib.pyplot as plt import matplotlib as mpl +from matplotlib.patches import Rectangle import seaborn as sns from .constants import REF_GENE_LENGTHS, NSP_LENGTHS, NSPS from .core import get_logo @@ -23,6 +24,11 @@ plt.rcdefaults() mpl.rcParams["pdf.fonttype"] = 42 +HEATMAP_MIN_FIG_WIDTH = 4.0 +HEATMAP_COL_WIDTH = 1.2 +HEATMAP_MIN_FIG_HEIGHT = 4.0 +HEATMAP_MIN_ROW_HEIGHT = 0.65 + def _ensure_joint_prefix(change_type: object) -> str: text = str(change_type or "").strip() @@ -43,6 +49,51 @@ def _match_any_pattern(value: object, patterns: Sequence[str]) -> bool: return any(fnmatch.fnmatch(text, pattern) for pattern in patterns) +def _normalise_all_samples_pass_qc(value: object) -> str: + text = str(value or "").strip().lower() + if text in {"pass", "true", "t", "1", "yes", "y"}: + return "true" + if text in {"fail", "false", "f", "0", "no", "n"}: + return "false" + return text + + +def _coerce_proportion_samples_passing_qc(row: object) -> float | None: + direct_value = getattr(row, "proportion_samples_passing_qc", None) + if direct_value not in {None, ""}: + try: + direct_text = str(direct_value).strip() + if direct_text: + return float(direct_text) + except (TypeError, ValueError): + pass + + per_sample = str(getattr(row, "per_sample_variant_qc", "")).strip() + if not per_sample: + return None + tokens = _parse_slash_separated_tokens(per_sample) + if not tokens: + return None + passing = sum(str(token).strip().upper() == "P" for token in tokens) + return passing / len(tokens) + + +def _parse_per_sample_qc_map(row: object) -> dict[str, str]: + sample_tokens = _parse_slash_separated_tokens(getattr(row, "samples", "")) + qc_tokens = _parse_slash_separated_tokens(getattr(row, "per_sample_variant_qc", "")) + return { + sample: str(token).strip().upper() + for sample, token in zip(sample_tokens, qc_tokens) + if str(sample).strip() + } + + +def _heatmap_figure_size(n_rows: int, n_cols: int) -> tuple[float, float]: + width = max(HEATMAP_MIN_FIG_WIDTH, HEATMAP_COL_WIDTH * n_cols) + height = max(HEATMAP_MIN_FIG_HEIGHT, HEATMAP_MIN_ROW_HEIGHT * n_rows) + return width, height + + def process_joint_variants(path): """ Process joint variants from bcftools csq output. @@ -523,6 +574,7 @@ def _prepare_variant_heatmap_matrix( gene_exclude: Sequence[str] | None = None, variant_type_include: Sequence[str] | None = None, qc_include: Sequence[str] | None = None, + min_prop_passing_qc: float | None = None, min_persistence: int | None = None, min_max_af: float | None = None, min_sample_af: float | None = None, @@ -610,6 +662,7 @@ def _prepare_variant_heatmap_matrix( records: List[Dict[str, Union[str, float, int]]] = [] seen_labels = set() + qc_maps: dict[str, dict[str, str]] = {} for row in table.itertuples(index=False): gene_value = getattr(row, "gene", "") @@ -657,10 +710,20 @@ def _prepare_variant_heatmap_matrix( fnmatch.fnmatch(variant_type, pattern) for pattern in variant_type_patterns ): continue - if qc_patterns and not _match_any_pattern( - getattr(row, "overall_variant_qc", ""), list(qc_patterns) + qc_value = _normalise_all_samples_pass_qc( + getattr(row, "all_samples_pass_qc", getattr(row, "overall_variant_qc", "")) + ) + if qc_patterns and not _match_any_pattern(qc_value, list(qc_patterns)): + continue + qc_proportion = _coerce_proportion_samples_passing_qc(row) + if ( + min_prop_passing_qc is not None + and qc_proportion is not None + and qc_proportion < min_prop_passing_qc ): continue + if min_prop_passing_qc is not None and qc_proportion is None: + continue sample_tokens = _parse_slash_separated_tokens(getattr(row, "samples", "")) freq_tokens = [ @@ -699,6 +762,7 @@ def _prepare_variant_heatmap_matrix( sample_presence_map.get(sample, "N") for sample in ordered_samples ] row_depths = [sample_depth_map.get(sample, 0.0) for sample in ordered_samples] + row_qc_map = _parse_per_sample_qc_map(row) max_freq = max(row_values) if row_values else 0.0 @@ -731,6 +795,9 @@ def _prepare_variant_heatmap_matrix( records.append(record) seen_labels.add(base_label) + qc_maps[display_label] = { + sample: row_qc_map.get(sample, "") for sample in ordered_samples + } if not records: return pd.DataFrame(columns=ordered_samples) @@ -751,6 +818,9 @@ def _prepare_variant_heatmap_matrix( matrix_df = matrix_df.reindex(columns=ordered_samples).fillna(0.0) matrix_df.attrs["base_labels"] = base_label_map matrix_df.attrs["canonical_labels"] = canonical_label_map + matrix_df.attrs["qc_by_label"] = { + label: qc_maps.get(label, {}) for label in matrix_df.index + } return matrix_df @@ -876,6 +946,7 @@ def _write_interactive_heatmap_html( label_map: Dict[str, str] = matrix.attrs.get("base_labels", {}) canonical_label_map: Dict[str, str] = matrix.attrs.get("canonical_labels", {}) + qc_by_label: Dict[str, Dict[str, str]] = matrix.attrs.get("qc_by_label", {}) if ( literature_df is None @@ -933,17 +1004,22 @@ def _frequency_to_color(value: float) -> tuple[str, str]: grid_cells.append(header_html) row_values = matrix.loc[label] + row_qc = qc_by_label.get(label, {}) for sample, value in zip(x_labels, row_values): freq = float(value) if value is not None else 0.0 color, text_value = _frequency_to_color(freq) + qc_value = row_qc.get(sample, "") + qc_failed = qc_value == "F" tooltip = html.escape( f"Variant: {label.replace(chr(10), ' ')} • " f"Sample: {sample} • " - f"Allele frequency: {freq:.3f}" + f"AF={freq:.2f}, QC={'FAIL' if qc_failed else 'PASS'}" ) classes = ["cell"] if not text_value: classes.append("cell-empty") + if qc_failed: + classes.append("cell-qc-fail") cell_html = ( f'
tuple[str, str]: .cell-empty {{ color: #1f2937; }} + .cell-qc-fail {{ + box-shadow: inset 0 0 0 2px #111827; + }} .cell:hover {{ transform: scale(1.03); z-index: 2; @@ -1237,6 +1316,7 @@ def generate_variant_heatmap( gene_exclude: Sequence[str] | None = None, variant_type_include: Sequence[str] | None = None, qc_include: Sequence[str] | None = None, + min_prop_passing_qc: float | None = None, min_persistence: int | None = None, min_max_af: float | None = None, min_sample_af: float | None = None, @@ -1265,6 +1345,7 @@ def generate_variant_heatmap( gene_exclude, variant_type_include, qc_include, + min_prop_passing_qc, min_persistence, min_max_af, min_sample_af, @@ -1279,10 +1360,14 @@ def generate_variant_heatmap( heatmap_path = os.path.join(outdir, "variant_allele_frequency_heatmap.pdf") - fig_width = max(4, 1.2 * len(heatmap_data.columns)) - fig_height = max(4, 0.4 * len(heatmap_data)) + fig_width, fig_height = _heatmap_figure_size( + len(heatmap_data.index), len(heatmap_data.columns) + ) fig, ax = plt.subplots(figsize=(fig_width, fig_height)) plot_data = heatmap_data.replace(0, np.nan) + qc_by_label: Dict[str, Dict[str, str]] = heatmap_data.attrs.get( + "qc_by_label", {} + ) cmap = sns.color_palette("viridis", as_cmap=True).copy() cmap.set_bad(color="#d9d9d9") sns.heatmap( @@ -1296,6 +1381,23 @@ def generate_variant_heatmap( ax=ax, ) + for row_idx, label in enumerate(heatmap_data.index): + row_qc = qc_by_label.get(label, {}) + for col_idx, sample in enumerate(heatmap_data.columns): + if row_qc.get(sample, "") != "F": + continue + ax.add_patch( + Rectangle( + (col_idx, row_idx), + 1, + 1, + fill=False, + edgecolor="black", + linewidth=1.4, + zorder=10, + ) + ) + heatmap_title = ( f"{project_name}: variant allele frequencies" if project_name @@ -1305,6 +1407,7 @@ def generate_variant_heatmap( tick_labels = list(sample_names) ax.set_xticklabels(tick_labels, rotation=45, ha="right") + ax.tick_params(axis="y", labelsize=10) ax.set_xlabel("Sample", fontweight="bold") ax.set_ylabel("Variant (5' → 3')", fontweight="bold") diff --git a/vartracker/main.py b/vartracker/main.py index 3dbf9d7..f48e20d 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -174,7 +174,17 @@ def _add_heatmap_option_arguments(group: argparse._ArgumentGroup) -> None: action="store", required=False, default="", - help="Comma-separated overall QC patterns to include (e.g. PASS)", + help=( + "Comma-separated all-samples QC patterns to include " + "(e.g. true,false,pass,fail)" + ), + ) + group.add_argument( + "--min-prop-passing-qc", + action="store", + type=float, + default=None, + help="Minimum proportion of samples that must pass per-sample QC (0-1)", ) group.add_argument( "--heatmap-min-persistence", @@ -240,6 +250,7 @@ def _collect_heatmap_kwargs(args) -> dict[str, object]: getattr(args, "heatmap_variant_type", "") ), "qc_include": _parse_csv_option_list(getattr(args, "heatmap_qc", "")), + "min_prop_passing_qc": getattr(args, "min_prop_passing_qc", None), "min_persistence": getattr(args, "heatmap_min_persistence", None), "min_max_af": getattr(args, "heatmap_min_max_af", None), "min_sample_af": getattr(args, "heatmap_min_sample_af", None), @@ -989,6 +1000,14 @@ def _normalise_rulegraph_path(path: str | None) -> str | None: return str(Path(resolved).resolve()) +def _drop_exact_duplicate_result_rows(table: pd.DataFrame) -> pd.DataFrame: + deduped = table.drop_duplicates().reset_index(drop=True) + removed = len(table) - len(deduped) + if removed: + print(f"Removed {removed} exact duplicate result rows.") + return deduped + + def main(sysargs=None): """Entry point for the vartracker CLI.""" if sysargs is None: @@ -1874,6 +1893,8 @@ def _process_files( else: pname = "" + table = _drop_exact_duplicate_result_rows(table) + # Write initial results outfile = os.path.join(args.outdir, args.filename) table.fillna("").to_csv(outfile, index=None) @@ -1884,6 +1905,9 @@ def _process_files( else: table = pd.read_csv(outfile, keep_default_na=False) + table = _drop_exact_duplicate_result_rows(table) + table.fillna("").to_csv(outfile, index=None) + literature_hits_df = None literature_full_csv_path = None if ( diff --git a/vartracker/schemas.py b/vartracker/schemas.py index 09e8edd..e0c7fb3 100644 --- a/vartracker/schemas.py +++ b/vartracker/schemas.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any, Iterable, Sequence -RESULTS_SCHEMA_VERSION = "1.0" +RESULTS_SCHEMA_VERSION = "1.1" RESULTS_SCHEMA: list[dict[str, str]] = [ { @@ -139,11 +139,18 @@ "values": "", }, { - "name": "overall_variant_qc", - "type": "string", - "description": "Aggregated QC status across samples.", + "name": "all_samples_pass_qc", + "type": "boolean", + "description": "True if every sample passes per-sample variant QC.", "units": "", - "values": "PASS, FAIL", + "values": "true, false", + }, + { + "name": "proportion_samples_passing_qc", + "type": "number", + "description": "Proportion of samples passing per-sample variant QC.", + "units": "fraction", + "values": "0-1", }, { "name": "per_sample_variant_qc", diff --git a/vartracker/test_data/precomputed/test_results.csv b/vartracker/test_data/precomputed/test_results.csv index e440247..0f86033 100644 --- a/vartracker/test_data/precomputed/test_results.csv +++ b/vartracker/test_data/precomputed/test_results.csv @@ -1,3 +1,3 @@ -chrom,start,end,gene,ref,alt,variant,amino_acid_consequence,nsp_aa_change,bcsq_nt_notation,bcsq_aa_notation,type_of_variant,type_of_change,variant_status,persistence_status,presence_absence,first_appearance,last_appearance,overall_variant_qc,per_sample_variant_qc,aa1_total_properties,aa2_total_properties,aa1_unique_properties,aa2_unique_properties,aa1_weight,aa2_weight,weight_difference,alt_freq,variant_depth,variant_site_depth,variant_window_depth,samples,total_genome_coverage -NC_045512.2,21563,21563,S,A,T,A21563T,S:D614V,NA,NA,NA,snp,missense,new,new_persistent,N / N / Y / Y / Y / Y / Y / Y / Y / Y,2,9,PASS,PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS,NA,NA,NA,NA,89.09,117.15,28.06,0.0 / 0.05 / 0.15 / 0.25 / 0.35 / 0.45 / 0.55 / 0.65 / 0.75 / 0.85,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 -NC_045512.2,23403,23403,S,A,G,A23403G,S:P681R,NA,NA,NA,snp,missense,new,new_persistent,N / Y / Y / Y / Y / Y / Y / Y / Y / Y,1,9,PASS,PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS / PASS,NA,NA,NA,NA,89.09,156.19,67.1,0.0 / 0.10 / 0.20 / 0.30 / 0.40 / 0.50 / 0.60 / 0.70 / 0.80 / 0.90,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 +chrom,start,end,gene,ref,alt,variant,amino_acid_consequence,nsp_aa_change,bcsq_nt_notation,bcsq_aa_notation,type_of_variant,type_of_change,variant_status,persistence_status,presence_absence,first_appearance,last_appearance,all_samples_pass_qc,proportion_samples_passing_qc,per_sample_variant_qc,aa1_total_properties,aa2_total_properties,aa1_unique_properties,aa2_unique_properties,aa1_weight,aa2_weight,weight_difference,alt_freq,variant_depth,variant_site_depth,variant_window_depth,samples,total_genome_coverage +NC_045512.2,21563,21563,S,A,T,A21563T,S:D614V,NA,NA,NA,snp,missense,new,new_persistent,N / N / Y / Y / Y / Y / Y / Y / Y / Y,2,9,True,1.0,P / P / P / P / P / P / P / P / P / P,NA,NA,NA,NA,89.09,117.15,28.06,0.0 / 0.05 / 0.15 / 0.25 / 0.35 / 0.45 / 0.55 / 0.65 / 0.75 / 0.85,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 +NC_045512.2,23403,23403,S,A,G,A23403G,S:P681R,NA,NA,NA,snp,missense,new,new_persistent,N / Y / Y / Y / Y / Y / Y / Y / Y / Y,1,9,True,1.0,P / P / P / P / P / P / P / P / P / P,NA,NA,NA,NA,89.09,156.19,67.1,0.0 / 0.10 / 0.20 / 0.30 / 0.40 / 0.50 / 0.60 / 0.70 / 0.80 / 0.90,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index 094bebb..f8a375a 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -581,7 +581,13 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): persistent_status = "unknown" depths_qc = calculate_variant_site_depths(cov_df, v, samples, min_depth) - overall_variant_qc = "FAIL" if "F" in depths_qc["variant_qc"] else "PASS" + all_samples_pass_qc = "F" not in depths_qc["variant_qc"] + proportion_samples_passing_qc = ( + sum(flag == "P" for flag in depths_qc["variant_qc"]) + / len(depths_qc["variant_qc"]) + if depths_qc["variant_qc"] + else 0.0 + ) first_appearance = ( samples[presence_absence.index("Y")] if "Y" in presence_absence else "None" ) @@ -604,7 +610,8 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): presence_absence, first_appearance, last_appearance, - overall_variant_qc, + all_samples_pass_qc, + proportion_samples_passing_qc, depths_qc, allele_freqs, samples, @@ -620,7 +627,8 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): presence_absence, first_appearance, last_appearance, - overall_variant_qc, + all_samples_pass_qc, + proportion_samples_passing_qc, depths_qc, allele_freqs, samples, @@ -639,7 +647,8 @@ def _process_annotation( presence_absence, first_appearance, last_appearance, - overall_variant_qc, + all_samples_pass_qc, + proportion_samples_passing_qc, depths_qc, allele_freqs, samples, @@ -667,7 +676,8 @@ def _process_annotation( "presence_absence": " / ".join(presence_absence), "first_appearance": first_appearance, "last_appearance": last_appearance, - "overall_variant_qc": overall_variant_qc, + "all_samples_pass_qc": all_samples_pass_qc, + "proportion_samples_passing_qc": proportion_samples_passing_qc, "per_sample_variant_qc": " / ".join(depths_qc["variant_qc"]), "aa1_total_properties": anno[0], "aa2_total_properties": anno[0], @@ -712,7 +722,8 @@ def _process_annotation( "presence_absence": " / ".join(presence_absence), "first_appearance": first_appearance, "last_appearance": last_appearance, - "overall_variant_qc": overall_variant_qc, + "all_samples_pass_qc": all_samples_pass_qc, + "proportion_samples_passing_qc": proportion_samples_passing_qc, "per_sample_variant_qc": " / ".join(depths_qc["variant_qc"]), "aa1_total_properties": ";".join(aa_exploration.aa1_total_properties), "aa2_total_properties": ";".join(aa_exploration.aa2_total_properties), @@ -737,7 +748,8 @@ def _create_unannotated_result( presence_absence, first_appearance, last_appearance, - overall_variant_qc, + all_samples_pass_qc, + proportion_samples_passing_qc, depths_qc, allele_freqs, samples, @@ -770,7 +782,8 @@ def _create_unannotated_result( "presence_absence": " / ".join(presence_absence), "first_appearance": first_appearance, "last_appearance": last_appearance, - "overall_variant_qc": overall_variant_qc, + "all_samples_pass_qc": all_samples_pass_qc, + "proportion_samples_passing_qc": proportion_samples_passing_qc, "per_sample_variant_qc": " / ".join(depths_qc["variant_qc"]), "aa1_total_properties": "None", "aa2_total_properties": "None", From 9b785b1fa8efd2b13f1cf4598e3c0629722d05b4 Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Thu, 26 Mar 2026 14:31:20 +1100 Subject: [PATCH 05/20] SOme new plot options added --- README.md | 47 ++ pytest_results.txt | 28 + tests/test_main.py | 167 +++++ tests/test_plotting.py | 166 +++++ vartracker/analysis.py | 4 +- vartracker/main.py | 468 +++++++++++++ vartracker/plotting.py | 651 ++++++++++++++++++ vartracker/schemas.py | 7 + .../test_data/precomputed/test_results.csv | 6 +- 9 files changed, 1540 insertions(+), 4 deletions(-) create mode 100644 pytest_results.txt create mode 100644 tests/test_plotting.py create mode 100644 vartracker/plotting.py diff --git a/README.md b/README.md index e50b859..74b6295 100755 --- a/README.md +++ b/README.md @@ -207,6 +207,18 @@ vartracker plot heatmap results/results.csv \ --heatmap-aa-exclude "*frameshift*" \ --outdir results/replots +# Plot whole-dataset turnover from an existing results file +vartracker plot turnover results/results.csv + +# Plot selected variant trajectories from an existing results file +vartracker plot trajectory results/results.csv \ + --variants "S:D614G,S:E484K,S:N501Y" + +# Plot takeover-style trajectories using AF thresholds +vartracker plot trajectory results/results.csv \ + --thresholds 0.5,0.9 \ + --crossing-only + # Generate a template spreadsheet for a directory of files vartracker prepare spreadsheet --mode e2e --dir data/passaging --out inputs.csv @@ -260,6 +272,9 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. - `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. - `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV. +- `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. +- `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. +- `vartracker plot lifespan` – plot first-to-last detection spans for a selected or auto-ranked subset of variants. Heatmap filtering: - By default, all consequence classes are included except joint variants. Use `--heatmap-include-joint` to show joint variants. @@ -278,6 +293,38 @@ Heatmap filtering: - `--heatmap-hide-singletons`: hide variants present in only one included sample. - `--heatmap-min-depth`: minimum site depth a variant must reach in at least one included sample. - Example: `--heatmap-aa-exclude "synonymous,*frameshift*,stop_gained"` + +Standalone plot filtering: +- `--gene`, `--effect`, `--min-af`, `--max-af`: restrict the plotted result set before ranking/selection. +- `--variants` or `--variant-file`: explicitly choose variants and preserve that order. +- `--sample-min`, `--sample-max`: restrict the passage/sample-number window. +- `--persistent-only` and `--new-only`: keep only persistent new variants or only variants with `variant_status == new`. +- `trajectory` and `lifespan` auto-select a limited subset by default (`--top-n`) to stay readable. +- `turnover` uses all filtered variants by default and is also written automatically during the main `vcf`/`bam`/`end-to-end` workflows as `variant_turnover_plot.pdf`. + +Standalone plot output: +- `--out`: write to an exact file path. +- `--outdir`: write beside the results CSV or into the chosen directory using deterministic names such as `variant_trajectory_plot.pdf`. +- `--format`: choose `pdf`, `png`, or `svg`. +- `--dpi`: set raster output resolution. + +Trajectory threshold mode: +- `--thresholds`: draw horizontal AF threshold lines, e.g. `0.5,0.9`. +- `--crossing-only`: keep only variants crossing at least one supplied threshold. +- `--label-threshold-crossers`: label only threshold-crossing variants to reduce clutter. +- `--crossing-rule`: choose whether threshold equality counts (`at_or_above`) or requires a strict exceedance (`strictly_above`). + +Standalone plot examples: +- `vartracker plot turnover results.csv` +- `vartracker plot trajectory results.csv --variants "S:D614G,S:E484K"` +- `vartracker plot trajectory results.csv --thresholds 0.5,0.9` +- `vartracker plot trajectory results.csv --thresholds 0.5,0.9 --crossing-only` +- `vartracker plot trajectory results.csv --thresholds 0.5,0.9 --crossing-only --label-threshold-crossers` +- `vartracker plot lifespan results.csv --top-n 20 --persistent-only` + +Note: +- The standalone `plot` commands require `results.csv` files written by current vartracker versions, which now include a slash-separated `sample_number` column for stable passage ordering. + - `vartracker prepare spreadsheet` – specify `--mode` (`vcf`, `bam`, or `e2e`), `--dir` to scan, `--out` for the CSV, and `--dry-run` to preview without writing a file. - `vartracker prepare reference` – build a merged FASTA/GFF3 bundle from GenBank nucleotide accessions. diff --git a/pytest_results.txt b/pytest_results.txt new file mode 100644 index 0000000..a7837ad --- /dev/null +++ b/pytest_results.txt @@ -0,0 +1,28 @@ +........................................................................ [ 87%] +.......... [100%] +================================ tests coverage ================================ +_______________ coverage: platform darwin, python 3.11.8-final-0 _______________ + +Name Stmts Miss Cover Missing +------------------------------------------------------------------- +vartracker/__init__.py 7 2 71% 17-19 +vartracker/_version.py 7 0 100% +vartracker/amino_acids.py 36 30 17% 20-94, 98-146, 150, 154 +vartracker/analysis.py 722 231 68% 36, 43, 50, 59, 70-71, 78, 111-112, 121, 158-160, 176, 189-234, 250-353, 365-400, 457, 462, 468, 476, 480, 486, 488, 491, 494-495, 521-526, 533-536, 554, 560, 590, 593, 597, 612, 616-625, 681, 694, 703, 708, 714, 719, 726, 728, 743, 759-760, 778, 783, 785, 787, 805, 834, 837, 843, 849, 859, 869, 889, 909, 919, 942, 947, 958-961, 1001, 1360-1361, 1433-1434, 1454-1457, 1466, 1476-1477, 1483, 1496-1501, 1504-1510, 1520-1521, 1526-1532, 1536-1560, 1565-1566, 1570-1594, 1599-1621, 1633-1634, 1644, 1650-1655, 1693 +vartracker/analysis_launcher.py 72 61 15% 24, 45-136, 140-155 +vartracker/annotation_processing.py 120 110 8% 10-14, 18-27, 33-42, 48-55, 68-200 +vartracker/constants.py 85 71 16% 149-176, 191-246, 260-280 +vartracker/core.py 62 10 84% 90, 95, 98-100, 113-118 +vartracker/data/__init__.py 8 0 100% +vartracker/data/parse_pokay.py 141 26 82% 69-70, 110, 127-130, 155-156, 162, 168, 234, 253-276, 280 +vartracker/generate.py 121 13 89% 54, 74, 84, 99, 115, 119, 132, 143, 156, 167, 202, 205, 207 +vartracker/main.py 939 277 71% 125, 404, 410, 422, 428-433, 439, 443, 476, 479-480, 484-489, 497, 515, 521-529, 534-572, 577-605, 616, 759-760, 796-802, 806-809, 960-963, 982-984, 1041-1044, 1090, 1140-1143, 1166-1168, 1209-1211, 1450-1453, 1467, 1500-1504, 1506-1507, 1528, 1543-1544, 1548-1549, 1560-1561, 1566-1571, 1588, 1592, 1608, 1612, 1617-1625, 1643-1666, 1719, 1721, 1731-1749, 1830-1834, 1851-1855, 1861-1862, 1867, 1888, 1906-1911, 1922, 1939-1952, 1972-1976, 1993-1997, 2001-2002, 2007, 2030, 2048-2053, 2064, 2081-2094, 2104, 2115, 2117, 2127, 2141, 2146-2150, 2170-2179, 2183-2227, 2254-2260, 2289, 2299-2304, 2366, 2438 +vartracker/plotting.py 297 54 82% 29-37, 43, 46, 52, 58, 62, 69, 71, 76-77, 85, 87, 97, 117, 120, 122, 163, 225, 241, 246, 251, 281, 318, 320, 334, 369-389, 407, 436, 456-457, 564, 572, 608, 626-627, 650 +vartracker/provenance.py 166 52 69% 45, 63-73, 77-87, 97-115, 127-128, 140, 158, 169-170, 173, 185-186, 245, 272-273, 280, 282, 307 +vartracker/reference_prepare.py 308 44 86% 82, 93, 104, 106, 114, 121-146, 155, 162-164, 170, 182-183, 204, 212, 218, 294, 353, 357, 364, 368, 376, 406, 414, 460, 486, 498, 517, 548, 591, 596, 601, 609-610 +vartracker/schemas.py 60 7 88% 270-298, 303 +vartracker/vcf_processing.py 337 308 9% 23-25, 33-112, 124, 136-140, 146-152, 180-353, 366-394, 409-484, 498-639, 658-706, 759-766 +------------------------------------------------------------------- +TOTAL 3488 1296 63% +Coverage XML written to file coverage.xml +82 passed in 42.73s diff --git a/tests/test_main.py b/tests/test_main.py index f05c9a4..d9bde75 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -461,6 +461,173 @@ def fake_heatmap(*args, **kwargs): assert recorded["kwargs"]["include_joint"] is True +def _write_plot_results_csv(path: Path, n_variants: int = 4) -> None: + rows = [ + "samples,sample_number,name,alt_freq,per_sample_variant_qc,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,amino_acid_consequence,nsp_aa_change,start,reference" + ] + for idx in range(n_variants): + gene = "S" if idx % 2 == 0 else "N" + aa = f"{gene}:V{idx + 1}A" + effect = "missense" if idx % 3 else "synonymous" + persistence = "new_persistent" if idx % 2 == 0 else "new_transient" + reference = "PMID123" if idx == 1 else "" + rows.append( + "P0 / P1 / P2 / P3," + "0 / 1 / 2 / 3," + "Example," + f"0.0 / 0.{idx + 2} / 0.{idx + 3} / 0.{idx + 4}," + "P / P / P / P," + "N / Y / Y / Y," + "new," + f"{persistence}," + "snp," + f"{effect}," + f"{gene}," + f"A{100 + idx}G," + f"{aa},," + f"{100 + idx}," + f"{reference}" + ) + path.write_text("\n".join(rows) + "\n", encoding="utf-8") + + +def test_plot_trajectory_variants_override_auto_selection(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + + recorded = {} + + def fake_plot(*args, **kwargs): + recorded["selected_variants"] = kwargs["selected_variants"] + + monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) + + exit_code = main_module.main( + [ + "plot", + "trajectory", + str(results_csv), + "--variants", + "S:V3A,N:V2A", + ] + ) + + assert exit_code == 0 + assert recorded["selected_variants"] == ["S:V3A", "N:V2A"] + + +def test_plot_trajectory_threshold_options_are_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + + recorded = {} + + def fake_plot(*args, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) + + exit_code = main_module.main( + [ + "plot", + "trajectory", + str(results_csv), + "--thresholds", + "0.5,0.9", + "--label-threshold-crossers", + "--crossing-rule", + "strictly_above", + ] + ) + + assert exit_code == 0 + assert recorded["thresholds"] == [0.5, 0.9] + assert recorded["label_threshold_crossers"] is True + assert recorded["crossing_rule"] == "strictly_above" + + +def test_plot_trajectory_crossing_only_requires_thresholds(tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + + exit_code = main_module.main( + ["plot", "trajectory", str(results_csv), "--crossing-only"] + ) + + assert exit_code == 1 + + +def test_plot_trajectory_crossing_only_filters_variants(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + + recorded = {} + + def fake_plot(summary, long_df, **kwargs): + recorded["selected_variants"] = kwargs["selected_variants"] + recorded["summary_variant_ids"] = list(summary["variant_id"]) + + monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) + + exit_code = main_module.main( + [ + "plot", + "trajectory", + str(results_csv), + "--thresholds", + "0.5", + "--crossing-only", + ] + ) + + assert exit_code == 0 + assert all( + variant in recorded["summary_variant_ids"] + for variant in recorded["selected_variants"] + ) + + +def test_plot_turnover_uses_all_filtered_variants_by_default(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=5) + + recorded = {} + + def fake_turnover(summary, long_df, **kwargs): + recorded["summary_count"] = len(summary) + recorded["variant_ids"] = list(summary["variant_id"]) + + monkeypatch.setattr(main_module, "plot_variant_turnover", fake_turnover) + + exit_code = main_module.main(["plot", "turnover", str(results_csv), "--gene", "S"]) + + assert exit_code == 0 + assert recorded["summary_count"] == 3 + assert all(variant.startswith("S:") for variant in recorded["variant_ids"]) + + +def test_plot_commands_auto_limit_variants_by_top_n(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=20) + + recorded = {} + + def fake_trajectory(*args, **kwargs): + recorded["trajectory"] = kwargs["selected_variants"] + + def fake_lifespan(*args, **kwargs): + recorded["lifespan"] = kwargs["selected_variants"] + + monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_trajectory) + monkeypatch.setattr(main_module, "plot_variant_lifespan", fake_lifespan) + + assert main_module.main(["plot", "trajectory", str(results_csv)]) == 0 + assert main_module.main(["plot", "lifespan", str(results_csv)]) == 0 + + assert len(recorded["trajectory"]) == main_module.DEFAULT_TRAJECTORY_TOP_N + assert len(recorded["lifespan"]) == main_module.DEFAULT_LIFESPAN_TOP_N + + def test_e2e_runs_snakemake_then_vcf(monkeypatch, tmp_path): updated_csv = tmp_path / "samples_updated.csv" vcf_out = tmp_path / "vcf.gz" diff --git a/tests/test_plotting.py b/tests/test_plotting.py new file mode 100644 index 0000000..1246889 --- /dev/null +++ b/tests/test_plotting.py @@ -0,0 +1,166 @@ +"""Tests for standalone plotting helpers.""" + +from __future__ import annotations + +import pandas as pd + +from vartracker.plotting import ( + apply_shared_plot_filters, + auto_select_variants, + get_threshold_crossing_variants, + plot_variant_lifespan, + plot_variant_trajectory, + plot_variant_turnover, + prepare_plot_inputs, + variant_crosses_thresholds, +) + + +def _results_table() -> pd.DataFrame: + return pd.DataFrame( + [ + { + "gene": "S", + "variant": "A23403G", + "amino_acid_consequence": "S:D614G", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "N / Y / Y / Y", + "alt_freq": "0.0 / 0.20 / 0.45 / 0.70", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "PMID123", + }, + { + "gene": "S", + "variant": "G23012A", + "amino_acid_consequence": "S:E484K", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_transient", + "presence_absence": "N / N / Y / N", + "alt_freq": "0.0 / 0.0 / 0.40 / 0.0", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + { + "gene": "S", + "variant": "A23063T", + "amino_acid_consequence": "S:N501Y", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "original", + "persistence_status": "original_retained", + "presence_absence": "Y / Y / Y / Y", + "alt_freq": "0.55 / 0.58 / 0.60 / 0.62", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + { + "gene": "N", + "variant": "C28977T", + "amino_acid_consequence": "N:S202=", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "synonymous", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "N / Y / Y / Y", + "alt_freq": "0.0 / 0.30 / 0.35 / 0.40", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + ] + ) + + +def test_prepare_plot_inputs_and_filters(): + summary, long_df, sample_names, sample_numbers = prepare_plot_inputs(_results_table()) + + assert sample_names == ["P0", "P1", "P2", "P3"] + assert sample_numbers == [0, 1, 2, 3] + assert set(summary["variant_id"]) >= {"S:D614G", "S:E484K", "S:N501Y", "N:N202="} + + filtered_summary, filtered_long = apply_shared_plot_filters( + summary, + long_df, + gene="S", + effects=["missense"], + min_af=0.5, + sample_min=1, + sample_max=3, + ) + + assert set(filtered_summary["variant_id"]) == {"S:D614G", "S:N501Y"} + assert filtered_long["sample_number"].min() == 1 + assert filtered_long["sample_number"].max() == 3 + + +def test_auto_select_variants_prefers_literature_and_persistent(): + summary, _, _, _ = prepare_plot_inputs(_results_table()) + + selected = auto_select_variants(summary, top_n=2) + + assert selected[0] == "S:D614G" + assert len(selected) == 2 + + +def test_plot_functions_create_output_files(tmp_path, monkeypatch): + mpl_dir = tmp_path / "mpl" + mpl_dir.mkdir() + monkeypatch.setenv("MPLCONFIGDIR", str(mpl_dir)) + + summary, long_df, _, _ = prepare_plot_inputs(_results_table()) + selected = ["S:D614G", "S:N501Y"] + + trajectory = tmp_path / "trajectory.pdf" + turnover = tmp_path / "turnover.pdf" + lifespan = tmp_path / "lifespan.pdf" + trajectory_threshold = tmp_path / "trajectory_threshold.pdf" + + plot_variant_trajectory( + summary, + long_df, + selected_variants=selected, + output_path=trajectory, + ) + plot_variant_turnover(summary, long_df, output_path=turnover) + plot_variant_lifespan( + summary, + selected_variants=selected, + output_path=lifespan, + ) + plot_variant_trajectory( + summary, + long_df, + selected_variants=selected, + output_path=trajectory_threshold, + thresholds=[0.5, 0.9], + ) + + assert trajectory.exists() + assert turnover.exists() + assert lifespan.exists() + assert trajectory_threshold.exists() + + +def test_threshold_crossing_helpers(): + assert variant_crosses_thresholds([0.1, 0.5], [0.5]) is True + assert variant_crosses_thresholds([0.1, 0.5], [0.5], crossing_rule="strictly_above") is False + + summary, long_df, _, _ = prepare_plot_inputs(_results_table()) + crossed = get_threshold_crossing_variants(summary, long_df, [0.5]) + assert set(crossed["variant_id"]) == {"S:D614G", "S:N501Y"} diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 5286625..1db9e75 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -39,7 +39,9 @@ def _ensure_joint_prefix(change_type: object) -> str: def _parse_slash_separated_tokens(value: object) -> list[str]: - return [token.strip() for token in str(value or "").split(" / ")] + if value is None: + return [] + return [token.strip() for token in str(value).split(" / ")] def _match_any_pattern(value: object, patterns: Sequence[str]) -> bool: diff --git a/vartracker/main.py b/vartracker/main.py index f48e20d..e8204b1 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -16,6 +16,7 @@ from contextlib import ExitStack from importlib import resources from pathlib import Path +from typing import Sequence import pandas as pd @@ -40,6 +41,22 @@ generate_variant_heatmap, search_literature, ) +from .plotting import ( + DEFAULT_LIFESPAN_TOP_N, + DEFAULT_TRAJECTORY_TOP_N, + apply_shared_plot_filters, + auto_select_variants, + collect_explicit_variants, + get_threshold_crossing_variants, + load_results_table, + parse_thresholds, + plot_variant_lifespan, + plot_variant_trajectory, + plot_variant_turnover, + prepare_plot_inputs, + _project_name_from_results, + resolve_plot_output_path, +) from ._version import __version__ from .provenance import ( RunManifest, @@ -262,6 +279,150 @@ def _collect_heatmap_kwargs(args) -> dict[str, object]: } +def _add_shared_plot_filter_arguments(group: argparse._ArgumentGroup) -> None: + group.add_argument("--gene", default=None, help="Restrict to a single gene") + group.add_argument( + "--effect", + default="", + help="Comma-separated effect classes to include (e.g. missense,synonymous)", + ) + group.add_argument( + "--min-af", + type=float, + default=None, + help="Minimum max allele frequency across included samples", + ) + group.add_argument( + "--max-af", + type=float, + default=None, + help="Maximum max allele frequency across included samples", + ) + group.add_argument( + "--variants", + default="", + help="Comma-separated variant identifiers to plot, in order", + ) + group.add_argument( + "--variant-file", + default=None, + help="Path to a file with one variant identifier per line", + ) + group.add_argument( + "--sample-min", + type=int, + default=None, + help="Minimum sample_number to include", + ) + group.add_argument( + "--sample-max", + type=int, + default=None, + help="Maximum sample_number to include", + ) + group.add_argument( + "--include-synonymous", + action="store_true", + default=True, + help="Preserve synonymous variants in standalone plot filtering", + ) + group.add_argument( + "--persistent-only", + action="store_true", + default=False, + help="Only include variants with persistence_status == new_persistent", + ) + group.add_argument( + "--new-only", + action="store_true", + default=False, + help="Only include variants with variant_status == new", + ) + + +def _add_plot_output_arguments(parser: argparse.ArgumentParser) -> None: + output_group = parser.add_argument_group("Output") + output_group.add_argument( + "--out", default=None, help="Write the plot to this exact path" + ) + output_group.add_argument( + "--outdir", + default=None, + help="Output directory for plot files (default: beside results.csv)", + ) + output_group.add_argument( + "--format", + choices=["pdf", "png", "svg"], + default="pdf", + help="Output format when using --outdir or default naming (default: pdf)", + ) + output_group.add_argument( + "--dpi", + type=int, + default=300, + help="Output DPI (default: 300)", + ) + + +def _collect_shared_plot_kwargs(args) -> dict[str, object]: + return { + "gene": getattr(args, "gene", None), + "effects": _parse_csv_option_list(getattr(args, "effect", "")), + "min_af": getattr(args, "min_af", None), + "max_af": getattr(args, "max_af", None), + "variants": collect_explicit_variants( + getattr(args, "variants", ""), getattr(args, "variant_file", None) + ), + "sample_min": getattr(args, "sample_min", None), + "sample_max": getattr(args, "sample_max", None), + "include_synonymous": getattr(args, "include_synonymous", True), + "persistent_only": getattr(args, "persistent_only", False), + "new_only": getattr(args, "new_only", False), + } + + +def _resolve_plot_variants( + summary: pd.DataFrame, + *, + explicit_variants: Sequence[str], + top_n: int, + prefer_crossing: bool = False, + thresholds: Sequence[float] | None = None, +) -> list[str]: + if explicit_variants: + selected = [] + seen = set() + for requested in explicit_variants: + match = summary[ + summary.apply( + lambda row: requested + in {row["variant_id"], row["variant_label"], row["variant_name"]}, + axis=1, + ) + ] + if match.empty: + continue + variant_id = str(match.iloc[0]["variant_id"]) + if variant_id not in seen: + selected.append(variant_id) + seen.add(variant_id) + if not selected: + raise ProcessingError( + "None of the requested variants were found in the filtered results" + ) + return selected + + selected = auto_select_variants( + summary, + top_n=top_n, + prefer_crossing=prefer_crossing, + thresholds=thresholds, + ) + if not selected: + raise ProcessingError("No variants were available for plotting") + return selected + + def _print_dependency_error(error: DependencyError) -> None: """Render a dependency error with optional remediation tips.""" message = str(error) @@ -883,6 +1044,173 @@ def _run_plot_command(args): return args.handler(args) +def _prepare_standalone_plot_inputs(args): + results_csv = Path(args.results_csv).expanduser().resolve() + table = load_results_table(results_csv) + summary, long_df, sample_names, sample_numbers = prepare_plot_inputs(table) + shared_kwargs = _collect_shared_plot_kwargs(args) + summary, long_df = apply_shared_plot_filters(summary, long_df, **shared_kwargs) + project_name = _project_name_from_results(table, getattr(args, "name", None)) + return ( + results_csv, + table, + summary, + long_df, + sample_names, + sample_numbers, + project_name, + shared_kwargs, + ) + + +def _run_plot_trajectory_command(args): + try: + ( + results_csv, + _table, + summary, + long_df, + _sample_names, + _sample_numbers, + project_name, + shared_kwargs, + ) = _prepare_standalone_plot_inputs(args) + thresholds = parse_thresholds(args.thresholds) + if args.crossing_only and not thresholds: + raise InputValidationError("--crossing-only requires --thresholds") + if args.crossing_only: + summary = get_threshold_crossing_variants( + summary, + long_df, + thresholds, + crossing_rule=args.crossing_rule, + ) + long_df = long_df[long_df["variant_id"].isin(summary["variant_id"])] + if summary.empty: + raise ProcessingError("No variants crossed the requested thresholds") + selected = _resolve_plot_variants( + summary, + explicit_variants=shared_kwargs["variants"], + top_n=args.top_n, + prefer_crossing=bool(thresholds), + thresholds=thresholds, + ) + output_path = resolve_plot_output_path( + results_csv, + out=args.out, + outdir=args.outdir, + fmt=args.format, + filename="variant_trajectory_plot", + ) + plot_variant_trajectory( + summary, + long_df, + selected_variants=selected, + output_path=output_path, + thresholds=thresholds, + title=args.title + or (f"{project_name}: variant trajectories" if project_name else None), + width=args.width, + height=args.height, + dpi=args.dpi, + label_lines=args.label_lines, + label_threshold_crossers=args.label_threshold_crossers, + crossing_rule=args.crossing_rule, + ) + print(f"\nFinished: wrote {output_path}\n") + return 0 + except (InputValidationError, ProcessingError) as exc: + print(f"\nERROR: {exc}\n") + return 1 + + +def _run_plot_turnover_command(args): + try: + ( + results_csv, + _table, + summary, + long_df, + _sample_names, + _sample_numbers, + project_name, + shared_kwargs, + ) = _prepare_standalone_plot_inputs(args) + if shared_kwargs["variants"]: + summary = summary[summary["variant_id"].isin(shared_kwargs["variants"])] + long_df = long_df[long_df["variant_id"].isin(summary["variant_id"])] + if summary.empty: + raise ProcessingError( + "None of the requested variants were found in the filtered results" + ) + output_path = resolve_plot_output_path( + results_csv, + out=args.out, + outdir=args.outdir, + fmt=args.format, + filename="variant_turnover_plot", + ) + plot_variant_turnover( + summary, + long_df, + output_path=output_path, + title=args.title + or (f"{project_name}: variant turnover" if project_name else None), + width=args.width, + height=args.height, + dpi=args.dpi, + count_mode=args.count_mode, + ) + print(f"\nFinished: wrote {output_path}\n") + return 0 + except (InputValidationError, ProcessingError) as exc: + print(f"\nERROR: {exc}\n") + return 1 + + +def _run_plot_lifespan_command(args): + try: + ( + results_csv, + _table, + summary, + _long_df, + _sample_names, + _sample_numbers, + project_name, + shared_kwargs, + ) = _prepare_standalone_plot_inputs(args) + selected = _resolve_plot_variants( + summary, + explicit_variants=shared_kwargs["variants"], + top_n=args.top_n, + ) + output_path = resolve_plot_output_path( + results_csv, + out=args.out, + outdir=args.outdir, + fmt=args.format, + filename="variant_lifespan_plot", + ) + plot_variant_lifespan( + summary, + selected_variants=selected, + output_path=output_path, + title=args.title + or (f"{project_name}: variant lifespan" if project_name else None), + width=args.width, + height=args.height, + dpi=args.dpi, + sort_by=args.sort_by, + annotate_class=args.annotate_class, + ) + print(f"\nFinished: wrote {output_path}\n") + return 0 + except (InputValidationError, ProcessingError) as exc: + print(f"\nERROR: {exc}\n") + return 1 + + def _add_plot_heatmap_subparser(subparsers): parser = subparsers.add_parser( "heatmap", @@ -930,6 +1258,128 @@ def _add_plot_heatmap_subparser(subparsers): parser.set_defaults(handler=_run_plot_heatmap_command) +def _add_plot_trajectory_subparser(subparsers): + parser = subparsers.add_parser( + "trajectory", + help="Plot selected variant trajectories from an existing results CSV", + description="Generate allele-frequency trajectory plots from a vartracker results CSV.", + formatter_class=HelpFormatter, + ) + parser.add_argument("results_csv", help="Path to a vartracker results CSV") + filter_group = parser.add_argument_group("Filtering") + _add_shared_plot_filter_arguments(filter_group) + parser.add_argument( + "--top-n", + type=int, + default=DEFAULT_TRAJECTORY_TOP_N, + help=f"Auto-select up to this many variants when --variants is not used (default: {DEFAULT_TRAJECTORY_TOP_N})", + ) + parser.add_argument( + "--label-lines", + action="store_true", + default=False, + help="Add end labels to plotted lines", + ) + parser.add_argument( + "--thresholds", + default="", + help="Comma-separated AF thresholds to draw, e.g. 0.5,0.9", + ) + parser.add_argument( + "--crossing-only", + action="store_true", + default=False, + help="Only keep variants crossing at least one requested threshold", + ) + parser.add_argument( + "--label-threshold-crossers", + action="store_true", + default=False, + help="Label variants that cross at least one supplied threshold", + ) + parser.add_argument( + "--crossing-rule", + choices=["at_or_above", "strictly_above"], + default="at_or_above", + help="How to evaluate exact threshold matches (default: at_or_above)", + ) + parser.add_argument("--title", default=None, help="Optional plot title") + parser.add_argument( + "--width", type=float, default=8.5, help="Figure width in inches" + ) + parser.add_argument( + "--height", type=float, default=5.5, help="Figure height in inches" + ) + _add_plot_output_arguments(parser) + parser.set_defaults(handler=_run_plot_trajectory_command) + + +def _add_plot_turnover_subparser(subparsers): + parser = subparsers.add_parser( + "turnover", + help="Plot longitudinal variant turnover from an existing results CSV", + description="Generate new-versus-lost turnover summaries from a vartracker results CSV.", + formatter_class=HelpFormatter, + ) + parser.add_argument("results_csv", help="Path to a vartracker results CSV") + filter_group = parser.add_argument_group("Filtering") + _add_shared_plot_filter_arguments(filter_group) + parser.add_argument( + "--count-mode", + choices=["count", "sum_af"], + default="count", + help="Aggregate turnover as counts or summed allele frequencies (default: count)", + ) + parser.add_argument("--title", default=None, help="Optional plot title") + parser.add_argument( + "--width", type=float, default=8.5, help="Figure width in inches" + ) + parser.add_argument( + "--height", type=float, default=5.0, help="Figure height in inches" + ) + _add_plot_output_arguments(parser) + parser.set_defaults(handler=_run_plot_turnover_command) + + +def _add_plot_lifespan_subparser(subparsers): + parser = subparsers.add_parser( + "lifespan", + help="Plot first-to-last detection spans from an existing results CSV", + description="Generate a horizontal lifespan plot from a vartracker results CSV.", + formatter_class=HelpFormatter, + ) + parser.add_argument("results_csv", help="Path to a vartracker results CSV") + filter_group = parser.add_argument_group("Filtering") + _add_shared_plot_filter_arguments(filter_group) + parser.add_argument( + "--top-n", + type=int, + default=DEFAULT_LIFESPAN_TOP_N, + help=f"Auto-select up to this many variants when --variants is not used (default: {DEFAULT_LIFESPAN_TOP_N})", + ) + parser.add_argument( + "--sort-by", + choices=["first_seen", "last_seen", "duration", "max_af"], + default="duration", + help="Sort plotted variants by this summary metric (default: duration)", + ) + parser.add_argument( + "--annotate-class", + action="store_true", + default=False, + help="Append variant/new and persistent/transient classes to labels", + ) + parser.add_argument("--title", default=None, help="Optional plot title") + parser.add_argument( + "--width", type=float, default=8.5, help="Figure width in inches" + ) + parser.add_argument( + "--height", type=float, default=6.0, help="Figure height in inches" + ) + _add_plot_output_arguments(parser) + parser.set_defaults(handler=_run_plot_lifespan_command) + + def _add_plot_subparser(subparsers): plot_parser = subparsers.add_parser( "plot", @@ -939,6 +1389,9 @@ def _add_plot_subparser(subparsers): ) plot_subparsers = plot_parser.add_subparsers(dest="plot_command") _add_plot_heatmap_subparser(plot_subparsers) + _add_plot_trajectory_subparser(plot_subparsers) + _add_plot_turnover_subparser(plot_subparsers) + _add_plot_lifespan_subparser(plot_subparsers) plot_parser.set_defaults(handler=_run_plot_command, _subparser=plot_parser) @@ -1242,6 +1695,9 @@ def resolve_path(value: str) -> str: "mutations_per_gene_plot": os.path.join( args.outdir, "mutations_per_gene.pdf" ), + "variant_turnover_plot": os.path.join( + args.outdir, "variant_turnover_plot.pdf" + ), "variant_allele_frequency_heatmap_html": os.path.join( args.outdir, "variant_allele_frequency_heatmap.html" ), @@ -1893,6 +2349,10 @@ def _process_files( else: pname = "" + table["sample_number"] = " / ".join( + [str(value) for value in list(input_file["sample_number"])] + ) + table = _drop_exact_duplicate_result_rows(table) # Write initial results @@ -1937,6 +2397,14 @@ def _process_files( gene_table = generate_gene_table(table, gene_lengths) plot_gene_table(gene_table, pname, args.outdir) + plot_variant_turnover( + *apply_shared_plot_filters( + *prepare_plot_inputs(table)[:2], + include_synonymous=True, + ), + output_path=os.path.join(args.outdir, "variant_turnover_plot.pdf"), + title=f"{pname}: variant turnover" if pname else None, + ) generate_variant_heatmap( table, sample_names, diff --git a/vartracker/plotting.py b/vartracker/plotting.py new file mode 100644 index 0000000..3669096 --- /dev/null +++ b/vartracker/plotting.py @@ -0,0 +1,651 @@ +"""Standalone plotting helpers for vartracker results tables.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterable, Sequence + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +from .analysis import _coerce_frequency, _resolve_variant_labels +from .core import InputValidationError, ProcessingError + +DEFAULT_TRAJECTORY_TOP_N = 12 +DEFAULT_LIFESPAN_TOP_N = 20 +DEFAULT_TRAJECTORY_THRESHOLDS = (0.5, 0.9) + + +def _parse_csv_option_list(value: str | None) -> list[str]: + if not value: + return [] + return [item.strip() for item in str(value).split(",") if item.strip()] + + +def _read_variant_list_file(path: str | None) -> list[str]: + if not path: + return [] + variant_path = Path(path).expanduser().resolve() + if not variant_path.exists(): + raise InputValidationError(f"Variant file not found: {variant_path}") + lines = [ + line.strip() + for line in variant_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + return lines + + +def load_results_table(results_csv: str | Path) -> pd.DataFrame: + path = Path(results_csv).expanduser().resolve() + if not path.exists(): + raise InputValidationError(f"Results CSV not found: {path}") + table = pd.read_csv(path, keep_default_na=False) + if table.empty: + raise InputValidationError("Results CSV is empty") + return table + + +def _parse_slash_tokens(value: object) -> list[str]: + if value is None: + return [] + return [token.strip() for token in str(value).split(" / ") if token.strip()] + + +def _get_sample_axis(table: pd.DataFrame) -> tuple[list[str], list[int]]: + if "samples" not in table.columns: + raise InputValidationError( + "Results CSV must contain a 'samples' column for standalone plotting" + ) + if "sample_number" not in table.columns: + raise InputValidationError( + "Results CSV must contain a 'sample_number' column for standalone plotting" + ) + + sample_names = _parse_slash_tokens(table.iloc[0]["samples"]) + sample_numbers_raw = _parse_slash_tokens(table.iloc[0]["sample_number"]) + if not sample_names or not sample_numbers_raw: + raise InputValidationError("Could not determine sample axis from results CSV") + if len(sample_names) != len(sample_numbers_raw): + raise InputValidationError( + "Results CSV has mismatched 'samples' and 'sample_number' columns" + ) + try: + sample_numbers = [int(float(value)) for value in sample_numbers_raw] + except ValueError as exc: + raise InputValidationError( + "Results CSV 'sample_number' values must be numeric" + ) from exc + return sample_names, sample_numbers + + +def _project_name_from_results(table: pd.DataFrame, explicit_name: str | None) -> str: + if explicit_name is not None: + return explicit_name + if "name" not in table.columns: + return "" + names = [str(value).strip() for value in table["name"].unique() if str(value).strip()] + return names[0] if len(names) == 1 else "" + + +def _sample_qc_flags(row: object, length: int) -> list[str]: + flags = _parse_slash_tokens(getattr(row, "per_sample_variant_qc", "")) + if not flags: + return [""] * length + if len(flags) < length: + flags.extend([""] * (length - len(flags))) + return [str(flag).strip().upper() for flag in flags[:length]] + + +def prepare_plot_inputs( + table: pd.DataFrame, +) -> tuple[pd.DataFrame, pd.DataFrame, list[str], list[int]]: + sample_names, sample_numbers = _get_sample_axis(table) + long_records: list[dict[str, object]] = [] + + for row in table.drop_duplicates().itertuples(index=False): + _, display_label, base_label = _resolve_variant_labels(row) + names = _parse_slash_tokens(getattr(row, "samples", "")) + numbers = _parse_slash_tokens(getattr(row, "sample_number", "")) + freqs = [_coerce_frequency(token) for token in _parse_slash_tokens(getattr(row, "alt_freq", ""))] + presence = _parse_slash_tokens(getattr(row, "presence_absence", "")) + qc_flags = _sample_qc_flags(row, len(names)) + + length = min(len(names), len(numbers), len(freqs)) if freqs else min(len(names), len(numbers)) + if length == 0: + continue + + if len(freqs) < length: + freqs.extend([0.0] * (length - len(freqs))) + if len(presence) < length: + presence.extend(["N"] * (length - len(presence))) + + for sample_name, sample_number, af, present, qc_flag in zip( + names[:length], + numbers[:length], + freqs[:length], + presence[:length], + qc_flags[:length], + ): + long_records.append( + { + "variant_id": base_label, + "variant_label": display_label.replace("\n", " "), + "variant_name": str(getattr(row, "variant", "")).strip(), + "gene": str(getattr(row, "gene", "")).strip(), + "type_of_change": str(getattr(row, "type_of_change", "")).strip(), + "type_of_variant": str(getattr(row, "type_of_variant", "")).strip(), + "variant_status": str(getattr(row, "variant_status", "")).strip(), + "persistence_status": str(getattr(row, "persistence_status", "")).strip(), + "sample_name": sample_name, + "sample_number": int(float(sample_number)), + "allele_frequency": float(af), + "present": str(present).strip().upper() == "Y" or float(af) > 0, + "sample_qc": qc_flag, + "sample_pass_qc": qc_flag != "F", + "has_literature": any( + str(getattr(row, column, "")).strip() + not in {"", "False", "false", "0", "None"} + for column in ( + "key_mutation", + "category", + "database_mutation_string", + "prior_information", + "reference", + ) + if hasattr(row, column) + ), + } + ) + + if not long_records: + raise ProcessingError("No plottable variant records found in results CSV") + + long_df = pd.DataFrame(long_records).drop_duplicates() + long_df = long_df.sort_values(["sample_number", "variant_id", "sample_name"]).reset_index(drop=True) + + summary = ( + long_df.groupby("variant_id", as_index=False) + .agg( + variant_label=("variant_label", "first"), + variant_name=("variant_name", "first"), + gene=("gene", "first"), + type_of_change=("type_of_change", "first"), + type_of_variant=("type_of_variant", "first"), + variant_status=("variant_status", "first"), + persistence_status=("persistence_status", "first"), + has_literature=("has_literature", "max"), + max_af=("allele_frequency", "max"), + ) + ) + + present_df = long_df[long_df["present"]] + first_seen = ( + present_df.groupby("variant_id")["sample_number"].min().rename("first_seen") + ) + last_seen = present_df.groupby("variant_id")["sample_number"].max().rename("last_seen") + summary = summary.merge(first_seen, on="variant_id", how="left") + summary = summary.merge(last_seen, on="variant_id", how="left") + summary["first_seen"] = summary["first_seen"].fillna(summary["max_af"].map(lambda _x: np.nan)) + summary["last_seen"] = summary["last_seen"].fillna(summary["first_seen"]) + summary["duration"] = ( + summary["last_seen"].fillna(0).astype(float) + - summary["first_seen"].fillna(0).astype(float) + ) + summary["is_persistent_new"] = summary["persistence_status"].eq("new_persistent") + summary["is_nonsynonymous"] = ~summary["type_of_change"].str.lower().str.contains( + "synonymous", na=False + ) + + return summary, long_df, sample_names, sample_numbers + + +def apply_shared_plot_filters( + summary: pd.DataFrame, + long_df: pd.DataFrame, + *, + gene: str | None = None, + effects: Sequence[str] | None = None, + min_af: float | None = None, + max_af: float | None = None, + variants: Sequence[str] | None = None, + sample_min: int | None = None, + sample_max: int | None = None, + include_synonymous: bool = True, + persistent_only: bool = False, + new_only: bool = False, +) -> tuple[pd.DataFrame, pd.DataFrame]: + filtered_long = long_df.copy() + if sample_min is not None: + filtered_long = filtered_long[filtered_long["sample_number"] >= sample_min] + if sample_max is not None: + filtered_long = filtered_long[filtered_long["sample_number"] <= sample_max] + if filtered_long.empty: + raise ProcessingError("No variant observations remain after sample-range filtering") + + filtered_summary = summary.copy() + + if gene: + filtered_summary = filtered_summary[ + filtered_summary["gene"].astype(str).str.lower() == str(gene).strip().lower() + ] + + if effects: + lowered = {str(effect).strip().lower() for effect in effects if str(effect).strip()} + filtered_summary = filtered_summary[ + filtered_summary["type_of_change"].astype(str).str.lower().isin(lowered) + ] + + if not include_synonymous: + filtered_summary = filtered_summary[ + ~filtered_summary["type_of_change"].astype(str).str.lower().str.contains("synonymous", na=False) + ] + + if persistent_only: + filtered_summary = filtered_summary[ + filtered_summary["persistence_status"].eq("new_persistent") + ] + + if new_only: + filtered_summary = filtered_summary[filtered_summary["variant_status"].eq("new")] + + if variants: + order_map = {name: idx for idx, name in enumerate(variants)} + filtered_summary = filtered_summary[ + filtered_summary.apply( + lambda row: any( + candidate in order_map + for candidate in ( + row["variant_id"], + row["variant_label"], + row["variant_name"], + ) + ), + axis=1, + ) + ] + filtered_summary["selection_order"] = filtered_summary.apply( + lambda row: min( + order_map[candidate] + for candidate in (row["variant_id"], row["variant_label"], row["variant_name"]) + if candidate in order_map + ), + axis=1, + ) + filtered_summary = filtered_summary.sort_values("selection_order").drop( + columns=["selection_order"] + ) + + if filtered_summary.empty: + raise ProcessingError("No variants remain after applying plot filters") + + filtered_long = filtered_long[filtered_long["variant_id"].isin(filtered_summary["variant_id"])] + recomputed = ( + filtered_long.groupby("variant_id", as_index=False) + .agg( + max_af=("allele_frequency", "max"), + ) + ) + filtered_summary = filtered_summary.drop(columns=["max_af"], errors="ignore").merge( + recomputed, on="variant_id", how="left" + ) + first_seen = ( + filtered_long[filtered_long["present"]] + .groupby("variant_id")["sample_number"] + .min() + .rename("first_seen") + ) + last_seen = ( + filtered_long[filtered_long["present"]] + .groupby("variant_id")["sample_number"] + .max() + .rename("last_seen") + ) + filtered_summary = filtered_summary.drop( + columns=["first_seen", "last_seen", "duration"], errors="ignore" + ) + filtered_summary = filtered_summary.merge(first_seen, on="variant_id", how="left") + filtered_summary = filtered_summary.merge(last_seen, on="variant_id", how="left") + filtered_summary["duration"] = ( + filtered_summary["last_seen"].fillna(0).astype(float) + - filtered_summary["first_seen"].fillna(0).astype(float) + ) + + if min_af is not None: + filtered_summary = filtered_summary[filtered_summary["max_af"] >= min_af] + if max_af is not None: + filtered_summary = filtered_summary[filtered_summary["max_af"] <= max_af] + if filtered_summary.empty: + raise ProcessingError("No variants remain after allele-frequency filtering") + + filtered_long = filtered_long[filtered_long["variant_id"].isin(filtered_summary["variant_id"])] + return filtered_summary.reset_index(drop=True), filtered_long.reset_index(drop=True) + + +def auto_select_variants( + summary: pd.DataFrame, + *, + top_n: int, + prefer_crossing: bool = False, + thresholds: Sequence[float] | None = None, +) -> list[str]: + if summary.empty: + return [] + + ranked = summary.copy() + threshold_values = [float(value) for value in (thresholds or [])] + if prefer_crossing and threshold_values: + ranked["crosses_threshold"] = ranked["max_af"].apply( + lambda value: any(float(value) >= threshold for threshold in threshold_values) + ) + else: + ranked["crosses_threshold"] = False + + ranked = ranked.sort_values( + by=[ + "crosses_threshold", + "has_literature", + "is_persistent_new", + "is_nonsynonymous", + "max_af", + "first_seen", + "variant_id", + ], + ascending=[False, False, False, False, False, True, True], + na_position="last", + ) + return ranked["variant_id"].head(max(1, int(top_n))).tolist() + + +def _resolve_variant_selection( + summary: pd.DataFrame, + *, + explicit_variants: Sequence[str], + top_n: int, + prefer_crossing: bool = False, + thresholds: Sequence[float] | None = None, +) -> list[str]: + if explicit_variants: + selected: list[str] = [] + seen: set[str] = set() + for requested in explicit_variants: + match = summary[ + summary.apply( + lambda row: requested + in {row["variant_id"], row["variant_label"], row["variant_name"]}, + axis=1, + ) + ] + if match.empty: + continue + variant_id = str(match.iloc[0]["variant_id"]) + if variant_id not in seen: + selected.append(variant_id) + seen.add(variant_id) + if not selected: + raise ProcessingError("None of the requested variants were found in the filtered results") + return selected + return auto_select_variants( + summary, + top_n=top_n, + prefer_crossing=prefer_crossing, + thresholds=thresholds, + ) + + +def resolve_plot_output_path( + results_csv: str | Path, + *, + out: str | None, + outdir: str | None, + fmt: str, + filename: str, +) -> Path: + results_path = Path(results_csv).expanduser().resolve() + if out: + return Path(out).expanduser().resolve() + destination_dir = ( + Path(outdir).expanduser().resolve() if outdir else results_path.parent + ) + destination_dir.mkdir(parents=True, exist_ok=True) + return destination_dir / f"{filename}.{fmt}" + + +def _apply_plot_title(ax, title: str | None, default_title: str) -> None: + ax.set_title(title or default_title, fontweight="bold") + + +def plot_variant_trajectory( + summary: pd.DataFrame, + long_df: pd.DataFrame, + *, + selected_variants: Sequence[str], + output_path: str | Path, + thresholds: Sequence[float] | None = None, + title: str | None = None, + width: float = 8.5, + height: float = 5.5, + dpi: int = 300, + label_lines: bool = False, + label_threshold_crossers: bool = False, + crossing_rule: str = "at_or_above", +) -> None: + plot_df = long_df[long_df["variant_id"].isin(selected_variants)].copy() + if plot_df.empty: + raise ProcessingError("No data available for the trajectory plot") + + fig, ax = plt.subplots(figsize=(width, height)) + threshold_values = [float(value) for value in (thresholds or [])] + for variant_id in selected_variants: + subset = plot_df[plot_df["variant_id"] == variant_id].sort_values("sample_number") + label = str(subset["variant_id"].iloc[0]) + ax.plot( + subset["sample_number"], + subset["allele_frequency"], + marker="o", + linewidth=1.8, + label=label, + ) + crosses_threshold = variant_crosses_thresholds( + subset["allele_frequency"].tolist(), + threshold_values, + crossing_rule=crossing_rule, + ) + if (label_lines or (label_threshold_crossers and crosses_threshold)) and not subset.empty: + last_row = subset.iloc[-1] + ax.text( + float(last_row["sample_number"]) + 0.1, + float(last_row["allele_frequency"]), + label, + fontsize=8, + va="center", + ) + + ax.set_xlabel("Sample") + ax.set_ylabel("Allele frequency") + ax.set_ylim(0, 1.02) + ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + for threshold in threshold_values: + ax.axhline(float(threshold), linestyle="--", linewidth=1.0, color="black", alpha=0.6) + xmax = max(plot_df["sample_number"]) if not plot_df.empty else 0 + ax.text( + float(xmax) + 0.1, + float(threshold), + f"{threshold:g}", + fontsize=8, + va="center", + ha="left", + color="black", + ) + _apply_plot_title(ax, title, "Variant trajectories") + if not label_lines and not label_threshold_crossers: + ax.legend(frameon=False, loc="best") + fig.tight_layout() + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def plot_variant_turnover( + summary: pd.DataFrame, + long_df: pd.DataFrame, + *, + output_path: str | Path, + title: str | None = None, + width: float = 8.5, + height: float = 5.0, + dpi: int = 300, + count_mode: str = "count", +) -> None: + plot_df = long_df.sort_values(["variant_id", "sample_number"]).copy() + records: list[dict[str, float | int]] = [] + + for _, variant_df in plot_df.groupby("variant_id"): + variant_df = variant_df.sort_values("sample_number").reset_index(drop=True) + for idx in range(1, len(variant_df)): + previous = variant_df.iloc[idx - 1] + current = variant_df.iloc[idx] + if not bool(previous["present"]) and bool(current["present"]): + value = 1.0 if count_mode == "count" else float(current["allele_frequency"]) + records.append( + { + "sample_number": int(current["sample_number"]), + "new": value, + "lost": 0.0, + } + ) + if bool(previous["present"]) and not bool(current["present"]): + value = 1.0 if count_mode == "count" else float(previous["allele_frequency"]) + records.append( + { + "sample_number": int(current["sample_number"]), + "new": 0.0, + "lost": value, + } + ) + + sample_numbers = sorted(plot_df["sample_number"].unique()) + turnover = pd.DataFrame({"sample_number": sample_numbers, "new": 0.0, "lost": 0.0}) + if records: + events = pd.DataFrame(records).groupby("sample_number", as_index=False).sum() + turnover = turnover.merge(events, on="sample_number", how="left", suffixes=("", "_event")) + turnover["new"] = turnover["new_event"].fillna(turnover["new"]) + turnover["lost"] = turnover["lost_event"].fillna(turnover["lost"]) + turnover = turnover.drop(columns=["new_event", "lost_event"]) + + fig, ax = plt.subplots(figsize=(width, height)) + ax.bar(turnover["sample_number"], turnover["new"], label="New") + ax.bar(turnover["sample_number"], -turnover["lost"], label="Lost") + ax.axhline(0, color="black", linewidth=0.8) + ax.set_xlabel("Sample") + ax.set_ylabel("Variant count" if count_mode == "count" else "Summed allele frequency") + ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + _apply_plot_title(ax, title, "Variant turnover") + ax.legend(frameon=False) + fig.tight_layout() + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def plot_variant_lifespan( + summary: pd.DataFrame, + *, + selected_variants: Sequence[str], + output_path: str | Path, + title: str | None = None, + width: float = 8.5, + height: float = 6.0, + dpi: int = 300, + sort_by: str = "duration", + annotate_class: bool = False, +) -> None: + plot_df = summary[summary["variant_id"].isin(selected_variants)].copy() + if plot_df.empty: + raise ProcessingError("No data available for the lifespan plot") + + ascending = {"first_seen": True, "last_seen": True, "duration": False, "max_af": False} + plot_df = plot_df.sort_values(sort_by, ascending=ascending.get(sort_by, False), na_position="last") + plot_df = plot_df.reset_index(drop=True) + + labels = plot_df["variant_id"].astype(str).tolist() + if annotate_class: + labels = [ + f"{label} [{status}; {persistence}]" + for label, status, persistence in zip( + plot_df["variant_id"], + plot_df["variant_status"], + plot_df["persistence_status"], + ) + ] + + fig, ax = plt.subplots(figsize=(width, height)) + y_positions = np.arange(len(plot_df)) + starts = plot_df["first_seen"].fillna(plot_df["last_seen"]).astype(float) + ends = plot_df["last_seen"].fillna(plot_df["first_seen"]).astype(float) + ax.hlines(y_positions, starts, ends, linewidth=2.2) + ax.scatter(starts, y_positions, s=24, zorder=3) + ax.scatter(ends, y_positions, s=24, zorder=3) + ax.set_yticks(y_positions) + ax.set_yticklabels(labels) + ax.invert_yaxis() + ax.set_xlabel("Sample") + ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + _apply_plot_title(ax, title, "Variant lifespan") + fig.tight_layout() + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + + +def get_threshold_crossing_variants( + summary: pd.DataFrame, + long_df: pd.DataFrame, + thresholds: Sequence[float], + *, + crossing_rule: str = "at_or_above", +) -> pd.DataFrame: + threshold_values = [float(value) for value in thresholds] + if not threshold_values: + return summary.copy() + crossed_ids: list[str] = [] + for variant_id, subset in long_df.groupby("variant_id"): + if variant_crosses_thresholds( + subset["allele_frequency"].tolist(), + threshold_values, + crossing_rule=crossing_rule, + ): + crossed_ids.append(str(variant_id)) + return summary[summary["variant_id"].isin(crossed_ids)].copy() + + +def parse_thresholds(value: str | None) -> list[float]: + thresholds = _parse_csv_option_list(value) + if not thresholds: + return [] + try: + return [float(item) for item in thresholds] + except ValueError as exc: + raise InputValidationError("Thresholds must be comma-separated numbers") from exc + + +def variant_crosses_thresholds( + values: Iterable[float], + thresholds: Sequence[float], + *, + crossing_rule: str = "at_or_above", +) -> bool: + if not thresholds: + return False + values_list = [float(value) for value in values] + if crossing_rule == "strictly_above": + return any(value > threshold for value in values_list for threshold in thresholds) + return any(value >= threshold for value in values_list for threshold in thresholds) + + +def collect_explicit_variants( + variants: str | None, variant_file: str | None +) -> list[str]: + requested = _parse_csv_option_list(variants) + file_variants = _read_variant_list_file(variant_file) + if requested and file_variants: + raise InputValidationError("Use either --variants or --variant-file, not both.") + return requested or file_variants diff --git a/vartracker/schemas.py b/vartracker/schemas.py index e0c7fb3..8ea37b2 100644 --- a/vartracker/schemas.py +++ b/vartracker/schemas.py @@ -243,6 +243,13 @@ "units": "", "values": "", }, + { + "name": "sample_number", + "type": "string (slash-separated)", + "description": "Sample ordering values corresponding to per-sample fields.", + "units": "", + "values": "integer-like sample numbers", + }, { "name": "total_genome_coverage", "type": "string (slash-separated)", diff --git a/vartracker/test_data/precomputed/test_results.csv b/vartracker/test_data/precomputed/test_results.csv index 0f86033..6b50551 100644 --- a/vartracker/test_data/precomputed/test_results.csv +++ b/vartracker/test_data/precomputed/test_results.csv @@ -1,3 +1,3 @@ -chrom,start,end,gene,ref,alt,variant,amino_acid_consequence,nsp_aa_change,bcsq_nt_notation,bcsq_aa_notation,type_of_variant,type_of_change,variant_status,persistence_status,presence_absence,first_appearance,last_appearance,all_samples_pass_qc,proportion_samples_passing_qc,per_sample_variant_qc,aa1_total_properties,aa2_total_properties,aa1_unique_properties,aa2_unique_properties,aa1_weight,aa2_weight,weight_difference,alt_freq,variant_depth,variant_site_depth,variant_window_depth,samples,total_genome_coverage -NC_045512.2,21563,21563,S,A,T,A21563T,S:D614V,NA,NA,NA,snp,missense,new,new_persistent,N / N / Y / Y / Y / Y / Y / Y / Y / Y,2,9,True,1.0,P / P / P / P / P / P / P / P / P / P,NA,NA,NA,NA,89.09,117.15,28.06,0.0 / 0.05 / 0.15 / 0.25 / 0.35 / 0.45 / 0.55 / 0.65 / 0.75 / 0.85,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 -NC_045512.2,23403,23403,S,A,G,A23403G,S:P681R,NA,NA,NA,snp,missense,new,new_persistent,N / Y / Y / Y / Y / Y / Y / Y / Y / Y,1,9,True,1.0,P / P / P / P / P / P / P / P / P / P,NA,NA,NA,NA,89.09,156.19,67.1,0.0 / 0.10 / 0.20 / 0.30 / 0.40 / 0.50 / 0.60 / 0.70 / 0.80 / 0.90,. / 20 / 22 / 24 / 26 / 28 / 30 / 32 / 34 / 36,. / 150 / 155 / 160 / 165 / 170 / 175 / 180 / 185 / 190,. / 200 / 205 / 210 / 215 / 220 / 225 / 230 / 235 / 240,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5 / Passage_6 / Passage_9 / Passage_12 / Passage_15,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 +chrom,start,end,gene,ref,alt,variant,amino_acid_consequence,nsp_aa_change,bcsq_nt_notation,bcsq_aa_notation,type_of_variant,type_of_change,variant_status,persistence_status,presence_absence,first_appearance,last_appearance,all_samples_pass_qc,proportion_samples_passing_qc,per_sample_variant_qc,aa1_total_properties,aa2_total_properties,aa1_unique_properties,aa2_unique_properties,aa1_weight,aa2_weight,weight_difference,alt_freq,variant_depth,variant_site_depth,variant_window_depth,samples,sample_number,total_genome_coverage +NC_045512.2,21563,21563,S,A,T,A21563T,S:D614V,NA,NA,NA,snp,missense,new,new_persistent,N / N / Y / Y / Y / Y,2,5,True,1.0,P / P / P / P / P / P,NA,NA,NA,NA,89.09,117.15,28.06,0.0 / 0.05 / 0.15 / 0.25 / 0.35 / 0.45,. / 20 / 22 / 24 / 26 / 28,. / 150 / 155 / 160 / 165 / 170,. / 200 / 205 / 210 / 215 / 220,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5,0 / 1 / 2 / 3 / 4 / 5,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 +NC_045512.2,23403,23403,S,A,G,A23403G,S:P681R,NA,NA,NA,snp,missense,new,new_persistent,N / Y / Y / Y / Y / Y,1,5,True,1.0,P / P / P / P / P / P,NA,NA,NA,NA,89.09,156.19,67.1,0.0 / 0.10 / 0.20 / 0.30 / 0.40 / 0.50,. / 20 / 22 / 24 / 26 / 28,. / 150 / 155 / 160 / 165 / 170,. / 200 / 205 / 210 / 215 / 220,Passage_0 / Passage_1 / Passage_2 / Passage_3 / Passage_4 / Passage_5,0 / 1 / 2 / 3 / 4 / 5,97.0 / 97.0 / 97.0 / 97.0 / 97.0 / 97.0 From d6352e5cb371fbb77d4f348bd854caf5de95db9b Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Fri, 27 Mar 2026 12:48:47 +1100 Subject: [PATCH 06/20] Plot improvements --- README.md | 22 +- pytest_results.txt | 363 ++++++++++++++++++- tests/test_main.py | 278 ++++++++++++++- tests/test_plotting.py | 237 ++++++++++++- vartracker/main.py | 171 ++++++++- vartracker/plotting.py | 783 +++++++++++++++++++++++++++++++++++++---- 6 files changed, 1747 insertions(+), 107 deletions(-) diff --git a/README.md b/README.md index 74b6295..2ea2207 100755 --- a/README.md +++ b/README.md @@ -210,6 +210,12 @@ vartracker plot heatmap results/results.csv \ # Plot whole-dataset turnover from an existing results file vartracker plot turnover results/results.csv +# Plot collapsed variant frequencies along the genome +vartracker plot genome results/results.csv + +# Zoom to a gene region, optionally using amino-acid coordinates +vartracker plot genome results/results.csv --gene F --aa-scale + # Plot selected variant trajectories from an existing results file vartracker plot trajectory results/results.csv \ --variants "S:D614G,S:E484K,S:N501Y" @@ -272,6 +278,7 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. - `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. - `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV. +- `vartracker plot genome` – plot SNP positions along the genome or a selected gene region using all observed allele-frequency values for each variant. - `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. - `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. - `vartracker plot lifespan` – plot first-to-last detection spans for a selected or auto-ranked subset of variants. @@ -301,13 +308,21 @@ Standalone plot filtering: - `--persistent-only` and `--new-only`: keep only persistent new variants or only variants with `variant_status == new`. - `trajectory` and `lifespan` auto-select a limited subset by default (`--top-n`) to stay readable. - `turnover` uses all filtered variants by default and is also written automatically during the main `vcf`/`bam`/`end-to-end` workflows as `variant_turnover_plot.pdf`. +- `genome` uses SNPs only by default, keeps all observed allele-frequency values for each plotted variant, and writes `variant_genome_plot.pdf` during the main workflows. Standalone plot output: - `--out`: write to an exact file path. -- `--outdir`: write beside the results CSV or into the chosen directory using deterministic names such as `variant_trajectory_plot.pdf`. +- `--outdir`: write beside the results CSV or into the chosen directory using deterministic names such as `variant_trajectory_plot.pdf` or `variant_genome_plot.pdf`. - `--format`: choose `pdf`, `png`, or `svg`. - `--dpi`: set raster output resolution. +Genome plot options: +- `--gene`: zoom to a single gene region. +- `--aa-scale`: with `--gene`, use amino-acid coordinates on the x-axis. +- `--focus-coords`: highlight nucleotide or amino-acid coordinate ranges, depending on the current x-axis mode. +- `--include-indels`: opt in to plotting indels too. This may be ambiguous or hard to interpret. +- The standalone genome plot auto-discovers `reference_features.json` beside `results.csv`; workflow runs generate this sidecar automatically. + Trajectory threshold mode: - `--thresholds`: draw horizontal AF threshold lines, e.g. `0.5,0.9`. - `--crossing-only`: keep only variants crossing at least one supplied threshold. @@ -315,6 +330,11 @@ Trajectory threshold mode: - `--crossing-rule`: choose whether threshold equality counts (`at_or_above`) or requires a strict exceedance (`strictly_above`). Standalone plot examples: +- `vartracker plot genome results.csv` +- `vartracker plot genome results.csv --gene F` +- `vartracker plot genome results.csv --gene F --aa-scale` +- `vartracker plot genome results.csv --focus-coords "150-300,900-1800"` +- `vartracker plot genome results.csv --gene F --aa-scale --focus-coords "50-120,180-220"` - `vartracker plot turnover results.csv` - `vartracker plot trajectory results.csv --variants "S:D614G,S:E484K"` - `vartracker plot trajectory results.csv --thresholds 0.5,0.9` diff --git a/pytest_results.txt b/pytest_results.txt index a7837ad..9412307 100644 --- a/pytest_results.txt +++ b/pytest_results.txt @@ -1,5 +1,350 @@ -........................................................................ [ 87%] -.......... [100%] +............................................FFF......................... [ 82%] +............... [100%] +=================================== FAILURES =================================== +______________________ test_main_resolves_relative_paths _______________________ + +tmp_path = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_main_resolves_relative_pa0') +monkeypatch = <_pytest.monkeypatch.MonkeyPatch object at 0x19bc11c90> +minimal_vcf = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_main_resolves_relative_pa0/sample.vcf') + + def test_main_resolves_relative_paths(tmp_path, monkeypatch, minimal_vcf): + coverage_path = tmp_path / "sample.cov.txt" + coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") + + csv_path = tmp_path / "inputs.csv" + csv_path.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + "Sample1,0,,,,sample.vcf,sample.cov.txt\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + + def fake_setup(args): + args.reference = "/tmp/mock_reference.fasta" + args.gff3 = "/tmp/mock_annotation.gff3" + return args + + monkeypatch.setattr(main_module, "setup_default_paths", fake_setup) + monkeypatch.setattr( + main_module, "validate_reference_and_annotation", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "generate_cumulative_lineplot", lambda *a, **k: None + ) + monkeypatch.setattr(main_module, "generate_variant_heatmap", lambda *a, **k: None) + monkeypatch.setattr( + main_module, "process_joint_variants", lambda path: pd.read_csv(path) + ) + monkeypatch.setattr( + main_module, "generate_gene_table", lambda table, *_a, **_k: table + ) + monkeypatch.setattr(main_module, "plot_gene_table", lambda *a, **k: None) + monkeypatch.setattr(main_module, "search_literature", lambda *a, **k: None) + + formatted_csq = tmp_path / "formatted.csq.vcf.gz" + monkeypatch.setattr( + main_module, + "format_vcf", + lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), + ) + monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr( + main_module, + "process_vcf", + lambda *a, **k: pd.DataFrame( + { + "gene": ["S"], + "variant": ["A266C"], + "amino_acid_consequence": ["S:A1C"], + "nsp_aa_change": [""], + "presence_absence": ["Y"], + "variant_status": ["new"], + "persistence_status": ["new_persistent"], + "samples": ["Sample1"], + "alt_freq": ["0.5"], + } + ), + ) + + exit_code = main_module.main( + [ + "vcf", + str(csv_path), + "--name", + "Example", + "--outdir", + str(tmp_path / "results"), + ] + ) + +> assert exit_code == 0 +E assert 1 == 0 + +tests/test_main.py:204: AssertionError +----------------------------- Captured stdout call ----------------------------- + +██ ██ █████ ██████ ████████ ██████ █████ ██████ ██ ██ ███████ ██████ +██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ +██ ██ ███████ ██████ ██ ██████ ███████ ██ █████ █████ ██████ + ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ + ████ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██████ ██ ██ ███████ ██ ██ + + + +Warning: vartracker was designed for longitudinal comparisons but only one input VCF was provided. Some results may not make sense. +Pre-processing VCF files for compatibility... +Annotating results... +Summarising results... +Plotting results... + +ERROR: GFF3 not found: /private/tmp/mock_annotation.gff3 + +________________ test_search_pokay_retains_parsed_database_csv _________________ + +tmp_path = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_search_pokay_retains_pars0') +monkeypatch = <_pytest.monkeypatch.MonkeyPatch object at 0x1a160e2d0> +minimal_vcf = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_search_pokay_retains_pars0/sample.vcf') + + def test_search_pokay_retains_parsed_database_csv(tmp_path, monkeypatch, minimal_vcf): + coverage_path = tmp_path / "sample.cov.txt" + coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") + + csv_path = tmp_path / "inputs.csv" + csv_path.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + "Sample1,0,,,,sample.vcf,sample.cov.txt\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + + def fake_setup(args): + args.reference = "/tmp/mock_reference.fasta" + args.gff3 = "/tmp/mock_annotation.gff3" + return args + + monkeypatch.setattr(main_module, "setup_default_paths", fake_setup) + monkeypatch.setattr( + main_module, "validate_reference_and_annotation", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "generate_cumulative_lineplot", lambda *a, **k: None + ) + monkeypatch.setattr(main_module, "generate_variant_heatmap", lambda *a, **k: None) + monkeypatch.setattr( + main_module, "process_joint_variants", lambda path: pd.read_csv(path) + ) + monkeypatch.setattr( + main_module, "generate_gene_table", lambda table, *_a, **_k: table + ) + monkeypatch.setattr(main_module, "plot_gene_table", lambda *a, **k: None) + monkeypatch.setattr( + main_module, "search_literature", lambda *a, **k: pd.DataFrame() + ) + + formatted_csq = tmp_path / "formatted.csq.vcf.gz" + monkeypatch.setattr( + main_module, + "format_vcf", + lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), + ) + monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr( + main_module, + "process_vcf", + lambda *a, **k: pd.DataFrame( + { + "gene": ["S"], + "variant": ["A266C"], + "amino_acid_consequence": ["S:A1C"], + "nsp_aa_change": [""], + "presence_absence": ["Y"], + "variant_status": ["new"], + "persistence_status": ["new_persistent"], + "samples": ["Sample1"], + "alt_freq": ["0.5"], + } + ), + ) + + def fake_parse_pokay(argv): + output_path = Path(argv[0]) + output_path.write_text( + "gene,mutation,information,reference\nS,S:A1C,Mock hit,PMID123\n", + encoding="utf-8", + ) + return 0 + + monkeypatch.setattr(main_module.parse_pokay_module, "main", fake_parse_pokay) + + outdir = tmp_path / "results" + exit_code = main_module.main( + [ + "vcf", + str(csv_path), + "--name", + "Example", + "--outdir", + str(outdir), + "--search-pokay", + ] + ) + +> assert exit_code == 0 +E assert 1 == 0 + +tests/test_main.py:292: AssertionError +----------------------------- Captured stdout call ----------------------------- + +██ ██ █████ ██████ ████████ ██████ █████ ██████ ██ ██ ███████ ██████ +██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ +██ ██ ███████ ██████ ██ ██████ ███████ ██ █████ █████ ██████ + ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ + ████ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██████ ██ ██ ███████ ██ ██ + + + +Warning: vartracker was designed for longitudinal comparisons but only one input VCF was provided. Some results may not make sense. +Pre-processing VCF files for compatibility... +Annotating results... +Summarising results... +Plotting results... + +ERROR: GFF3 not found: /private/tmp/mock_annotation.gff3 + +______________ test_vcf_heatmap_options_are_forwarded_to_heatmap _______________ + +monkeypatch = <_pytest.monkeypatch.MonkeyPatch object at 0x1a13ed750> +tmp_path = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_vcf_heatmap_options_are_f0') +minimal_vcf = PosixPath('/private/var/folders/gq/5znbqdq51158c_xf1391nj1m0000gn/T/pytest-of-cfos/pytest-78/test_vcf_heatmap_options_are_f0/sample.vcf') + + def test_vcf_heatmap_options_are_forwarded_to_heatmap( + monkeypatch, tmp_path, minimal_vcf + ): + coverage_path = tmp_path / "sample.cov.txt" + coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") + + csv_path = tmp_path / "inputs.csv" + csv_path.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + "Sample1,0,,,,sample.vcf,sample.cov.txt\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + + def fake_setup(args): + args.reference = "/tmp/mock_reference.fasta" + args.gff3 = "/tmp/mock_annotation.gff3" + return args + + monkeypatch.setattr(main_module, "setup_default_paths", fake_setup) + monkeypatch.setattr( + main_module, "validate_reference_and_annotation", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "generate_cumulative_lineplot", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "process_joint_variants", lambda path: pd.read_csv(path) + ) + monkeypatch.setattr( + main_module, "generate_gene_table", lambda table, *_a, **_k: table + ) + monkeypatch.setattr(main_module, "plot_gene_table", lambda *a, **k: None) + monkeypatch.setattr(main_module, "search_literature", lambda *a, **k: None) + + recorded = {} + + def fake_heatmap(*args, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "generate_variant_heatmap", fake_heatmap) + + formatted_csq = tmp_path / "formatted.csq.vcf.gz" + monkeypatch.setattr( + main_module, + "format_vcf", + lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), + ) + monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr( + main_module, + "process_vcf", + lambda *a, **k: pd.DataFrame( + { + "gene": ["S"], + "variant": ["A266C"], + "amino_acid_consequence": ["S:A1C"], + "nsp_aa_change": [""], + "presence_absence": ["Y"], + "variant_status": ["new"], + "persistence_status": ["new_persistent"], + "samples": ["Sample1"], + "alt_freq": ["0.5"], + } + ), + ) + + exit_code = main_module.main( + [ + "vcf", + str(csv_path), + "--outdir", + str(tmp_path / "results"), + "--heatmap-aa-exclude", + "synonymous,*frameshift*,stop_gained", + "--heatmap-aa-include", + "missense", + "--heatmap-only-persistent", + "--heatmap-only-new", + "--heatmap-gene-include", + "S", + "--heatmap-gene-exclude", + "N", + "--heatmap-variant-type", + "snp", + "--heatmap-qc", + "PASS", + "--min-prop-passing-qc", + "0.75", + "--heatmap-min-persistence", + "2", + "--heatmap-min-max-af", + "0.4", + "--heatmap-min-sample-af", + "0.3", + "--heatmap-sample-subset", + "Sample1", + "--heatmap-hide-singletons", + "--heatmap-min-depth", + "20", + ] + ) + +> assert exit_code == 0 +E assert 1 == 0 + +tests/test_main.py:400: AssertionError +----------------------------- Captured stdout call ----------------------------- + +██ ██ █████ ██████ ████████ ██████ █████ ██████ ██ ██ ███████ ██████ +██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ +██ ██ ███████ ██████ ██ ██████ ███████ ██ █████ █████ ██████ + ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ + ████ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██████ ██ ██ ███████ ██ ██ + + + +Warning: vartracker was designed for longitudinal comparisons but only one input VCF was provided. Some results may not make sense. +Pre-processing VCF files for compatibility... +Annotating results... +Summarising results... +Plotting results... + +ERROR: GFF3 not found: /private/tmp/mock_annotation.gff3 + ================================ tests coverage ================================ _______________ coverage: platform darwin, python 3.11.8-final-0 _______________ @@ -16,13 +361,17 @@ vartracker/core.py 62 10 84% 90, 95, 98-100, 113- vartracker/data/__init__.py 8 0 100% vartracker/data/parse_pokay.py 141 26 82% 69-70, 110, 127-130, 155-156, 162, 168, 234, 253-276, 280 vartracker/generate.py 121 13 89% 54, 74, 84, 99, 115, 119, 132, 143, 156, 167, 202, 205, 207 -vartracker/main.py 939 277 71% 125, 404, 410, 422, 428-433, 439, 443, 476, 479-480, 484-489, 497, 515, 521-529, 534-572, 577-605, 616, 759-760, 796-802, 806-809, 960-963, 982-984, 1041-1044, 1090, 1140-1143, 1166-1168, 1209-1211, 1450-1453, 1467, 1500-1504, 1506-1507, 1528, 1543-1544, 1548-1549, 1560-1561, 1566-1571, 1588, 1592, 1608, 1612, 1617-1625, 1643-1666, 1719, 1721, 1731-1749, 1830-1834, 1851-1855, 1861-1862, 1867, 1888, 1906-1911, 1922, 1939-1952, 1972-1976, 1993-1997, 2001-2002, 2007, 2030, 2048-2053, 2064, 2081-2094, 2104, 2115, 2117, 2127, 2141, 2146-2150, 2170-2179, 2183-2227, 2254-2260, 2289, 2299-2304, 2366, 2438 -vartracker/plotting.py 297 54 82% 29-37, 43, 46, 52, 58, 62, 69, 71, 76-77, 85, 87, 97, 117, 120, 122, 163, 225, 241, 246, 251, 281, 318, 320, 334, 369-389, 407, 436, 456-457, 564, 572, 608, 626-627, 650 -vartracker/provenance.py 166 52 69% 45, 63-73, 77-87, 97-115, 127-128, 140, 158, 169-170, 173, 185-186, 245, 272-273, 280, 282, 307 +vartracker/main.py 972 279 71% 129, 408, 414, 426, 432-437, 443, 447, 480, 483-484, 488-493, 501, 519, 525-533, 538-576, 581-609, 620, 763-764, 800-806, 810-813, 964-967, 986-988, 1045-1048, 1094, 1144-1147, 1170-1172, 1213-1215, 1556-1559, 1573, 1606-1610, 1612-1613, 1634, 1649-1650, 1654-1655, 1666-1667, 1672-1677, 1694, 1698, 1714, 1718, 1723-1731, 1749-1772, 1822-1836, 1844-1847, 1853-1861, 1942-1946, 1963-1967, 1973-1974, 1979, 2000, 2018-2023, 2034, 2051-2064, 2084-2088, 2105-2109, 2113-2114, 2119, 2142, 2160-2165, 2176, 2193-2206, 2216, 2227, 2229, 2239, 2253, 2258-2262, 2282-2291, 2295-2339, 2366-2372, 2401, 2411-2416, 2478, 2557 +vartracker/plotting.py 533 84 84% 33-41, 47, 50, 56, 62, 66, 73, 75, 80-81, 89, 91, 101, 121, 124, 126, 167, 229, 245, 250, 255, 285, 322, 324, 338, 373-393, 411, 440, 460-461, 568, 576, 612, 630-631, 654, 664, 671-672, 676, 685, 707, 730, 740, 769, 785-787, 794, 827, 835, 861-862, 864, 866, 868, 870, 872, 886, 891, 966, 970, 973, 999, 1002, 1080 +vartracker/provenance.py 166 51 69% 45, 63-73, 77-87, 97-115, 127-128, 140, 158, 169-170, 173, 185-186, 245, 272-273, 282, 307 vartracker/reference_prepare.py 308 44 86% 82, 93, 104, 106, 114, 121-146, 155, 162-164, 170, 182-183, 204, 212, 218, 294, 353, 357, 364, 368, 376, 406, 414, 460, 486, 498, 517, 548, 591, 596, 601, 609-610 vartracker/schemas.py 60 7 88% 270-298, 303 vartracker/vcf_processing.py 337 308 9% 23-25, 33-112, 124, 136-140, 146-152, 180-353, 366-394, 409-484, 498-639, 658-706, 759-766 ------------------------------------------------------------------- -TOTAL 3488 1296 63% +TOTAL 3757 1327 65% Coverage XML written to file coverage.xml -82 passed in 42.73s +=========================== short test summary info ============================ +FAILED tests/test_main.py::test_main_resolves_relative_paths - assert 1 == 0 +FAILED tests/test_main.py::test_search_pokay_retains_parsed_database_csv - as... +FAILED tests/test_main.py::test_vcf_heatmap_options_are_forwarded_to_heatmap +3 failed, 84 passed in 41.09s diff --git a/tests/test_main.py b/tests/test_main.py index d9bde75..a2627f9 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -463,15 +463,19 @@ def fake_heatmap(*args, **kwargs): def _write_plot_results_csv(path: Path, n_variants: int = 4) -> None: rows = [ - "samples,sample_number,name,alt_freq,per_sample_variant_qc,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,amino_acid_consequence,nsp_aa_change,start,reference" + "chrom,start,end,samples,sample_number,name,alt_freq,per_sample_variant_qc,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,amino_acid_consequence,nsp_aa_change,reference" ] for idx in range(n_variants): gene = "S" if idx % 2 == 0 else "N" + chrom = "segA" if gene == "S" else "segB" aa = f"{gene}:V{idx + 1}A" effect = "missense" if idx % 3 else "synonymous" persistence = "new_persistent" if idx % 2 == 0 else "new_transient" reference = "PMID123" if idx == 1 else "" rows.append( + f"{chrom}," + f"{100 + idx}," + f"{100 + idx}," "P0 / P1 / P2 / P3," "0 / 1 / 2 / 3," "Example," @@ -485,20 +489,41 @@ def _write_plot_results_csv(path: Path, n_variants: int = 4) -> None: f"{gene}," f"A{100 + idx}G," f"{aa},," - f"{100 + idx}," f"{reference}" ) path.write_text("\n".join(rows) + "\n", encoding="utf-8") +def _write_reference_features_json(path: Path) -> None: + path.write_text( + """{ + "contig_lengths": {"segA": 30000, "segB": 15000}, + "contig_order": ["segA", "segB"], + "features": [ + {"aa_length": 1200, "contig": "segA", "end": 25000, "name": "S", "start": 21000, "strand": "+", "type": "mRNA"}, + {"aa_length": 400, "contig": "segB", "end": 12000, "name": "N", "start": 9000, "strand": "+", "type": "mRNA"} + ], + "source_gff3": "/tmp/reference.gff3" +} +""", + encoding="utf-8", + ) + + def test_plot_trajectory_variants_override_auto_selection(monkeypatch, tmp_path): results_csv = tmp_path / "results.csv" _write_plot_results_csv(results_csv, n_variants=6) recorded = {} - def fake_plot(*args, **kwargs): - recorded["selected_variants"] = kwargs["selected_variants"] + def fake_plot(summary, _long_df, **kwargs): + selected_ids = kwargs["selected_variants"] + recorded["selected_labels"] = ( + summary.set_index("variant_id") + .loc[selected_ids, "variant_label"] + .str.replace(r" \(.*\)$", "", regex=True) + .tolist() + ) monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) @@ -513,7 +538,7 @@ def fake_plot(*args, **kwargs): ) assert exit_code == 0 - assert recorded["selected_variants"] == ["S:V3A", "N:V2A"] + assert recorded["selected_labels"] == ["S:V3A", "N:V2A"] def test_plot_trajectory_threshold_options_are_forwarded(monkeypatch, tmp_path): @@ -546,6 +571,63 @@ def fake_plot(*args, **kwargs): assert recorded["crossing_rule"] == "strictly_above" +def test_plot_trajectory_label_mode_is_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + + recorded = {} + + def fake_plot(*args, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) + + exit_code = main_module.main( + [ + "plot", + "trajectory", + str(results_csv), + "--label-mode", + "nt", + ] + ) + + assert exit_code == 0 + assert recorded["label_mode"] == "nt" + + +def test_plot_sample_axis_mode_is_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + + recorded = {} + + def fake_turnover(*args, **kwargs): + recorded["turnover"] = kwargs["sample_axis_mode"] + + def fake_lifespan(*args, **kwargs): + recorded["lifespan"] = kwargs["sample_axis_mode"] + + monkeypatch.setattr(main_module, "plot_variant_turnover", fake_turnover) + monkeypatch.setattr(main_module, "plot_variant_lifespan", fake_lifespan) + + assert ( + main_module.main( + ["plot", "turnover", str(results_csv), "--sample-axis", "name"] + ) + == 0 + ) + assert ( + main_module.main( + ["plot", "lifespan", str(results_csv), "--sample-axis", "name"] + ) + == 0 + ) + + assert recorded["turnover"] == "name" + assert recorded["lifespan"] == "name" + + def test_plot_trajectory_crossing_only_requires_thresholds(tmp_path): results_csv = tmp_path / "results.csv" _write_plot_results_csv(results_csv, n_variants=4) @@ -563,8 +645,14 @@ def test_plot_trajectory_crossing_only_filters_variants(monkeypatch, tmp_path): recorded = {} - def fake_plot(summary, long_df, **kwargs): - recorded["selected_variants"] = kwargs["selected_variants"] + def fake_plot(summary, _long_df, **kwargs): + selected_ids = kwargs["selected_variants"] + recorded["selected_labels"] = ( + summary.set_index("variant_id") + .loc[selected_ids, "variant_label"] + .str.replace(r" \(.*\)$", "", regex=True) + .tolist() + ) recorded["summary_variant_ids"] = list(summary["variant_id"]) monkeypatch.setattr(main_module, "plot_variant_trajectory", fake_plot) @@ -581,10 +669,8 @@ def fake_plot(summary, long_df, **kwargs): ) assert exit_code == 0 - assert all( - variant in recorded["summary_variant_ids"] - for variant in recorded["selected_variants"] - ) + assert recorded["selected_labels"] + assert len(recorded["selected_labels"]) <= main_module.DEFAULT_TRAJECTORY_TOP_N def test_plot_turnover_uses_all_filtered_variants_by_default(monkeypatch, tmp_path): @@ -595,7 +681,9 @@ def test_plot_turnover_uses_all_filtered_variants_by_default(monkeypatch, tmp_pa def fake_turnover(summary, long_df, **kwargs): recorded["summary_count"] = len(summary) - recorded["variant_ids"] = list(summary["variant_id"]) + recorded["variant_labels"] = ( + summary["variant_label"].str.replace(r" \(.*\)$", "", regex=True).tolist() + ) monkeypatch.setattr(main_module, "plot_variant_turnover", fake_turnover) @@ -603,7 +691,7 @@ def fake_turnover(summary, long_df, **kwargs): assert exit_code == 0 assert recorded["summary_count"] == 3 - assert all(variant.startswith("S:") for variant in recorded["variant_ids"]) + assert all(label.startswith("S:") for label in recorded["variant_labels"]) def test_plot_commands_auto_limit_variants_by_top_n(monkeypatch, tmp_path): @@ -628,6 +716,170 @@ def fake_lifespan(*args, **kwargs): assert len(recorded["lifespan"]) == main_module.DEFAULT_LIFESPAN_TOP_N +def test_plot_genome_uses_sidecar_metadata(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=6) + _write_reference_features_json(tmp_path / "reference_features.json") + + recorded = {} + + def fake_genome(table, metadata, **kwargs): + recorded["rows"] = len(table) + recorded["metadata"] = metadata + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + + exit_code = main_module.main( + [ + "plot", + "genome", + str(results_csv), + "--gene", + "S", + "--aa-scale", + "--focus-coords", + "50-120,180-220", + "--min-af", + "0.2", + ] + ) + + assert exit_code == 0 + assert recorded["rows"] == 6 + assert recorded["gene"] == "S" + assert recorded["aa_scale"] is True + assert recorded["focus_ranges"] == [(50.0, 120.0), (180.0, 220.0)] + assert recorded["min_af"] == 0.2 + assert recorded["include_indels"] is False + assert recorded["metadata"]["contig_order"] == ["segA", "segB"] + + +def test_plot_genome_include_indels_is_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + _write_reference_features_json(tmp_path / "reference_features.json") + + recorded = {} + + def fake_genome(_table, _metadata, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + + exit_code = main_module.main( + ["plot", "genome", str(results_csv), "--include-indels"] + ) + + assert exit_code == 0 + assert recorded["include_indels"] is True + + +def test_plot_genome_aa_scale_requires_gene(tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + _write_reference_features_json(tmp_path / "reference_features.json") + + exit_code = main_module.main(["plot", "genome", str(results_csv), "--aa-scale"]) + + assert exit_code == 1 + + +def test_vcf_workflow_writes_default_genome_plot(monkeypatch, tmp_path, minimal_vcf): + coverage_path = tmp_path / "sample.cov.txt" + coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") + + csv_path = tmp_path / "inputs.csv" + csv_path.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + "Sample1,0,,,,sample.vcf,sample.cov.txt\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + + def fake_setup(args): + args.reference = "/tmp/mock_reference.fasta" + args.gff3 = "/tmp/mock_annotation.gff3" + return args + + monkeypatch.setattr(main_module, "setup_default_paths", fake_setup) + monkeypatch.setattr( + main_module, "validate_reference_and_annotation", lambda *a, **k: None + ) + monkeypatch.setattr( + main_module, "generate_cumulative_lineplot", lambda *a, **k: None + ) + monkeypatch.setattr(main_module, "generate_variant_heatmap", lambda *a, **k: None) + monkeypatch.setattr( + main_module, "process_joint_variants", lambda path: pd.read_csv(path) + ) + monkeypatch.setattr( + main_module, "generate_gene_table", lambda table, *_a, **_k: table + ) + monkeypatch.setattr(main_module, "plot_gene_table", lambda *a, **k: None) + monkeypatch.setattr(main_module, "search_literature", lambda *a, **k: None) + monkeypatch.setattr( + main_module, "write_reference_feature_metadata", lambda *a, **k: None + ) + + recorded = {} + + def fake_genome(*args, **kwargs): + recorded["output_path"] = kwargs["output_path"] + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + monkeypatch.setattr( + main_module, + "load_reference_feature_metadata", + lambda *_a, **_k: { + "contig_order": ["NC_045512.2"], + "contig_lengths": {"NC_045512.2": 29903}, + "features": [], + }, + ) + + formatted_csq = tmp_path / "formatted.csq.vcf.gz" + monkeypatch.setattr( + main_module, + "format_vcf", + lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), + ) + monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr( + main_module, + "process_vcf", + lambda *a, **k: pd.DataFrame( + { + "chrom": ["NC_045512.2"], + "start": [266], + "end": [266], + "gene": ["S"], + "variant": ["A266C"], + "amino_acid_consequence": ["S:A1C"], + "nsp_aa_change": [""], + "presence_absence": ["Y"], + "variant_status": ["new"], + "persistence_status": ["new_persistent"], + "samples": ["Sample1"], + "alt_freq": ["0.5"], + } + ), + ) + + exit_code = main_module.main( + [ + "vcf", + str(csv_path), + "--outdir", + str(tmp_path / "results"), + ] + ) + + assert exit_code == 0 + assert recorded["output_path"].endswith("variant_genome_plot.pdf") + + def test_e2e_runs_snakemake_then_vcf(monkeypatch, tmp_path): updated_csv = tmp_path / "samples_updated.csv" vcf_out = tmp_path / "vcf.gz" diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 1246889..c97d146 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -3,11 +3,18 @@ from __future__ import annotations import pandas as pd +import pytest +from vartracker.core import InputValidationError from vartracker.plotting import ( + _collapse_variants_for_genome_plot, + _parse_focus_ranges, + _subset_collapsed_variants, + build_reference_feature_metadata, apply_shared_plot_filters, auto_select_variants, get_threshold_crossing_variants, + plot_variant_genome, plot_variant_lifespan, plot_variant_trajectory, plot_variant_turnover, @@ -20,6 +27,9 @@ def _results_table() -> pd.DataFrame: return pd.DataFrame( [ { + "chrom": "segA", + "start": 23403, + "end": 23403, "gene": "S", "variant": "A23403G", "amino_acid_consequence": "S:D614G", @@ -36,6 +46,9 @@ def _results_table() -> pd.DataFrame: "reference": "PMID123", }, { + "chrom": "segA", + "start": 23012, + "end": 23012, "gene": "S", "variant": "G23012A", "amino_acid_consequence": "S:E484K", @@ -52,6 +65,9 @@ def _results_table() -> pd.DataFrame: "reference": "", }, { + "chrom": "segA", + "start": 23063, + "end": 23063, "gene": "S", "variant": "A23063T", "amino_acid_consequence": "S:N501Y", @@ -68,6 +84,9 @@ def _results_table() -> pd.DataFrame: "reference": "", }, { + "chrom": "segB", + "start": 28977, + "end": 28977, "gene": "N", "variant": "C28977T", "amino_acid_consequence": "N:S202=", @@ -88,11 +107,18 @@ def _results_table() -> pd.DataFrame: def test_prepare_plot_inputs_and_filters(): - summary, long_df, sample_names, sample_numbers = prepare_plot_inputs(_results_table()) + summary, long_df, sample_names, sample_numbers = prepare_plot_inputs( + _results_table() + ) assert sample_names == ["P0", "P1", "P2", "P3"] assert sample_numbers == [0, 1, 2, 3] - assert set(summary["variant_id"]) >= {"S:D614G", "S:E484K", "S:N501Y", "N:N202="} + assert {label.split(" (", 1)[0] for label in summary["variant_label"]} >= { + "S:D614G", + "S:E484K", + "S:N501Y", + "N:N202=", + } filtered_summary, filtered_long = apply_shared_plot_filters( summary, @@ -104,7 +130,10 @@ def test_prepare_plot_inputs_and_filters(): sample_max=3, ) - assert set(filtered_summary["variant_id"]) == {"S:D614G", "S:N501Y"} + assert {label.split(" (", 1)[0] for label in filtered_summary["variant_label"]} == { + "S:D614G", + "S:N501Y", + } assert filtered_long["sample_number"].min() == 1 assert filtered_long["sample_number"].max() == 3 @@ -114,7 +143,10 @@ def test_auto_select_variants_prefers_literature_and_persistent(): selected = auto_select_variants(summary, top_n=2) - assert selected[0] == "S:D614G" + first_label = summary.loc[ + summary["variant_id"] == selected[0], "variant_label" + ].iloc[0] + assert first_label.split(" (", 1)[0] == "S:D614G" assert len(selected) == 2 @@ -124,7 +156,12 @@ def test_plot_functions_create_output_files(tmp_path, monkeypatch): monkeypatch.setenv("MPLCONFIGDIR", str(mpl_dir)) summary, long_df, _, _ = prepare_plot_inputs(_results_table()) - selected = ["S:D614G", "S:N501Y"] + selected = summary.loc[ + summary["variant_label"] + .str.replace(r" \(.*\)$", "", regex=True) + .isin(["S:D614G", "S:N501Y"]), + "variant_id", + ].tolist() trajectory = tmp_path / "trajectory.pdf" turnover = tmp_path / "turnover.pdf" @@ -136,12 +173,18 @@ def test_plot_functions_create_output_files(tmp_path, monkeypatch): long_df, selected_variants=selected, output_path=trajectory, + sample_axis_mode="name", + ) + plot_variant_turnover( + summary, long_df, output_path=turnover, sample_axis_mode="name" ) - plot_variant_turnover(summary, long_df, output_path=turnover) plot_variant_lifespan( summary, selected_variants=selected, output_path=lifespan, + sample_axis_mode="name", + sample_names=["P0", "P1", "P2", "P3"], + sample_numbers=[0, 1, 2, 3], ) plot_variant_trajectory( summary, @@ -159,8 +202,186 @@ def test_plot_functions_create_output_files(tmp_path, monkeypatch): def test_threshold_crossing_helpers(): assert variant_crosses_thresholds([0.1, 0.5], [0.5]) is True - assert variant_crosses_thresholds([0.1, 0.5], [0.5], crossing_rule="strictly_above") is False + assert ( + variant_crosses_thresholds([0.1, 0.5], [0.5], crossing_rule="strictly_above") + is False + ) summary, long_df, _, _ = prepare_plot_inputs(_results_table()) crossed = get_threshold_crossing_variants(summary, long_df, [0.5]) - assert set(crossed["variant_id"]) == {"S:D614G", "S:N501Y"} + assert {label.split(" (", 1)[0] for label in crossed["variant_label"]} == { + "S:D614G", + "S:N501Y", + } + + +def test_genome_plot_collapses_variants_by_max_af(): + table = pd.concat([_results_table(), _results_table().iloc[[0]]], ignore_index=True) + + collapsed = _collapse_variants_for_genome_plot(table) + + assert len(collapsed) == 4 + d614g = collapsed.loc[ + collapsed["variant_label"].str.replace(r" \(.*\)$", "", regex=True) == "S:D614G" + ].iloc[0] + assert d614g["summary_af"] == 0.7 + assert d614g["af_values"] == [0.0, 0.2, 0.45, 0.7] + assert d614g["chrom"] == "segA" + assert d614g["start"] == 23403 + + +def test_prepare_plot_inputs_uses_unique_internal_variant_keys(): + table = pd.DataFrame( + [ + { + "chrom": "segA", + "start": 100, + "end": 100, + "ref": "A", + "alt": "G", + "gene": "INTERGENIC", + "variant": "A100G", + "amino_acid_consequence": "None", + "nsp_aa_change": "None", + "type_of_variant": "snp", + "type_of_change": "intergenic", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "Y / N", + "alt_freq": "0.8 / 0.0", + "samples": "P1 / P2", + "sample_number": "1 / 2", + "per_sample_variant_qc": "P / P", + }, + { + "chrom": "segA", + "start": 200, + "end": 200, + "ref": "C", + "alt": "T", + "gene": "INTERGENIC", + "variant": "C200T", + "amino_acid_consequence": "None", + "nsp_aa_change": "None", + "type_of_variant": "snp", + "type_of_change": "intergenic", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "N / Y", + "alt_freq": "0.0 / 0.7", + "samples": "P1 / P2", + "sample_number": "1 / 2", + "per_sample_variant_qc": "P / P", + }, + ] + ) + + summary, long_df, _, _ = prepare_plot_inputs(table) + + assert summary["variant_id"].nunique() == 2 + assert {label.split(" (", 1)[0] for label in summary["variant_label"]} == { + "INTERGENIC:None" + } + assert ( + long_df.loc[ + long_df["variant_label"].str.startswith("INTERGENIC:None"), "variant_id" + ].nunique() + == 2 + ) + + +def test_genome_plot_defaults_to_snps_only(capsys): + table = pd.DataFrame( + [ + { + "chrom": "segA", + "start": 100, + "end": 100, + "gene": "S", + "variant": "A100G", + "amino_acid_consequence": "S:V1A", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "Y / Y", + "alt_freq": "0.1 / 0.4", + }, + { + "chrom": "segA", + "start": 101, + "end": 101, + "gene": "S", + "variant": "A101AG", + "amino_acid_consequence": "S:V2del", + "nsp_aa_change": "", + "type_of_variant": "indel", + "type_of_change": "frameshift", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "Y / Y", + "alt_freq": "0.2 / 0.5", + }, + ] + ) + + collapsed = _subset_collapsed_variants(_collapse_variants_for_genome_plot(table)) + + assert list(collapsed["type_of_variant"]) == ["snp"] + assert capsys.readouterr().out == "" + + +def test_genome_plot_builds_outputs_and_supports_gene_modes(tmp_path, monkeypatch): + mpl_dir = tmp_path / "mpl" + mpl_dir.mkdir() + monkeypatch.setenv("MPLCONFIGDIR", str(mpl_dir)) + + gff = tmp_path / "reference.gff3" + gff.write_text( + "\n".join( + [ + "##gff-version 3", + "##sequence-region segA 1 30000", + "##sequence-region segB 1 15000", + "segA\tRefSeq\tgene\t21563\t25384\t.\t+\t.\tID=gene:S;Name=S", + "segA\tRefSeq\tmRNA\t21563\t25384\t.\t+\t.\tID=transcript:S;Name=S", + "segB\tRefSeq\tgene\t28274\t29533\t.\t+\t.\tID=gene:N;Name=N", + "segB\tRefSeq\tmRNA\t28274\t29533\t.\t+\t.\tID=transcript:N;Name=N", + ] + ) + + "\n", + encoding="utf-8", + ) + metadata = build_reference_feature_metadata(gff) + table = _results_table() + + genome_out = tmp_path / "genome.pdf" + focus_out = tmp_path / "genome_focus.pdf" + aa_out = tmp_path / "genome_gene_aa.pdf" + + plot_variant_genome(table, metadata, output_path=genome_out) + plot_variant_genome( + table, + metadata, + output_path=focus_out, + focus_ranges=_parse_focus_ranges("23000-24000,50000-51000"), + gene="S", + ) + plot_variant_genome( + table, + metadata, + output_path=aa_out, + gene="S", + aa_scale=True, + focus_ranges=_parse_focus_ranges("450-650"), + ) + + assert genome_out.exists() + assert focus_out.exists() + assert aa_out.exists() + + +def test_parse_focus_ranges_caps_number_of_regions(): + with pytest.raises(InputValidationError): + _parse_focus_ranges("1-2,3-4,5-6,7-8,9-10,11-12,13-14") diff --git a/vartracker/main.py b/vartracker/main.py index e8204b1..6734bbc 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -44,18 +44,22 @@ from .plotting import ( DEFAULT_LIFESPAN_TOP_N, DEFAULT_TRAJECTORY_TOP_N, + _parse_focus_ranges, apply_shared_plot_filters, auto_select_variants, + load_reference_feature_metadata, collect_explicit_variants, get_threshold_crossing_variants, load_results_table, parse_thresholds, + plot_variant_genome, plot_variant_lifespan, plot_variant_trajectory, plot_variant_turnover, prepare_plot_inputs, _project_name_from_results, resolve_plot_output_path, + write_reference_feature_metadata, ) from ._version import __version__ from .provenance import ( @@ -389,6 +393,9 @@ def _resolve_plot_variants( prefer_crossing: bool = False, thresholds: Sequence[float] | None = None, ) -> list[str]: + def _label_prefix(value: object) -> str: + return str(value).split(" (", 1)[0].strip() + if explicit_variants: selected = [] seen = set() @@ -396,7 +403,12 @@ def _resolve_plot_variants( match = summary[ summary.apply( lambda row: requested - in {row["variant_id"], row["variant_label"], row["variant_name"]}, + in { + row["variant_id"], + row["variant_label"], + _label_prefix(row["variant_label"]), + row["variant_name"], + }, axis=1, ) ] @@ -1116,6 +1128,8 @@ def _run_plot_trajectory_command(args): label_lines=args.label_lines, label_threshold_crossers=args.label_threshold_crossers, crossing_rule=args.crossing_rule, + label_mode=args.label_mode, + sample_axis_mode=args.sample_axis, ) print(f"\nFinished: wrote {output_path}\n") return 0 @@ -1160,6 +1174,7 @@ def _run_plot_turnover_command(args): height=args.height, dpi=args.dpi, count_mode=args.count_mode, + sample_axis_mode=args.sample_axis, ) print(f"\nFinished: wrote {output_path}\n") return 0 @@ -1203,6 +1218,50 @@ def _run_plot_lifespan_command(args): dpi=args.dpi, sort_by=args.sort_by, annotate_class=args.annotate_class, + sample_axis_mode=args.sample_axis, + sample_names=_sample_names, + sample_numbers=_sample_numbers, + ) + print(f"\nFinished: wrote {output_path}\n") + return 0 + except (InputValidationError, ProcessingError) as exc: + print(f"\nERROR: {exc}\n") + return 1 + + +def _run_plot_genome_command(args): + try: + results_csv = Path(args.results_csv).expanduser().resolve() + table = load_results_table(results_csv) + project_name = _project_name_from_results(table, getattr(args, "name", None)) + metadata = load_reference_feature_metadata(results_csv) + output_path = resolve_plot_output_path( + results_csv, + out=args.out, + outdir=args.outdir, + fmt=args.format, + filename="variant_genome_plot", + ) + plot_variant_genome( + table, + metadata, + output_path=output_path, + gene=args.gene, + aa_scale=args.aa_scale, + focus_ranges=_parse_focus_ranges(args.focus_coords), + min_af=args.min_af, + max_af=args.max_af, + effects=_parse_csv_option_list(args.effect), + persistent_only=args.persistent_only, + new_only=args.new_only, + include_indels=args.include_indels, + title=args.title + or ( + f"{project_name}: genome variant distribution" if project_name else None + ), + width=args.width, + height=args.height, + dpi=args.dpi, ) print(f"\nFinished: wrote {output_path}\n") return 0 @@ -1280,6 +1339,18 @@ def _add_plot_trajectory_subparser(subparsers): default=False, help="Add end labels to plotted lines", ) + parser.add_argument( + "--label-mode", + choices=["aa", "nt"], + default="aa", + help="Use amino-acid-style labels or nucleotide-level variant labels (default: aa)", + ) + parser.add_argument( + "--sample-axis", + choices=["number", "name"], + default="number", + help="Use sample numbers or sample names on the x-axis (default: number)", + ) parser.add_argument( "--thresholds", default="", @@ -1330,6 +1401,12 @@ def _add_plot_turnover_subparser(subparsers): default="count", help="Aggregate turnover as counts or summed allele frequencies (default: count)", ) + parser.add_argument( + "--sample-axis", + choices=["number", "name"], + default="number", + help="Use sample numbers or sample names on the x-axis (default: number)", + ) parser.add_argument("--title", default=None, help="Optional plot title") parser.add_argument( "--width", type=float, default=8.5, help="Figure width in inches" @@ -1369,6 +1446,12 @@ def _add_plot_lifespan_subparser(subparsers): default=False, help="Append variant/new and persistent/transient classes to labels", ) + parser.add_argument( + "--sample-axis", + choices=["number", "name"], + default="number", + help="Use sample numbers or sample names on the x-axis (default: number)", + ) parser.add_argument("--title", default=None, help="Optional plot title") parser.add_argument( "--width", type=float, default=8.5, help="Figure width in inches" @@ -1380,6 +1463,75 @@ def _add_plot_lifespan_subparser(subparsers): parser.set_defaults(handler=_run_plot_lifespan_command) +def _add_plot_genome_subparser(subparsers): + parser = subparsers.add_parser( + "genome", + help="Plot collapsed variant allele frequencies along the genome", + description=( + "Generate a genome-position summary plot from a vartracker results CSV." + ), + formatter_class=HelpFormatter, + ) + parser.add_argument("results_csv", help="Path to a vartracker results CSV") + filter_group = parser.add_argument_group("Filtering") + filter_group.add_argument("--gene", default=None, help="Restrict to a single gene") + filter_group.add_argument( + "--effect", + default="", + help="Comma-separated effect classes to include (e.g. missense,synonymous)", + ) + filter_group.add_argument( + "--min-af", + type=float, + default=None, + help="Minimum collapsed allele frequency to include", + ) + filter_group.add_argument( + "--max-af", + type=float, + default=None, + help="Maximum collapsed allele frequency to include", + ) + filter_group.add_argument( + "--persistent-only", + action="store_true", + default=False, + help="Only include variants with persistence_status == new_persistent", + ) + filter_group.add_argument( + "--new-only", + action="store_true", + default=False, + help="Only include variants with variant_status == new", + ) + filter_group.add_argument( + "--include-indels", + action="store_true", + default=False, + help="Include indels as well as SNPs in the genome plot (default: SNPs only)", + ) + parser.add_argument( + "--aa-scale", + action="store_true", + default=False, + help="With --gene, plot x-axis in amino-acid coordinates for that gene", + ) + parser.add_argument( + "--focus-coords", + default="", + help="Comma-separated coordinate ranges to highlight, e.g. 150-300,900-1800", + ) + parser.add_argument("--title", default=None, help="Optional plot title") + parser.add_argument( + "--width", type=float, default=10.0, help="Figure width in inches" + ) + parser.add_argument( + "--height", type=float, default=6.5, help="Figure height in inches" + ) + _add_plot_output_arguments(parser) + parser.set_defaults(handler=_run_plot_genome_command) + + def _add_plot_subparser(subparsers): plot_parser = subparsers.add_parser( "plot", @@ -1389,6 +1541,7 @@ def _add_plot_subparser(subparsers): ) plot_subparsers = plot_parser.add_subparsers(dest="plot_command") _add_plot_heatmap_subparser(plot_subparsers) + _add_plot_genome_subparser(plot_subparsers) _add_plot_trajectory_subparser(plot_subparsers) _add_plot_turnover_subparser(plot_subparsers) _add_plot_lifespan_subparser(plot_subparsers) @@ -1685,6 +1838,9 @@ def resolve_path(value: str) -> str: outputs = { "results_csv": os.path.join(args.outdir, args.filename), "results_metadata": str(results_metadata_path), + "reference_features": os.path.join( + args.outdir, "reference_features.json" + ), "new_mutations_csv": os.path.join(args.outdir, "new_mutations.csv"), "persistent_new_mutations_csv": os.path.join( args.outdir, "persistent_new_mutations.csv" @@ -1698,6 +1854,9 @@ def resolve_path(value: str) -> str: "variant_turnover_plot": os.path.join( args.outdir, "variant_turnover_plot.pdf" ), + "variant_genome_plot": os.path.join( + args.outdir, "variant_genome_plot.pdf" + ), "variant_allele_frequency_heatmap_html": os.path.join( args.outdir, "variant_allele_frequency_heatmap.html" ), @@ -2405,6 +2564,16 @@ def _process_files( output_path=os.path.join(args.outdir, "variant_turnover_plot.pdf"), title=f"{pname}: variant turnover" if pname else None, ) + try: + write_reference_feature_metadata(args.gff3, args.outdir) + plot_variant_genome( + table, + load_reference_feature_metadata(outfile), + output_path=os.path.join(args.outdir, "variant_genome_plot.pdf"), + title=f"{pname}: genome variant distribution" if pname else None, + ) + except (InputValidationError, ProcessingError) as exc: + print(f"Warning: skipped genome plot generation ({exc})") generate_variant_heatmap( table, sample_names, diff --git a/vartracker/plotting.py b/vartracker/plotting.py index 3669096..d19650b 100644 --- a/vartracker/plotting.py +++ b/vartracker/plotting.py @@ -2,19 +2,36 @@ from __future__ import annotations +import json from pathlib import Path -from typing import Iterable, Sequence +from typing import Iterable, Sequence, cast +from urllib.parse import unquote import matplotlib.pyplot as plt import numpy as np import pandas as pd +from matplotlib.gridspec import GridSpec -from .analysis import _coerce_frequency, _resolve_variant_labels +from .analysis import ( + _coerce_frequency, + _extract_numeric_position, + _resolve_variant_labels, +) from .core import InputValidationError, ProcessingError DEFAULT_TRAJECTORY_TOP_N = 12 DEFAULT_LIFESPAN_TOP_N = 20 DEFAULT_TRAJECTORY_THRESHOLDS = (0.5, 0.9) +MAX_FOCUS_RANGES = 6 +FOCUS_RANGE_COLORS = [ + "#f8bbd0", + "#c5cae9", + "#c8e6c9", + "#ffe0b2", + "#d1c4e9", + "#b2dfdb", +] +FeatureRecord = dict[str, object] def _parse_csv_option_list(value: str | None) -> list[str]: @@ -85,10 +102,32 @@ def _project_name_from_results(table: pd.DataFrame, explicit_name: str | None) - return explicit_name if "name" not in table.columns: return "" - names = [str(value).strip() for value in table["name"].unique() if str(value).strip()] + names = [ + str(value).strip() for value in table["name"].unique() if str(value).strip() + ] return names[0] if len(names) == 1 else "" +def _build_variant_key(row: object, fallback_label: str) -> str: + chrom = str(getattr(row, "chrom", "")).strip() + start = str(getattr(row, "start", "")).strip() + ref = str(getattr(row, "ref", "")).strip() + alt = str(getattr(row, "alt", "")).strip() + variant_name = str(getattr(row, "variant", "")).strip() + + if chrom and start and ref and alt: + return f"{chrom}:{start}:{ref}>{alt}" + if chrom and start and variant_name: + return f"{chrom}:{start}:{variant_name}" + if variant_name: + return variant_name + return fallback_label + + +def _display_label_prefix(label: object) -> str: + return str(label).split(" (", 1)[0].strip() + + def _sample_qc_flags(row: object, length: int) -> list[str]: flags = _parse_slash_tokens(getattr(row, "per_sample_variant_qc", "")) if not flags: @@ -106,13 +145,21 @@ def prepare_plot_inputs( for row in table.drop_duplicates().itertuples(index=False): _, display_label, base_label = _resolve_variant_labels(row) + variant_key = _build_variant_key(row, base_label) names = _parse_slash_tokens(getattr(row, "samples", "")) numbers = _parse_slash_tokens(getattr(row, "sample_number", "")) - freqs = [_coerce_frequency(token) for token in _parse_slash_tokens(getattr(row, "alt_freq", ""))] + freqs = [ + _coerce_frequency(token) + for token in _parse_slash_tokens(getattr(row, "alt_freq", "")) + ] presence = _parse_slash_tokens(getattr(row, "presence_absence", "")) qc_flags = _sample_qc_flags(row, len(names)) - length = min(len(names), len(numbers), len(freqs)) if freqs else min(len(names), len(numbers)) + length = ( + min(len(names), len(numbers), len(freqs)) + if freqs + else min(len(names), len(numbers)) + ) if length == 0: continue @@ -130,14 +177,16 @@ def prepare_plot_inputs( ): long_records.append( { - "variant_id": base_label, + "variant_id": variant_key, "variant_label": display_label.replace("\n", " "), "variant_name": str(getattr(row, "variant", "")).strip(), "gene": str(getattr(row, "gene", "")).strip(), "type_of_change": str(getattr(row, "type_of_change", "")).strip(), "type_of_variant": str(getattr(row, "type_of_variant", "")).strip(), "variant_status": str(getattr(row, "variant_status", "")).strip(), - "persistence_status": str(getattr(row, "persistence_status", "")).strip(), + "persistence_status": str( + getattr(row, "persistence_status", "") + ).strip(), "sample_name": sample_name, "sample_number": int(float(sample_number)), "allele_frequency": float(af), @@ -163,39 +212,41 @@ def prepare_plot_inputs( raise ProcessingError("No plottable variant records found in results CSV") long_df = pd.DataFrame(long_records).drop_duplicates() - long_df = long_df.sort_values(["sample_number", "variant_id", "sample_name"]).reset_index(drop=True) - - summary = ( - long_df.groupby("variant_id", as_index=False) - .agg( - variant_label=("variant_label", "first"), - variant_name=("variant_name", "first"), - gene=("gene", "first"), - type_of_change=("type_of_change", "first"), - type_of_variant=("type_of_variant", "first"), - variant_status=("variant_status", "first"), - persistence_status=("persistence_status", "first"), - has_literature=("has_literature", "max"), - max_af=("allele_frequency", "max"), - ) + long_df = long_df.sort_values( + ["sample_number", "variant_id", "sample_name"] + ).reset_index(drop=True) + + summary = long_df.groupby("variant_id", as_index=False).agg( + variant_label=("variant_label", "first"), + variant_name=("variant_name", "first"), + gene=("gene", "first"), + type_of_change=("type_of_change", "first"), + type_of_variant=("type_of_variant", "first"), + variant_status=("variant_status", "first"), + persistence_status=("persistence_status", "first"), + has_literature=("has_literature", "max"), + max_af=("allele_frequency", "max"), ) present_df = long_df[long_df["present"]] first_seen = ( present_df.groupby("variant_id")["sample_number"].min().rename("first_seen") ) - last_seen = present_df.groupby("variant_id")["sample_number"].max().rename("last_seen") + last_seen = ( + present_df.groupby("variant_id")["sample_number"].max().rename("last_seen") + ) summary = summary.merge(first_seen, on="variant_id", how="left") summary = summary.merge(last_seen, on="variant_id", how="left") - summary["first_seen"] = summary["first_seen"].fillna(summary["max_af"].map(lambda _x: np.nan)) - summary["last_seen"] = summary["last_seen"].fillna(summary["first_seen"]) - summary["duration"] = ( - summary["last_seen"].fillna(0).astype(float) - - summary["first_seen"].fillna(0).astype(float) + summary["first_seen"] = summary["first_seen"].fillna( + summary["max_af"].map(lambda _x: np.nan) ) + summary["last_seen"] = summary["last_seen"].fillna(summary["first_seen"]) + summary["duration"] = summary["last_seen"].fillna(0).astype(float) - summary[ + "first_seen" + ].fillna(0).astype(float) summary["is_persistent_new"] = summary["persistence_status"].eq("new_persistent") - summary["is_nonsynonymous"] = ~summary["type_of_change"].str.lower().str.contains( - "synonymous", na=False + summary["is_nonsynonymous"] = ( + ~summary["type_of_change"].str.lower().str.contains("synonymous", na=False) ) return summary, long_df, sample_names, sample_numbers @@ -222,24 +273,32 @@ def apply_shared_plot_filters( if sample_max is not None: filtered_long = filtered_long[filtered_long["sample_number"] <= sample_max] if filtered_long.empty: - raise ProcessingError("No variant observations remain after sample-range filtering") + raise ProcessingError( + "No variant observations remain after sample-range filtering" + ) filtered_summary = summary.copy() if gene: filtered_summary = filtered_summary[ - filtered_summary["gene"].astype(str).str.lower() == str(gene).strip().lower() + filtered_summary["gene"].astype(str).str.lower() + == str(gene).strip().lower() ] if effects: - lowered = {str(effect).strip().lower() for effect in effects if str(effect).strip()} + lowered = { + str(effect).strip().lower() for effect in effects if str(effect).strip() + } filtered_summary = filtered_summary[ filtered_summary["type_of_change"].astype(str).str.lower().isin(lowered) ] if not include_synonymous: filtered_summary = filtered_summary[ - ~filtered_summary["type_of_change"].astype(str).str.lower().str.contains("synonymous", na=False) + ~filtered_summary["type_of_change"] + .astype(str) + .str.lower() + .str.contains("synonymous", na=False) ] if persistent_only: @@ -248,7 +307,9 @@ def apply_shared_plot_filters( ] if new_only: - filtered_summary = filtered_summary[filtered_summary["variant_status"].eq("new")] + filtered_summary = filtered_summary[ + filtered_summary["variant_status"].eq("new") + ] if variants: order_map = {name: idx for idx, name in enumerate(variants)} @@ -259,16 +320,24 @@ def apply_shared_plot_filters( for candidate in ( row["variant_id"], row["variant_label"], + _display_label_prefix(row["variant_label"]), row["variant_name"], ) ), axis=1, ) ] + if filtered_summary.empty: + raise ProcessingError("No variants remain after applying plot filters") filtered_summary["selection_order"] = filtered_summary.apply( lambda row: min( order_map[candidate] - for candidate in (row["variant_id"], row["variant_label"], row["variant_name"]) + for candidate in ( + row["variant_id"], + row["variant_label"], + _display_label_prefix(row["variant_label"]), + row["variant_name"], + ) if candidate in order_map ), axis=1, @@ -280,12 +349,11 @@ def apply_shared_plot_filters( if filtered_summary.empty: raise ProcessingError("No variants remain after applying plot filters") - filtered_long = filtered_long[filtered_long["variant_id"].isin(filtered_summary["variant_id"])] - recomputed = ( - filtered_long.groupby("variant_id", as_index=False) - .agg( - max_af=("allele_frequency", "max"), - ) + filtered_long = filtered_long[ + filtered_long["variant_id"].isin(filtered_summary["variant_id"]) + ] + recomputed = filtered_long.groupby("variant_id", as_index=False).agg( + max_af=("allele_frequency", "max"), ) filtered_summary = filtered_summary.drop(columns=["max_af"], errors="ignore").merge( recomputed, on="variant_id", how="left" @@ -307,10 +375,9 @@ def apply_shared_plot_filters( ) filtered_summary = filtered_summary.merge(first_seen, on="variant_id", how="left") filtered_summary = filtered_summary.merge(last_seen, on="variant_id", how="left") - filtered_summary["duration"] = ( - filtered_summary["last_seen"].fillna(0).astype(float) - - filtered_summary["first_seen"].fillna(0).astype(float) - ) + filtered_summary["duration"] = filtered_summary["last_seen"].fillna(0).astype( + float + ) - filtered_summary["first_seen"].fillna(0).astype(float) if min_af is not None: filtered_summary = filtered_summary[filtered_summary["max_af"] >= min_af] @@ -319,7 +386,9 @@ def apply_shared_plot_filters( if filtered_summary.empty: raise ProcessingError("No variants remain after allele-frequency filtering") - filtered_long = filtered_long[filtered_long["variant_id"].isin(filtered_summary["variant_id"])] + filtered_long = filtered_long[ + filtered_long["variant_id"].isin(filtered_summary["variant_id"]) + ] return filtered_summary.reset_index(drop=True), filtered_long.reset_index(drop=True) @@ -337,7 +406,9 @@ def auto_select_variants( threshold_values = [float(value) for value in (thresholds or [])] if prefer_crossing and threshold_values: ranked["crosses_threshold"] = ranked["max_af"].apply( - lambda value: any(float(value) >= threshold for threshold in threshold_values) + lambda value: any( + float(value) >= threshold for threshold in threshold_values + ) ) else: ranked["crosses_threshold"] = False @@ -373,7 +444,12 @@ def _resolve_variant_selection( match = summary[ summary.apply( lambda row: requested - in {row["variant_id"], row["variant_label"], row["variant_name"]}, + in { + row["variant_id"], + row["variant_label"], + _display_label_prefix(row["variant_label"]), + row["variant_name"], + }, axis=1, ) ] @@ -384,7 +460,9 @@ def _resolve_variant_selection( selected.append(variant_id) seen.add(variant_id) if not selected: - raise ProcessingError("None of the requested variants were found in the filtered results") + raise ProcessingError( + "None of the requested variants were found in the filtered results" + ) return selected return auto_select_variants( summary, @@ -416,6 +494,36 @@ def _apply_plot_title(ax, title: str | None, default_title: str) -> None: ax.set_title(title or default_title, fontweight="bold") +def _sample_axis_ticks_from_long_df( + long_df: pd.DataFrame, +) -> tuple[list[int], list[str]]: + sample_axis = ( + long_df[["sample_number", "sample_name"]] + .drop_duplicates() + .sort_values(["sample_number", "sample_name"]) + ) + return ( + sample_axis["sample_number"].astype(int).tolist(), + sample_axis["sample_name"].astype(str).tolist(), + ) + + +def _apply_sample_axis( + ax, + *, + sample_numbers: Sequence[int], + sample_names: Sequence[str], + sample_axis_mode: str, +) -> None: + ax.set_xticks(list(sample_numbers)) + if sample_axis_mode == "name": + ax.set_xlabel("Sample Name") + ax.set_xticklabels(list(sample_names), rotation=45, ha="right") + else: + ax.set_xlabel("Sample Number") + ax.set_xticklabels([str(value) for value in sample_numbers]) + + def plot_variant_trajectory( summary: pd.DataFrame, long_df: pd.DataFrame, @@ -430,6 +538,8 @@ def plot_variant_trajectory( label_lines: bool = False, label_threshold_crossers: bool = False, crossing_rule: str = "at_or_above", + label_mode: str = "aa", + sample_axis_mode: str = "number", ) -> None: plot_df = long_df[long_df["variant_id"].isin(selected_variants)].copy() if plot_df.empty: @@ -438,8 +548,16 @@ def plot_variant_trajectory( fig, ax = plt.subplots(figsize=(width, height)) threshold_values = [float(value) for value in (thresholds or [])] for variant_id in selected_variants: - subset = plot_df[plot_df["variant_id"] == variant_id].sort_values("sample_number") - label = str(subset["variant_id"].iloc[0]) + subset = plot_df[plot_df["variant_id"] == variant_id].sort_values( + "sample_number" + ) + if label_mode == "nt": + label = ( + str(subset["variant_name"].iloc[0]).strip() + or str(subset["variant_label"].iloc[0]).split(" (", 1)[0] + ) + else: + label = str(subset["variant_label"].iloc[0]).split(" (", 1)[0] ax.plot( subset["sample_number"], subset["allele_frequency"], @@ -452,7 +570,9 @@ def plot_variant_trajectory( threshold_values, crossing_rule=crossing_rule, ) - if (label_lines or (label_threshold_crossers and crosses_threshold)) and not subset.empty: + if ( + label_lines or (label_threshold_crossers and crosses_threshold) + ) and not subset.empty: last_row = subset.iloc[-1] ax.text( float(last_row["sample_number"]) + 0.1, @@ -462,21 +582,18 @@ def plot_variant_trajectory( va="center", ) - ax.set_xlabel("Sample") ax.set_ylabel("Allele frequency") ax.set_ylim(0, 1.02) - ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + axis_numbers, axis_names = _sample_axis_ticks_from_long_df(plot_df) + _apply_sample_axis( + ax, + sample_numbers=axis_numbers, + sample_names=axis_names, + sample_axis_mode=sample_axis_mode, + ) for threshold in threshold_values: - ax.axhline(float(threshold), linestyle="--", linewidth=1.0, color="black", alpha=0.6) - xmax = max(plot_df["sample_number"]) if not plot_df.empty else 0 - ax.text( - float(xmax) + 0.1, - float(threshold), - f"{threshold:g}", - fontsize=8, - va="center", - ha="left", - color="black", + ax.axhline( + float(threshold), linestyle="--", linewidth=1.0, color="black", alpha=0.6 ) _apply_plot_title(ax, title, "Variant trajectories") if not label_lines and not label_threshold_crossers: @@ -496,6 +613,7 @@ def plot_variant_turnover( height: float = 5.0, dpi: int = 300, count_mode: str = "count", + sample_axis_mode: str = "number", ) -> None: plot_df = long_df.sort_values(["variant_id", "sample_number"]).copy() records: list[dict[str, float | int]] = [] @@ -506,7 +624,9 @@ def plot_variant_turnover( previous = variant_df.iloc[idx - 1] current = variant_df.iloc[idx] if not bool(previous["present"]) and bool(current["present"]): - value = 1.0 if count_mode == "count" else float(current["allele_frequency"]) + value = ( + 1.0 if count_mode == "count" else float(current["allele_frequency"]) + ) records.append( { "sample_number": int(current["sample_number"]), @@ -515,7 +635,11 @@ def plot_variant_turnover( } ) if bool(previous["present"]) and not bool(current["present"]): - value = 1.0 if count_mode == "count" else float(previous["allele_frequency"]) + value = ( + 1.0 + if count_mode == "count" + else float(previous["allele_frequency"]) + ) records.append( { "sample_number": int(current["sample_number"]), @@ -528,7 +652,9 @@ def plot_variant_turnover( turnover = pd.DataFrame({"sample_number": sample_numbers, "new": 0.0, "lost": 0.0}) if records: events = pd.DataFrame(records).groupby("sample_number", as_index=False).sum() - turnover = turnover.merge(events, on="sample_number", how="left", suffixes=("", "_event")) + turnover = turnover.merge( + events, on="sample_number", how="left", suffixes=("", "_event") + ) turnover["new"] = turnover["new_event"].fillna(turnover["new"]) turnover["lost"] = turnover["lost_event"].fillna(turnover["lost"]) turnover = turnover.drop(columns=["new_event", "lost_event"]) @@ -537,9 +663,16 @@ def plot_variant_turnover( ax.bar(turnover["sample_number"], turnover["new"], label="New") ax.bar(turnover["sample_number"], -turnover["lost"], label="Lost") ax.axhline(0, color="black", linewidth=0.8) - ax.set_xlabel("Sample") - ax.set_ylabel("Variant count" if count_mode == "count" else "Summed allele frequency") - ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + ax.set_ylabel( + "Variant count" if count_mode == "count" else "Summed allele frequency" + ) + axis_numbers, axis_names = _sample_axis_ticks_from_long_df(plot_df) + _apply_sample_axis( + ax, + sample_numbers=axis_numbers, + sample_names=axis_names, + sample_axis_mode=sample_axis_mode, + ) _apply_plot_title(ax, title, "Variant turnover") ax.legend(frameon=False) fig.tight_layout() @@ -558,13 +691,23 @@ def plot_variant_lifespan( dpi: int = 300, sort_by: str = "duration", annotate_class: bool = False, + sample_axis_mode: str = "number", + sample_names: Sequence[str] | None = None, + sample_numbers: Sequence[int] | None = None, ) -> None: plot_df = summary[summary["variant_id"].isin(selected_variants)].copy() if plot_df.empty: raise ProcessingError("No data available for the lifespan plot") - ascending = {"first_seen": True, "last_seen": True, "duration": False, "max_af": False} - plot_df = plot_df.sort_values(sort_by, ascending=ascending.get(sort_by, False), na_position="last") + ascending = { + "first_seen": True, + "last_seen": True, + "duration": False, + "max_af": False, + } + plot_df = plot_df.sort_values( + sort_by, ascending=ascending.get(sort_by, False), na_position="last" + ) plot_df = plot_df.reset_index(drop=True) labels = plot_df["variant_id"].astype(str).tolist() @@ -588,8 +731,30 @@ def plot_variant_lifespan( ax.set_yticks(y_positions) ax.set_yticklabels(labels) ax.invert_yaxis() - ax.set_xlabel("Sample") - ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True)) + if sample_numbers is None: + inferred_numbers = sorted( + { + int(value) + for value in pd.concat([starts, ends]) + .dropna() + .astype(float) + .astype(int) + .tolist() + } + ) + else: + inferred_numbers = [int(value) for value in sample_numbers] + inferred_names = ( + [str(value) for value in sample_names] + if sample_names is not None + else [str(value) for value in inferred_numbers] + ) + _apply_sample_axis( + ax, + sample_numbers=inferred_numbers, + sample_names=inferred_names, + sample_axis_mode=sample_axis_mode, + ) _apply_plot_title(ax, title, "Variant lifespan") fig.tight_layout() fig.savefig(output_path, dpi=dpi, bbox_inches="tight") @@ -624,7 +789,9 @@ def parse_thresholds(value: str | None) -> list[float]: try: return [float(item) for item in thresholds] except ValueError as exc: - raise InputValidationError("Thresholds must be comma-separated numbers") from exc + raise InputValidationError( + "Thresholds must be comma-separated numbers" + ) from exc def variant_crosses_thresholds( @@ -637,7 +804,9 @@ def variant_crosses_thresholds( return False values_list = [float(value) for value in values] if crossing_rule == "strictly_above": - return any(value > threshold for value in values_list for threshold in thresholds) + return any( + value > threshold for value in values_list for threshold in thresholds + ) return any(value >= threshold for value in values_list for threshold in thresholds) @@ -649,3 +818,463 @@ def collect_explicit_variants( if requested and file_variants: raise InputValidationError("Use either --variants or --variant-file, not both.") return requested or file_variants + + +def _parse_focus_ranges(value: str | None) -> list[tuple[float, float]]: + if not value: + return [] + ranges: list[tuple[float, float]] = [] + for part in _parse_csv_option_list(value): + if "-" not in part: + raise InputValidationError( + "Focus ranges must use start-end syntax, e.g. 150-300,900-1800" + ) + start_text, end_text = part.split("-", 1) + try: + start = float(start_text) + end = float(end_text) + except ValueError as exc: + raise InputValidationError( + "Focus ranges must use numeric start-end coordinates" + ) from exc + if end < start: + start, end = end, start + ranges.append((start, end)) + if len(ranges) > MAX_FOCUS_RANGES: + raise InputValidationError( + f"Use at most {MAX_FOCUS_RANGES} focus ranges so each can retain a distinct highlight color" + ) + return ranges + + +def _parse_gff_attrs(attr_text: str) -> dict[str, str]: + attrs: dict[str, str] = {} + for part in attr_text.split(";"): + if not part: + continue + if "=" in part: + key, value = part.split("=", 1) + attrs[key] = unquote(value) + return attrs + + +def build_reference_feature_metadata(gff_path: str | Path) -> dict[str, object]: + path = Path(gff_path).expanduser().resolve() + if not path.exists(): + raise InputValidationError(f"GFF3 not found: {path}") + + contig_lengths: dict[str, int] = {} + contig_order: list[str] = [] + transcript_features: list[FeatureRecord] = [] + gene_features: list[FeatureRecord] = [] + max_end_by_contig: dict[str, int] = {} + + with path.open(encoding="utf-8") as handle: + for raw in handle: + line = raw.strip() + if not line: + continue + if line.startswith("##sequence-region"): + parts = line.split() + if len(parts) >= 4: + contig = parts[1] + seq_start = int(parts[2]) + seq_end = int(parts[3]) + if contig not in contig_order: + contig_order.append(contig) + contig_lengths[contig] = max( + contig_lengths.get(contig, 0), seq_end - seq_start + 1 + ) + continue + if line.startswith("#"): + continue + + seqid, _source, feature_type, start, end, _score, strand, _phase, attrs = ( + line.split("\t") + ) + start_i = int(start) + end_i = int(end) + max_end_by_contig[seqid] = max(max_end_by_contig.get(seqid, 0), end_i) + if seqid not in contig_order: + contig_order.append(seqid) + + parsed = _parse_gff_attrs(attrs) + name = ( + parsed.get("Name") + or parsed.get("gene") + or parsed.get("product") + or parsed.get("ID", "").split(":")[-1] + ) + if not name: + continue + record = { + "contig": seqid, + "start": start_i, + "end": end_i, + "name": str(name), + "strand": strand, + "type": feature_type, + "aa_length": max(1, int((end_i - start_i + 1) / 3)), + } + if feature_type in {"mRNA", "transcript"}: + transcript_features.append(record) + elif feature_type == "gene": + gene_features.append(record) + + for contig, max_end in max_end_by_contig.items(): + contig_lengths.setdefault(contig, max_end) + + features = transcript_features or gene_features + deduped_features: list[FeatureRecord] = [] + seen = set() + for feature in features: + key = ( + feature["contig"], + feature["start"], + feature["end"], + feature["name"], + ) + if key in seen: + continue + seen.add(key) + deduped_features.append(feature) + + return { + "source_gff3": str(path), + "contig_order": contig_order, + "contig_lengths": contig_lengths, + "features": deduped_features, + } + + +def write_reference_feature_metadata(gff_path: str | Path, outdir: str | Path) -> Path: + outdir_path = Path(outdir).expanduser().resolve() + outdir_path.mkdir(parents=True, exist_ok=True) + metadata = build_reference_feature_metadata(gff_path) + destination = outdir_path / "reference_features.json" + destination.write_text( + json.dumps(metadata, indent=2, sort_keys=True), encoding="utf-8" + ) + return destination + + +def load_reference_feature_metadata(results_csv: str | Path) -> dict[str, object]: + results_path = Path(results_csv).expanduser().resolve() + sidecar = results_path.parent / "reference_features.json" + if not sidecar.exists(): + raise InputValidationError( + f"Reference feature metadata not found beside results CSV: {sidecar}" + ) + return cast(dict[str, object], json.loads(sidecar.read_text(encoding="utf-8"))) + + +def _collapse_variants_for_genome_plot( + table: pd.DataFrame, +) -> pd.DataFrame: + records: list[dict[str, object]] = [] + for row in table.drop_duplicates().itertuples(index=False): + gene_label, display_label, base_label = _resolve_variant_labels(row) + variant_key = _build_variant_key(row, base_label) + af_values = [ + _coerce_frequency(token) + for token in _parse_slash_tokens(getattr(row, "alt_freq", "")) + ] + aa_position = _extract_numeric_position( + base_label.split(":", 1)[1] if ":" in base_label else base_label + ) + records.append( + { + "variant_id": variant_key, + "variant_label": display_label.replace("\n", " "), + "chrom": str(getattr(row, "chrom", "")).strip(), + "start": int(getattr(row, "start", 0)), + "end": int(getattr(row, "end", getattr(row, "start", 0))), + "plot_gene": str(gene_label).strip(), + "raw_gene": str(getattr(row, "gene", "")).strip(), + "type_of_change": str(getattr(row, "type_of_change", "")).strip(), + "type_of_variant": str(getattr(row, "type_of_variant", "")).strip(), + "variant_status": str(getattr(row, "variant_status", "")).strip(), + "persistence_status": str( + getattr(row, "persistence_status", "") + ).strip(), + "af_values": af_values, + "summary_af": max(af_values) if af_values else 0.0, + "aa_position": aa_position, + } + ) + + collapsed = pd.DataFrame(records) + if collapsed.empty: + raise ProcessingError("No variants available for genome plotting") + + grouped_rows: list[dict[str, object]] = [] + skipped_variants = 0 + for variant_id, group in collapsed.groupby("variant_id", dropna=False): + chroms = group["chrom"].dropna().unique().tolist() + starts = group["start"].dropna().unique().tolist() + ends = group["end"].dropna().unique().tolist() + if len(chroms) > 1 or len(starts) > 1 or len(ends) > 1: + skipped_variants += 1 + print( + f"Warning: skipped variant '{variant_id}' due to inconsistent positional metadata" + ) + continue + first = group.iloc[0].to_dict() + merged_af_values: list[float] = [] + for value in group["af_values"].tolist(): + merged_af_values.extend([float(item) for item in cast(list[float], value)]) + first["af_values"] = merged_af_values + first["summary_af"] = float(group["summary_af"].max()) + grouped_rows.append(first) + + if skipped_variants: + print( + f"Warning: skipped {skipped_variants} variant(s) with inconsistent positional metadata" + ) + if not grouped_rows: + raise ProcessingError("No variants available for genome plotting") + + return ( + pd.DataFrame(grouped_rows) + .sort_values(["chrom", "start", "variant_id"]) + .reset_index(drop=True) + ) + + +def _subset_collapsed_variants( + collapsed: pd.DataFrame, + *, + gene: str | None = None, + effects: Sequence[str] | None = None, + min_af: float | None = None, + max_af: float | None = None, + persistent_only: bool = False, + new_only: bool = False, + include_indels: bool = False, +) -> pd.DataFrame: + subset = collapsed.copy() + if not include_indels and "type_of_variant" in subset.columns: + subset = subset[subset["type_of_variant"].astype(str).str.lower() == "snp"] + if gene: + subset = subset[ + subset["plot_gene"].astype(str).str.lower() == str(gene).strip().lower() + ] + if effects: + lowered = { + str(effect).strip().lower() for effect in effects if str(effect).strip() + } + subset = subset[subset["type_of_change"].astype(str).str.lower().isin(lowered)] + if min_af is not None: + subset = subset[subset["summary_af"] >= min_af] + if max_af is not None: + subset = subset[subset["summary_af"] <= max_af] + if persistent_only: + subset = subset[subset["persistence_status"].eq("new_persistent")] + if new_only: + subset = subset[subset["variant_status"].eq("new")] + if subset.empty: + raise ProcessingError("No variants remain for the genome plot after filtering") + return subset + + +def _find_gene_feature(metadata: dict[str, object], gene_name: str) -> FeatureRecord: + features = cast(list[FeatureRecord], metadata.get("features", [])) + matches = [ + feature + for feature in features + if str(feature.get("name", "")).strip().lower() == gene_name.strip().lower() + ] + if not matches: + raise InputValidationError(f"Gene not found in reference features: {gene_name}") + unique_regions = { + (match["contig"], match["start"], match["end"], match["name"]) + for match in matches + } + if len(unique_regions) != 1: + raise InputValidationError( + f"Gene '{gene_name}' maps to multiple regions; --aa-scale requires a unique gene region" + ) + return matches[0] + + +def _clip_focus_ranges( + ranges: Sequence[tuple[float, float]], xmin: float, xmax: float +) -> tuple[list[tuple[float, float]], int]: + clipped: list[tuple[float, float]] = [] + ignored = 0 + for start, end in ranges: + if end < xmin or start > xmax: + ignored += 1 + continue + clipped.append((max(start, xmin), min(end, xmax))) + return clipped, ignored + + +def plot_variant_genome( + table: pd.DataFrame, + reference_metadata: dict[str, object], + *, + output_path: str | Path, + gene: str | None = None, + aa_scale: bool = False, + focus_ranges: Sequence[tuple[float, float]] | None = None, + min_af: float | None = None, + max_af: float | None = None, + effects: Sequence[str] | None = None, + persistent_only: bool = False, + new_only: bool = False, + include_indels: bool = False, + title: str | None = None, + width: float = 10.0, + height: float = 6.5, + dpi: int = 300, +) -> None: + if aa_scale and not gene: + raise InputValidationError("--aa-scale requires --gene") + if include_indels: + print( + "Warning: including indels in the genome plot may produce ambiguous or hard-to-interpret positions" + ) + + collapsed = _collapse_variants_for_genome_plot(table) + collapsed = _subset_collapsed_variants( + collapsed, + gene=gene, + effects=effects, + min_af=min_af, + max_af=max_af, + persistent_only=persistent_only, + new_only=new_only, + include_indels=include_indels, + ) + + contig_order = cast(list[str], reference_metadata.get("contig_order", [])) + contig_lengths = { + str(key): int(cast(int | float | str, value)) + for key, value in cast( + dict[str, object], reference_metadata.get("contig_lengths", {}) + ).items() + } + + region_by_contig: dict[str, tuple[float, float]] = {} + plot_df = collapsed.copy() + + if gene: + gene_feature = _find_gene_feature(reference_metadata, gene) + contig = str(gene_feature["contig"]) + start = float(cast(int | float | str, gene_feature["start"])) + end = float(cast(int | float | str, gene_feature["end"])) + plot_df = plot_df[plot_df["chrom"] == contig] + plot_df = plot_df[(plot_df["start"] >= start) & (plot_df["start"] <= end)] + if plot_df.empty: + raise ProcessingError(f"No variants remain in gene region '{gene}'") + if aa_scale: + omitted = plot_df["aa_position"].isna().sum() + if omitted: + print( + f"Warning: omitted {omitted} variants without amino-acid coordinates" + ) + plot_df = plot_df.dropna(subset=["aa_position"]).copy() + if plot_df.empty: + raise ProcessingError( + "No variants could be mapped to amino-acid coordinates" + ) + plot_df["x_coord"] = plot_df["aa_position"].astype(float) + region_by_contig[contig] = ( + 1.0, + float(cast(int | float | str, gene_feature.get("aa_length", 1))), + ) + else: + plot_df["x_coord"] = plot_df["start"].astype(float) + region_by_contig[contig] = (start, end) + else: + plot_df["x_coord"] = plot_df["start"].astype(float) + for contig in plot_df["chrom"].unique(): + region_by_contig[str(contig)] = ( + 1.0, + float( + contig_lengths.get( + str(contig), + cast( + int | float, + plot_df.loc[plot_df["chrom"] == contig, "end"].max(), + ), + ) + ), + ) + + if not contig_order: + contig_order = sorted(plot_df["chrom"].astype(str).unique().tolist()) + selected_contigs = [ + contig + for contig in contig_order + if contig in plot_df["chrom"].astype(str).unique().tolist() + ] + if not selected_contigs: + raise ProcessingError("No contigs remain for genome plotting") + + n_panels = len(selected_contigs) + fig = plt.figure( + figsize=(width, max(height, 2.1 * n_panels)), + constrained_layout=True, + ) + grid = GridSpec( + n_panels, + 1, + figure=fig, + hspace=0.35, + ) + + focus = list(focus_ranges or []) + for idx, contig in enumerate(selected_contigs): + ax = fig.add_subplot(grid[idx, 0]) + subset = plot_df[plot_df["chrom"].astype(str) == contig].sort_values("x_coord") + xmin, xmax = region_by_contig[contig] + clipped_focus, ignored_focus = _clip_focus_ranges(focus, xmin, xmax) + for focus_index, (start, end) in enumerate(clipped_focus): + ax.axvspan( + start, + end, + color=FOCUS_RANGE_COLORS[focus_index % len(FOCUS_RANGE_COLORS)], + alpha=0.25, + zorder=0, + ) + if ignored_focus: + print( + f"Warning: ignored {ignored_focus} focus range(s) outside the plotted region" + ) + + for row in subset.itertuples(index=False): + af_values = [float(value) for value in cast(list[float], row.af_values)] + if af_values: + ax.vlines( + float(row.x_coord), + min(af_values), + max(af_values), + color="#607d8b", + linewidth=0.9, + alpha=0.55, + zorder=1, + ) + ax.scatter( + [float(row.x_coord)] * len(af_values), + af_values, + s=18, + alpha=0.45, + color="#1f77b4", + edgecolors="none", + zorder=2, + ) + ax.axhline(0.5, color="black", linestyle="--", linewidth=0.8, alpha=0.4) + ax.set_ylim(0, 1.02) + ax.set_ylabel("Allele frequency") + ax.set_xlim(xmin, xmax) + ax.set_title(str(contig), loc="left", fontsize=10, fontweight="bold") + ax.set_xlabel("Amino-acid position" if aa_scale else "Genomic position") + for spine in ("top", "right"): + ax.spines[spine].set_visible(False) + + if title: + fig.suptitle(title, fontweight="bold") + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) From 796cb14a2955ab7a5cc83d4e4ce2a8e2ca8e9599 Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Tue, 7 Apr 2026 10:55:42 +1000 Subject: [PATCH 07/20] Updated plotting capabilities --- README.md | 18 +- tests/test_main.py | 163 +++++++++----- tests/test_plotting.py | 99 ++++++++- vartracker/main.py | 40 +++- vartracker/plotting.py | 470 ++++++++++++++++++++++++++++++++++++----- 5 files changed, 676 insertions(+), 114 deletions(-) diff --git a/README.md b/README.md index 2ea2207..1cc5615 100755 --- a/README.md +++ b/README.md @@ -271,19 +271,20 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test ### Mode-specific options -- `vartracker vcf` – accepts plotting and filtering options such as `--min-snv-freq`, `--min-indel-freq`, - `--allele-frequency-tag`, `--heatmap-aa-exclude`, `--heatmap-aa-include`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls +- `vartracker vcf` – accepts core analysis options such as `--min-snv-freq`, `--min-indel-freq`, + `--allele-frequency-tag`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. - `vartracker bam` – everything from `vcf`, plus Snakemake options: `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. - `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. -- `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV. +- `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV, including all heatmap customization filters. - `vartracker plot genome` – plot SNP positions along the genome or a selected gene region using all observed allele-frequency values for each variant. - `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. - `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. - `vartracker plot lifespan` – plot first-to-last detection spans for a selected or auto-ranked subset of variants. Heatmap filtering: +- `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customize heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. - By default, all consequence classes are included except joint variants. Use `--heatmap-include-joint` to show joint variants. - `--heatmap-aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. - `--heatmap-aa-include`: comma-separated `type_of_change` patterns to include. @@ -319,7 +320,11 @@ Standalone plot output: Genome plot options: - `--gene`: zoom to a single gene region. - `--aa-scale`: with `--gene`, use amino-acid coordinates on the x-axis. -- `--focus-coords`: highlight nucleotide or amino-acid coordinate ranges, depending on the current x-axis mode. +- `--cds-scale`: with `--gene`, use CDS-relative nucleotide coordinates on the x-axis. +- `--focus-coords`: highlight nucleotide or amino-acid coordinate ranges, depending on the current x-axis mode. Separate color groups with `;`, ranges within a group with `,`, and optionally prefix a group with `Name:`. +- `--focus-region-file`: read named focus region groups from a `.json`, `.csv`, or `.tsv` file for an inset legend. +- `--show-intersections`: add a compact `Region | Variant` table below the genome plot for highlighted-region hits. +- In the genome plot, undetected samples are rendered at the detection threshold rather than zero; by default this floor is `0.03`, or `--min-af` if supplied, and the dashed guide line follows that same threshold. - `--include-indels`: opt in to plotting indels too. This may be ambiguous or hard to interpret. - The standalone genome plot auto-discovers `reference_features.json` beside `results.csv`; workflow runs generate this sidecar automatically. @@ -333,7 +338,12 @@ Standalone plot examples: - `vartracker plot genome results.csv` - `vartracker plot genome results.csv --gene F` - `vartracker plot genome results.csv --gene F --aa-scale` +- `vartracker plot genome results.csv --gene F --cds-scale --focus-coords "184-210,586-630"` - `vartracker plot genome results.csv --focus-coords "150-300,900-1800"` +- `vartracker plot genome results.csv --focus-coords "62-69,196-210;31-42,323-332,379-399;254-277"` +- `vartracker plot genome results.csv --gene F --aa-scale --focus-coords "Ø:62-69,196-210;I:31-42,323-332,379-399;II:254-277"` +- `vartracker plot genome results.csv --gene F --aa-scale --focus-coords "Ø:62-69,196-210;I:31-42,323-332,379-399" --show-intersections` +- `vartracker plot genome results.csv --focus-region-file fusion_regions.json` - `vartracker plot genome results.csv --gene F --aa-scale --focus-coords "50-120,180-220"` - `vartracker plot turnover results.csv` - `vartracker plot trajectory results.csv --variants "S:D614G,S:E484K"` diff --git a/tests/test_main.py b/tests/test_main.py index a2627f9..15be8aa 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -293,9 +293,7 @@ def fake_parse_pokay(argv): assert (outdir / "literature_database.csv").exists() -def test_vcf_heatmap_options_are_forwarded_to_heatmap( - monkeypatch, tmp_path, minimal_vcf -): +def test_vcf_workflow_uses_default_heatmap_options(monkeypatch, tmp_path, minimal_vcf): coverage_path = tmp_path / "sample.cov.txt" coverage_path.write_text("NC_045512.2\t266\t100\n", encoding="utf-8") @@ -367,57 +365,29 @@ def fake_heatmap(*args, **kwargs): str(csv_path), "--outdir", str(tmp_path / "results"), - "--heatmap-aa-exclude", - "synonymous,*frameshift*,stop_gained", - "--heatmap-aa-include", - "missense", - "--heatmap-only-persistent", - "--heatmap-only-new", - "--heatmap-gene-include", - "S", - "--heatmap-gene-exclude", - "N", - "--heatmap-variant-type", - "snp", - "--heatmap-qc", - "PASS", - "--min-prop-passing-qc", - "0.75", - "--heatmap-min-persistence", - "2", - "--heatmap-min-max-af", - "0.4", - "--heatmap-min-sample-af", - "0.3", - "--heatmap-sample-subset", - "Sample1", - "--heatmap-hide-singletons", - "--heatmap-min-depth", - "20", ] ) assert exit_code == 0 - assert recorded["excluded_consequence_types"] == [ - "synonymous", - "*frameshift*", - "stop_gained", - ] - assert recorded["included_consequence_types"] == ["missense"] - assert recorded["include_joint"] is False - assert recorded["only_persistent"] is True - assert recorded["only_new"] is True - assert recorded["gene_include"] == ["S"] - assert recorded["gene_exclude"] == ["N"] - assert recorded["variant_type_include"] == ["snp"] - assert recorded["qc_include"] == ["PASS"] - assert recorded["min_prop_passing_qc"] == 0.75 - assert recorded["min_persistence"] == 2 - assert recorded["min_max_af"] == 0.4 - assert recorded["min_sample_af"] == 0.3 - assert recorded["sample_subset"] == ["Sample1"] - assert recorded["hide_singletons"] is True - assert recorded["min_depth"] == 20 + for key in ( + "excluded_consequence_types", + "included_consequence_types", + "include_joint", + "only_persistent", + "only_new", + "gene_include", + "gene_exclude", + "variant_type_include", + "qc_include", + "min_prop_passing_qc", + "min_persistence", + "min_max_af", + "min_sample_af", + "sample_subset", + "hide_singletons", + "min_depth", + ): + assert key not in recorded def test_plot_heatmap_replots_results_csv(monkeypatch, tmp_path): @@ -749,12 +719,103 @@ def fake_genome(table, metadata, **kwargs): assert recorded["rows"] == 6 assert recorded["gene"] == "S" assert recorded["aa_scale"] is True - assert recorded["focus_ranges"] == [(50.0, 120.0), (180.0, 220.0)] + assert recorded["focus_ranges"] == [(50.0, 120.0, 0), (180.0, 220.0, 1)] + assert recorded["focus_labels"] == [None, None] assert recorded["min_af"] == 0.2 assert recorded["include_indels"] is False + assert recorded["show_intersections"] is False assert recorded["metadata"]["contig_order"] == ["segA", "segB"] +def test_plot_genome_named_focus_regions_are_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + _write_reference_features_json(tmp_path / "reference_features.json") + + recorded = {} + + def fake_genome(_table, _metadata, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + + exit_code = main_module.main( + [ + "plot", + "genome", + str(results_csv), + "--focus-coords", + "I:50-120,180-220;II:300-330", + ] + ) + + assert exit_code == 0 + assert recorded["focus_ranges"] == [ + (50.0, 120.0, 0), + (180.0, 220.0, 0), + (300.0, 330.0, 1), + ] + assert recorded["focus_labels"] == ["I", "II"] + + +def test_plot_genome_cds_scale_is_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + _write_reference_features_json(tmp_path / "reference_features.json") + + recorded = {} + + def fake_genome(_table, _metadata, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + + exit_code = main_module.main( + [ + "plot", + "genome", + str(results_csv), + "--gene", + "S", + "--cds-scale", + "--focus-coords", + "10-30", + ] + ) + + assert exit_code == 0 + assert recorded["gene"] == "S" + assert recorded["cds_scale"] is True + assert recorded["focus_ranges"] == [(10.0, 30.0, 0)] + + +def test_plot_genome_show_intersections_is_forwarded(monkeypatch, tmp_path): + results_csv = tmp_path / "results.csv" + _write_plot_results_csv(results_csv, n_variants=4) + _write_reference_features_json(tmp_path / "reference_features.json") + + recorded = {} + + def fake_genome(_table, _metadata, **kwargs): + recorded.update(kwargs) + + monkeypatch.setattr(main_module, "plot_variant_genome", fake_genome) + + exit_code = main_module.main( + [ + "plot", + "genome", + str(results_csv), + "--focus-coords", + "I:50-120,180-220", + "--show-intersections", + ] + ) + + assert exit_code == 0 + assert recorded["show_intersections"] is True + + def test_plot_genome_include_indels_is_forwarded(monkeypatch, tmp_path): results_csv = tmp_path / "results.csv" _write_plot_results_csv(results_csv, n_variants=4) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index c97d146..fb3b200 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -2,12 +2,15 @@ from __future__ import annotations +import json + import pandas as pd import pytest from vartracker.core import InputValidationError from vartracker.plotting import ( _collapse_variants_for_genome_plot, + _genome_axis_label, _parse_focus_ranges, _subset_collapsed_variants, build_reference_feature_metadata, @@ -19,6 +22,7 @@ plot_variant_trajectory, plot_variant_turnover, prepare_plot_inputs, + resolve_focus_regions, variant_crosses_thresholds, ) @@ -150,6 +154,13 @@ def test_auto_select_variants_prefers_literature_and_persistent(): assert len(selected) == 2 +def test_genome_axis_label_uses_gene_name_for_cds_scale(): + assert ( + _genome_axis_label(gene="F", aa_scale=False, cds_scale=True) + == "CDS position in F" + ) + + def test_plot_functions_create_output_files(tmp_path, monkeypatch): mpl_dir = tmp_path / "mpl" mpl_dir.mkdir() @@ -359,14 +370,21 @@ def test_genome_plot_builds_outputs_and_supports_gene_modes(tmp_path, monkeypatc genome_out = tmp_path / "genome.pdf" focus_out = tmp_path / "genome_focus.pdf" aa_out = tmp_path / "genome_gene_aa.pdf" + cds_out = tmp_path / "genome_gene_cds.pdf" plot_variant_genome(table, metadata, output_path=genome_out) plot_variant_genome( table, metadata, output_path=focus_out, - focus_ranges=_parse_focus_ranges("23000-24000,50000-51000"), + focus_ranges=resolve_focus_regions( + "Fusion peptide:23000-24000;Heptad repeat:50000-51000", "" + )[0], + focus_labels=resolve_focus_regions( + "Fusion peptide:23000-24000;Heptad repeat:50000-51000", "" + )[1], gene="S", + show_intersections=True, ) plot_variant_genome( table, @@ -376,12 +394,91 @@ def test_genome_plot_builds_outputs_and_supports_gene_modes(tmp_path, monkeypatc aa_scale=True, focus_ranges=_parse_focus_ranges("450-650"), ) + plot_variant_genome( + table, + metadata, + output_path=cds_out, + gene="S", + cds_scale=True, + focus_ranges=_parse_focus_ranges("440-500"), + ) assert genome_out.exists() assert focus_out.exists() assert aa_out.exists() + assert cds_out.exists() def test_parse_focus_ranges_caps_number_of_regions(): with pytest.raises(InputValidationError): _parse_focus_ranges("1-2,3-4,5-6,7-8,9-10,11-12,13-14") + + +def test_parse_focus_ranges_supports_semicolon_color_groups(): + parsed = _parse_focus_ranges( + "62-69,196-210;31-42,323-332,379-399;254-277;50-50,305-310;422-438;163-181" + ) + + assert parsed == [ + (62.0, 69.0, 0), + (196.0, 210.0, 0), + (31.0, 42.0, 1), + (323.0, 332.0, 1), + (379.0, 399.0, 1), + (254.0, 277.0, 2), + (50.0, 50.0, 3), + (305.0, 310.0, 3), + (422.0, 438.0, 4), + (163.0, 181.0, 5), + ] + + +def test_parse_focus_ranges_caps_number_of_groups(): + with pytest.raises(InputValidationError): + _parse_focus_ranges("1-2;3-4;5-6;7-8;9-10;11-12;13-14") + + +def test_resolve_focus_regions_supports_named_inline_groups(): + ranges, labels = resolve_focus_regions( + "Ø:62-69,196-210;I:31-42,323-332,379-399;II:254-277", "" + ) + + assert labels == ["Ø", "I", "II"] + assert ranges == [ + (62.0, 69.0, 0), + (196.0, 210.0, 0), + (31.0, 42.0, 1), + (323.0, 332.0, 1), + (379.0, 399.0, 1), + (254.0, 277.0, 2), + ] + + +def test_resolve_focus_regions_supports_json_and_csv_files(tmp_path): + json_path = tmp_path / "regions.json" + json_path.write_text( + json.dumps({"Ø": ["62-69", "196-210"], "I": ["31-42", "323-332"]}), + encoding="utf-8", + ) + csv_path = tmp_path / "regions.csv" + csv_path.write_text( + "label,start,end\nII,254,277\nIII,50,50\nIII,305,310\n", + encoding="utf-8", + ) + + json_ranges, json_labels = resolve_focus_regions("", str(json_path)) + csv_ranges, csv_labels = resolve_focus_regions("", str(csv_path)) + + assert json_labels == ["Ø", "I"] + assert json_ranges == [ + (62.0, 69.0, 0), + (196.0, 210.0, 0), + (31.0, 42.0, 1), + (323.0, 332.0, 1), + ] + assert csv_labels == ["II", "III"] + assert csv_ranges == [ + (254.0, 277.0, 0), + (50.0, 50.0, 1), + (305.0, 310.0, 1), + ] diff --git a/vartracker/main.py b/vartracker/main.py index 6734bbc..c1083f2 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -44,7 +44,6 @@ from .plotting import ( DEFAULT_LIFESPAN_TOP_N, DEFAULT_TRAJECTORY_TOP_N, - _parse_focus_ranges, apply_shared_plot_filters, auto_select_variants, load_reference_feature_metadata, @@ -58,6 +57,7 @@ plot_variant_turnover, prepare_plot_inputs, _project_name_from_results, + resolve_focus_regions, resolve_plot_output_path, write_reference_feature_metadata, ) @@ -631,7 +631,6 @@ def _configure_vcf_parser( analysis_group = parser.add_argument_group("Vartracker Analysis Options") output_group = parser.add_argument_group("Vartracker Output Options") - heatmap_group = parser.add_argument_group("Heatmap") analysis_group.add_argument( "-r", @@ -754,8 +753,6 @@ def _configure_vcf_parser( default=False, ) - _add_heatmap_option_arguments(heatmap_group) - def _move_action_group_after( parser: argparse.ArgumentParser, group_title: str, anchor_title: str @@ -1235,6 +1232,9 @@ def _run_plot_genome_command(args): table = load_results_table(results_csv) project_name = _project_name_from_results(table, getattr(args, "name", None)) metadata = load_reference_feature_metadata(results_csv) + focus_ranges, focus_labels = resolve_focus_regions( + args.focus_coords, args.focus_region_file + ) output_path = resolve_plot_output_path( results_csv, out=args.out, @@ -1248,13 +1248,16 @@ def _run_plot_genome_command(args): output_path=output_path, gene=args.gene, aa_scale=args.aa_scale, - focus_ranges=_parse_focus_ranges(args.focus_coords), + cds_scale=args.cds_scale, + focus_ranges=focus_ranges, + focus_labels=focus_labels, min_af=args.min_af, max_af=args.max_af, effects=_parse_csv_option_list(args.effect), persistent_only=args.persistent_only, new_only=args.new_only, include_indels=args.include_indels, + show_intersections=args.show_intersections, title=args.title or ( f"{project_name}: genome variant distribution" if project_name else None @@ -1516,10 +1519,34 @@ def _add_plot_genome_subparser(subparsers): default=False, help="With --gene, plot x-axis in amino-acid coordinates for that gene", ) + parser.add_argument( + "--cds-scale", + action="store_true", + default=False, + help="With --gene, plot x-axis in CDS-relative nucleotide coordinates for that gene", + ) parser.add_argument( "--focus-coords", default="", - help="Comma-separated coordinate ranges to highlight, e.g. 150-300,900-1800", + help=( + "Coordinate ranges to highlight. Use commas for separate ranges or " + "semicolons to group ranges with the same color, e.g. " + "150-300,900-1800;50-120 or Name:150-300,900-1800;Other:50-120" + ), + ) + parser.add_argument( + "--focus-region-file", + default="", + help=( + "Optional JSON/CSV/TSV file defining named focus region groups for the " + "genome plot" + ), + ) + parser.add_argument( + "--show-intersections", + action="store_true", + default=False, + help="Show a compact table of variants intersecting the highlighted focus regions", ) parser.add_argument("--title", default=None, help="Optional plot title") parser.add_argument( @@ -2586,7 +2613,6 @@ def _process_files( literature_hits=literature_hits_df, literature_table_path=literature_full_csv_path, cli_command=cli_command, - **_collect_heatmap_kwargs(args), ) # Write specialized tables diff --git a/vartracker/plotting.py b/vartracker/plotting.py index d19650b..0ae53b0 100644 --- a/vartracker/plotting.py +++ b/vartracker/plotting.py @@ -2,8 +2,10 @@ from __future__ import annotations +import csv import json from pathlib import Path +from textwrap import wrap from typing import Iterable, Sequence, cast from urllib.parse import unquote @@ -11,6 +13,8 @@ import numpy as np import pandas as pd from matplotlib.gridspec import GridSpec +from matplotlib.patches import Patch +from matplotlib.transforms import Bbox from .analysis import ( _coerce_frequency, @@ -31,6 +35,7 @@ "#d1c4e9", "#b2dfdb", ] +FocusRange = tuple[float, float, int] FeatureRecord = dict[str, object] @@ -545,7 +550,25 @@ def plot_variant_trajectory( if plot_df.empty: raise ProcessingError("No data available for the trajectory plot") - fig, ax = plt.subplots(figsize=(width, height)) + show_legend = not label_lines and not label_threshold_crossers + if show_legend: + fig = plt.figure( + figsize=(width + 1.4, height), + constrained_layout=True, + ) + grid = GridSpec( + 1, + 2, + figure=fig, + width_ratios=[1.0, 0.2], + wspace=0.03, + ) + ax = fig.add_subplot(grid[0, 0]) + legend_ax = fig.add_subplot(grid[0, 1]) + legend_ax.axis("off") + else: + fig, ax = plt.subplots(figsize=(width, height)) + legend_ax = None threshold_values = [float(value) for value in (thresholds or [])] for variant_id in selected_variants: subset = plot_df[plot_df["variant_id"] == variant_id].sort_values( @@ -596,9 +619,21 @@ def plot_variant_trajectory( float(threshold), linestyle="--", linewidth=1.0, color="black", alpha=0.6 ) _apply_plot_title(ax, title, "Variant trajectories") - if not label_lines and not label_threshold_crossers: - ax.legend(frameon=False, loc="best") - fig.tight_layout() + if show_legend and legend_ax is not None: + handles, labels = ax.get_legend_handles_labels() + legend = legend_ax.legend( + handles, + labels, + title="Variants", + loc="upper left", + frameon=False, + borderaxespad=0.0, + handlelength=1.8, + labelspacing=0.35, + ) + legend.get_title().set_fontweight("bold") + elif hasattr(fig, "tight_layout"): + fig.tight_layout() fig.savefig(output_path, dpi=dpi, bbox_inches="tight") plt.close(fig) @@ -820,31 +855,168 @@ def collect_explicit_variants( return requested or file_variants -def _parse_focus_ranges(value: str | None) -> list[tuple[float, float]]: +def _parse_single_focus_range(part: str) -> tuple[float, float]: + if "-" not in part: + raise InputValidationError( + "Focus ranges must use start-end syntax, e.g. 150-300,900-1800 or Name:150-300,900-1800;Other:50-120" + ) + start_text, end_text = part.split("-", 1) + try: + start = float(start_text) + end = float(end_text) + except ValueError as exc: + raise InputValidationError( + "Focus ranges must use numeric start-end coordinates" + ) from exc + if end < start: + start, end = end, start + return start, end + + +def _parse_focus_region_groups( + value: str | None, +) -> tuple[list[FocusRange], list[str | None]]: if not value: - return [] - ranges: list[tuple[float, float]] = [] - for part in _parse_csv_option_list(value): - if "-" not in part: + return [], [] + + ranges: list[FocusRange] = [] + labels: list[str | None] = [] + group_texts = [part.strip() for part in str(value).split(";") if part.strip()] + if not group_texts: + return [], [] + + group_specs: list[tuple[str | None, list[str]]] + # Preserve legacy behavior for purely comma-separated input: one range per color. + if len(group_texts) == 1 and ":" not in group_texts[0]: + raw_ranges = _parse_csv_option_list(group_texts[0]) + group_specs = [(None, [part]) for part in raw_ranges] + else: + group_specs = [] + for group_text in group_texts: + label: str | None = None + range_text = group_text + if ":" in group_text: + maybe_label, maybe_ranges = group_text.split(":", 1) + if maybe_label.strip(): + label = maybe_label.strip() + range_text = maybe_ranges + group_parts = _parse_csv_option_list(range_text) + if not group_parts: + raise InputValidationError( + "Each focus region group must include at least one range" + ) + group_specs.append((label, group_parts)) + + if len(group_specs) > MAX_FOCUS_RANGES: + raise InputValidationError( + f"Use at most {MAX_FOCUS_RANGES} focus region groups so each can retain a distinct highlight color" + ) + + for group_index, (label, group_parts) in enumerate(group_specs): + labels.append(label) + for part in group_parts: + start, end = _parse_single_focus_range(part) + ranges.append((start, end, group_index)) + return ranges, labels + + +def _parse_focus_ranges(value: str | None) -> list[FocusRange]: + return _parse_focus_region_groups(value)[0] + + +def _load_focus_region_file( + path: str | None, +) -> tuple[list[FocusRange], list[str | None]]: + if not path: + return [], [] + focus_path = Path(path).expanduser().resolve() + if not focus_path.exists(): + raise InputValidationError(f"Focus region file not found: {focus_path}") + + suffix = focus_path.suffix.lower() + if suffix == ".json": + payload = json.loads(focus_path.read_text(encoding="utf-8")) + if isinstance(payload, dict): + group_specs = [ + (str(label).strip(), value) for label, value in payload.items() + ] + elif isinstance(payload, list): + group_specs = [] + for entry in payload: + if not isinstance(entry, dict) or "label" not in entry: + raise InputValidationError( + "JSON focus region entries must be objects with 'label' and 'ranges'" + ) + group_specs.append( + (str(entry["label"]).strip(), entry.get("ranges", [])) + ) + else: raise InputValidationError( - "Focus ranges must use start-end syntax, e.g. 150-300,900-1800" + "JSON focus region file must contain an object or list of labeled ranges" ) - start_text, end_text = part.split("-", 1) - try: - start = float(start_text) - end = float(end_text) - except ValueError as exc: - raise InputValidationError( - "Focus ranges must use numeric start-end coordinates" - ) from exc - if end < start: - start, end = end, start - ranges.append((start, end)) - if len(ranges) > MAX_FOCUS_RANGES: + + specs: list[str] = [] + for label, value in group_specs: + if isinstance(value, str): + ranges_text = value + elif isinstance(value, list): + ranges_text = ",".join( + str(item).strip() for item in value if str(item).strip() + ) + else: + raise InputValidationError( + "JSON focus region values must be a range string or list of range strings" + ) + specs.append(f"{label}:{ranges_text}") + return _parse_focus_region_groups(";".join(specs)) + + if suffix in {".csv", ".tsv"}: + delimiter = "\t" if suffix == ".tsv" else "," + with focus_path.open("r", encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle, delimiter=delimiter) + required = {"label", "start", "end"} + if reader.fieldnames is None or not required.issubset( + {field.strip().lower() for field in reader.fieldnames} + ): + raise InputValidationError( + "Focus region CSV/TSV must contain label,start,end columns" + ) + grouped: dict[str, list[str]] = {} + for row in reader: + label = str(row.get("label", "")).strip() + start = str(row.get("start", "")).strip() + end = str(row.get("end", "")).strip() + if not label or not start or not end: + raise InputValidationError( + "Focus region CSV/TSV rows must include label,start,end values" + ) + grouped.setdefault(label, []).append(f"{start}-{end}") + specs = [f"{label}:{','.join(ranges)}" for label, ranges in grouped.items()] + return _parse_focus_region_groups(";".join(specs)) + + raise InputValidationError("Focus region file must be .json, .csv, or .tsv") + + +def resolve_focus_regions( + focus_coords: str | None, focus_region_file: str | None +) -> tuple[list[FocusRange], list[str | None]]: + cli_ranges, cli_labels = _parse_focus_region_groups(focus_coords) + file_ranges, file_labels = _load_focus_region_file(focus_region_file) + if not cli_ranges and not file_ranges: + return [], [] + + offset = len(file_labels) + merged_ranges = list(file_ranges) + merged_labels = list(file_labels) + merged_labels.extend(cli_labels) + for start, end, group_index in cli_ranges: + merged_ranges.append((start, end, group_index + offset)) + + if len(merged_labels) > MAX_FOCUS_RANGES: raise InputValidationError( - f"Use at most {MAX_FOCUS_RANGES} focus ranges so each can retain a distinct highlight color" + f"Use at most {MAX_FOCUS_RANGES} focus region groups so each can retain a distinct highlight color" ) - return ranges + return merged_ranges, merged_labels def _parse_gff_attrs(attr_text: str) -> dict[str, str]: @@ -986,6 +1158,7 @@ def _collapse_variants_for_genome_plot( { "variant_id": variant_key, "variant_label": display_label.replace("\n", " "), + "variant_name": str(getattr(row, "variant", "")).strip(), "chrom": str(getattr(row, "chrom", "")).strip(), "start": int(getattr(row, "start", 0)), "end": int(getattr(row, "end", getattr(row, "start", 0))), @@ -1098,18 +1271,48 @@ def _find_gene_feature(metadata: dict[str, object], gene_name: str) -> FeatureRe def _clip_focus_ranges( - ranges: Sequence[tuple[float, float]], xmin: float, xmax: float -) -> tuple[list[tuple[float, float]], int]: - clipped: list[tuple[float, float]] = [] + ranges: Sequence[FocusRange], xmin: float, xmax: float +) -> tuple[list[FocusRange], int]: + clipped: list[FocusRange] = [] ignored = 0 - for start, end in ranges: + for start, end, group_index in ranges: if end < xmin or start > xmax: ignored += 1 continue - clipped.append((max(start, xmin), min(end, xmax))) + clipped.append((max(start, xmin), min(end, xmax), group_index)) return clipped, ignored +def _focus_group_name(label: str | None, group_index: int) -> str: + return ( + str(label).strip() + if label and str(label).strip() + else f"Region {group_index + 1}" + ) + + +def _format_focus_variants(variants: Sequence[str]) -> str: + joined = ", ".join(variants) + return "\n".join(wrap(joined, width=54)) if joined else "" + + +def _format_focus_table_column(values: Sequence[str], width: int) -> str: + lines: list[str] = [] + for value in values: + wrapped = wrap(str(value), width=width) or [str(value)] + lines.extend(wrapped) + return "\n".join(lines) + + +def _genome_axis_label(*, gene: str | None, aa_scale: bool, cds_scale: bool) -> str: + if aa_scale: + return "Amino-acid position" + if cds_scale: + gene_label = str(gene).strip() if gene else "gene" + return f"CDS position in {gene_label}" + return "Genomic position" + + def plot_variant_genome( table: pd.DataFrame, reference_metadata: dict[str, object], @@ -1117,13 +1320,16 @@ def plot_variant_genome( output_path: str | Path, gene: str | None = None, aa_scale: bool = False, - focus_ranges: Sequence[tuple[float, float]] | None = None, + cds_scale: bool = False, + focus_ranges: Sequence[FocusRange] | None = None, + focus_labels: Sequence[str | None] | None = None, min_af: float | None = None, max_af: float | None = None, effects: Sequence[str] | None = None, persistent_only: bool = False, new_only: bool = False, include_indels: bool = False, + show_intersections: bool = False, title: str | None = None, width: float = 10.0, height: float = 6.5, @@ -1131,10 +1337,15 @@ def plot_variant_genome( ) -> None: if aa_scale and not gene: raise InputValidationError("--aa-scale requires --gene") + if cds_scale and not gene: + raise InputValidationError("--cds-scale requires --gene") + if aa_scale and cds_scale: + raise InputValidationError("Use either --aa-scale or --cds-scale, not both") if include_indels: print( "Warning: including indels in the genome plot may produce ambiguous or hard-to-interpret positions" ) + detection_floor = float(min_af) if min_af is not None else 0.03 collapsed = _collapse_variants_for_genome_plot(table) collapsed = _subset_collapsed_variants( @@ -1184,6 +1395,9 @@ def plot_variant_genome( 1.0, float(cast(int | float | str, gene_feature.get("aa_length", 1))), ) + elif cds_scale: + plot_df["x_coord"] = plot_df["start"].astype(float) - start + 1.0 + region_by_contig[contig] = (1.0, end - start + 1.0) else: plot_df["x_coord"] = plot_df["start"].astype(float) region_by_contig[contig] = (start, end) @@ -1213,66 +1427,220 @@ def plot_variant_genome( if not selected_contigs: raise ProcessingError("No contigs remain for genome plotting") + focus = list(focus_ranges or []) + focus_group_labels = list(focus_labels or []) + focus_hits: dict[int, set[tuple[str, str]]] = {} + ignored_focus_total = 0 + if focus: + ignored_focus_total = sum( + 1 + for start, end, _group_index in focus + if not any( + not ( + end < region_by_contig[contig][0] + or start > region_by_contig[contig][1] + ) + for contig in selected_contigs + ) + ) + for contig in selected_contigs: + subset = plot_df[plot_df["chrom"].astype(str) == contig].sort_values("x_coord") + xmin, xmax = region_by_contig[contig] + clipped_focus, _ignored_focus = _clip_focus_ranges(focus, xmin, xmax) + for start, end, group_index in clipped_focus: + matching = subset.loc[ + (subset["x_coord"] >= start) & (subset["x_coord"] <= end), + ["variant_name", "variant_label"], + ] + if matching.empty: + continue + focus_hits.setdefault(group_index, set()).update( + ( + str(nt_change).strip(), + _display_label_prefix(aa_change), + ) + for nt_change, aa_change in matching.itertuples(index=False, name=None) + ) + if ignored_focus_total: + print( + f"Warning: ignored {ignored_focus_total} focus range(s) outside the plotted region" + ) + focus_table_rows = [] + if show_intersections: + focus_table_rows = [ + ( + _focus_group_name( + ( + focus_group_labels[group_index] + if group_index < len(focus_group_labels) + else None + ), + group_index, + ), + _format_focus_table_column( + [nt_change for nt_change, _aa_change in sorted(variants)], + width=20, + ), + _format_focus_table_column( + [aa_change for _nt_change, aa_change in sorted(variants)], + width=24, + ), + ) + for group_index, variants in sorted(focus_hits.items()) + if variants + ] + show_focus_legend = any( + label and str(label).strip() for label in focus_group_labels + ) + show_focus_table = bool(focus_table_rows) + n_panels = len(selected_contigs) + n_columns = 2 if show_focus_legend else 1 + n_rows = n_panels + (1 if show_focus_table else 0) fig = plt.figure( - figsize=(width, max(height, 2.1 * n_panels)), + figsize=( + width + (1.7 if show_focus_legend else 0.0), + max(height, 2.1 * n_panels + (1.2 if show_focus_table else 0.0)), + ), constrained_layout=True, ) grid = GridSpec( - n_panels, - 1, + n_rows, + n_columns, figure=fig, - hspace=0.35, + hspace=0.08, + width_ratios=[1.0, 0.18] if show_focus_legend else None, + height_ratios=[1.0] * n_panels + ([0.34] if show_focus_table else []), ) - focus = list(focus_ranges or []) for idx, contig in enumerate(selected_contigs): ax = fig.add_subplot(grid[idx, 0]) subset = plot_df[plot_df["chrom"].astype(str) == contig].sort_values("x_coord") xmin, xmax = region_by_contig[contig] clipped_focus, ignored_focus = _clip_focus_ranges(focus, xmin, xmax) - for focus_index, (start, end) in enumerate(clipped_focus): + for start, end, group_index in clipped_focus: ax.axvspan( start, end, - color=FOCUS_RANGE_COLORS[focus_index % len(FOCUS_RANGE_COLORS)], - alpha=0.25, + color=FOCUS_RANGE_COLORS[group_index % len(FOCUS_RANGE_COLORS)], + alpha=0.34, zorder=0, ) - if ignored_focus: - print( - f"Warning: ignored {ignored_focus} focus range(s) outside the plotted region" - ) for row in subset.itertuples(index=False): af_values = [float(value) for value in cast(list[float], row.af_values)] + plotted_af_values = [ + value if value > 0 else detection_floor for value in af_values + ] if af_values: ax.vlines( float(row.x_coord), - min(af_values), - max(af_values), + min(plotted_af_values), + max(plotted_af_values), color="#607d8b", linewidth=0.9, alpha=0.55, zorder=1, ) ax.scatter( - [float(row.x_coord)] * len(af_values), - af_values, + [float(row.x_coord)] * len(plotted_af_values), + plotted_af_values, s=18, alpha=0.45, color="#1f77b4", edgecolors="none", zorder=2, ) - ax.axhline(0.5, color="black", linestyle="--", linewidth=0.8, alpha=0.4) + ax.axhline( + detection_floor, + color="black", + linestyle="--", + linewidth=0.8, + alpha=0.35, + ) ax.set_ylim(0, 1.02) ax.set_ylabel("Allele frequency") ax.set_xlim(xmin, xmax) - ax.set_title(str(contig), loc="left", fontsize=10, fontweight="bold") - ax.set_xlabel("Amino-acid position" if aa_scale else "Genomic position") - for spine in ("top", "right"): - ax.spines[spine].set_visible(False) + ax.set_title(f"Reference: {contig}", loc="left", fontsize=10, fontweight="bold") + ax.set_xlabel( + _genome_axis_label(gene=gene, aa_scale=aa_scale, cds_scale=cds_scale) + ) + for spine in ax.spines.values(): + spine.set_visible(True) + spine.set_color("black") + spine.set_linewidth(1.0) + + if show_focus_legend: + legend_ax = fig.add_subplot(grid[:n_panels, 1]) + legend_ax.axis("off") + handles = [ + Patch( + facecolor=FOCUS_RANGE_COLORS[group_index % len(FOCUS_RANGE_COLORS)], + edgecolor="none", + alpha=0.35, + label=_focus_group_name(label, group_index), + ) + for group_index, label in enumerate(focus_group_labels) + if label and str(label).strip() + ] + if handles: + legend_ax.legend( + handles=handles, + title="Focus regions", + loc="upper left", + frameon=True, + framealpha=0.9, + fontsize=8, + title_fontsize=9, + borderpad=0.5, + handlelength=1.6, + handletextpad=0.6, + labelspacing=0.35, + ) + + if show_focus_table: + table_ax = fig.add_subplot(grid[n_panels, 0]) + table_ax.axis("off") + table_ax.text( + 0.0, + 0.94, + "Intersecting Variants", + transform=table_ax.transAxes, + ha="left", + va="top", + fontsize=9, + fontweight="bold", + ) + max_nt_chars = max( + len(nt_values.replace("\n", ", ")) for _, nt_values, _ in focus_table_rows + ) + max_aa_chars = max( + len(aa_values.replace("\n", ", ")) for _, _, aa_values in focus_table_rows + ) + table_width = min( + 0.96, max(0.62, 0.22 + (max_nt_chars / 120.0) + (max_aa_chars / 110.0)) + ) + table = table_ax.table( + cellText=[ + [region, nt_change, aa_change] + for region, nt_change, aa_change in focus_table_rows + ], + colLabels=["Region", "NT Change", "AA Change"], + colLoc="left", + cellLoc="left", + colWidths=[0.14, 0.26, 0.60], + bbox=Bbox.from_bounds(0.0, 0.10, table_width, 0.74), + ) + table.auto_set_font_size(False) + table.set_fontsize(8) + table.scale(1.0, 1.2) + for (row_idx, col_idx), cell in table.get_celld().items(): + cell.set_edgecolor("#d0d0d0") + if row_idx == 0: + cell.set_facecolor("#f5f5f5") + cell.set_text_props(weight="bold") + elif col_idx == 0: + cell.set_facecolor("#fafafa") if title: fig.suptitle(title, fontweight="bold") From 49c40edfecc06b5d24312b3d783496368d661c5a Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Thu, 23 Apr 2026 09:54:05 +1000 Subject: [PATCH 08/20] Added ability to generate consensus genomes + bug huntin --- README.md | 58 ++++-- docs/DEPENDENCIES.md | 3 + tests/test_analysis.py | 142 ++++++++++++++ tests/test_analysis_launcher.py | 37 ++++ tests/test_cli.py | 20 ++ tests/test_consensus.py | 48 +++++ tests/test_main.py | 64 ++++++- tests/test_snakemake_workflow.py | 14 ++ tests/test_vcf_processing.py | 267 ++++++++++++++++++++++++++- vartracker/Snakefile | 305 ++++++++++++++++++++++++++++++- vartracker/analysis.py | 62 ++++++- vartracker/analysis_launcher.py | 55 ++++++ vartracker/consensus.py | 108 +++++++++++ vartracker/main.py | 258 ++++++++++++++++++-------- vartracker/vcf_processing.py | 234 ++++++++++++++++++------ 15 files changed, 1502 insertions(+), 173 deletions(-) create mode 100644 tests/test_analysis_launcher.py create mode 100644 tests/test_consensus.py create mode 100644 vartracker/consensus.py diff --git a/README.md b/README.md index 1cc5615..8ba282d 100755 --- a/README.md +++ b/README.md @@ -99,6 +99,7 @@ Minimum tested versions are tracked in `docs/DEPENDENCIES.md`. - **samtools**, **lofreq**, **fastp**, **bwa**, and **snakemake** – required for the `bam` and `end-to-end` Snakemake workflows If you only plan to run `vartracker vcf` against pre-generated VCFs, the first pair is sufficient. The additional tools are needed whenever you ask vartracker to align reads or call variants for you. +Consensus genome generation in the `bam` and `end-to-end` workflows uses `bcftools` and `samtools`; it does not require `bedtools`. Note: the pinned micromamba environment installs `tabix`/`bgzip` via `htslib`. @@ -204,8 +205,10 @@ vartracker end-to-end path/to/read_inputs.csv \ # Re-plot a heatmap from an existing vartracker results file vartracker plot heatmap results/results.csv \ - --heatmap-aa-exclude "*frameshift*" \ - --outdir results/replots + --aa-exclude "*frameshift*" \ + --x-labels sample-number \ + --literature-csv results/sample.literature_database_hits.full.csv \ + --title "Variant allele frequencies" # Plot whole-dataset turnover from an existing results file vartracker plot turnover results/results.csv @@ -262,6 +265,17 @@ Mode-specific expectations: - **BAM mode** requires `bam` and will fill `vcf` + `coverage` during the workflow. - **End-to-end mode** requires `reads1` (and optionally `reads2`); remaining fields are generated. +The `bam` and `end-to-end` workflows also write two consensus FASTA columns to +the updated Snakemake spreadsheet: `consensus` for a simple consensus and +`iupac_consensus` for an IUPAC-aware consensus. SNPs below +`--consensus-snp-min-af` are ignored, SNPs from `--consensus-snp-min-af` up to +`--consensus-snp-thresh` stay as reference bases in the simple consensus +and become REF+ALT ambiguity codes in the IUPAC consensus, and SNPs at or above +`--consensus-snp-thresh` become ALT bases. Indels are controlled +separately by `--consensus-indel-thresh` in both consensus modes. Low-depth bases +are masked as `N`, except for called deletion intervals so true deletions are not +converted to low-depth masks. + Relative paths are resolved with respect to the CSV location, so you can store the sheet alongside your sequencing artefacts. The `prepare spreadsheet` subcommand can scaffold a CSV and highlight missing files. @@ -275,8 +289,11 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test `--allele-frequency-tag`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. - `vartracker bam` – everything from `vcf`, plus Snakemake options: - `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`. -- `vartracker end-to-end` – similar to `bam`, with an optional `--primer-bed` for amplicon clipping. + `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, + `--rulegraph`, `--consensus-snp-min-af`, `--consensus-snp-thresh`, and + `--consensus-indel-thresh`. +- `vartracker end-to-end` – similar to `bam`, with optional amplicon clipping controls: + `--primer-bed` and `--ampliconclip-tolerance` (default: `1`). - `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV, including all heatmap customization filters. - `vartracker plot genome` – plot SNP positions along the genome or a selected gene region using all observed allele-frequency values for each variant. - `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. @@ -285,22 +302,25 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test Heatmap filtering: - `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customize heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. -- By default, all consequence classes are included except joint variants. Use `--heatmap-include-joint` to show joint variants. -- `--heatmap-aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. -- `--heatmap-aa-include`: comma-separated `type_of_change` patterns to include. -- `--heatmap-only-persistent`: only include `new_persistent` variants. -- `--heatmap-only-new`: only include variants with `variant_status == new`. -- `--heatmap-gene-include` and `--heatmap-gene-exclude`: comma-separated gene patterns. -- `--heatmap-variant-type`: comma-separated variant-type patterns such as `snp` or `indel`. -- `--heatmap-qc`: comma-separated `all_samples_pass_qc` patterns to include. Accepted values include `true`, `false`, `pass`, and `fail`. +- By default, all consequence classes are included except joint variants. Use `--include-joint` to show joint variants. +- `--aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. +- `--aa-include`: comma-separated `type_of_change` patterns to include. +- `--only-persistent`: only include `new_persistent` variants. +- `--only-new`: only include variants with `variant_status == new`. +- `--gene-include` and `--gene-exclude`: comma-separated gene patterns. +- `--variant-type`: comma-separated variant-type patterns such as `snp` or `indel`. +- `--qc`: comma-separated `all_samples_pass_qc` patterns to include. Accepted values include `true`, `false`, `pass`, and `fail`. - `--min-prop-passing-qc`: minimum fraction of samples that must pass per-sample QC. -- `--heatmap-min-persistence`: minimum number of included samples in which the variant must be present. -- `--heatmap-min-max-af`: minimum maximum allele frequency across included samples. -- `--heatmap-min-sample-af`: minimum allele frequency that must be reached in at least one included sample. -- `--heatmap-sample-subset`: comma-separated sample-name patterns to plot. -- `--heatmap-hide-singletons`: hide variants present in only one included sample. -- `--heatmap-min-depth`: minimum site depth a variant must reach in at least one included sample. -- Example: `--heatmap-aa-exclude "synonymous,*frameshift*,stop_gained"` +- `--min-persistence`: minimum number of included samples in which the variant must be present. +- `--min-max-af`: minimum maximum allele frequency across included samples. +- `--min-sample-af`: minimum allele frequency that must be reached in at least one included sample. +- `--sample-subset`: comma-separated sample-name patterns to plot. +- `--hide-singletons`: hide variants present in only one included sample. +- `--min-depth`: minimum site depth a variant must reach in at least one included sample. +- `--x-labels sample-number`: label heatmap x-axis columns by `sample_number` instead of sample name. +- `--title`: set the heatmap plot title. The default is `Variant allele frequencies`. +- `--literature-csv`: include literature links in the interactive HTML heatmap using a literature hits CSV. +- Example: `--aa-exclude "synonymous,*frameshift*,stop_gained"` Standalone plot filtering: - `--gene`, `--effect`, `--min-af`, `--max-af`: restrict the plotted result set before ranking/selection. diff --git a/docs/DEPENDENCIES.md b/docs/DEPENDENCIES.md index 02a56e0..474724b 100644 --- a/docs/DEPENDENCIES.md +++ b/docs/DEPENDENCIES.md @@ -14,6 +14,9 @@ The recommended micromamba environment pins the same versions for reproducibilit | snakemake | 9.0.1 | 9.0.1 | Snakemake API used for workflows. | | bgzip | 1.21 | 1.21 (via htslib) | Provided by htslib; used for VCF compression. | +Consensus FASTA outputs in `bam` and `end-to-end` mode are produced with +`bcftools consensus` and `samtools depth`; `bedtools` is not required. + If you validate compatibility with older tool versions, update the minimum tested version here and consider expanding CI coverage. diff --git a/tests/test_analysis.py b/tests/test_analysis.py index b52a732..b24193b 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -356,6 +356,148 @@ def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): assert result.loc[1, "type_of_change"] == "joint_frameshift" +def test_process_joint_variants_matches_main_row_by_presence_pattern(tmp_path): + csv_path = tmp_path / "results.csv" + pd.DataFrame( + [ + { + "start": 100, + "gene": "S", + "amino_acid_consequence": "K2A", + "nsp_aa_change": "", + "bcsq_nt_notation": "4A>G+5A>C", + "bcsq_aa_notation": "2K>2A", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "missense", + "presence_absence": "Y / N", + }, + { + "start": 100, + "gene": "S", + "amino_acid_consequence": "K2V", + "nsp_aa_change": "", + "bcsq_nt_notation": "4A>T+5A>G", + "bcsq_aa_notation": "2K>2V", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "missense", + "presence_absence": "N / Y", + }, + { + "start": 101, + "gene": "", + "amino_acid_consequence": "", + "nsp_aa_change": "", + "bcsq_nt_notation": "", + "bcsq_aa_notation": "@100", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "@100", + "presence_absence": "N / Y", + }, + ] + ).to_csv(csv_path, index=False) + + result = process_joint_variants(str(csv_path)) + + assert result.loc[2, "amino_acid_consequence"] == "K2V" + assert result.loc[2, "type_of_change"] == "joint_missense" + + +@pytest.mark.xfail( + strict=True, + reason=( + "process_joint_variants currently resolves tied main-row candidates by row " + "order, which is unsafe for overlapping genes" + ), +) +def test_process_joint_variants_is_order_invariant_for_overlapping_gene_rows(tmp_path): + shared_rows = [ + { + "start": 25470, + "gene": "ORF3a", + "amino_acid_consequence": "ORF3a:A10V", + "nsp_aa_change": "", + "bcsq_nt_notation": "c.30C>T", + "bcsq_aa_notation": "p.A10V", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "missense", + "presence_absence": "Y / N / Y", + }, + { + "start": 25470, + "gene": "ORF3c", + "amino_acid_consequence": "ORF3c:M5I", + "nsp_aa_change": "", + "bcsq_nt_notation": "c.15G>A", + "bcsq_aa_notation": "p.M5I", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "missense", + "presence_absence": "Y / N / Y", + }, + { + "start": 25471, + "gene": "", + "amino_acid_consequence": "", + "nsp_aa_change": "", + "bcsq_nt_notation": "", + "bcsq_aa_notation": "@25470", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "@25470", + "presence_absence": "Y / N / Y", + }, + ] + + first_csv = tmp_path / "overlap_first.csv" + second_csv = tmp_path / "overlap_second.csv" + pd.DataFrame(shared_rows).to_csv(first_csv, index=False) + pd.DataFrame([shared_rows[1], shared_rows[0], shared_rows[2]]).to_csv( + second_csv, index=False + ) + + first_result = process_joint_variants(str(first_csv)) + second_result = process_joint_variants(str(second_csv)) + + first_joint = first_result.loc[2, ["gene", "amino_acid_consequence"]].to_dict() + second_joint = second_result.loc[2, ["gene", "amino_acid_consequence"]].to_dict() + + assert first_joint == second_joint + + def test_generate_variant_heatmap_creates_interactive_html(tmp_path, monkeypatch): mpl_dir = tmp_path / "mpl" mpl_dir.mkdir() diff --git a/tests/test_analysis_launcher.py b/tests/test_analysis_launcher.py new file mode 100644 index 0000000..004c12d --- /dev/null +++ b/tests/test_analysis_launcher.py @@ -0,0 +1,37 @@ +import pytest + +from vartracker.analysis_launcher import _validate_primer_bed_reference + + +def test_validate_primer_bed_reference_accepts_matching_contig(tmp_path): + reference = tmp_path / "reference.fasta" + reference.write_text(">NC_045512.2 reference\nACGT\n", encoding="utf-8") + primer_bed = tmp_path / "primers.bed" + primer_bed.write_text( + "track name=primers\n" + "NC_045512.2\t0\t20\tprimer_1\n" + "NC_045512.2\t20\t40\tprimer_2\n", + encoding="utf-8", + ) + + _validate_primer_bed_reference(primer_bed, reference) + + +def test_validate_primer_bed_reference_rejects_mismatched_contig(tmp_path): + reference = tmp_path / "reference.fasta" + reference.write_text(">NC_045512.2 reference\nACGT\n", encoding="utf-8") + primer_bed = tmp_path / "primers.bed" + primer_bed.write_text("MN908947.3\t0\t20\tprimer_1\n", encoding="utf-8") + + with pytest.raises(ValueError, match="do not match reference FASTA ID"): + _validate_primer_bed_reference(primer_bed, reference) + + +def test_validate_primer_bed_reference_rejects_empty_bed(tmp_path): + reference = tmp_path / "reference.fasta" + reference.write_text(">NC_045512.2 reference\nACGT\n", encoding="utf-8") + primer_bed = tmp_path / "primers.bed" + primer_bed.write_text("# no intervals\n", encoding="utf-8") + + with pytest.raises(ValueError, match="contains no intervals"): + _validate_primer_bed_reference(primer_bed, reference) diff --git a/tests/test_cli.py b/tests/test_cli.py index 2418817..3f29f0c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -20,6 +20,26 @@ def test_cli_help_smoke(): assert "usage" in out.lower() or "help" in out.lower() +def test_plot_heatmap_help_uses_standalone_option_names(): + exit_code, out = _run_main("plot", "heatmap", "-h") + assert exit_code == 0 + assert "--aa-exclude" in out + assert "--include-joint" in out + assert "--title" in out + assert "--literature-csv" in out + assert "--x-labels" in out + heatmap_options = out.split("Heatmap options:", 1)[1] + assert "--title" in heatmap_options + assert "--literature-csv" in heatmap_options + assert "--x-labels" in heatmap_options + assert "--heatmap-aa-exclude" not in out + assert "--heatmap-include-joint" not in out + assert "--outdir" not in out + assert "--name" not in out + assert "--min-snv-freq" not in out + assert "--min-indel-freq" not in out + + def test_cli_version_option(): exit_code, out = _run_main("--version") assert exit_code == 0 diff --git a/tests/test_consensus.py b/tests/test_consensus.py new file mode 100644 index 0000000..5fe13c9 --- /dev/null +++ b/tests/test_consensus.py @@ -0,0 +1,48 @@ +import pytest + +from vartracker.consensus import iupac_consensus_base, simple_consensus_base + + +@pytest.mark.parametrize( + ("af", "simple", "iupac"), + [ + (0.20, "A", "A"), + (0.30, "A", "R"), + (0.50, "A", "R"), + (0.80, "G", "G"), + ], +) +def test_snp_consensus_three_tier_thresholds(af, simple, iupac): + kwargs = { + "depth": 100, + "af": af, + "min_depth": 10, + "consensus_snp_min_af": 0.25, + "consensus_snp_thresh": 0.75, + "consensus_indel_thresh": 0.75, + } + + assert simple_consensus_base("A", "G", **kwargs) == simple + assert iupac_consensus_base("A", "G", **kwargs) == iupac + + +@pytest.mark.parametrize( + ("af", "simple"), + [ + (0.74, "AT"), + (0.75, "A"), + (0.80, "A"), + ], +) +def test_indel_af_cutoff_is_separate_from_snp_ambiguity(af, simple): + kwargs = { + "depth": 100, + "af": af, + "min_depth": 10, + "consensus_snp_min_af": 0.25, + "consensus_snp_thresh": 0.75, + "consensus_indel_thresh": 0.75, + } + + assert simple_consensus_base("AT", "A", **kwargs) == simple + assert iupac_consensus_base("AT", "A", **kwargs) == simple diff --git a/tests/test_main.py b/tests/test_main.py index 15be8aa..340c90c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -47,7 +47,7 @@ def test_drop_exact_duplicate_result_rows_removes_only_exact_duplicates(capsys): {"variant": "A1C", "gene": "S", "amino_acid_consequence": "S:A1C"}, {"variant": "A1C", "gene": "N", "amino_acid_consequence": "N:A1C"}, ] - assert "Removed 1 exact duplicate result rows." in capsys.readouterr().out + assert capsys.readouterr().out == "" def test_prepare_reference_command_invokes_bundle(monkeypatch, tmp_path): @@ -172,6 +172,7 @@ def fake_setup(args): lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), ) monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr(main_module, "annotate_vcf", lambda *a, **k: None) monkeypatch.setattr( main_module, "process_vcf", @@ -248,6 +249,7 @@ def fake_setup(args): lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), ) monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr(main_module, "annotate_vcf", lambda *a, **k: None) monkeypatch.setattr( main_module, "process_vcf", @@ -341,6 +343,7 @@ def fake_heatmap(*args, **kwargs): lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), ) monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr(main_module, "annotate_vcf", lambda *a, **k: None) monkeypatch.setattr( main_module, "process_vcf", @@ -393,42 +396,61 @@ def fake_heatmap(*args, **kwargs): def test_plot_heatmap_replots_results_csv(monkeypatch, tmp_path): results_csv = tmp_path / "results.csv" results_csv.write_text( - "samples,name,alt_freq,variant_site_depth,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,start\n" - "P0 / P1,Example,0.0 / 0.5,0 / 100,N / Y,new,new_persistent,snp,missense,S,A266C,266\n", + "samples,sample_number,name,alt_freq,variant_site_depth,presence_absence,variant_status,persistence_status,type_of_variant,type_of_change,gene,variant,start\n" + "P0 / P1,0 / 1,Example,0.0 / 0.5,0 / 100,N / Y,new,new_persistent,snp,missense,S,A266C,266\n", encoding="utf-8", ) recorded = {} + literature_csv = tmp_path / "literature_hits.csv" + literature_csv.write_text( + "gene,amino_acid_consequence,information,reference\n" + "S,S:A266C,Mock evidence,PMID123\n", + encoding="utf-8", + ) def fake_heatmap(*args, **kwargs): recorded["table"] = args[0] recorded["sample_names"] = args[1] + recorded["sample_numbers"] = args[2] recorded["outdir"] = args[3] recorded["project_name"] = args[4] + recorded["min_snv_freq"] = args[5] + recorded["min_indel_freq"] = args[6] recorded["kwargs"] = kwargs monkeypatch.setattr(main_module, "generate_variant_heatmap", fake_heatmap) - outdir = tmp_path / "plots" exit_code = main_module.main( [ "plot", "heatmap", str(results_csv), - "--outdir", - str(outdir), - "--heatmap-aa-exclude", + "--aa-exclude", "*frameshift*", - "--heatmap-include-joint", + "--include-joint", + "--title", + "Custom heatmap", + "--x-labels", + "sample-number", + "--literature-csv", + str(literature_csv), ] ) assert exit_code == 0 assert list(recorded["sample_names"]) == ["P0", "P1"] - assert recorded["outdir"] == str(outdir) - assert recorded["project_name"] == "Example" + assert list(recorded["sample_numbers"]) == ["0", "1"] + assert recorded["outdir"] == str(tmp_path) + assert recorded["project_name"] == "" + assert recorded["min_snv_freq"] == 0.03 + assert recorded["min_indel_freq"] == 0.1 assert recorded["kwargs"]["excluded_consequence_types"] == ["*frameshift*"] assert recorded["kwargs"]["include_joint"] is True + assert recorded["kwargs"]["x_tick_labels"] == ["0", "1"] + assert recorded["kwargs"]["plot_title"] == "Custom heatmap" + assert recorded["kwargs"]["literature_table_path"] == str(literature_csv.resolve()) + assert recorded["kwargs"]["literature_hits"]["gene"].tolist() == ["S"] def _write_plot_results_csv(path: Path, n_variants: int = 4) -> None: @@ -907,6 +929,7 @@ def fake_genome(*args, **kwargs): lambda *a, **k: (str(tmp_path / "formatted.vcf.gz"), str(formatted_csq)), ) monkeypatch.setattr(main_module, "merge_consequences", lambda *a, **k: None) + monkeypatch.setattr(main_module, "annotate_vcf", lambda *a, **k: None) monkeypatch.setattr( main_module, "process_vcf", @@ -993,6 +1016,14 @@ def fake_vcf(args): "ref.fasta", "--snakemake-outdir", str(tmp_path / "snakemake"), + "--consensus-snp-min-af", + "0.20", + "--consensus-snp-thresh", + "0.80", + "--consensus-indel-thresh", + "0.70", + "--ampliconclip-tolerance", + "2", "--outdir", str(tmp_path / "results"), ] @@ -1003,6 +1034,11 @@ def fake_vcf(args): assert recorded["workflow_kwargs"]["force_all"] is False assert recorded["workflow_kwargs"]["quiet"] is True assert recorded["workflow_kwargs"]["mode"] == "reads" + assert recorded["workflow_kwargs"]["min_depth"] == 10 + assert recorded["workflow_kwargs"]["consensus_snp_min_af"] == 0.20 + assert recorded["workflow_kwargs"]["consensus_snp_thresh"] == 0.80 + assert recorded["workflow_kwargs"]["consensus_indel_thresh"] == 0.70 + assert recorded["workflow_kwargs"]["ampliconclip_tolerance"] == 2 assert recorded["vcf_input"] == str(updated_csv) assert modes_checked == ["e2e"] @@ -1170,6 +1206,10 @@ def fake_vcf(args): str(samples_csv), "--reference", "ref.fasta", + "--min-snv-freq", + "0.05", + "--min-depth", + "12", "--outdir", str(tmp_path / "results"), ] @@ -1180,6 +1220,10 @@ def fake_vcf(args): assert recorded["force_all"] is False assert recorded["quiet"] is True assert recorded["mode"] == "bam" + assert recorded["min_depth"] == 12 + assert recorded["consensus_snp_min_af"] == 0.25 + assert recorded["consensus_snp_thresh"] == 0.75 + assert recorded["consensus_indel_thresh"] == 0.75 assert recorded["vcf_input"] == str(updated_csv) assert recorded["suppress_logo"] is True assert modes_checked == ["bam"] diff --git a/tests/test_snakemake_workflow.py b/tests/test_snakemake_workflow.py index 184d7a4..acc5dea 100644 --- a/tests/test_snakemake_workflow.py +++ b/tests/test_snakemake_workflow.py @@ -17,8 +17,22 @@ def test_snakemake_rules_write_logs_under_outdir(): 'f"{OUTDIR}/{{sample}}/logs/lofreq_indelqual.log"', 'f"{OUTDIR}/{{sample}}/logs/samtools_depth.log"', 'f"{OUTDIR}/{{sample}}/logs/lofreq_call.log"', + 'f"{OUTDIR}/{{sample}}/logs/deletion_variants_bed.log"', + 'f"{OUTDIR}/{{sample}}/logs/depth_mask.log"', + 'f"{OUTDIR}/{{sample}}/logs/iupac_genotyped_vcf.log"', + 'f"{OUTDIR}/{{sample}}/logs/bcftools_consensus.log"', + 'f"{OUTDIR}/{{sample}}/logs/bcftools_iupac_consensus.log"', 'f"{OUTDIR}/logs/update_csv.log"', ] for expected in expected_logs: assert expected in snakefile + + assert "bedtools" not in snakefile + assert "--both-ends" not in snakefile + assert "--tolerance {params.tolerance}" in snakefile + assert "_validate_primer_bed_reference(PRIMER_BED, REF)" in snakefile + assert "_consensus.fasta" in snakefile + assert "_iupac_consensus.fasta" in snakefile + assert "df['consensus']" in snakefile + assert "df['iupac_consensus']" in snakefile diff --git a/tests/test_vcf_processing.py b/tests/test_vcf_processing.py index c461300..eb3bbb3 100644 --- a/tests/test_vcf_processing.py +++ b/tests/test_vcf_processing.py @@ -3,8 +3,22 @@ from __future__ import annotations import os +import shutil +import subprocess +from pathlib import Path -from vartracker.vcf_processing import _derive_vcf_output_paths +import pandas as pd +import pytest +from cyvcf2 import VCF + +from vartracker.analysis import process_joint_variants +from vartracker.vcf_processing import ( + _derive_vcf_output_paths, + annotate_vcf, + format_vcf, + merge_consequences, + process_vcf, +) def test_derive_vcf_output_paths_handles_gz(tmp_path): @@ -25,3 +39,254 @@ def test_derive_vcf_output_paths_handles_plain_vcf(tmp_path): assert csq.endswith("alpha.csq.vcf.gz") assert log.endswith("alpha.log") assert os.path.dirname(out) == str(tmp_path) + + +def _write_depth_file(path: Path, *, seqid: str = "chr1", length: int = 9) -> None: + with open(path, "w", encoding="utf-8") as handle: + for pos in range(1, length + 1): + handle.write(f"{seqid}\t{pos}\t100\n") + + +def _write_minimal_reference_bundle(tmp_path: Path) -> tuple[Path, Path]: + ref = tmp_path / "ref.fa" + gff = tmp_path / "ref.gff3" + ref.write_text(">chr1\nATGAAATAA\n", encoding="utf-8") + gff.write_text( + "##gff-version 3\n" + "chr1\tvartracker_test\tgene\t1\t9\t.\t+\t.\tID=gene:chr1_1;Name=GENE1;biotype=protein_coding\n" + "chr1\tvartracker_test\tmRNA\t1\t9\t.\t+\t.\tID=transcript:chr1_1;Parent=gene:chr1_1;Name=GENE1;biotype=protein_coding\n" + "chr1\tvartracker_test\tCDS\t1\t9\t.\t+\t0\tID=cds:chr1_1;Parent=transcript:chr1_1;gene=GENE1\n", + encoding="utf-8", + ) + return ref, gff + + +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +@pytest.mark.xfail( + strict=True, + reason="format_vcf currently deduplicates records by POS and can drop a real ALT allele", +) +def test_format_vcf_preserves_distinct_alt_records_at_same_position(tmp_path): + ref, gff = _write_minimal_reference_bundle(tmp_path) + vcf_path = tmp_path / "same_pos_two_alts.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\n" + "chr1\t5\t.\tA\tC\t.\tPASS\tAF=0.20;DP=100\tGT:DP:AF\t1:100:0.20\n" + "chr1\t5\t.\tA\tG\t.\tPASS\tAF=0.30;DP=100\tGT:DP:AF\t1:100:0.30\n", + encoding="utf-8", + ) + + formatted_vcf, _ = format_vcf( + str(vcf_path), + "s1", + str(tmp_path), + 0.0, + 0.0, + str(ref), + str(gff), + False, + ) + + records = list(VCF(formatted_vcf)) + + assert len(records) == 2 + assert {record.ALT[0] for record in records} == {"C", "G"} + + +def test_process_vcf_splits_sample_specific_bcsq_annotations(tmp_path): + vcf_path = tmp_path / "annotated.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\ts2\n" + "chr1\t4\t.\tA\tG\t.\tPASS\tDP=100;AF=0.1;BCSQ=missense|GENE1|tx|protein_coding|+|2K>2A|4A>G+5A>C\tGT:DP:AF:BCSQ\t1:100:0.1:1\t0:.:.:0\n" + "chr1\t5\t.\tA\tC\t.\tPASS\tDP=200;AF=0.1;BCSQ=@4,missense|GENE1|tx|protein_coding|+|2K>2T|5A>C\tGT:DP:AF:BCSQ\t1:100:0.1:1\t1:100:0.1:4\n", + encoding="utf-8", + ) + + cov1 = tmp_path / "s1.depth.txt" + cov2 = tmp_path / "s2.depth.txt" + _write_depth_file(cov1) + _write_depth_file(cov2) + + table = process_vcf(str(vcf_path), [str(cov1), str(cov2)], 10, ["s1", "s2"]) + observed = ( + table[ + [ + "start", + "amino_acid_consequence", + "bcsq_aa_notation", + "presence_absence", + "alt_freq", + ] + ] + .sort_values(["start", "amino_acid_consequence"]) + .reset_index(drop=True) + ) + + expected = pd.DataFrame( + [ + { + "start": 4, + "amino_acid_consequence": "K2A", + "bcsq_aa_notation": "2K>2A", + "presence_absence": "Y / N", + "alt_freq": "0.100 / .", + }, + { + "start": 5, + "amino_acid_consequence": "@4", + "bcsq_aa_notation": "@4", + "presence_absence": "Y / N", + "alt_freq": "0.100 / .", + }, + { + "start": 5, + "amino_acid_consequence": "K2T", + "bcsq_aa_notation": "2K>2T", + "presence_absence": "N / Y", + "alt_freq": ". / 0.100", + }, + ] + ) + + pd.testing.assert_frame_equal(observed, expected) + + +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +def test_merge_then_annotate_preserves_joint_annotations_across_samples(tmp_path): + ref, gff = _write_minimal_reference_bundle(tmp_path) + + sample1_vcf = tmp_path / "s1.vcf" + sample2_vcf = tmp_path / "s2.vcf" + header = ( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t{sample}\n" + ) + sample1_vcf.write_text( + header.format(sample="s1") + + "chr1\t4\t.\tA\tG\t.\tPASS\tAF=0.10;DP=100\tGT:DP:AF\t1:100:0.10\n" + + "chr1\t5\t.\tA\tC\t.\tPASS\tAF=0.10;DP=100\tGT:DP:AF\t1:100:0.10\n", + encoding="utf-8", + ) + sample2_vcf.write_text( + header.format(sample="s2") + + "chr1\t5\t.\tA\tC\t.\tPASS\tAF=0.10;DP=100\tGT:DP:AF\t1:100:0.10\n", + encoding="utf-8", + ) + + for path in (sample1_vcf, sample2_vcf): + gz_path = Path(f"{path}.gz") + subprocess.run( + ["bcftools", "view", "-Oz", "-o", str(gz_path), str(path)], + check=True, + capture_output=True, + text=True, + ) + subprocess.run( + ["bcftools", "index", "-f", str(gz_path)], + check=True, + capture_output=True, + text=True, + ) + + vcf_list = tmp_path / "vcf_list.txt" + vcf_list.write_text( + f"{sample1_vcf}.gz\n{sample2_vcf}.gz\n", + encoding="utf-8", + ) + sample_names = tmp_path / "sample_names.txt" + sample_names.write_text("s1\ns2\n", encoding="utf-8") + + merged_vcf = tmp_path / "merged.vcf" + merge_consequences(str(tmp_path), str(merged_vcf), str(sample_names), debug=False) + + annotated_vcf = tmp_path / "annotated.vcf" + annotate_vcf( + str(merged_vcf), + str(annotated_vcf), + str(ref), + str(gff), + debug=False, + ) + + cov1 = tmp_path / "s1.depth.txt" + cov2 = tmp_path / "s2.depth.txt" + _write_depth_file(cov1) + _write_depth_file(cov2) + + table = process_vcf( + str(annotated_vcf), + [str(cov1), str(cov2)], + 10, + ["s1", "s2"], + ) + results_csv = tmp_path / "results.csv" + table.to_csv(results_csv, index=False) + processed = process_joint_variants(str(results_csv)) + + observed = ( + processed[ + [ + "start", + "variant", + "amino_acid_consequence", + "presence_absence", + "type_of_change", + "joint_variant", + ] + ] + .sort_values(["start", "amino_acid_consequence", "presence_absence"]) + .reset_index(drop=True) + ) + + expected = pd.DataFrame( + [ + { + "start": 4, + "variant": "A4G", + "amino_acid_consequence": "K2A", + "presence_absence": "Y / N", + "type_of_change": "joint_missense", + "joint_variant": True, + }, + { + "start": 5, + "variant": "A5C", + "amino_acid_consequence": "K2A", + "presence_absence": "Y / N", + "type_of_change": "joint_missense", + "joint_variant": True, + }, + { + "start": 5, + "variant": "A5C", + "amino_acid_consequence": "K2T", + "presence_absence": "N / Y", + "type_of_change": "missense", + "joint_variant": False, + }, + ] + ) + + pd.testing.assert_frame_equal(observed, expected) diff --git a/vartracker/Snakefile b/vartracker/Snakefile index 8aae8ac..b1bc34b 100644 --- a/vartracker/Snakefile +++ b/vartracker/Snakefile @@ -1,5 +1,9 @@ import pandas as pd from pathlib import Path +try: + from vartracker.consensus import consensus_genotype_for_variant +except ModuleNotFoundError: + from consensus import consensus_genotype_for_variant # Load samples from CSV samples_df = pd.read_csv(config["samples_csv"]) @@ -9,7 +13,12 @@ SAMPLES = samples_df.set_index("sample_name", drop=False).to_dict(orient="index" REF = config["reference"] OUTDIR = config.get("outdir", "results") PRIMER_BED = config.get("primer_bed", None) +AMPLICONCLIP_TOLERANCE = int(config.get("ampliconclip_tolerance", 1)) MODE = config.get("mode", "reads") +MIN_DEPTH = int(config.get("min_depth", 10)) +CONSENSUS_SNP_MIN_AF = float(config.get("consensus_snp_min_af", 0.25)) +CONSENSUS_SNP_THRESH = float(config.get("consensus_snp_thresh", 0.75)) +CONSENSUS_INDEL_THRESH = float(config.get("consensus_indel_thresh", 0.75)) def _value_or_blank(sample, key, allow_blank=False): value = SAMPLES[sample].get(key) @@ -28,6 +37,51 @@ def _require_value(sample, key): return value +def _first_fasta_record_id(fasta_path): + with open(fasta_path, "r", encoding="utf-8") as handle: + for line in handle: + if line.startswith(">"): + return line[1:].split()[0].strip() + return None + + +def _primer_bed_contigs(primer_bed_path): + contigs = set() + with open(primer_bed_path, "r", encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("track ") or line.startswith("browser "): + continue + contigs.add(line.split()[0]) + return contigs + + +def _validate_primer_bed_reference(primer_bed_path, reference_path): + reference_id = _first_fasta_record_id(reference_path) + if not reference_id: + raise ValueError( + f"Unable to determine reference FASTA ID from: {reference_path}" + ) + + bed_contigs = _primer_bed_contigs(primer_bed_path) + if not bed_contigs: + raise ValueError(f"Primer BED file contains no intervals: {primer_bed_path}") + + mismatched = sorted(contig for contig in bed_contigs if contig != reference_id) + if mismatched: + joined = ", ".join(mismatched) + raise ValueError( + f"Primer BED contig(s) {joined!r} do not match reference FASTA ID " + f"{reference_id!r}" + ) + + +if PRIMER_BED: + _validate_primer_bed_reference(PRIMER_BED, REF) + + def _initial_bam(wildcards): if MODE == "reads": return f"{OUTDIR}/{wildcards.sample}/aligned.clipped.bam" @@ -43,6 +97,8 @@ rule all: expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_consensus.fasta", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_iupac_consensus.fasta", sample=SAMPLES.keys()), f"{OUTDIR}/vartracker_execution_spreadsheet.csv" if MODE == "reads": @@ -143,7 +199,8 @@ if MODE == "reads": bam = temp(f"{OUTDIR}/{{sample}}/aligned.clipped.bam"), bai = temp(f"{OUTDIR}/{{sample}}/aligned.clipped.bam.bai") params: - bed = PRIMER_BED if PRIMER_BED else "" + bed = PRIMER_BED if PRIMER_BED else "", + tolerance = AMPLICONCLIP_TOLERANCE log: f"{OUTDIR}/{{sample}}/logs/ampliconclip.log" threads: max(1, int(workflow.cores * 0.5)) @@ -152,7 +209,7 @@ if MODE == "reads": shell(""" mkdir -p $(dirname {log}) samtools ampliconclip -b {params.bed} -@ {threads} \ - --strand --both-ends -o - {input.bam} 2> {log} \ + --strand --tolerance {params.tolerance} -o - {input.bam} 2> {log} \ | samtools sort -@ {threads} -o {output.bam} - 2>> {log} samtools index {output.bam} 2>> {log} """) @@ -214,11 +271,247 @@ rule lofreq_call: bcftools index {output.vcf} 2>> {log} """ +rule deletion_variants_bed: + input: + vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz" + output: + bed = temp(f"{OUTDIR}/{{sample}}/{{sample}}_deletions.bed") + params: + min_depth = MIN_DEPTH, + consensus_indel_thresh = CONSENSUS_INDEL_THRESH + log: + f"{OUTDIR}/{{sample}}/logs/deletion_variants_bed.log" + shell: + """ + mkdir -p $(dirname {log}) + bcftools query \ + -i 'TYPE="indel" && strlen(REF)>strlen(ALT) && INFO/DP >= {params.min_depth} && INFO/AF >= {params.consensus_indel_thresh}' \ + -f'%CHROM\\t%POS0\\t%END\\n' {input.vcf} > {output.bed} 2> {log} + """ + +rule depth_mask: + input: + depth = f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt", + deletions = rules.deletion_variants_bed.output.bed + output: + mask = f"{OUTDIR}/{{sample}}/{{sample}}_lowdepth.bed" + params: + min_depth = MIN_DEPTH + log: + f"{OUTDIR}/{{sample}}/logs/depth_mask.log" + run: + from collections import defaultdict + + log_path = Path(str(log[0])) + log_path.parent.mkdir(parents=True, exist_ok=True) + + deletion_intervals = defaultdict(list) + with open(input.deletions, "r", encoding="utf-8") as handle: + for line in handle: + parts = line.rstrip("\n").split("\t") + if len(parts) < 3: + continue + chrom, start, end = parts[:3] + deletion_intervals[chrom].append((int(start), int(end))) + + for intervals in deletion_intervals.values(): + intervals.sort() + + def overlaps_deletion(chrom, pos0): + intervals = deletion_intervals.get(chrom, []) + for start, end in intervals: + if pos0 < start: + return False + if start <= pos0 < end: + return True + return False + + mask_path = Path(str(output.mask)) + mask_path.parent.mkdir(parents=True, exist_ok=True) + + intervals_written = 0 + current_chrom = None + current_start = None + current_end = None + with open(input.depth, "r", encoding="utf-8") as depth_handle, open( + mask_path, "w", encoding="utf-8" + ) as mask_handle: + for line in depth_handle: + parts = line.rstrip("\n").split("\t") + if len(parts) < 3: + continue + chrom, pos_text, depth_text = parts[:3] + pos0 = int(pos_text) - 1 + depth = int(depth_text) + should_mask = depth < params.min_depth and not overlaps_deletion( + chrom, pos0 + ) + + if not should_mask: + if current_chrom is not None: + mask_handle.write( + f"{current_chrom}\t{current_start}\t{current_end}\n" + ) + intervals_written += 1 + current_chrom = current_start = current_end = None + continue + + if current_chrom == chrom and current_end == pos0: + current_end = pos0 + 1 + else: + if current_chrom is not None: + mask_handle.write( + f"{current_chrom}\t{current_start}\t{current_end}\n" + ) + intervals_written += 1 + current_chrom = chrom + current_start = pos0 + current_end = pos0 + 1 + + if current_chrom is not None: + mask_handle.write(f"{current_chrom}\t{current_start}\t{current_end}\n") + intervals_written += 1 + + log_path.write_text( + f"Wrote {intervals_written} low-depth mask intervals to {output.mask}\n", + encoding="utf-8", + ) + +rule iupac_genotyped_vcf: + input: + vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz" + output: + vcf = temp(f"{OUTDIR}/{{sample}}/{{sample}}_iupac_genotyped.vcf.gz"), + csi = temp(f"{OUTDIR}/{{sample}}/{{sample}}_iupac_genotyped.vcf.gz.csi") + params: + sample = "{sample}", + min_depth = MIN_DEPTH, + consensus_snp_min_af = CONSENSUS_SNP_MIN_AF, + consensus_snp_thresh = CONSENSUS_SNP_THRESH, + consensus_indel_thresh = CONSENSUS_INDEL_THRESH + log: + f"{OUTDIR}/{{sample}}/logs/iupac_genotyped_vcf.log" + run: + import gzip + import os + + log_path = Path(str(log[0])) + log_path.parent.mkdir(parents=True, exist_ok=True) + + def open_text(path): + if str(path).endswith(".gz"): + return gzip.open(path, "rt", encoding="utf-8") + return open(path, "r", encoding="utf-8") + + def info_value(info, key, default=None): + for item in info.split(";"): + if item.startswith(f"{key}="): + return item.split("=", 1)[1] + return default + + def genotype_for(ref, alt, info): + try: + depth = int(float(info_value(info, "DP", "0"))) + af = float(info_value(info, "AF", "0").split(",", 1)[0]) + except ValueError: + return "./." + + return consensus_genotype_for_variant( + ref, + alt, + depth=depth, + af=af, + min_depth=params.min_depth, + consensus_snp_min_af=params.consensus_snp_min_af, + consensus_snp_thresh=params.consensus_snp_thresh, + consensus_indel_thresh=params.consensus_indel_thresh, + ) + + tmp_vcf = f"{output.vcf}.tmp.vcf" + with open_text(input.vcf) as src, open( + tmp_vcf, "w", encoding="utf-8" + ) as dst: + for line in src: + if line.startswith("##"): + dst.write(line) + continue + if line.startswith("#CHROM"): + dst.write( + '##FORMAT=\n' + ) + dst.write(line.rstrip("\n") + f"\tFORMAT\t{params.sample}\n") + continue + parts = line.rstrip("\n").split("\t") + if len(parts) < 8: + continue + gt = genotype_for(parts[3], parts[4], parts[7]) + dst.write("\t".join(parts[:8] + ["GT", gt]) + "\n") + + try: + shell("bgzip -c {tmp_vcf} > {output.vcf} 2> {log}") + shell("bcftools index -f -c {output.vcf} 2>> {log}") + finally: + if os.path.exists(tmp_vcf): + os.remove(tmp_vcf) + +rule consensus: + input: + ref = REF, + mask = rules.depth_mask.output.mask, + vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz" + output: + fasta = f"{OUTDIR}/{{sample}}/{{sample}}_consensus.fasta" + params: + prefix = "{sample}", + min_depth = MIN_DEPTH, + consensus_snp_thresh = CONSENSUS_SNP_THRESH, + consensus_indel_thresh = CONSENSUS_INDEL_THRESH + log: + f"{OUTDIR}/{{sample}}/logs/bcftools_consensus.log" + shell: + """ + mkdir -p $(dirname {log}) + bcftools consensus -p "{params.prefix} " \ + -f {input.ref} \ + --mark-del '-' \ + -m {input.mask} \ + -i 'INFO/DP >= {params.min_depth} && ((TYPE="snp" && INFO/AF >= {params.consensus_snp_thresh}) || (TYPE!="snp" && INFO/AF >= {params.consensus_indel_thresh}))' \ + {input.vcf} -o {output.fasta} > {log} 2>&1 + """ + +rule iupac_consensus: + input: + ref = REF, + mask = rules.depth_mask.output.mask, + vcf = rules.iupac_genotyped_vcf.output.vcf, + csi = rules.iupac_genotyped_vcf.output.csi + output: + fasta = f"{OUTDIR}/{{sample}}/{{sample}}_iupac_consensus.fasta" + params: + prefix = "{sample}", + min_depth = MIN_DEPTH + log: + f"{OUTDIR}/{{sample}}/logs/bcftools_iupac_consensus.log" + shell: + """ + mkdir -p $(dirname {log}) + bcftools consensus -p "{params.prefix} IUPAC " \ + -f {input.ref} \ + --mark-del '-' \ + -m {input.mask} \ + -s "{params.prefix}" \ + -H I \ + -i 'INFO/DP >= {params.min_depth} && GT!="mis"' \ + {input.vcf} -o {output.fasta} > {log} 2>&1 + """ + rule update_csv: input: vcfs = expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", sample=SAMPLES.keys()), depths = expand(f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt", sample=SAMPLES.keys()), - bams = expand(f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam", sample=SAMPLES.keys()) + bams = expand(f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam", sample=SAMPLES.keys()), + consensus = expand(f"{OUTDIR}/{{sample}}/{{sample}}_consensus.fasta", sample=SAMPLES.keys()), + iupac_consensus = expand(f"{OUTDIR}/{{sample}}/{{sample}}_iupac_consensus.fasta", sample=SAMPLES.keys()) output: csv = f"{OUTDIR}/vartracker_execution_spreadsheet.csv" params: @@ -246,6 +539,12 @@ rule update_csv: df['coverage'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_depth.txt") ) + df['consensus'] = df['sample_name'].apply( + lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_consensus.fasta") + ) + df['iupac_consensus'] = df['sample_name'].apply( + lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_iupac_consensus.fasta") + ) # Write updated CSV df.to_csv(output.csv, index=False) diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 1db9e75..0ecb115 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -96,6 +96,34 @@ def _heatmap_figure_size(n_rows: int, n_cols: int) -> tuple[float, float]: return width, height +def _presence_vector(value: object) -> tuple[str, ...]: + return tuple(_parse_slash_separated_tokens(value)) + + +def _find_joint_main_index(tab: pd.DataFrame, joint_index: int, main_pos: int) -> int: + candidates = tab[ + (tab["start"] == main_pos) + & ~tab["bcsq_aa_notation"].astype(str).str.startswith("@", na=False) + ].index.tolist() + if not candidates: + raise IndexError(f"No main annotation found for joint position {main_pos}") + if len(candidates) == 1: + return candidates[0] + + joint_presence = _presence_vector(tab.at[joint_index, "presence_absence"]) + + def candidate_key(candidate_index: int) -> tuple[int, int, int]: + candidate_presence = _presence_vector(tab.at[candidate_index, "presence_absence"]) + exact_match = int(candidate_presence == joint_presence) + shared_present = sum( + lhs == rhs == "Y" for lhs, rhs in zip(joint_presence, candidate_presence) + ) + total_present = sum(token == "Y" for token in candidate_presence) + return (exact_match, shared_present, total_present) + + return max(candidates, key=candidate_key) + + def process_joint_variants(path): """ Process joint variants from bcftools csq output. @@ -124,7 +152,7 @@ def process_joint_variants(path): try: # Find the main pos and main index number main_pos = int(tab.loc[i]["bcsq_aa_notation"].replace("@", "")) - j = tab[tab["start"] == main_pos].index[0] + j = _find_joint_main_index(tab, i, main_pos) # Update the joint variant key tab.at[i, "joint_variant"] = True @@ -937,13 +965,16 @@ def _write_interactive_heatmap_html( literature_df: Optional[pd.DataFrame], literature_table_path: Optional[str], cli_command: Optional[str], + sample_labels: Sequence[str] | None = None, + plot_title: str | None = None, ) -> None: if matrix.empty: return - x_labels = [str(name) for name in sample_names] + sample_keys = [str(name) for name in matrix.columns] + x_labels = [str(name) for name in (sample_labels or sample_names)] y_labels = list(matrix.index) - if not x_labels or not y_labels: + if not sample_keys or not x_labels or not y_labels: return label_map: Dict[str, str] = matrix.attrs.get("base_labels", {}) @@ -1007,14 +1038,14 @@ def _frequency_to_color(value: float) -> tuple[str, str]: row_values = matrix.loc[label] row_qc = qc_by_label.get(label, {}) - for sample, value in zip(x_labels, row_values): + for sample_key, sample_label, value in zip(sample_keys, x_labels, row_values): freq = float(value) if value is not None else 0.0 color, text_value = _frequency_to_color(freq) - qc_value = row_qc.get(sample, "") + qc_value = row_qc.get(sample_key, "") qc_failed = qc_value == "F" tooltip = html.escape( f"Variant: {label.replace(chr(10), ' ')} • " - f"Sample: {sample} • " + f"Sample: {sample_label} • " f"AF={freq:.2f}, QC={'FAIL' if qc_failed else 'PASS'}" ) classes = ["cell"] @@ -1037,7 +1068,7 @@ def _frequency_to_color(value: float) -> tuple[str, str]: ) heatmap_scroll_html = f'
{heatmap_grid_html}
' - heatmap_title = ( + heatmap_title = plot_title or ( f"{project_name}: variant allele frequencies" if project_name else "Variant allele frequencies" @@ -1329,6 +1360,8 @@ def generate_variant_heatmap( literature_hits: Optional[pd.DataFrame] = None, literature_table_path: Optional[str] = None, cli_command: Optional[str] = None, + x_tick_labels: Sequence[str] | None = None, + plot_title: str | None = None, ): """Generate a heatmap of variant allele frequencies across passages.""" @@ -1400,14 +1433,23 @@ def generate_variant_heatmap( ) ) - heatmap_title = ( + heatmap_title = plot_title or ( f"{project_name}: variant allele frequencies" if project_name else "Variant allele frequencies" ) ax.set_title(heatmap_title, weight="bold") - tick_labels = list(sample_names) + display_label_by_sample = dict( + zip( + [str(name) for name in sample_names], + [str(label) for label in (x_tick_labels or sample_names)], + ) + ) + tick_labels = [ + display_label_by_sample.get(str(sample), str(sample)) + for sample in heatmap_data.columns + ] ax.set_xticklabels(tick_labels, rotation=45, ha="right") ax.tick_params(axis="y", labelsize=10) ax.set_xlabel("Sample", fontweight="bold") @@ -1426,6 +1468,8 @@ def generate_variant_heatmap( literature_hits, literature_table_path, cli_command, + sample_labels=tick_labels, + plot_title=heatmap_title, ) except Exception as html_exc: # pragma: no cover - best-effort UX print(f"Warning: failed to generate interactive heatmap: {html_exc}") diff --git a/vartracker/analysis_launcher.py b/vartracker/analysis_launcher.py index 5aa7de4..a5fda82 100644 --- a/vartracker/analysis_launcher.py +++ b/vartracker/analysis_launcher.py @@ -24,12 +24,60 @@ def _normalise_path(value: str | Path) -> str: return str(Path(value).expanduser().resolve()) +def _first_fasta_record_id(fasta_path: str | Path) -> str | None: + with Path(fasta_path).open("r", encoding="utf-8") as handle: + for line in handle: + if line.startswith(">"): + return line[1:].split()[0].strip() + return None + + +def _primer_bed_contigs(primer_bed_path: str | Path) -> set[str]: + contigs = set() + with Path(primer_bed_path).open("r", encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("track ") or line.startswith("browser "): + continue + contigs.add(line.split()[0]) + return contigs + + +def _validate_primer_bed_reference( + primer_bed_path: str | Path, reference_path: str | Path +) -> None: + reference_id = _first_fasta_record_id(reference_path) + if not reference_id: + raise ValueError( + f"Unable to determine reference FASTA ID from: {reference_path}" + ) + + bed_contigs = _primer_bed_contigs(primer_bed_path) + if not bed_contigs: + raise ValueError(f"Primer BED file contains no intervals: {primer_bed_path}") + + mismatched = sorted(contig for contig in bed_contigs if contig != reference_id) + if mismatched: + joined = ", ".join(mismatched) + raise ValueError( + f"Primer BED contig(s) {joined!r} do not match reference FASTA ID " + f"{reference_id!r}" + ) + + def run_workflow( samples_csv: str | Path, reference: str | Path, outdir: str | Path = "results", cores: int = 8, primer_bed: Optional[str | Path] = None, + ampliconclip_tolerance: int = 1, + min_depth: int = 10, + consensus_snp_min_af: float = 0.25, + consensus_snp_thresh: float = 0.75, + consensus_indel_thresh: float = 0.75, dryrun: bool = False, force_all: bool = False, quiet: bool = True, @@ -57,6 +105,8 @@ def run_workflow( raise FileNotFoundError(f"Reference genome not found: {reference}") if primer_bed_path and not Path(primer_bed_path).exists(): raise FileNotFoundError(f"Primer BED file not found: {primer_bed_path}") + if primer_bed_path: + _validate_primer_bed_reference(primer_bed_path, reference) Path(outdir).mkdir(parents=True, exist_ok=True) @@ -66,6 +116,11 @@ def run_workflow( "reference": reference, "outdir": outdir, "mode": mode, + "ampliconclip_tolerance": ampliconclip_tolerance, + "min_depth": min_depth, + "consensus_snp_min_af": consensus_snp_min_af, + "consensus_snp_thresh": consensus_snp_thresh, + "consensus_indel_thresh": consensus_indel_thresh, } if primer_bed_path: config_dict["primer_bed"] = primer_bed_path diff --git a/vartracker/consensus.py b/vartracker/consensus.py new file mode 100644 index 0000000..33d55ae --- /dev/null +++ b/vartracker/consensus.py @@ -0,0 +1,108 @@ +"""Consensus allele/genotype helpers for the Snakemake workflow.""" + +from __future__ import annotations + +IUPAC_CODES: dict[frozenset[str], str] = { + frozenset({"A", "G"}): "R", + frozenset({"C", "T"}): "Y", + frozenset({"G", "C"}): "S", + frozenset({"A", "T"}): "W", + frozenset({"G", "T"}): "K", + frozenset({"A", "C"}): "M", +} + + +def first_alt_allele(alt: str) -> str: + """Return the first ALT allele from a possibly multi-allelic VCF field.""" + return str(alt).split(",", 1)[0] + + +def is_snp_variant(ref: str, alt: str) -> bool: + """Return True when the first ALT allele represents a SNP.""" + return len(str(ref)) == 1 and len(first_alt_allele(alt)) == 1 + + +def consensus_genotype_for_variant( + ref: str, + alt: str, + *, + depth: int, + af: float, + min_depth: int, + consensus_snp_min_af: float, + consensus_snp_thresh: float, + consensus_indel_thresh: float, +) -> str: + """Return the genotype bcftools consensus should use for IUPAC output.""" + if depth < min_depth: + return "./." + + if is_snp_variant(ref, alt): + if af >= consensus_snp_thresh: + return "1/1" + if af >= consensus_snp_min_af: + return "0/1" + return "./." + + if af >= consensus_indel_thresh: + return "1/1" + return "./." + + +def simple_consensus_base( + ref: str, + alt: str, + *, + depth: int, + af: float, + min_depth: int, + consensus_snp_min_af: float, + consensus_snp_thresh: float, + consensus_indel_thresh: float, +) -> str: + """Return REF or ALT for simple consensus tests.""" + if depth < min_depth: + return ref + + if is_snp_variant(ref, alt): + if af >= consensus_snp_thresh: + return first_alt_allele(alt) + return ref + + if af >= consensus_indel_thresh: + return first_alt_allele(alt) + return ref + + +def iupac_consensus_base( + ref: str, + alt: str, + *, + depth: int, + af: float, + min_depth: int, + consensus_snp_min_af: float, + consensus_snp_thresh: float, + consensus_indel_thresh: float, +) -> str: + """Return the expected single-site IUPAC consensus symbol for tests.""" + genotype = consensus_genotype_for_variant( + ref, + alt, + depth=depth, + af=af, + min_depth=min_depth, + consensus_snp_min_af=consensus_snp_min_af, + consensus_snp_thresh=consensus_snp_thresh, + consensus_indel_thresh=consensus_indel_thresh, + ) + if genotype == "./.": + return ref + + first_alt = first_alt_allele(alt) + if genotype == "1/1": + return first_alt + + if is_snp_variant(ref, alt): + return IUPAC_CODES.get(frozenset({ref.upper(), first_alt.upper()}), "N") + return ref diff --git a/vartracker/main.py b/vartracker/main.py index c1083f2..190a46d 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -32,7 +32,7 @@ ProcessingError, FILE_COLUMNS, ) -from .vcf_processing import format_vcf, merge_consequences, process_vcf +from .vcf_processing import annotate_vcf, format_vcf, merge_consequences, process_vcf from .analysis import ( process_joint_variants, generate_cumulative_lineplot, @@ -130,76 +130,114 @@ def _parse_csv_option_list(value: str | None) -> list[str]: return [item.strip() for item in str(value).split(",") if item.strip()] -def _add_heatmap_option_arguments(group: argparse._ArgumentGroup) -> None: +def _add_heatmap_option_arguments( + group: argparse._ArgumentGroup, + *, + prefix: str = "heatmap", + add_legacy_prefixed_aliases: bool = False, +) -> None: + option_prefix = f"{prefix}-" if prefix else "" + + def option(name: str) -> str: + return f"--{option_prefix}{name}" + + def add_legacy(name: str, dest: str, **kwargs) -> None: + if add_legacy_prefixed_aliases and not prefix: + kwargs.setdefault("default", argparse.SUPPRESS) + group.add_argument( + f"--heatmap-{name}", + dest=dest, + help=argparse.SUPPRESS, + **kwargs, + ) + group.add_argument( - "--heatmap-aa-exclude", + option("aa-exclude"), action="store", required=False, default="", + dest="heatmap_aa_exclude", help=( "Comma-separated `type_of_change` patterns to exclude from heatmaps " "(wildcards supported, e.g. synonymous,*frameshift*)" ), ) + add_legacy("aa-exclude", "heatmap_aa_exclude", action="store") group.add_argument( - "--heatmap-aa-include", + option("aa-include"), action="store", required=False, default="", + dest="heatmap_aa_include", help=( "Comma-separated `type_of_change` patterns to include in heatmaps " "(wildcards supported)" ), ) + add_legacy("aa-include", "heatmap_aa_include", action="store") group.add_argument( - "--heatmap-include-joint", + option("include-joint"), action="store_true", default=False, + dest="heatmap_include_joint", help="Include joint variants in heatmaps (default: exclude them)", ) + add_legacy("include-joint", "heatmap_include_joint", action="store_true") group.add_argument( - "--heatmap-only-persistent", + option("only-persistent"), action="store_true", default=False, + dest="heatmap_only_persistent", help="Only include variants with persistence_status == new_persistent", ) + add_legacy("only-persistent", "heatmap_only_persistent", action="store_true") group.add_argument( - "--heatmap-only-new", + option("only-new"), action="store_true", default=False, + dest="heatmap_only_new", help="Only include variants with variant_status == new", ) + add_legacy("only-new", "heatmap_only_new", action="store_true") group.add_argument( - "--heatmap-gene-include", + option("gene-include"), action="store", required=False, default="", + dest="heatmap_gene_include", help="Comma-separated gene patterns to include in heatmaps", ) + add_legacy("gene-include", "heatmap_gene_include", action="store") group.add_argument( - "--heatmap-gene-exclude", + option("gene-exclude"), action="store", required=False, default="", + dest="heatmap_gene_exclude", help="Comma-separated gene patterns to exclude from heatmaps", ) + add_legacy("gene-exclude", "heatmap_gene_exclude", action="store") group.add_argument( - "--heatmap-variant-type", + option("variant-type"), action="store", required=False, default="", + dest="heatmap_variant_type", help="Comma-separated variant type patterns to include (e.g. snp,indel)", ) + add_legacy("variant-type", "heatmap_variant_type", action="store") group.add_argument( - "--heatmap-qc", + option("qc"), action="store", required=False, default="", + dest="heatmap_qc", help=( "Comma-separated all-samples QC patterns to include " "(e.g. true,false,pass,fail)" ), ) + add_legacy("qc", "heatmap_qc", action="store") group.add_argument( "--min-prop-passing-qc", action="store", @@ -208,46 +246,58 @@ def _add_heatmap_option_arguments(group: argparse._ArgumentGroup) -> None: help="Minimum proportion of samples that must pass per-sample QC (0-1)", ) group.add_argument( - "--heatmap-min-persistence", + option("min-persistence"), action="store", type=int, default=None, + dest="heatmap_min_persistence", help="Minimum number of samples in which a variant must be present", ) + add_legacy("min-persistence", "heatmap_min_persistence", action="store", type=int) group.add_argument( - "--heatmap-min-max-af", + option("min-max-af"), action="store", type=float, default=None, + dest="heatmap_min_max_af", help="Minimum maximum allele frequency across included samples", ) + add_legacy("min-max-af", "heatmap_min_max_af", action="store", type=float) group.add_argument( - "--heatmap-min-sample-af", + option("min-sample-af"), action="store", type=float, default=None, + dest="heatmap_min_sample_af", help="Minimum allele frequency that must be reached in at least one included sample", ) + add_legacy("min-sample-af", "heatmap_min_sample_af", action="store", type=float) group.add_argument( - "--heatmap-sample-subset", + option("sample-subset"), action="store", required=False, default="", + dest="heatmap_sample_subset", help="Comma-separated sample name patterns to plot", ) + add_legacy("sample-subset", "heatmap_sample_subset", action="store") group.add_argument( - "--heatmap-hide-singletons", + option("hide-singletons"), action="store_true", default=False, + dest="heatmap_hide_singletons", help="Hide variants present in only one included sample", ) + add_legacy("hide-singletons", "heatmap_hide_singletons", action="store_true") group.add_argument( - "--heatmap-min-depth", + option("min-depth"), action="store", type=int, default=None, + dest="heatmap_min_depth", help="Minimum site depth a variant must reach in at least one included sample", ) + add_legacy("min-depth", "heatmap_min_depth", action="store", type=int) def _collect_heatmap_kwargs(args) -> dict[str, object]: @@ -622,6 +672,7 @@ def _configure_vcf_parser( *, include_input_csv: bool, input_csv_required: bool = False, + include_consensus_options: bool = False, ) -> None: if include_input_csv: if input_csv_required: @@ -682,6 +733,37 @@ def _configure_vcf_parser( default=10, help="Minimum depth threshold for variant QC (default: 10)", ) + if include_consensus_options: + analysis_group.add_argument( + "--consensus-snp-min-af", + action="store", + required=False, + type=float, + default=0.25, + help=( + "Minimum SNP allele frequency required before the site is " + "considered for consensus (default: 0.25)" + ), + ) + analysis_group.add_argument( + "--consensus-snp-thresh", + action="store", + required=False, + type=float, + default=0.75, + help=( + "Minimum SNP allele frequency required for ALT to become the " + "consensus base (default: 0.75)" + ), + ) + analysis_group.add_argument( + "--consensus-indel-thresh", + action="store", + required=False, + type=float, + default=0.75, + help="Minimum indel allele frequency required for consensus (default: 0.75)", + ) analysis_group.add_argument( "--sample-cap", action="store", @@ -873,7 +955,12 @@ def _add_bam_subparser(subparsers): default=False, ) - _configure_vcf_parser(bam_parser, include_input_csv=True, input_csv_required=False) + _configure_vcf_parser( + bam_parser, + include_input_csv=True, + input_csv_required=False, + include_consensus_options=True, + ) _move_action_group_after( bam_parser, "Snakemake options", "Vartracker Analysis Options" ) @@ -1282,41 +1369,26 @@ def _add_plot_heatmap_subparser(subparsers): formatter_class=HelpFormatter, ) parser.add_argument("results_csv", help="Path to a vartracker results CSV") - parser.add_argument( - "--outdir", - default=None, - help="Output directory for regenerated heatmap files (default: results CSV directory)", - ) - parser.add_argument( - "--name", + heatmap_group = parser.add_argument_group("Heatmap options") + heatmap_group.add_argument( + "--title", default=None, - help="Optional plot title prefix (default: use the `name` column if present)", + help="Plot title for the heatmap (default: Variant allele frequencies)", ) - parser.add_argument( + heatmap_group.add_argument( "--literature-csv", default=None, help="Optional literature hits CSV to link from the interactive heatmap", ) - parser.add_argument( - "-m", - "--min-snv-freq", - action="store", - required=False, - type=float, - default=0.03, - help="Minimum allele frequency of SNV variants to keep (default: 0.03)", + heatmap_group.add_argument( + "--x-labels", + choices=["sample-name", "sample-number"], + default="sample-name", + help="X-axis labels to use for the heatmap (default: sample-name)", ) - parser.add_argument( - "-M", - "--min-indel-freq", - action="store", - required=False, - type=float, - default=0.1, - help="Minimum allele frequency of indel variants to keep (default: 0.1)", + _add_heatmap_option_arguments( + heatmap_group, prefix="", add_legacy_prefixed_aliases=True ) - heatmap_group = parser.add_argument_group("Heatmap options") - _add_heatmap_option_arguments(heatmap_group) parser.set_defaults(handler=_run_plot_heatmap_command) @@ -1635,9 +1707,6 @@ def _normalise_rulegraph_path(path: str | None) -> str | None: def _drop_exact_duplicate_result_rows(table: pd.DataFrame) -> pd.DataFrame: deduped = table.drop_duplicates().reset_index(drop=True) - removed = len(table) - len(deduped) - if removed: - print(f"Removed {removed} exact duplicate result rows.") return deduped @@ -1967,6 +2036,12 @@ def _add_e2e_subparser(subparsers): "--primer-bed", help="Optional primer BED file for amplicon clipping in Snakemake", ) + snk_group.add_argument( + "--ampliconclip-tolerance", + type=int, + default=1, + help="Tolerance for samtools ampliconclip primer matching (default: 1)", + ) snk_group.add_argument( "--snakemake-dryrun", action="store_true", @@ -1991,7 +2066,11 @@ def _add_e2e_subparser(subparsers): default=False, ) - _configure_vcf_parser(e2e_parser, include_input_csv=False) + _configure_vcf_parser( + e2e_parser, + include_input_csv=False, + include_consensus_options=True, + ) _move_action_group_after( e2e_parser, "Snakemake options", "Vartracker Analysis Options" ) @@ -2081,6 +2160,11 @@ def _run_e2e_command(args): outdir=snakemake_outdir, cores=args.cores, primer_bed=args.primer_bed, + ampliconclip_tolerance=args.ampliconclip_tolerance, + min_depth=args.min_depth, + consensus_snp_min_af=args.consensus_snp_min_af, + consensus_snp_thresh=args.consensus_snp_thresh, + consensus_indel_thresh=args.consensus_indel_thresh, dryrun=args.snakemake_dryrun, force_all=args.redo, quiet=not args.verbose, @@ -2223,6 +2307,10 @@ def _run_bam_command(args): outdir=snakemake_outdir, cores=args.cores, primer_bed=None, + min_depth=args.min_depth, + consensus_snp_min_af=args.consensus_snp_min_af, + consensus_snp_thresh=args.consensus_snp_thresh, + consensus_indel_thresh=args.consensus_indel_thresh, dryrun=args.snakemake_dryrun, force_all=args.redo, quiet=not args.verbose, @@ -2289,11 +2377,7 @@ def _run_plot_heatmap_command(args): if not results_csv.exists(): raise InputValidationError(f"Results CSV not found: {results_csv}") - outdir = ( - Path(args.outdir).expanduser().resolve() - if args.outdir - else results_csv.parent - ) + outdir = results_csv.parent outdir.mkdir(parents=True, exist_ok=True) table = pd.read_csv(results_csv, keep_default_na=False) @@ -2313,23 +2397,33 @@ def _run_plot_heatmap_command(args): raise InputValidationError( "Could not determine sample names from the results CSV" ) - - project_name = args.name - if project_name is None and "name" in table.columns: - names = [ - str(value).strip() - for value in table["name"].unique() - if str(value).strip() - ] - if len(names) == 1: - project_name = names[0] - if project_name is None: - project_name = "" + sample_numbers = [ + token.strip() + for token in str(table.iloc[0].get("sample_number", "")).split(" / ") + if token.strip() + ] + if args.x_labels == "sample-number": + if not sample_numbers: + raise InputValidationError( + "Results CSV must contain a 'sample_number' column to use " + "--x-labels sample-number" + ) + if len(sample_numbers) != len(sample_names): + raise InputValidationError( + "Results CSV has mismatched 'samples' and 'sample_number' columns" + ) + x_tick_labels = sample_numbers + else: + x_tick_labels = sample_names literature_df = None literature_path = None if args.literature_csv: literature_path = str(Path(args.literature_csv).expanduser().resolve()) + if not Path(literature_path).exists(): + raise InputValidationError( + f"Literature CSV not found: {literature_path}" + ) try: literature_df = pd.read_csv(literature_path) except Exception as exc: @@ -2341,14 +2435,16 @@ def _run_plot_heatmap_command(args): generate_variant_heatmap( table, sample_names, - sample_names, + sample_numbers or sample_names, str(outdir), - project_name, - args.min_snv_freq, - args.min_indel_freq, + "", + 0.03, + 0.1, + cli_command=cli_command, + x_tick_labels=x_tick_labels, + plot_title=args.title, literature_hits=literature_df, literature_table_path=literature_path, - cli_command=cli_command, **_collect_heatmap_kwargs(args), ) print(f"\nFinished: find results in {outdir}\n") @@ -2490,9 +2586,9 @@ def _process_files( table = pd.read_csv(precomputed_path, keep_default_na=False) else: # Format VCF files - csq_paths = [] + formatted_vcfs = [] for vcf, sample in zip(vcfs, sample_names): - _, csq_path = format_vcf( + formatted_vcf, _ = format_vcf( vcf, sample, tempdir, @@ -2503,10 +2599,10 @@ def _process_files( args.debug, args.allele_frequency_tag, ) - csq_paths.append(csq_path) + formatted_vcfs.append(formatted_vcf) # Prepare file lists for merging - new_vcfs = csq_paths + new_vcfs = formatted_vcfs # Write VCF list file with open(os.path.join(tempdir, "vcf_list.txt"), "w") as f: @@ -2519,10 +2615,18 @@ def _process_files( for sample_name in sample_names: f.write(f"{sample_name}\n") - # Merge VCF files + # Merge VCF files and annotate once across the merged longitudinal set. print("Annotating results...") + merged_vcf = os.path.join(tempdir, "vcf_merged.vcf") + merge_consequences(tempdir, merged_vcf, sample_names_file, args.debug) csq_file = os.path.join(tempdir, "vcf_annotated.vcf") - merge_consequences(tempdir, csq_file, sample_names_file, args.debug) + annotate_vcf( + merged_vcf, + csq_file, + args.reference, + args.gff3, + args.debug, + ) # Process VCF and extract variants print("Summarising results...") diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index f8a375a..afcd69c 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -164,7 +164,7 @@ def format_vcf( allele_frequency_tag="AF", ): """ - Format VCF file for compatibility and add consequences. + Format VCF file for compatibility and filter variants before merging. Args: vcf (str): Path to input VCF file @@ -178,6 +178,7 @@ def format_vcf( allele_frequency_tag (str): INFO tag name for allele frequency (default: AF) """ out, csq_file, log = _derive_vcf_output_paths(vcf, tempdir, sample) + raw_out = os.path.join(tempdir, f"{Path(out).stem}.raw.vcf.gz") try: prepared_vcf = _ensure_format_and_sample(vcf, sample, tempdir, debug) @@ -251,7 +252,7 @@ def format_vcf( else: header_lines.append(normalized_line) header_str = "\n".join(header_lines) + "\n" - w = Writer.from_string(out, header_str, mode="wz") + w = Writer.from_string(raw_out, header_str, mode="wz") variants = {} def _scalar(value): @@ -313,21 +314,28 @@ def _scalar(value): w.write_record(variant) w.close() - # Index the output file - cmd = f"bcftools index -f {out}" + # Index the intermediate output file + cmd = f"bcftools index -f {raw_out}" subprocess.run(cmd, shell=True, check=True) - # Run csq before merging for better AF filtering + # Filter variants before the cross-sample merge. Consequence annotation is + # run only after merging so BCSQ reflects the sample-specific haplotype at + # each longitudinal timepoint. cmd = ( f'bcftools view -i \'(INFO/AF >= {min_snv_freq} & INFO/TYPE == "SNP") | ' - f'(INFO/AF >= {min_indel_freq} & INFO/TYPE == "INDEL")\' {out} | ' - f"bcftools csq -f {reference} -g {annotation} --force -Oz -o {csq_file}; " - f"tabix -f -p vcf {csq_file}" + f'(INFO/AF >= {min_indel_freq} & INFO/TYPE == "INDEL")\' {raw_out} ' + f"-Oz -o {out}" ) try: with open(log, "a", encoding="utf-8") as err: subprocess.run(cmd, shell=True, stderr=err, check=True) + subprocess.run( + f"bcftools index -f {out}", + shell=True, + stderr=err, + check=True, + ) except subprocess.CalledProcessError as exc: log_tail = "" if os.path.exists(log): @@ -355,7 +363,7 @@ def _scalar(value): def merge_consequences(tempdir, csq_file, sample_names, debug): """ - Merge VCF files with consequences. + Merge per-sample VCF files before consequence annotation. Args: tempdir (str): Temporary directory path @@ -394,6 +402,22 @@ def merge_consequences(tempdir, csq_file, sample_names, debug): raise RuntimeError(f"Error merging VCF files: {str(e)}") +def annotate_vcf(vcf_file, output_file, reference, annotation, debug): + """Annotate a merged multi-sample VCF with bcftools csq.""" + cmd = ( + f"bcftools csq -f {reference} -g {annotation} --force " + f"-Ov -o {output_file} {vcf_file}" + ) + + if debug: + print(f"Command: {cmd}") + + try: + subprocess.run(cmd, shell=True, check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Error annotating merged VCF with bcftools csq: {str(e)}") + + def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): """ Calculate depth metrics for variant sites. @@ -484,6 +508,95 @@ def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): return result +def _summarise_sample_trajectory(allele_freqs, samples): + """Derive trajectory metadata from per-sample allele frequencies.""" + presence_absence = ["N" if x == "." else "Y" for x in allele_freqs] + variant_status = "new" if allele_freqs[0] == "." else "original" + + if allele_freqs[0] != "." and allele_freqs[-1] == ".": + persistent_status = "original_lost" + elif allele_freqs[0] != "." and allele_freqs[-1] != ".": + persistent_status = "original_retained" + elif allele_freqs[0] == "." and allele_freqs[-1] != ".": + persistent_status = "new_persistent" + elif allele_freqs[0] == "." and allele_freqs[-1] == ".": + persistent_status = "new_transient" + else: + persistent_status = "unknown" + + first_appearance = ( + samples[presence_absence.index("Y")] if "Y" in presence_absence else "None" + ) + last_appearance = ( + samples[rindex(presence_absence, "Y")] if "Y" in presence_absence else "None" + ) + + return { + "presence_absence": presence_absence, + "variant_status": variant_status, + "persistent_status": persistent_status, + "first_appearance": first_appearance, + "last_appearance": last_appearance, + } + + +def _extract_scalar_mask(value) -> int: + if isinstance(value, np.ndarray): + if value.size == 0: + return 0 + value = value.flat[0] + elif isinstance(value, (list, tuple)): + if not value: + return 0 + value = value[0] + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + +def _decode_sample_bcsq_annotations(v, samples, annotations): + """Decode FORMAT/BCSQ bitmasks into sample-specific annotation assignments.""" + if not annotations: + return {} + + try: + mask_values = v.format("BCSQ") + except (KeyError, AttributeError, ValueError, RuntimeError): + return {} + + if mask_values is None: + return {} + + decoded = {} + for sample, sample_mask in zip(samples, mask_values): + bitmask = _extract_scalar_mask(sample_mask) + matches = [] + for idx, annotation in enumerate(annotations): + first_haplotype_bit = 1 << (2 * idx) + second_haplotype_bit = 1 << (2 * idx + 1) + if bitmask & first_haplotype_bit or bitmask & second_haplotype_bit: + matches.append(annotation) + decoded[sample] = matches + + return decoded + + +def _mask_allele_frequencies_for_annotation( + annotation, allele_freqs, samples, sample_bcsq_map +): + """Keep allele frequencies only for samples where the annotation applies.""" + masked = [] + for sample, allele_freq in zip(samples, allele_freqs): + if allele_freq == ".": + masked.append(".") + continue + masked.append( + allele_freq if annotation in sample_bcsq_map.get(sample, []) else "." + ) + return masked + + def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): """ Process VCF file and extract variant information. @@ -565,20 +678,7 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): allele_freqs.append(".") if not allele_freqs: allele_freqs = ["."] * max(1, len(samples)) - presence_absence = ["N" if x == "." else "Y" for x in allele_freqs] - variant_status = "new" if allele_freqs[0] == "." else "original" - - # Calculate persistence status - if allele_freqs[0] != "." and allele_freqs[-1] == ".": - persistent_status = "original_lost" - elif allele_freqs[0] != "." and allele_freqs[-1] != ".": - persistent_status = "original_retained" - elif allele_freqs[0] == "." and allele_freqs[-1] != ".": - persistent_status = "new_persistent" - elif allele_freqs[0] == "." and allele_freqs[-1] == ".": - persistent_status = "new_transient" - else: - persistent_status = "unknown" + trajectory = _summarise_sample_trajectory(allele_freqs, samples) depths_qc = calculate_variant_site_depths(cov_df, v, samples, min_depth) all_samples_pass_qc = "F" not in depths_qc["variant_qc"] @@ -588,45 +688,71 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): if depths_qc["variant_qc"] else 0.0 ) - first_appearance = ( - samples[presence_absence.index("Y")] if "Y" in presence_absence else "None" - ) - last_appearance = ( - samples[rindex(presence_absence, "Y")] - if "Y" in presence_absence - else "None" - ) # Process annotations if "BCSQ" in info: annotations = v.INFO["BCSQ"].split(",") - for annot in annotations: - anno = annot.split("|") - result = _process_annotation( - v, - anno, - variant_status, - persistent_status, - presence_absence, - first_appearance, - last_appearance, - all_samples_pass_qc, - proportion_samples_passing_qc, - depths_qc, - allele_freqs, - samples, - total_cov_list, - ) - results.append(result) + sample_bcsq_map = _decode_sample_bcsq_annotations(v, samples, annotations) + produced_annotation_specific_row = False + + if sample_bcsq_map: + for annot in annotations: + masked_allele_freqs = _mask_allele_frequencies_for_annotation( + annot, allele_freqs, samples, sample_bcsq_map + ) + masked_trajectory = _summarise_sample_trajectory( + masked_allele_freqs, samples + ) + if "Y" not in masked_trajectory["presence_absence"]: + continue + + anno = annot.split("|") + result = _process_annotation( + v, + anno, + masked_trajectory["variant_status"], + masked_trajectory["persistent_status"], + masked_trajectory["presence_absence"], + masked_trajectory["first_appearance"], + masked_trajectory["last_appearance"], + all_samples_pass_qc, + proportion_samples_passing_qc, + depths_qc, + masked_allele_freqs, + samples, + total_cov_list, + ) + results.append(result) + produced_annotation_specific_row = True + + if not produced_annotation_specific_row: + for annot in annotations: + anno = annot.split("|") + result = _process_annotation( + v, + anno, + trajectory["variant_status"], + trajectory["persistent_status"], + trajectory["presence_absence"], + trajectory["first_appearance"], + trajectory["last_appearance"], + all_samples_pass_qc, + proportion_samples_passing_qc, + depths_qc, + allele_freqs, + samples, + total_cov_list, + ) + results.append(result) else: # Handle variants without annotations result = _create_unannotated_result( v, - variant_status, - persistent_status, - presence_absence, - first_appearance, - last_appearance, + trajectory["variant_status"], + trajectory["persistent_status"], + trajectory["presence_absence"], + trajectory["first_appearance"], + trajectory["last_appearance"], all_samples_pass_qc, proportion_samples_passing_qc, depths_qc, From a8a458564363b044a249e5a38d34f9bc44c6da3f Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Tue, 5 May 2026 10:40:41 +1000 Subject: [PATCH 09/20] Fixed treatment of low frequency multiallelics --- CITATION.cff | 12 +- README.md | 54 ++- pyproject.toml | 2 +- tests/test_analysis.py | 7 - tests/test_analysis_launcher.py | 31 +- tests/test_lofreq_primer_rescue.py | 66 ++++ tests/test_main.py | 35 ++ tests/test_snakemake_workflow.py | 6 + tests/test_vcf_processing.py | 251 +++++++++++- vartracker/Snakefile | 82 +++- vartracker/_version.py | 2 +- vartracker/analysis.py | 25 +- vartracker/analysis_launcher.py | 53 +++ vartracker/lofreq_primer_rescue.py | 470 +++++++++++++++++++++++ vartracker/main.py | 86 ++++- vartracker/vcf_processing.py | 588 ++++++++++++++++++++++++----- 16 files changed, 1628 insertions(+), 142 deletions(-) create mode 100644 tests/test_lofreq_primer_rescue.py create mode 100644 vartracker/lofreq_primer_rescue.py diff --git a/CITATION.cff b/CITATION.cff index 54d1801..b303859 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -1,12 +1,12 @@ cff-version: 1.2.0 message: "If you use this software, please cite it as below." title: "vartracker" -version: "2.1.0" -date-released: "2026-02-11" +version: "2.2.0" +date-released: "2026-05-05" license: MIT repository-code: "https://github.com/charlesfoster/vartracker" url: "https://github.com/charlesfoster/vartracker" -doi: "10.5281/zenodo.XXXXXXX" +doi: "10.5281/zenodo.18452274" keywords: - bioinformatics - genomics @@ -22,7 +22,7 @@ preferred-citation: - family-names: Foster given-names: Charles title: "vartracker" - version: "2.1.0" - doi: "10.5281/zenodo.XXXXXXX" + version: "2.2.0" + doi: "10.5281/zenodo.18452274" url: "https://github.com/charlesfoster/vartracker" - date-released: "2026-02-11" + date-released: "2026-05-05" diff --git a/README.md b/README.md index 8ba282d..b92c28b 100755 --- a/README.md +++ b/README.md @@ -248,6 +248,24 @@ Temporary LoFreq note: - This is a temporary workaround for an older Bioconda LoFreq build that can fail during `call-parallel` final filtering when many shards produce an excessively long merged VCF header. - The cap will be revisited once an updated LoFreq build is available through Bioconda. +LoFreq primer-overlap rescue: +- Amplicon schemes can create a specific LoFreq false-negative mode: after primer clipping, reads from one strand may be soft clipped at primer-overlap sites, so a genuine near-fixed variant can fail LoFreq's default strand-bias filter. +- `bam` and `end-to-end` therefore run LoFreq with `--no-default-filter`, then apply the normal `lofreq filter` step so standard LoFreq PASS calls are unchanged. +- With the default `--lofreq-primer-rescue auto`, the rescue step runs only when `--primer-bed` is supplied. In other words, `auto` means "use primer rescue when an amplicon primer scheme has been explicitly provided." +- In `end-to-end` mode, the same `--primer-bed` is used for `samtools ampliconclip` and for rescue. In `bam` mode, vartracker does not clip the input BAMs; the primer BED is used only to identify primer-overlap sites for rescue. +- Rescue candidates must be single-ALT SNPs that overlap a primer interval, fail the default LoFreq filter, and pass conservative near-fixed thresholds (`AF>=0.95`, `DP>=100`, `DP4 alt count>=95`, `QUAL>=100`, `DP4 ref count<=20`) with one-sided alternate-strand support. Indels, multi-ALT records, lower-frequency variants, and non-primer-overlap variants are not rescued by this rule. +- Rescued variants are marked with `FILTER=RESCUED_PRIMER_OVERLAP`, `INFO/PRIMER_OVERLAP`, and `INFO/RESCUED_BY=overlap_primer_interval`; per-sample details are written to `_variants.rescued.tsv` and listed in the updated spreadsheet as `lofreq_rescued_tsv`. +- Use `--lofreq-primer-rescue off` to disable rescue even when a primer BED is supplied, or `--lofreq-primer-rescue on` to require rescue and fail if `--primer-bed` is missing. The rescue thresholds can be adjusted with the `--lofreq-rescue-*` options. + +Example amplicon run with primer rescue: + +```bash +vartracker end-to-end inputs.csv \ + --primer-bed primers.bed \ + --ampliconclip-tolerance 1 \ + --outdir results/e2e_amplicon +``` + ### Input Spreadsheets Every CLI mode reads the same canonical columns: @@ -266,8 +284,10 @@ Mode-specific expectations: - **End-to-end mode** requires `reads1` (and optionally `reads2`); remaining fields are generated. The `bam` and `end-to-end` workflows also write two consensus FASTA columns to -the updated Snakemake spreadsheet: `consensus` for a simple consensus and -`iupac_consensus` for an IUPAC-aware consensus. SNPs below +the updated Snakemake spreadsheet, plus the LoFreq rescue audit column: +`consensus` for a simple consensus, `iupac_consensus` for an IUPAC-aware +consensus, and `lofreq_rescued_tsv` for the per-sample primer-overlap rescue +table. SNPs below `--consensus-snp-min-af` are ignored, SNPs from `--consensus-snp-min-af` up to `--consensus-snp-thresh` stay as reference bases in the simple consensus and become REF+ALT ambiguity codes in the IUPAC consensus, and SNPs at or above @@ -286,20 +306,27 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test ### Mode-specific options - `vartracker vcf` – accepts core analysis options such as `--min-snv-freq`, `--min-indel-freq`, - `--allele-frequency-tag`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls + `--allele-frequency-tag`, `--multiallelic-overflow`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. - `vartracker bam` – everything from `vcf`, plus Snakemake options: `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, - `--rulegraph`, `--consensus-snp-min-af`, `--consensus-snp-thresh`, and - `--consensus-indel-thresh`. + `--rulegraph`, `--primer-bed`, `--lofreq-primer-rescue`, `--consensus-snp-min-af`, + `--consensus-snp-thresh`, and `--consensus-indel-thresh`. - `vartracker end-to-end` – similar to `bam`, with optional amplicon clipping controls: - `--primer-bed` and `--ampliconclip-tolerance` (default: `1`). + `--primer-bed` and `--ampliconclip-tolerance` (default: `1`). Supplying + `--primer-bed` also enables LoFreq primer-overlap rescue by default. - `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV, including all heatmap customization filters. - `vartracker plot genome` – plot SNP positions along the genome or a selected gene region using all observed allele-frequency values for each variant. - `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. - `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. - `vartracker plot lifespan` – plot first-to-last detection spans for a selected or auto-ranked subset of variants. +Consequence-calling note: +- Vartracker keeps distinct ALT alleles at the same position separate during preprocessing, then rejoins them immediately before `bcftools csq` so codon-level consequences can still be inferred correctly. +- If more than two ALT alleles remain present in a single sample at one genomic position after frequency filtering, vartracker defaults to stopping with an informative error before `bcftools csq`. This is the safest behavior and the default `--multiallelic-overflow error` mode. +- `--multiallelic-overflow drop-lowest-af` continues by removing the lowest-frequency retained ALT allele(s) for the affected sample before `bcftools csq`, and prints a warning describing the site and the dropped allele(s). +- `--multiallelic-overflow skip-site` continues by skipping consequence calling for the affected site entirely, leaving those variants in the results as unannotated rows and printing a warning describing the site. + Heatmap filtering: - `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customize heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. - By default, all consequence classes are included except joint variants. Use `--include-joint` to show joint variants. @@ -482,6 +509,7 @@ vartracker produces several output files: - **results.csv**: Comprehensive variant analysis with all metrics - **results_metadata.json**: Output schema version and results metadata +- **`_variants.rescued.tsv`** (`bam`/`end-to-end`): LoFreq primer-overlap rescue audit table, empty when rescue is disabled or no variants are rescued - **new_mutations.csv**: Mutations not present in the first sample - **persistent_new_mutations.csv**: New mutations that persist to the final sample - **cumulative_mutations.pdf**: Plot showing mutation accumulation over time @@ -519,9 +547,9 @@ vartracker schema literature The pipeline performs the following analysis: -1. **VCF Standardization**: Normalizes and standardizes input VCF files -2. **Annotation**: Adds amino acid consequences using `bcftools csq` -3. **Variant Merging**: Combines all longitudinal samples +1. **VCF Standardization**: Normalizes and standardizes input VCF files, preserving distinct ALT alleles at the same genomic position +2. **Variant Merging**: Combines all longitudinal samples +3. **Annotation**: Adds amino acid consequences using `bcftools csq` on the merged VCF so sample-specific joint consequences are inferred from each sample's surviving ALT combination 4. **Comprehensive Analysis**: For each variant, determines: - Gene location and amino acid consequences - Variant type (SNP/indel) and change type (synonymous/missense/etc.) @@ -536,13 +564,11 @@ The pipeline performs the following analysis: ## Citation When using vartracker, please cite the software release you used. Citation metadata is provided -in `CITATION.cff`, and GitHub releases are archived on Zenodo (DOI will appear here once minted). +in `CITATION.cff`, and GitHub releases are archived on Zenodo. -If you use vartracker, please cite the software record on Zenodo: +- Foster, C. (2026). *vartracker* (Version 2.2.0). Zenodo. https://doi.org/10.5281/zenodo.18452274 -- Foster, C. (2026). *vartracker* (Version x.y.z). Zenodo. https://doi.org/10.5281/zenodo.XXXXX -- Concept DOI (all versions): https://doi.org/10.5281/zenodo.18452274 -Note: a version-specific DOI is minted by Zenodo after each GitHub release. +Note: the DOI above is the Zenodo concept DOI for all versions; a version-specific DOI is minted by Zenodo after each GitHub release. Also cite relevant methods or data sources, for example: diff --git a/pyproject.toml b/pyproject.toml index eb8aa07..49a348b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "vartracker" -version = "2.1.1" +version = "2.2.0" authors = [ {name = "Dr Charles Foster"}, ] diff --git a/tests/test_analysis.py b/tests/test_analysis.py index b24193b..d7b886e 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -420,13 +420,6 @@ def test_process_joint_variants_matches_main_row_by_presence_pattern(tmp_path): assert result.loc[2, "type_of_change"] == "joint_missense" -@pytest.mark.xfail( - strict=True, - reason=( - "process_joint_variants currently resolves tied main-row candidates by row " - "order, which is unsafe for overlapping genes" - ), -) def test_process_joint_variants_is_order_invariant_for_overlapping_gene_rows(tmp_path): shared_rows = [ { diff --git a/tests/test_analysis_launcher.py b/tests/test_analysis_launcher.py index 004c12d..977f832 100644 --- a/tests/test_analysis_launcher.py +++ b/tests/test_analysis_launcher.py @@ -1,6 +1,9 @@ import pytest -from vartracker.analysis_launcher import _validate_primer_bed_reference +from vartracker.analysis_launcher import ( + _validate_lofreq_primer_rescue, + _validate_primer_bed_reference, +) def test_validate_primer_bed_reference_accepts_matching_contig(tmp_path): @@ -35,3 +38,29 @@ def test_validate_primer_bed_reference_rejects_empty_bed(tmp_path): with pytest.raises(ValueError, match="contains no intervals"): _validate_primer_bed_reference(primer_bed, reference) + + +def test_validate_lofreq_primer_rescue_on_requires_primer_bed(): + with pytest.raises(ValueError, match="requires --primer-bed"): + _validate_lofreq_primer_rescue( + "on", + None, + min_af=0.95, + min_dp=100, + min_alt_count=95, + min_qual=100, + max_ref_count=20, + ) + + +def test_validate_lofreq_primer_rescue_rejects_invalid_threshold(): + with pytest.raises(ValueError, match="between 0 and 1"): + _validate_lofreq_primer_rescue( + "auto", + None, + min_af=1.5, + min_dp=100, + min_alt_count=95, + min_qual=100, + max_ref_count=20, + ) diff --git a/tests/test_lofreq_primer_rescue.py b/tests/test_lofreq_primer_rescue.py new file mode 100644 index 0000000..6822785 --- /dev/null +++ b/tests/test_lofreq_primer_rescue.py @@ -0,0 +1,66 @@ +import subprocess +from pathlib import Path + +from vartracker.lofreq_primer_rescue import rescue_lofreq_primer_variants + + +def test_rescue_lofreq_primer_variants_keeps_pass_and_rescues_overlap( + monkeypatch, tmp_path +): + raw_vcf = tmp_path / "raw.vcf" + raw_vcf.write_text( + "##fileformat=VCFv4.2\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" + "chr1\t10\t.\tA\tG\t200\t.\tAF=0.98;DP=120;DP4=1,0,118,0\n" + "chr1\t30\t.\tC\tT\t150\t.\tAF=0.40;DP=120;DP4=30,30,30,30\n", + encoding="utf-8", + ) + primers = tmp_path / "primers.bed" + primers.write_text("chr1\t0\t20\tprimer_1\n", encoding="utf-8") + default_filtered_vcf = ( + "##fileformat=VCFv4.2\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" + "chr1\t10\t.\tA\tG\t200\tstrandbias\tAF=0.98;DP=120;DP4=1,0,118,0\n" + "chr1\t30\t.\tC\tT\t150\tPASS\tAF=0.40;DP=120;DP4=30,30,30,30\n" + ) + + def fake_run(cmd, check): + assert cmd[:4] == ["lofreq", "filter", "-i", str(raw_vcf)] + assert check is True + output_path = Path(cmd[cmd.index("-o") + 1]) + output_path.write_text(default_filtered_vcf, encoding="utf-8") + return subprocess.CompletedProcess(cmd, 0) + + monkeypatch.setattr("vartracker.lofreq_primer_rescue.subprocess.run", fake_run) + + output_vcf = tmp_path / "final.vcf" + rescued_tsv = tmp_path / "rescued.tsv" + + result = rescue_lofreq_primer_variants( + raw_vcf=raw_vcf, + primers_bed=primers, + output_vcf=output_vcf, + rescued_tsv=rescued_tsv, + ) + + assert result.normal_passed == 1 + assert result.rescued == 1 + assert result.discarded == 0 + + output_text = output_vcf.read_text(encoding="utf-8") + assert "##FILTER= tuple[Path, Path]: return ref, gff +def _prepare_merged_multiallelic_vcf( + tmp_path: Path, record_line: str, sample_name: str = "s1" +) -> tuple[Path, Path, Path]: + ref, gff = _write_minimal_reference_bundle(tmp_path) + sample_vcf = tmp_path / f"{sample_name}.vcf" + sample_vcf.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + f"#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t{sample_name}\n" + f"{record_line}\n", + encoding="utf-8", + ) + + formatted_vcf, _ = format_vcf( + str(sample_vcf), + sample_name, + str(tmp_path), + 0.0, + 0.0, + str(ref), + str(gff), + False, + ) + + vcf_list = tmp_path / "vcf_list.txt" + vcf_list.write_text(f"{formatted_vcf}\n", encoding="utf-8") + sample_names = tmp_path / "sample_names.txt" + sample_names.write_text(f"{sample_name}\n", encoding="utf-8") + + merged_vcf = tmp_path / "merged.vcf" + merge_consequences(str(tmp_path), str(merged_vcf), str(sample_names), debug=False) + return ref, gff, merged_vcf + + @pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") -@pytest.mark.xfail( - strict=True, - reason="format_vcf currently deduplicates records by POS and can drop a real ALT allele", -) def test_format_vcf_preserves_distinct_alt_records_at_same_position(tmp_path): ref, gff = _write_minimal_reference_bundle(tmp_path) vcf_path = tmp_path / "same_pos_two_alts.vcf" @@ -100,6 +135,214 @@ def test_format_vcf_preserves_distinct_alt_records_at_same_position(tmp_path): assert {record.ALT[0] for record in records} == {"C", "G"} +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +def test_multiallelic_input_survives_merge_and_csq_annotation(tmp_path): + ref, gff = _write_minimal_reference_bundle(tmp_path) + sample_vcf = tmp_path / "multiallelic.vcf" + sample_vcf.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\n" + "chr1\t5\t.\tA\tC,G\t.\tPASS\tAF=0.20,0.30;DP=100\tGT:DP:AF\t1/2:100:0.20,0.30\n", + encoding="utf-8", + ) + + formatted_vcf, _ = format_vcf( + str(sample_vcf), + "s1", + str(tmp_path), + 0.0, + 0.0, + str(ref), + str(gff), + False, + ) + + vcf_list = tmp_path / "vcf_list.txt" + vcf_list.write_text(f"{formatted_vcf}\n", encoding="utf-8") + sample_names = tmp_path / "sample_names.txt" + sample_names.write_text("s1\n", encoding="utf-8") + + merged_vcf = tmp_path / "merged.vcf" + merge_consequences(str(tmp_path), str(merged_vcf), str(sample_names), debug=False) + + annotated_vcf = tmp_path / "annotated.vcf" + annotate_vcf( + str(merged_vcf), + str(annotated_vcf), + str(ref), + str(gff), + debug=False, + ) + + depth = tmp_path / "s1.depth.txt" + _write_depth_file(depth) + + table = process_vcf(str(annotated_vcf), [str(depth)], 10, ["s1"]) + + observed = ( + table[["variant", "amino_acid_consequence", "bcsq_aa_notation", "alt_freq"]] + .sort_values("variant") + .reset_index(drop=True) + ) + + expected = pd.DataFrame( + [ + { + "variant": "A5C", + "amino_acid_consequence": "K2T", + "bcsq_aa_notation": "2K>2T", + "alt_freq": "0.200", + }, + { + "variant": "A5G", + "amino_acid_consequence": "K2R", + "bcsq_aa_notation": "2K>2R", + "alt_freq": "0.300", + }, + ] + ) + + pd.testing.assert_frame_equal(observed, expected) + + +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +def test_annotate_vcf_reports_informative_error_for_multiallelic_overflow(tmp_path): + ref, gff, merged_vcf = _prepare_merged_multiallelic_vcf( + tmp_path, + "chr1\t5\t.\tA\tG,T,C\t.\tPASS\tAF=0.04,0.34,0.12;DP=100\tGT:DP:AF\t1/2:100:0.04,0.34,0.12", + ) + + with pytest.raises(RuntimeError) as excinfo: + annotate_vcf( + str(merged_vcf), + str(tmp_path / "annotated.vcf"), + str(ref), + str(gff), + debug=False, + multiallelic_overflow="error", + ) + + message = str(excinfo.value) + assert "chr1:5" in message + assert "Sample s1:" in message + assert "A>G AF=0.040" in message + assert "A>T AF=0.340" in message + assert "A>C AF=0.120" in message + assert "--min-snv-freq above 0.040" in message + + +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +def test_annotate_vcf_drop_lowest_af_continues_with_warning(tmp_path, capsys): + ref, gff, merged_vcf = _prepare_merged_multiallelic_vcf( + tmp_path, + "chr1\t5\t.\tA\tG,T,C\t.\tPASS\tAF=0.04,0.34,0.12;DP=100\tGT:DP:AF\t1/2:100:0.04,0.34,0.12", + ) + + annotated_vcf = tmp_path / "annotated.vcf" + annotate_vcf( + str(merged_vcf), + str(annotated_vcf), + str(ref), + str(gff), + debug=False, + multiallelic_overflow="drop-lowest-af", + ) + + stdout = capsys.readouterr().out + assert "Warning:" in stdout + assert "Dropping the lowest-frequency ALT allele(s) for this sample" in stdout + assert "A>G AF=0.040" in stdout + assert "--min-snv-freq above 0.040" in stdout + + depth = tmp_path / "s1.depth.txt" + _write_depth_file(depth) + table = process_vcf(str(annotated_vcf), [str(depth)], 10, ["s1"]) + + observed = ( + table[["variant", "amino_acid_consequence", "alt_freq"]] + .sort_values("variant") + .reset_index(drop=True) + ) + + expected = pd.DataFrame( + [ + { + "variant": "A5C", + "amino_acid_consequence": "K2T", + "alt_freq": "0.120", + }, + { + "variant": "A5T", + "amino_acid_consequence": "K2I", + "alt_freq": "0.340", + }, + ] + ) + + pd.testing.assert_frame_equal(observed, expected) + + +@pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") +def test_annotate_vcf_skip_site_keeps_unannotated_rows(tmp_path, capsys): + ref, gff, merged_vcf = _prepare_merged_multiallelic_vcf( + tmp_path, + "chr1\t5\t.\tA\tG,T,C\t.\tPASS\tAF=0.04,0.34,0.12;DP=100\tGT:DP:AF\t1/2:100:0.04,0.34,0.12", + ) + + annotated_vcf = tmp_path / "annotated.vcf" + annotate_vcf( + str(merged_vcf), + str(annotated_vcf), + str(ref), + str(gff), + debug=False, + multiallelic_overflow="skip-site", + ) + + stdout = capsys.readouterr().out + assert "Warning:" in stdout + assert "Skipping consequence calling for this site" in stdout + assert "--min-snv-freq above 0.040" in stdout + + depth = tmp_path / "s1.depth.txt" + _write_depth_file(depth) + table = process_vcf(str(annotated_vcf), [str(depth)], 10, ["s1"]) + + observed = ( + table[["variant", "amino_acid_consequence", "alt_freq"]] + .sort_values("variant") + .reset_index(drop=True) + ) + + expected = pd.DataFrame( + [ + { + "variant": "A5C", + "amino_acid_consequence": "None", + "alt_freq": "0.120", + }, + { + "variant": "A5G", + "amino_acid_consequence": "None", + "alt_freq": "0.040", + }, + { + "variant": "A5T", + "amino_acid_consequence": "None", + "alt_freq": "0.340", + }, + ] + ) + + pd.testing.assert_frame_equal(observed, expected) + + def test_process_vcf_splits_sample_specific_bcsq_annotations(tmp_path): vcf_path = tmp_path / "annotated.vcf" vcf_path.write_text( diff --git a/vartracker/Snakefile b/vartracker/Snakefile index b1bc34b..1da3c2e 100644 --- a/vartracker/Snakefile +++ b/vartracker/Snakefile @@ -2,8 +2,16 @@ import pandas as pd from pathlib import Path try: from vartracker.consensus import consensus_genotype_for_variant + from vartracker.lofreq_primer_rescue import ( + PrimerRescueThresholds, + rescue_lofreq_primer_variants, + ) except ModuleNotFoundError: from consensus import consensus_genotype_for_variant + from lofreq_primer_rescue import ( + PrimerRescueThresholds, + rescue_lofreq_primer_variants, + ) # Load samples from CSV samples_df = pd.read_csv(config["samples_csv"]) @@ -19,6 +27,22 @@ MIN_DEPTH = int(config.get("min_depth", 10)) CONSENSUS_SNP_MIN_AF = float(config.get("consensus_snp_min_af", 0.25)) CONSENSUS_SNP_THRESH = float(config.get("consensus_snp_thresh", 0.75)) CONSENSUS_INDEL_THRESH = float(config.get("consensus_indel_thresh", 0.75)) +LOFREQ_PRIMER_RESCUE = str(config.get("lofreq_primer_rescue", "auto")).lower() +LOFREQ_RESCUE_MIN_AF = float(config.get("lofreq_rescue_min_af", 0.95)) +LOFREQ_RESCUE_MIN_DP = int(config.get("lofreq_rescue_min_dp", 100)) +LOFREQ_RESCUE_MIN_ALT_COUNT = int(config.get("lofreq_rescue_min_alt_count", 95)) +LOFREQ_RESCUE_MIN_QUAL = float(config.get("lofreq_rescue_min_qual", 100.0)) +LOFREQ_RESCUE_MAX_REF_COUNT = int(config.get("lofreq_rescue_max_ref_count", 20)) + +if LOFREQ_PRIMER_RESCUE not in {"auto", "on", "off"}: + raise ValueError("lofreq_primer_rescue must be one of: auto, on, off") +if LOFREQ_PRIMER_RESCUE == "on" and not PRIMER_BED: + raise ValueError("lofreq_primer_rescue='on' requires primer_bed") + +LOFREQ_PRIMER_RESCUE_ENABLED = ( + LOFREQ_PRIMER_RESCUE == "on" + or (LOFREQ_PRIMER_RESCUE == "auto" and bool(PRIMER_BED)) +) def _value_or_blank(sample, key, allow_blank=False): value = SAMPLES[sample].get(key) @@ -255,21 +279,58 @@ rule lofreq_call: bam = f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam", ref = REF output: - vcf_raw = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf"), + vcf_raw = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz"), + vcf_filtered = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.filtered.vcf"), + rescued_tsv = f"{OUTDIR}/{{sample}}/{{sample}}_variants.rescued.tsv", vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", csi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi" log: f"{OUTDIR}/{{sample}}/logs/lofreq_call.log" threads: min(8, max(1, workflow.cores)) - shell: - """ - mkdir -p $(dirname {log}) - lofreq call-parallel --no-baq --call-indels --pp-threads {threads} \ - -f {input.ref} -o {output.vcf_raw} {input.bam} 2> {log} + run: + log_path = Path(str(log[0])) + log_path.parent.mkdir(parents=True, exist_ok=True) - bgzip -c {output.vcf_raw} > {output.vcf} 2>> {log} - bcftools index {output.vcf} 2>> {log} - """ + shell(""" + lofreq call-parallel --no-baq --call-indels --no-default-filter \ + --pp-threads {threads} -f {input.ref} -o {output.vcf_raw} \ + {input.bam} 2> {log} + """) + + if LOFREQ_PRIMER_RESCUE_ENABLED: + thresholds = PrimerRescueThresholds( + min_af=LOFREQ_RESCUE_MIN_AF, + min_dp=LOFREQ_RESCUE_MIN_DP, + min_alt_count=LOFREQ_RESCUE_MIN_ALT_COUNT, + min_qual=LOFREQ_RESCUE_MIN_QUAL, + max_ref_count=LOFREQ_RESCUE_MAX_REF_COUNT, + ) + result = rescue_lofreq_primer_variants( + raw_vcf=str(output.vcf_raw), + primers_bed=PRIMER_BED, + output_vcf=str(output.vcf_filtered), + rescued_tsv=str(output.rescued_tsv), + tmp_dir=str(log_path.parent), + thresholds=thresholds, + ) + with log_path.open("a", encoding="utf-8") as handle: + handle.write( + "lofreq primer rescue: " + f"normal_passed={result.normal_passed} " + f"rescued={result.rescued} discarded={result.discarded}\n" + ) + else: + shell( + "lofreq filter -i {output.vcf_raw} " + "-o {output.vcf_filtered} 2>> {log}" + ) + with open(output.rescued_tsv, "w", encoding="utf-8") as handle: + handle.write("variant\treason_filtered\treason_rescued\tmetrics\n") + with log_path.open("a", encoding="utf-8") as handle: + handle.write("lofreq primer rescue: disabled\n") + + shell("bgzip -c {output.vcf_filtered} > {output.vcf} 2>> {log}") + shell("bcftools index -f {output.vcf} 2>> {log}") rule deletion_variants_bed: input: @@ -539,6 +600,9 @@ rule update_csv: df['coverage'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_depth.txt") ) + df['lofreq_rescued_tsv'] = df['sample_name'].apply( + lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_variants.rescued.tsv") + ) df['consensus'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_consensus.fasta") ) diff --git a/vartracker/_version.py b/vartracker/_version.py index fa9afad..4e0c032 100644 --- a/vartracker/_version.py +++ b/vartracker/_version.py @@ -9,7 +9,7 @@ # NOTE: When bumping the project version remember to update this fallback value # alongside the version declared in pyproject.toml. -_FALLBACK_VERSION = "2.1.1" +_FALLBACK_VERSION = "2.2.0" try: __version__ = metadata.version("vartracker") diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 0ecb115..9a395ad 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -112,8 +112,10 @@ def _find_joint_main_index(tab: pd.DataFrame, joint_index: int, main_pos: int) - joint_presence = _presence_vector(tab.at[joint_index, "presence_absence"]) - def candidate_key(candidate_index: int) -> tuple[int, int, int]: - candidate_presence = _presence_vector(tab.at[candidate_index, "presence_absence"]) + def candidate_score(candidate_index: int) -> tuple[int, int, int]: + candidate_presence = _presence_vector( + tab.at[candidate_index, "presence_absence"] + ) exact_match = int(candidate_presence == joint_presence) shared_present = sum( lhs == rhs == "Y" for lhs, rhs in zip(joint_presence, candidate_presence) @@ -121,7 +123,24 @@ def candidate_key(candidate_index: int) -> tuple[int, int, int]: total_present = sum(token == "Y" for token in candidate_presence) return (exact_match, shared_present, total_present) - return max(candidates, key=candidate_key) + best_score = max(candidate_score(candidate_index) for candidate_index in candidates) + tied_candidates = [ + candidate_index + for candidate_index in candidates + if candidate_score(candidate_index) == best_score + ] + if len(tied_candidates) == 1: + return tied_candidates[0] + + def tie_break_key(candidate_index: int) -> tuple[str, str, str, str]: + return ( + str(tab.at[candidate_index, "gene"]), + str(tab.at[candidate_index, "amino_acid_consequence"]), + str(tab.at[candidate_index, "bcsq_nt_notation"]), + str(tab.at[candidate_index, "bcsq_aa_notation"]), + ) + + return min(tied_candidates, key=tie_break_key) def process_joint_variants(path): diff --git a/vartracker/analysis_launcher.py b/vartracker/analysis_launcher.py index a5fda82..791fe68 100644 --- a/vartracker/analysis_launcher.py +++ b/vartracker/analysis_launcher.py @@ -67,6 +67,31 @@ def _validate_primer_bed_reference( ) +def _validate_lofreq_primer_rescue( + mode: str, + primer_bed_path: Optional[str], + min_af: float, + min_dp: int, + min_alt_count: int, + min_qual: float, + max_ref_count: int, +) -> None: + if mode not in {"auto", "on", "off"}: + raise ValueError("lofreq_primer_rescue must be one of: auto, on, off") + if mode == "on" and not primer_bed_path: + raise ValueError("--lofreq-primer-rescue on requires --primer-bed") + if not 0 <= min_af <= 1: + raise ValueError("lofreq_rescue_min_af must be between 0 and 1") + if min_dp < 0: + raise ValueError("lofreq_rescue_min_dp must be >= 0") + if min_alt_count < 0: + raise ValueError("lofreq_rescue_min_alt_count must be >= 0") + if min_qual < 0: + raise ValueError("lofreq_rescue_min_qual must be >= 0") + if max_ref_count < 0: + raise ValueError("lofreq_rescue_max_ref_count must be >= 0") + + def run_workflow( samples_csv: str | Path, reference: str | Path, @@ -83,6 +108,12 @@ def run_workflow( quiet: bool = True, mode: str = "reads", rulegraph_path: Optional[str | Path] = None, + lofreq_primer_rescue: str = "auto", + lofreq_rescue_min_af: float = 0.95, + lofreq_rescue_min_dp: int = 100, + lofreq_rescue_min_alt_count: int = 95, + lofreq_rescue_min_qual: float = 100.0, + lofreq_rescue_max_ref_count: int = 20, ) -> Optional[str]: """Run the lofreq variant calling workflow via the Snakemake API. @@ -107,6 +138,15 @@ def run_workflow( raise FileNotFoundError(f"Primer BED file not found: {primer_bed_path}") if primer_bed_path: _validate_primer_bed_reference(primer_bed_path, reference) + _validate_lofreq_primer_rescue( + lofreq_primer_rescue, + primer_bed_path, + lofreq_rescue_min_af, + lofreq_rescue_min_dp, + lofreq_rescue_min_alt_count, + lofreq_rescue_min_qual, + lofreq_rescue_max_ref_count, + ) Path(outdir).mkdir(parents=True, exist_ok=True) @@ -121,6 +161,12 @@ def run_workflow( "consensus_snp_min_af": consensus_snp_min_af, "consensus_snp_thresh": consensus_snp_thresh, "consensus_indel_thresh": consensus_indel_thresh, + "lofreq_primer_rescue": lofreq_primer_rescue, + "lofreq_rescue_min_af": lofreq_rescue_min_af, + "lofreq_rescue_min_dp": lofreq_rescue_min_dp, + "lofreq_rescue_min_alt_count": lofreq_rescue_min_alt_count, + "lofreq_rescue_min_qual": lofreq_rescue_min_qual, + "lofreq_rescue_max_ref_count": lofreq_rescue_max_ref_count, } if primer_bed_path: config_dict["primer_bed"] = primer_bed_path @@ -202,6 +248,12 @@ def run_workflow( parser.add_argument( "--primer-bed", help="Optional primer BED file for amplicon clipping" ) + parser.add_argument( + "--lofreq-primer-rescue", + choices=("auto", "on", "off"), + default="auto", + help="Primer-overlap rescue mode for LoFreq calls (default: auto)", + ) parser.add_argument("--dryrun", action="store_true", help="Perform dry run") args = parser.parse_args() @@ -213,6 +265,7 @@ def run_workflow( outdir=args.outdir, cores=args.cores, primer_bed=args.primer_bed, + lofreq_primer_rescue=args.lofreq_primer_rescue, dryrun=args.dryrun, ) except Exception as exc: # pragma: no cover - CLI surface diff --git a/vartracker/lofreq_primer_rescue.py b/vartracker/lofreq_primer_rescue.py new file mode 100644 index 0000000..f95701a --- /dev/null +++ b/vartracker/lofreq_primer_rescue.py @@ -0,0 +1,470 @@ +"""Rescue high-confidence LoFreq SNPs filtered at primer-overlap sites.""" + +from __future__ import annotations + +import argparse +import gzip +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import TextIO, cast + + +RESCUE_FILTER = ( + '##FILTER=' +) +PRIMER_INFO = ( + '##INFO=' +) +RESCUED_BY_INFO = ( + '##INFO=' +) + + +@dataclass(frozen=True) +class PrimerRescueThresholds: + """Thresholds applied only to primer-overlap rescue candidates.""" + + min_af: float = 0.95 + min_dp: int = 100 + min_alt_count: int = 95 + min_qual: float = 100.0 + max_ref_count: int = 20 + + +@dataclass(frozen=True) +class PrimerRescueResult: + """Summary of a LoFreq primer rescue run.""" + + normal_passed: int + rescued: int + discarded: int + rescued_tsv: str + + +def _open_text(path: str | Path, mode: str = "rt") -> TextIO: + text_path = str(path) + if text_path.endswith(".gz"): + return cast(TextIO, gzip.open(text_path, mode, encoding="utf-8")) + return cast(TextIO, open(text_path, mode, encoding="utf-8")) + + +def load_primers(path: str | Path) -> dict[str, list[tuple[int, int]]]: + """Load a primer BED file as 0-based half-open intervals.""" + + intervals: dict[str, list[tuple[int, int]]] = {} + with Path(path).open("r", encoding="utf-8") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("track ") or line.startswith("browser "): + continue + fields = line.split() + if len(fields) < 3: + continue + chrom = fields[0] + start = int(fields[1]) + end = int(fields[2]) + intervals.setdefault(chrom, []).append((start, end)) + return intervals + + +def overlaps_primer( + chrom: str, start: int, end: int, intervals: dict[str, list[tuple[int, int]]] +) -> bool: + """Return whether a half-open interval overlaps a primer interval.""" + + for primer_start, primer_end in intervals.get(chrom, []): + if primer_start < end and start < primer_end: + return True + return False + + +def parse_info(info_text: str) -> dict[str, str | bool]: + """Parse a VCF INFO string into a simple mapping.""" + + info: dict[str, str | bool] = {} + if info_text in ("", "."): + return info + for item in info_text.split(";"): + if not item: + continue + if "=" in item: + key, value = item.split("=", 1) + info[key] = value + else: + info[item] = True + return info + + +def _first_float(value: str | bool) -> float: + return float(str(value).split(",", 1)[0]) + + +def _first_int(value: str | bool) -> int: + return int(str(value).split(",", 1)[0]) + + +def _dp4_counts(value: str | bool) -> list[int] | None: + parts = str(value).split(",") + if len(parts) != 4: + return None + return [int(part) for part in parts] + + +def _record_key(fields: list[str]) -> tuple[str, str, str, str]: + return (fields[0], fields[1], fields[3], fields[4]) + + +def _filter_is_pass(fields: list[str]) -> bool: + return fields[6] in ("PASS", ".") + + +def read_default_filter_results( + path: str | Path, +) -> tuple[set[tuple[str, str, str, str]], dict[tuple[str, str, str, str], str]]: + """Read records emitted by ``lofreq filter`` and return PASS keys/reasons.""" + + pass_keys: set[tuple[str, str, str, str]] = set() + filter_reasons: dict[tuple[str, str, str, str], str] = {} + with _open_text(path) as handle: + for line in handle: + if line.startswith("#"): + continue + fields = line.rstrip("\n").split("\t") + if len(fields) >= 8: + key = _record_key(fields) + filter_reasons[key] = fields[6] + if _filter_is_pass(fields): + pass_keys.add(key) + return pass_keys, filter_reasons + + +def rescue_metrics( + fields: list[str], + primers: dict[str, list[tuple[int, int]]], + thresholds: PrimerRescueThresholds, +) -> str | None: + """Return rescue metrics for a candidate record, or ``None`` if it fails.""" + + chrom = fields[0] + pos = int(fields[1]) + ref = fields[3] + alt = fields[4] + qual_text = fields[5] + + if "," in alt: + return None + if len(ref) != 1 or len(alt) != 1: + return None + + start = pos - 1 + end = pos + if not overlaps_primer(chrom, start, end, primers): + return None + + if qual_text == ".": + return None + + info = parse_info(fields[7]) + if "AF" not in info or "DP" not in info or "DP4" not in info: + return None + + try: + af = _first_float(info["AF"]) + dp = _first_int(info["DP"]) + qual = float(qual_text) + counts = _dp4_counts(info["DP4"]) + except ValueError: + return None + + if counts is None: + return None + + ref_fwd, ref_rev, alt_fwd, alt_rev = counts + ref_count = ref_fwd + ref_rev + alt_count = alt_fwd + alt_rev + dp4_total = ref_count + alt_count + + if dp4_total == 0: + return None + if af < thresholds.min_af: + return None + if dp < thresholds.min_dp: + return None + if qual < thresholds.min_qual: + return None + if alt_count < thresholds.min_alt_count: + return None + if alt_count / dp4_total < thresholds.min_af: + return None + if ref_count > thresholds.max_ref_count: + return None + if not (alt_fwd == 0 or alt_rev == 0): + return None + + dp4_af = alt_count / dp4_total + return ( + f"AF={af:g},DP={dp},QUAL={qual:g}," + f"DP4={ref_fwd}/{ref_rev}/{alt_fwd}/{alt_rev}," + f"alt_count={alt_count},ref_count={ref_count},dp4_af={dp4_af:g}" + ) + + +def _add_info_flag(info_text: str, flag: str) -> str: + if info_text in ("", "."): + return flag + items = info_text.split(";") + if flag not in items: + items.append(flag) + return ";".join(items) + + +def _add_info_value(info_text: str, key: str, value: str) -> str: + item = f"{key}={value}" + if info_text in ("", "."): + return item + + items = [] + replaced = False + for existing in info_text.split(";"): + if existing == key or existing.startswith(key + "="): + items.append(item) + replaced = True + else: + items.append(existing) + if not replaced: + items.append(item) + return ";".join(items) + + +def _rescued_fields(fields: list[str]) -> list[str]: + out = fields[:] + out[6] = "RESCUED_PRIMER_OVERLAP" + out[7] = _add_info_flag(out[7], "PRIMER_OVERLAP") + out[7] = _add_info_value(out[7], "RESCUED_BY", "overlap_primer_interval") + return out + + +def _pass_fields(fields: list[str]) -> list[str]: + out = fields[:] + out[6] = "PASS" + return out + + +def _variant_name(fields: list[str]) -> str: + return f"{fields[3]}{fields[1]}{fields[4]}" + + +def _write_headers(headers: list[str], output: TextIO) -> None: + has_rescue_filter = any( + line.startswith("##FILTER= PrimerRescueResult: + """Write the default-filtered VCF plus any rescued primer-overlap records.""" + + normal_passed = 0 + rescued = 0 + discarded = 0 + headers: list[str] = [] + + with ( + _open_text(raw_vcf) as raw, + _open_text(output_vcf, "wt") as output, + Path(rescued_tsv_path).open("w", encoding="utf-8") as rescued_tsv, + ): + rescued_tsv.write("variant\treason_filtered\treason_rescued\tmetrics\n") + + for line in raw: + if line.startswith("#"): + headers.append(line) + continue + + if headers: + _write_headers(headers, output) + headers = [] + + fields = line.rstrip("\n").split("\t") + if len(fields) < 8: + discarded += 1 + continue + + key = _record_key(fields) + if key in pass_keys: + output.write("\t".join(_pass_fields(fields)) + "\n") + normal_passed += 1 + else: + metrics = rescue_metrics(fields, primers, thresholds) + if metrics is None: + discarded += 1 + continue + + output.write("\t".join(_rescued_fields(fields)) + "\n") + reason_filtered = filter_reasons.get( + key, "not_passed_default_lofreq_filter" + ) + rescued_tsv.write( + f"{_variant_name(fields)}\t{reason_filtered}\t" + f"overlap_primer_interval\t{metrics}\n" + ) + rescued += 1 + + if headers: + _write_headers(headers, output) + + return PrimerRescueResult( + normal_passed=normal_passed, + rescued=rescued, + discarded=discarded, + rescued_tsv=str(rescued_tsv_path), + ) + + +def rescue_lofreq_primer_variants( + raw_vcf: str | Path, + primers_bed: str | Path, + output_vcf: str | Path, + rescued_tsv: str | Path | None = None, + tmp_dir: str | Path | None = None, + lofreq: str = "lofreq", + thresholds: PrimerRescueThresholds | None = None, +) -> PrimerRescueResult: + """Apply default LoFreq filtering plus the primer-overlap rescue rule.""" + + thresholds = thresholds or PrimerRescueThresholds() + if rescued_tsv is None: + rescued_tsv = f"{output_vcf}.rescued.tsv" + + primers = load_primers(primers_bed) + + with tempfile.TemporaryDirectory(dir=tmp_dir) as run_tmp_dir: + filtered_vcf = os.path.join(run_tmp_dir, "lofreq.default_filtered.vcf") + subprocess.run( + [lofreq, "filter", "-i", str(raw_vcf), "-o", filtered_vcf], + check=True, + ) + pass_keys, filter_reasons = read_default_filter_results(filtered_vcf) + + return write_final_vcf( + raw_vcf, + output_vcf, + rescued_tsv, + pass_keys, + filter_reasons, + primers, + thresholds, + ) + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Preserve default LoFreq filtering while rescuing near-fixed " + "primer-overlap SNPs as RESCUED_PRIMER_OVERLAP." + ) + ) + parser.add_argument( + "--raw-vcf", required=True, help="Raw LoFreq VCF made with --no-default-filter" + ) + parser.add_argument("--primers-bed", required=True, help="Primer BED file") + parser.add_argument("--output-vcf", required=True, help="Final filtered VCF") + parser.add_argument( + "--rescued-tsv", + help="TSV of rescued variants; default: OUTPUT_VCF.rescued.tsv", + ) + parser.add_argument("--tmp-dir", help="Optional temporary directory") + parser.add_argument("--lofreq", default="lofreq", help="lofreq executable") + parser.add_argument( + "--min-af", + type=float, + default=PrimerRescueThresholds.min_af, + help="Minimum INFO/AF for rescue candidates only (default: 0.95)", + ) + parser.add_argument( + "--min-dp", + type=int, + default=PrimerRescueThresholds.min_dp, + help="Minimum INFO/DP for rescue candidates only (default: 100)", + ) + parser.add_argument( + "--min-alt-count", + type=int, + default=PrimerRescueThresholds.min_alt_count, + help="Minimum DP4 alt count for rescue candidates only (default: 95)", + ) + parser.add_argument( + "--min-qual", + type=float, + default=PrimerRescueThresholds.min_qual, + help="Minimum QUAL for rescue candidates only (default: 100)", + ) + parser.add_argument( + "--max-ref-count", + type=int, + default=PrimerRescueThresholds.max_ref_count, + help="Maximum DP4 ref count for rescue candidates only (default: 20)", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + thresholds = PrimerRescueThresholds( + min_af=args.min_af, + min_dp=args.min_dp, + min_alt_count=args.min_alt_count, + min_qual=args.min_qual, + max_ref_count=args.max_ref_count, + ) + result = rescue_lofreq_primer_variants( + raw_vcf=args.raw_vcf, + primers_bed=args.primers_bed, + output_vcf=args.output_vcf, + rescued_tsv=args.rescued_tsv, + tmp_dir=args.tmp_dir, + lofreq=args.lofreq, + thresholds=thresholds, + ) + print( + f"normal_passed={result.normal_passed} rescued={result.rescued} " + f"discarded={result.discarded}", + file=sys.stderr, + ) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/vartracker/main.py b/vartracker/main.py index 190a46d..c9fb2c3 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -79,6 +79,7 @@ validate_reference_and_annotation, ) from .reference_prepare import parse_accessions, prepare_reference_bundle +from .lofreq_primer_rescue import PrimerRescueThresholds _RED = "\033[91m" _YELLOW = "\033[93m" @@ -724,6 +725,17 @@ def _configure_vcf_parser( default=0.1, help="Minimum allele frequency of indel variants to keep (default: 0.1)", ) + analysis_group.add_argument( + "--multiallelic-overflow", + action="store", + required=False, + choices=("error", "drop-lowest-af", "skip-site"), + default="error", + help=( + "How to handle sites where more than two ALT alleles remain present in " + "one sample after filtering (default: error)" + ), + ) analysis_group.add_argument( "-d", "--min-depth", @@ -836,6 +848,51 @@ def _configure_vcf_parser( ) +def _add_lofreq_primer_rescue_arguments(group: argparse._ArgumentGroup) -> None: + group.add_argument( + "--lofreq-primer-rescue", + choices=("auto", "on", "off"), + default="auto", + help=( + "Primer-overlap rescue mode for LoFreq calls. 'auto' runs rescue when " + "--primer-bed is supplied; 'off' disables it (default: auto)" + ), + ) + group.add_argument( + "--lofreq-rescue-min-af", + type=float, + default=PrimerRescueThresholds.min_af, + help="Minimum INFO/AF for primer rescue candidates only (default: 0.95)", + ) + group.add_argument( + "--lofreq-rescue-min-dp", + type=int, + default=PrimerRescueThresholds.min_dp, + help="Minimum INFO/DP for primer rescue candidates only (default: 100)", + ) + group.add_argument( + "--lofreq-rescue-min-alt-count", + type=int, + default=PrimerRescueThresholds.min_alt_count, + help=( + "Minimum DP4 alternate count for primer rescue candidates only " + "(default: 95)" + ), + ) + group.add_argument( + "--lofreq-rescue-min-qual", + type=float, + default=PrimerRescueThresholds.min_qual, + help="Minimum QUAL for primer rescue candidates only (default: 100)", + ) + group.add_argument( + "--lofreq-rescue-max-ref-count", + type=int, + default=PrimerRescueThresholds.max_ref_count, + help="Maximum DP4 reference count for primer rescue candidates only (default: 20)", + ) + + def _move_action_group_after( parser: argparse.ArgumentParser, group_title: str, anchor_title: str ) -> None: @@ -931,6 +988,14 @@ def _add_bam_subparser(subparsers): default=8, help="Number of cores for Snakemake execution (default: 8)", ) + snk_group.add_argument( + "--primer-bed", + help=( + "Optional primer BED file for LoFreq primer-overlap rescue in BAM mode " + "(BAMs are not amplicon-clipped by vartracker)" + ), + ) + _add_lofreq_primer_rescue_arguments(snk_group) snk_group.add_argument( "--snakemake-dryrun", action="store_true", @@ -2034,7 +2099,10 @@ def _add_e2e_subparser(subparsers): ) snk_group.add_argument( "--primer-bed", - help="Optional primer BED file for amplicon clipping in Snakemake", + help=( + "Optional primer BED file for amplicon clipping and LoFreq " + "primer-overlap rescue" + ), ) snk_group.add_argument( "--ampliconclip-tolerance", @@ -2042,6 +2110,7 @@ def _add_e2e_subparser(subparsers): default=1, help="Tolerance for samtools ampliconclip primer matching (default: 1)", ) + _add_lofreq_primer_rescue_arguments(snk_group) snk_group.add_argument( "--snakemake-dryrun", action="store_true", @@ -2170,6 +2239,12 @@ def _run_e2e_command(args): quiet=not args.verbose, mode="reads", rulegraph_path=rulegraph_path, + lofreq_primer_rescue=args.lofreq_primer_rescue, + lofreq_rescue_min_af=args.lofreq_rescue_min_af, + lofreq_rescue_min_dp=args.lofreq_rescue_min_dp, + lofreq_rescue_min_alt_count=args.lofreq_rescue_min_alt_count, + lofreq_rescue_min_qual=args.lofreq_rescue_min_qual, + lofreq_rescue_max_ref_count=args.lofreq_rescue_max_ref_count, ) if rulegraph_path: @@ -2306,7 +2381,7 @@ def _run_bam_command(args): reference=args.reference, outdir=snakemake_outdir, cores=args.cores, - primer_bed=None, + primer_bed=args.primer_bed, min_depth=args.min_depth, consensus_snp_min_af=args.consensus_snp_min_af, consensus_snp_thresh=args.consensus_snp_thresh, @@ -2316,6 +2391,12 @@ def _run_bam_command(args): quiet=not args.verbose, mode="bam", rulegraph_path=rulegraph_path, + lofreq_primer_rescue=args.lofreq_primer_rescue, + lofreq_rescue_min_af=args.lofreq_rescue_min_af, + lofreq_rescue_min_dp=args.lofreq_rescue_min_dp, + lofreq_rescue_min_alt_count=args.lofreq_rescue_min_alt_count, + lofreq_rescue_min_qual=args.lofreq_rescue_min_qual, + lofreq_rescue_max_ref_count=args.lofreq_rescue_max_ref_count, ) if rulegraph_path: @@ -2626,6 +2707,7 @@ def _process_files( args.reference, args.gff3, args.debug, + args.multiallelic_overflow, ) # Process VCF and extract variants diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index afcd69c..bf1c3e0 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -8,7 +8,9 @@ import re import subprocess import gzip +import shutil from pathlib import Path +from typing import TypedDict import pandas as pd import numpy as np @@ -19,6 +21,12 @@ from .amino_acids import AminoAcidChange +class _AltEntry(TypedDict): + index: int + alt: str + af: float + + def _open_vcf(path: str, mode: str): if path.endswith(".gz"): return gzip.open(path, mode, encoding="utf-8") @@ -140,6 +148,48 @@ def _normalize_header_line(line: str) -> str: return line +def _run_logged_command(cmd: str, log_file: str, context: str) -> None: + """Run a shell command and surface the tail of the log on failure.""" + + try: + with open(log_file, "a", encoding="utf-8") as err: + subprocess.run(cmd, shell=True, stderr=err, check=True) + except subprocess.CalledProcessError as exc: + log_tail = "" + if os.path.exists(log_file): + try: + with open( + log_file, "r", encoding="utf-8", errors="replace" + ) as err_file: + tail_lines = err_file.readlines()[-20:] + log_tail = "".join(tail_lines).strip() + except OSError: + log_tail = "" + + details = log_tail if log_tail else exc.stderr or "" + message = ( + f"{context} failed with exit status {exc.returncode} for command '{cmd}'." + ) + if details: + message = f"{message}\nLast bcftools log lines:\n{details}" + raise RuntimeError(message) from exc + + +def _split_multiallelic_records( + vcf_path: str, tempdir: str, sample: str, log_file: str, debug: bool = False +) -> str: + """Split multi-ALT records so downstream processing can track each allele.""" + + split_path = os.path.join(tempdir, f"{Path(vcf_path).stem}.{sample}.split.vcf") + cmd = f"bcftools norm -m -any -Ov -o {split_path} {vcf_path}" + _run_logged_command(cmd, log_file, "bcftools norm") + + if debug: + print(f"Command: {cmd}") + + return split_path + + def rindex(lst, item): """Find the last occurrence of an item in a list.""" @@ -182,8 +232,11 @@ def format_vcf( try: prepared_vcf = _ensure_format_and_sample(vcf, sample, tempdir, debug) + split_vcf = _split_multiallelic_records( + prepared_vcf, tempdir, sample, log, debug + ) - vcf_mod = VCF(prepared_vcf, strict_gt=True) + vcf_mod = VCF(split_vcf, strict_gt=True) existing_samples = list(vcf_mod.samples) sample_count = len(existing_samples) if existing_samples else 1 @@ -261,7 +314,10 @@ def _scalar(value): return value for v in vcf_mod: - pos_key = str(v.POS) + if not v.ALT: + continue + + variant_key = (v.CHROM, int(v.POS), v.REF, v.ALT[0]) # Get allele frequency from the specified tag if allele_frequency_tag != "AF": @@ -285,9 +341,9 @@ def _scalar(value): round(x, 6) if isinstance(x, float) else x for x in af_value ] - if pos_key in variants: - if variants[pos_key].INFO["AF"] < v.INFO["AF"]: - del variants[pos_key] + if variant_key in variants: + if variants[variant_key].INFO["AF"] < v.INFO["AF"]: + del variants[variant_key] else: continue @@ -308,7 +364,7 @@ def _scalar(value): v.set_format("DP", np.array(dp_array, dtype=int)) v.set_format("AF", np.array(af_array, dtype=float)) v.genotypes = [[1, True] for _ in range(sample_count)] - variants[pos_key] = v + variants[variant_key] = v for variant in variants.values(): w.write_record(variant) @@ -327,30 +383,8 @@ def _scalar(value): f"-Oz -o {out}" ) - try: - with open(log, "a", encoding="utf-8") as err: - subprocess.run(cmd, shell=True, stderr=err, check=True) - subprocess.run( - f"bcftools index -f {out}", - shell=True, - stderr=err, - check=True, - ) - except subprocess.CalledProcessError as exc: - log_tail = "" - if os.path.exists(log): - try: - with open(log, "r", encoding="utf-8", errors="replace") as err_file: - tail_lines = err_file.readlines()[-20:] - log_tail = "".join(tail_lines).strip() - except OSError: - log_tail = "" - - details = log_tail if log_tail else exc.stderr or "" - message = f"Command '{cmd}' returned non-zero exit status {exc.returncode}." - if details: - message = f"{message}\nLast bcftools log lines:\n{details}" - raise RuntimeError(message) from exc + _run_logged_command(cmd, log, "bcftools view filter") + _run_logged_command(f"bcftools index -f {out}", log, "bcftools index") if debug: print(f"Command: {cmd}") @@ -402,11 +436,263 @@ def merge_consequences(tempdir, csq_file, sample_names, debug): raise RuntimeError(f"Error merging VCF files: {str(e)}") -def annotate_vcf(vcf_file, output_file, reference, annotation, debug): +def _variant_threshold_flag(v) -> str: + ref = str(v.REF or "") + alts = [str(alt) for alt in (v.ALT or [])] + is_indel = any(len(alt) != len(ref) for alt in alts) + return "--min-indel-freq" if is_indel else "--min-snv-freq" + + +def _normalise_af_row(sample_values, alt_count: int) -> list[float]: + values = np.atleast_1d(sample_values).tolist() + if len(values) < alt_count: + values.extend([np.nan] * (alt_count - len(values))) + return values[:alt_count] + + +def _present_alt_entries(v, sample_values) -> list[_AltEntry]: + entries: list[_AltEntry] = [] + values = _normalise_af_row(sample_values, len(v.ALT or [])) + for alt_index, raw_value in enumerate(values, start=1): + try: + numeric_value = float(raw_value) + except (TypeError, ValueError): + continue + if np.isnan(numeric_value) or numeric_value <= 0: + continue + alt = str(v.ALT[alt_index - 1]) + entries.append({"index": alt_index, "alt": alt, "af": numeric_value}) + return entries + + +def _describe_alt_entries(v, entries: list[_AltEntry]) -> str: + if not entries: + return "None" + ref = str(v.REF or "") + return ", ".join(f"{ref}>{entry['alt']} AF={entry['af']:.3f}" for entry in entries) + + +def _build_overflow_events(v, sample_names, af_matrix) -> list[dict[str, object]]: + events = [] + for sample_name, sample_values in zip(sample_names, af_matrix.tolist()): + present_entries = _present_alt_entries(v, sample_values) + if len(present_entries) <= 2: + continue + + sorted_entries = sorted( + present_entries, key=lambda entry: (entry["af"], entry["index"]) + ) + drop_count = len(sorted_entries) - 2 + dropped_entries = sorted_entries[:drop_count] + kept_entries = sorted_entries[drop_count:] + threshold_suggestion = max(float(entry["af"]) for entry in dropped_entries) + + events.append( + { + "sample": sample_name, + "present": present_entries, + "dropped": dropped_entries, + "kept": kept_entries, + "suggestion": threshold_suggestion, + } + ) + + return events + + +def _format_overflow_message(v, overflow_events, *, action: str) -> str: + threshold_flag = _variant_threshold_flag(v) + location = f"{v.CHROM}:{v.POS}" + + lines = [ + f"bcftools csq input preparation found more than two ALT alleles present at {location}.", + f"Action: {action}.", + ] + for event in overflow_events: + lines.append(f"Sample {event['sample']}:") + lines.append( + f"Retained ALT alleles after filtering: {_describe_alt_entries(v, event['present'])}." + ) + lines.append( + f"To avoid this at the current site, increase {threshold_flag} above {event['suggestion']:.3f}." + ) + if event["dropped"]: + lines.append( + f"Lowest-frequency ALT alleles affected: {_describe_alt_entries(v, event['dropped'])}." + ) + lines.append("bcftools csq can represent at most two ALT alleles per sample/site.") + return "\n".join(lines) + + +def _vcf_record_sort_key(line: str) -> tuple[str, int, str, str]: + fields = line.rstrip("\n").split("\t") + chrom = fields[0] if len(fields) > 0 else "" + try: + pos = int(fields[1]) if len(fields) > 1 else 0 + except ValueError: + pos = 0 + ref = fields[3] if len(fields) > 3 else "" + alt = fields[4] if len(fields) > 4 else "" + return chrom, pos, ref, alt + + +def _merge_vcf_outputs(primary_vcf: str, secondary_vcf: str, output_file: str) -> None: + with open(primary_vcf, "r", encoding="utf-8") as primary_handle: + primary_lines = primary_handle.readlines() + with open(secondary_vcf, "r", encoding="utf-8") as secondary_handle: + secondary_lines = secondary_handle.readlines() + + header_lines = [line for line in primary_lines if line.startswith("#")] + record_lines = [ + line for line in primary_lines if line.strip() and not line.startswith("#") + ] + record_lines.extend( + line for line in secondary_lines if line.strip() and not line.startswith("#") + ) + record_lines.sort(key=_vcf_record_sort_key) + + with open(output_file, "w", encoding="utf-8") as out_handle: + out_handle.writelines(header_lines) + out_handle.writelines(record_lines) + + +def annotate_vcf( + vcf_file, + output_file, + reference, + annotation, + debug, + multiallelic_overflow="error", +): """Annotate a merged multi-sample VCF with bcftools csq.""" + if multiallelic_overflow not in {"error", "drop-lowest-af", "skip-site"}: + raise RuntimeError( + "multiallelic_overflow must be one of: error, drop-lowest-af, skip-site" + ) + + output_dir = os.path.dirname(output_file) or "." + prefix = Path(output_file).stem + joined_vcf = os.path.join(output_dir, f"{prefix}.joined_for_csq.vcf") + prepared_vcf = os.path.join(output_dir, f"{prefix}.prepared_for_csq.vcf") + skipped_vcf = os.path.join(output_dir, f"{prefix}.skipped_from_csq.vcf") + annotated_only_vcf = os.path.join(output_dir, f"{prefix}.annotated_only.vcf") + + join_cmd = f"bcftools norm -m +any -Ov -o {joined_vcf} {vcf_file}" + if debug: + print(f"Command: {join_cmd}") + try: + subprocess.run(join_cmd, shell=True, check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Error preparing merged VCF for bcftools csq with multiallelic joins: {str(e)}" + ) + + joined = VCF(joined_vcf, strict_gt=True) + sample_names = list(joined.samples) + writer = Writer.from_string(prepared_vcf, joined.raw_header, mode="w") + skipped_writer = Writer.from_string(skipped_vcf, joined.raw_header, mode="w") + records_for_csq = 0 + skipped_records = 0 + + try: + for variant in joined: + try: + af_matrix = variant.format("AF") + except (KeyError, AttributeError, ValueError, RuntimeError): + af_matrix = None + + if af_matrix is None: + writer.write_record(variant) + records_for_csq += 1 + continue + + overflow_events = _build_overflow_events(variant, sample_names, af_matrix) + if overflow_events and multiallelic_overflow == "error": + raise RuntimeError( + _format_overflow_message( + variant, + overflow_events, + action="Stopping analysis before bcftools csq", + ) + ) + + if overflow_events and multiallelic_overflow == "skip-site": + print( + "Warning: " + + _format_overflow_message( + variant, + overflow_events, + action="Skipping consequence calling for this site", + ) + ) + skipped_writer.write_record(variant) + skipped_records += 1 + continue + + updated_af_rows = [] + genotypes = [] + for sample_name, sample_values in zip(sample_names, af_matrix.tolist()): + row_values = _normalise_af_row(sample_values, len(variant.ALT or [])) + present_entries = _present_alt_entries(variant, row_values) + + if len(present_entries) > 2: + sorted_entries = sorted( + present_entries, + key=lambda entry: (entry["af"], entry["index"]), + ) + dropped_entries = sorted_entries[: len(sorted_entries) - 2] + kept_entries = sorted_entries[len(sorted_entries) - 2 :] + dropped_indices = {entry["index"] for entry in dropped_entries} + + for alt_index in dropped_indices: + row_values[alt_index - 1] = np.nan + + print( + "Warning: " + + _format_overflow_message( + variant, + [ + { + "sample": sample_name, + "present": present_entries, + "dropped": dropped_entries, + "kept": kept_entries, + "suggestion": max( + float(entry["af"]) for entry in dropped_entries + ), + } + ], + action="Dropping the lowest-frequency ALT allele(s) for this sample", + ) + ) + present = sorted(entry["index"] for entry in kept_entries) + else: + present = sorted(entry["index"] for entry in present_entries) + + updated_af_rows.append(row_values) + if not present: + genotypes.append([0, 0, False]) + elif len(present) == 1: + genotypes.append([present[0], present[0], False]) + else: + genotypes.append([present[0], present[1], False]) + + variant.set_format("AF", np.asarray(updated_af_rows, dtype=float)) + variant.genotypes = genotypes + writer.write_record(variant) + records_for_csq += 1 + finally: + writer.close() + skipped_writer.close() + joined.close() + + if records_for_csq == 0: + shutil.copyfile(skipped_vcf, output_file) + return + cmd = ( - f"bcftools csq -f {reference} -g {annotation} --force " - f"-Ov -o {output_file} {vcf_file}" + f"bcftools csq -p R -f {reference} -g {annotation} --force " + f"-Ov -o {annotated_only_vcf} {prepared_vcf}" ) if debug: @@ -417,6 +703,11 @@ def annotate_vcf(vcf_file, output_file, reference, annotation, debug): except subprocess.CalledProcessError as e: raise RuntimeError(f"Error annotating merged VCF with bcftools csq: {str(e)}") + if skipped_records: + _merge_vcf_outputs(annotated_only_vcf, skipped_vcf, output_file) + else: + shutil.copyfile(annotated_only_vcf, output_file) + def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): """ @@ -597,6 +888,98 @@ def _mask_allele_frequencies_for_annotation( return masked +def _format_frequency_token(value) -> str: + try: + numeric = float(value) + except (TypeError, ValueError): + return "." + if np.isnan(numeric): + return "." + return "{:.3f}".format(numeric) + + +def _extract_alt_frequency_map(v, samples): + """Return per-ALT allele-frequency trajectories for a VCF record.""" + + alt_map = {str(alt): ["."] * len(samples) for alt in (v.ALT or [])} + if not alt_map: + return alt_map + + try: + allele_freq_arrays = v.format("AF") + except (KeyError, AttributeError, ValueError, RuntimeError): + allele_freq_arrays = None + + if allele_freq_arrays is not None: + for sample_index, sample_values in enumerate(allele_freq_arrays.tolist()): + values = np.atleast_1d(sample_values).tolist() + for alt_index, alt in enumerate(v.ALT): + raw_value = values[alt_index] if alt_index < len(values) else None + alt_map[str(alt)][sample_index] = _format_frequency_token(raw_value) + return alt_map + + info_af = v.INFO.get("AF") or v.INFO.get("VAF") + if isinstance(info_af, (list, tuple)): + values = list(info_af) + elif info_af is None: + values = [] + else: + values = [info_af] + + for alt_index, alt in enumerate(v.ALT): + raw_value = values[alt_index] if alt_index < len(values) else None + token = _format_frequency_token(raw_value) + alt_map[str(alt)] = [token] * len(samples) + + return alt_map + + +def _collapse_alt_frequency_map(alt_frequency_map, sample_count: int): + """Collapse per-ALT frequencies into a site-level trajectory using max AF.""" + + collapsed = [] + for sample_index in range(sample_count): + observed = [] + for allele_freqs in alt_frequency_map.values(): + if sample_index >= len(allele_freqs): + continue + token = allele_freqs[sample_index] + if token == ".": + continue + try: + observed.append(float(token)) + except (TypeError, ValueError): + continue + collapsed.append("{:.3f}".format(max(observed)) if observed else ".") + return collapsed + + +def _annotation_alt_for_record(v, anno): + """Infer which ALT allele a BCSQ annotation belongs to for the current record.""" + + if not v.ALT: + return None + if len(v.ALT) == 1: + return str(v.ALT[0]) + if len(anno) <= 6: + return None + + dna_change = str(anno[6] or "").strip() + if not dna_change: + return None + + for token in dna_change.split("+"): + match = re.match(r"^(\d+)([^>]+)>([^>]+)$", token.strip()) + if not match: + continue + pos = int(match.group(1)) + alt = match.group(3) + if pos == int(v.POS) and alt in v.ALT: + return alt + + return None + + def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): """ Process VCF file and extract variant information. @@ -646,38 +1029,10 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): # Process variants for v in vcf: info = dict(list(v.INFO)) - try: - allele_freq_arrays = v.format("AF").tolist() - except (KeyError, AttributeError, ValueError, RuntimeError): - allele_freq_arrays = None - - if allele_freq_arrays: - allele_freqs = [ - "{:.3f}".format(x[0]).replace("nan", ".") for x in allele_freq_arrays - ] - else: - info_af = v.INFO.get("AF") or v.INFO.get("VAF") - if isinstance(info_af, (list, tuple)): - values = list(info_af) - elif info_af is None: - values = [None] * len(samples) - else: - values = [info_af] - if not values: - values = [None] - if len(values) < len(samples): - values.extend([values[-1]] * (len(samples) - len(values))) - allele_freqs = [] - for val in values[: len(samples)]: - if val is None: - allele_freqs.append(".") - continue - try: - allele_freqs.append("{:.3f}".format(float(val))) - except (TypeError, ValueError): - allele_freqs.append(".") - if not allele_freqs: - allele_freqs = ["."] * max(1, len(samples)) + alt_frequency_map = _extract_alt_frequency_map(v, samples) + allele_freqs = _collapse_alt_frequency_map(alt_frequency_map, len(samples)) + if not allele_freqs: + allele_freqs = ["."] * max(1, len(samples)) trajectory = _summarise_sample_trajectory(allele_freqs, samples) depths_qc = calculate_variant_site_depths(cov_df, v, samples, min_depth) @@ -697,8 +1052,15 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): if sample_bcsq_map: for annot in annotations: + anno = annot.split("|") + annotation_alt = _annotation_alt_for_record(v, anno) + annotation_allele_freqs = ( + alt_frequency_map.get(annotation_alt, allele_freqs) + if annotation_alt is not None + else allele_freqs + ) masked_allele_freqs = _mask_allele_frequencies_for_annotation( - annot, allele_freqs, samples, sample_bcsq_map + annot, annotation_allele_freqs, samples, sample_bcsq_map ) masked_trajectory = _summarise_sample_trajectory( masked_allele_freqs, samples @@ -706,7 +1068,6 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): if "Y" not in masked_trajectory["presence_absence"]: continue - anno = annot.split("|") result = _process_annotation( v, anno, @@ -721,6 +1082,7 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): masked_allele_freqs, samples, total_cov_list, + annotation_alt, ) results.append(result) produced_annotation_specific_row = True @@ -728,39 +1090,73 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): if not produced_annotation_specific_row: for annot in annotations: anno = annot.split("|") + annotation_alt = _annotation_alt_for_record(v, anno) + annotation_allele_freqs = ( + alt_frequency_map.get(annotation_alt, allele_freqs) + if annotation_alt is not None + else allele_freqs + ) + annotation_trajectory = _summarise_sample_trajectory( + annotation_allele_freqs, samples + ) result = _process_annotation( v, anno, - trajectory["variant_status"], - trajectory["persistent_status"], - trajectory["presence_absence"], - trajectory["first_appearance"], - trajectory["last_appearance"], + annotation_trajectory["variant_status"], + annotation_trajectory["persistent_status"], + annotation_trajectory["presence_absence"], + annotation_trajectory["first_appearance"], + annotation_trajectory["last_appearance"], all_samples_pass_qc, proportion_samples_passing_qc, depths_qc, - allele_freqs, + annotation_allele_freqs, samples, total_cov_list, + annotation_alt, ) results.append(result) else: # Handle variants without annotations - result = _create_unannotated_result( - v, - trajectory["variant_status"], - trajectory["persistent_status"], - trajectory["presence_absence"], - trajectory["first_appearance"], - trajectory["last_appearance"], - all_samples_pass_qc, - proportion_samples_passing_qc, - depths_qc, - allele_freqs, - samples, - total_cov_list, - ) - results.append(result) + if len(alt_frequency_map) > 1: + for alt, alt_allele_freqs in alt_frequency_map.items(): + alt_trajectory = _summarise_sample_trajectory( + alt_allele_freqs, samples + ) + if "Y" not in alt_trajectory["presence_absence"]: + continue + result = _create_unannotated_result( + v, + alt_trajectory["variant_status"], + alt_trajectory["persistent_status"], + alt_trajectory["presence_absence"], + alt_trajectory["first_appearance"], + alt_trajectory["last_appearance"], + all_samples_pass_qc, + proportion_samples_passing_qc, + depths_qc, + alt_allele_freqs, + samples, + total_cov_list, + alt, + ) + results.append(result) + else: + result = _create_unannotated_result( + v, + trajectory["variant_status"], + trajectory["persistent_status"], + trajectory["presence_absence"], + trajectory["first_appearance"], + trajectory["last_appearance"], + all_samples_pass_qc, + proportion_samples_passing_qc, + depths_qc, + allele_freqs, + samples, + total_cov_list, + ) + results.append(result) return pd.DataFrame(results) @@ -779,8 +1175,10 @@ def _process_annotation( allele_freqs, samples, total_cov_list, + alt_allele=None, ): """Process a single annotation from bcftools csq.""" + selected_alt = alt_allele or (v.ALT[0] if v.ALT else "") if len(anno) == 1 and anno[0].startswith("@"): # Joint variant annotation return { @@ -789,8 +1187,8 @@ def _process_annotation( "end": v.end, "gene": anno[0], "ref": v.REF, - "alt": v.ALT[0], - "variant": v.REF + str(v.POS) + v.ALT[0], + "alt": selected_alt, + "variant": v.REF + str(v.POS) + selected_alt, "amino_acid_consequence": anno[0], "nsp_aa_change": anno[0], "bcsq_nt_notation": anno[0], @@ -835,8 +1233,8 @@ def _process_annotation( "end": v.end, "gene": anno[1], "ref": v.REF, - "alt": v.ALT[0], - "variant": v.REF + str(v.POS) + v.ALT[0], + "alt": selected_alt, + "variant": v.REF + str(v.POS) + selected_alt, "amino_acid_consequence": reformatted_aa[0], "nsp_aa_change": reformatted_aa[1], "bcsq_nt_notation": anno[6] if len(anno) > 5 else "", @@ -880,8 +1278,10 @@ def _create_unannotated_result( allele_freqs, samples, total_cov_list, + alt_allele=None, ): """Create result for variants without annotations.""" + selected_alt = alt_allele or (v.ALT[0] if v.ALT else "") if v.start + 1 < 266: gene = "5' UTR" elif v.start + 1 > 29674: @@ -895,8 +1295,8 @@ def _create_unannotated_result( "end": v.end, "gene": gene, "ref": v.REF, - "alt": v.ALT[0], - "variant": v.REF + str(v.POS) + v.ALT[0], + "alt": selected_alt, + "variant": v.REF + str(v.POS) + selected_alt, "amino_acid_consequence": "None", "nsp_aa_change": "None", "bcsq_nt_notation": "None", From 1fb3d8cfb5475ee08e8db898754b0e6ce42c1032 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Tue, 5 May 2026 12:51:48 +1000 Subject: [PATCH 10/20] Script lint fix --- vartracker/lofreq_primer_rescue.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vartracker/lofreq_primer_rescue.py b/vartracker/lofreq_primer_rescue.py index f95701a..d0a80cb 100644 --- a/vartracker/lofreq_primer_rescue.py +++ b/vartracker/lofreq_primer_rescue.py @@ -12,7 +12,6 @@ from pathlib import Path from typing import TextIO, cast - RESCUE_FILTER = ( '##FILTER=' From 7a21a7d70196975e8f4ed13c47c1035a4727a5bc Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Tue, 5 May 2026 13:38:57 +1000 Subject: [PATCH 11/20] Fixed mypy pin to pass CI --- .pre-commit-config.yaml | 3 ++- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 32da42f..0bb39c0 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -32,6 +32,7 @@ repos: types: [python] - id: mypy name: mypy - entry: mypy + entry: mypy vartracker tests language: system + pass_filenames: false types: [python] diff --git a/pyproject.toml b/pyproject.toml index 49a348b..1b2c97a 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ dev = [ "pytest-cov", "black", "flake8", - "mypy", + "mypy>=1.19,<1.20", "pre-commit>=3.6", ] From 0cc7d22e097954d7e6e57269d4b7149231a8403c Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Tue, 5 May 2026 14:13:10 +1000 Subject: [PATCH 12/20] Fixed schema to pass CI --- docs/OUTPUT_SCHEMA.md | 6 ++++-- tests/test_schema.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/docs/OUTPUT_SCHEMA.md b/docs/OUTPUT_SCHEMA.md index f74bd2a..578c9d1 100644 --- a/docs/OUTPUT_SCHEMA.md +++ b/docs/OUTPUT_SCHEMA.md @@ -1,6 +1,6 @@ # Output schema -Schema version: `1.0` +Schema version: `1.1` Generated from `vartracker.schemas.RESULTS_SCHEMA`. @@ -28,7 +28,8 @@ Columns that encode per-sample values are slash-separated and ordered by the inp | presence_absence | string (slash-separated) | Per-sample presence (Y) or absence (N), ordered by input. | | Y/N | | first_appearance | string | Sample name where the variant first appears. | | | | last_appearance | string | Sample name where the variant last appears. | | | -| overall_variant_qc | string | Aggregated QC status across samples. | | PASS, FAIL | +| all_samples_pass_qc | boolean | True if every sample passes per-sample variant QC. | | true, false | +| proportion_samples_passing_qc | number | Proportion of samples passing per-sample variant QC. | fraction | 0-1 | | per_sample_variant_qc | string (slash-separated) | Per-sample QC flags (P/F) ordered by input. | | P, F | | aa1_total_properties | string | Physicochemical properties for the reference amino acid. | | semicolon-separated properties | | aa2_total_properties | string | Physicochemical properties for the alternate amino acid. | | semicolon-separated properties | @@ -42,4 +43,5 @@ Columns that encode per-sample values are slash-separated and ordered by the inp | variant_site_depth | string (slash-separated) | Total read depth at the variant site per sample. | reads | | | variant_window_depth | string (slash-separated) | Mean read depth in the variant window per sample. | reads | | | samples | string (slash-separated) | Sample names corresponding to per-sample fields. | | | +| sample_number | string (slash-separated) | Sample ordering values corresponding to per-sample fields. | | integer-like sample numbers | | total_genome_coverage | string (slash-separated) | Total genome coverage (bases covered) per sample. | bases | | diff --git a/tests/test_schema.py b/tests/test_schema.py index 23529f5..4a89baa 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -1,6 +1,8 @@ """Tests for output schema metadata.""" from pathlib import Path +import subprocess +import sys import importlib.util import pandas as pd @@ -22,3 +24,14 @@ def test_results_schema_matches_precomputed_columns(): Path("vartracker") / "test_data" / "precomputed" / "test_results.csv" ) assert schema_columns == list(precomputed.columns) + + +def test_output_schema_documentation_is_current(): + result = subprocess.run( + [sys.executable, "scripts/check_output_schema.py"], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stdout + result.stderr From 146481101c3683d4b434adbd6127347042d5a1ba Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Wed, 6 May 2026 10:34:46 +1000 Subject: [PATCH 13/20] Fixing some variant rescue and longitudinal tracking quirks --- CITATION.cff | 8 +- README.md | 18 +- pyproject.toml | 2 +- tests/test_analysis.py | 65 +++++++- tests/test_analysis_launcher.py | 14 ++ tests/test_constants.py | 12 ++ tests/test_lofreq_primer_rescue.py | 82 ++++++++- tests/test_main.py | 4 + tests/test_snakemake_workflow.py | 18 ++ tests/test_vcf_processing.py | 85 +++++++++- vartracker/Snakefile | 38 ++++- vartracker/_version.py | 2 +- vartracker/analysis.py | 84 +++++++--- vartracker/analysis_launcher.py | 6 + vartracker/constants.py | 37 ++-- vartracker/lofreq_primer_rescue.py | 260 +++++++++++++++++++++++++---- vartracker/main.py | 15 ++ vartracker/vcf_processing.py | 97 +++++++++-- 18 files changed, 729 insertions(+), 118 deletions(-) create mode 100644 tests/test_constants.py diff --git a/CITATION.cff b/CITATION.cff index b303859..55d2fd2 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -1,8 +1,8 @@ cff-version: 1.2.0 message: "If you use this software, please cite it as below." title: "vartracker" -version: "2.2.0" -date-released: "2026-05-05" +version: "2.2.1" +date-released: "2026-05-06" license: MIT repository-code: "https://github.com/charlesfoster/vartracker" url: "https://github.com/charlesfoster/vartracker" @@ -22,7 +22,7 @@ preferred-citation: - family-names: Foster given-names: Charles title: "vartracker" - version: "2.2.0" + version: "2.2.1" doi: "10.5281/zenodo.18452274" url: "https://github.com/charlesfoster/vartracker" - date-released: "2026-05-05" + date-released: "2026-05-06" diff --git a/README.md b/README.md index b92c28b..9923e8a 100755 --- a/README.md +++ b/README.md @@ -167,7 +167,8 @@ Docker is a self-contained reproducible option. If you publish the image, record set it when running to include it in the run manifest: ```bash -export VARTRACKER_CONTAINER_IMAGE=ghcr.io/your-org/vartracker:2.0.0 +export +.2.1 export VARTRACKER_CONTAINER_DIGEST=sha256:... ``` @@ -253,8 +254,10 @@ LoFreq primer-overlap rescue: - `bam` and `end-to-end` therefore run LoFreq with `--no-default-filter`, then apply the normal `lofreq filter` step so standard LoFreq PASS calls are unchanged. - With the default `--lofreq-primer-rescue auto`, the rescue step runs only when `--primer-bed` is supplied. In other words, `auto` means "use primer rescue when an amplicon primer scheme has been explicitly provided." - In `end-to-end` mode, the same `--primer-bed` is used for `samtools ampliconclip` and for rescue. In `bam` mode, vartracker does not clip the input BAMs; the primer BED is used only to identify primer-overlap sites for rescue. -- Rescue candidates must be single-ALT SNPs that overlap a primer interval, fail the default LoFreq filter, and pass conservative near-fixed thresholds (`AF>=0.95`, `DP>=100`, `DP4 alt count>=95`, `QUAL>=100`, `DP4 ref count<=20`) with one-sided alternate-strand support. Indels, multi-ALT records, lower-frequency variants, and non-primer-overlap variants are not rescued by this rule. +- Rescue candidates must be single-ALT SNPs that overlap a primer interval, fail LoFreq's default strand-bias filtering, and pass conservative near-fixed thresholds (`AF>=0.95`, `DP>=100`, `DP4 alt count>=95`, `QUAL>=100`, `DP4 ref count<=20`, minor ALT strand fraction `<=0.05`). Indels, multi-ALT records, lower-frequency variants, non-primer-overlap variants, and variants filtered for non-strand-bias reasons are not rescued by this rule. +- The raw LoFreq calls are retained as `_variants.raw.vcf.gz` and listed in the updated spreadsheet as `raw_vcf`. - Rescued variants are marked with `FILTER=RESCUED_PRIMER_OVERLAP`, `INFO/PRIMER_OVERLAP`, and `INFO/RESCUED_BY=overlap_primer_interval`; per-sample details are written to `_variants.rescued.tsv` and listed in the updated spreadsheet as `lofreq_rescued_tsv`. +- Variants called by raw LoFreq but filtered out of the final VCF are written to `_variants.filtered_out.tsv` with the LoFreq filter reason and core metrics. This is useful for auditing high-frequency calls that fail strand-bias or other LoFreq filters. - Use `--lofreq-primer-rescue off` to disable rescue even when a primer BED is supplied, or `--lofreq-primer-rescue on` to require rescue and fail if `--primer-bed` is missing. The rescue thresholds can be adjusted with the `--lofreq-rescue-*` options. Example amplicon run with primer rescue: @@ -284,10 +287,11 @@ Mode-specific expectations: - **End-to-end mode** requires `reads1` (and optionally `reads2`); remaining fields are generated. The `bam` and `end-to-end` workflows also write two consensus FASTA columns to -the updated Snakemake spreadsheet, plus the LoFreq rescue audit column: +the updated Snakemake spreadsheet, plus LoFreq audit columns: `consensus` for a simple consensus, `iupac_consensus` for an IUPAC-aware -consensus, and `lofreq_rescued_tsv` for the per-sample primer-overlap rescue -table. SNPs below +consensus, `raw_vcf` for raw LoFreq calls, `lofreq_rescued_tsv` for the +per-sample primer-overlap rescue table, and `lofreq_filtered_out_tsv` for raw +LoFreq records excluded from the final VCF. SNPs below `--consensus-snp-min-af` are ignored, SNPs from `--consensus-snp-min-af` up to `--consensus-snp-thresh` stay as reference bases in the simple consensus and become REF+ALT ambiguity codes in the IUPAC consensus, and SNPs at or above @@ -509,7 +513,9 @@ vartracker produces several output files: - **results.csv**: Comprehensive variant analysis with all metrics - **results_metadata.json**: Output schema version and results metadata +- **`_variants.raw.vcf.gz`** (`bam`/`end-to-end`): Raw LoFreq calls before default filtering and primer-overlap rescue - **`_variants.rescued.tsv`** (`bam`/`end-to-end`): LoFreq primer-overlap rescue audit table, empty when rescue is disabled or no variants are rescued +- **`_variants.filtered_out.tsv`** (`bam`/`end-to-end`): Raw LoFreq calls excluded from the final VCF, including filter reason and call metrics - **new_mutations.csv**: Mutations not present in the first sample - **persistent_new_mutations.csv**: New mutations that persist to the final sample - **cumulative_mutations.pdf**: Plot showing mutation accumulation over time @@ -566,7 +572,7 @@ The pipeline performs the following analysis: When using vartracker, please cite the software release you used. Citation metadata is provided in `CITATION.cff`, and GitHub releases are archived on Zenodo. -- Foster, C. (2026). *vartracker* (Version 2.2.0). Zenodo. https://doi.org/10.5281/zenodo.18452274 +- Foster, C. (2026). *vartracker* (Version 2.2.1). Zenodo. https://doi.org/10.5281/zenodo.18452274 Note: the DOI above is the Zenodo concept DOI for all versions; a version-specific DOI is minted by Zenodo after each GitHub release. diff --git a/pyproject.toml b/pyproject.toml index 1b2c97a..790e3cc 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "vartracker" -version = "2.2.0" +version = "2.2.1" authors = [ {name = "Dr Charles Foster"}, ] diff --git a/tests/test_analysis.py b/tests/test_analysis.py index d7b886e..a97c007 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -133,18 +133,75 @@ def test_prepare_variant_heatmap_matrix_orders_variants_by_genome(): expected_long_label = f"nsp2:{head}+{middle}{tail}\n(G1946GT)" expected_index = [ - "nsp2:T809=\n(C1059T)", + "nsp2:T629=\n(C1059T)", expected_long_label, "S:D215G\n(A22206G)", ] assert list(matrix.index) == expected_index - assert matrix.loc["nsp2:T809=\n(C1059T)", "P0"] == 0.0 - assert matrix.loc["nsp2:T809=\n(C1059T)", "P1"] == 0.5 + assert matrix.loc["nsp2:T629=\n(C1059T)", "P0"] == 0.0 + assert matrix.loc["nsp2:T629=\n(C1059T)", "P1"] == 0.5 assert matrix.loc[expected_long_label, "P1"] == 0.8 assert matrix.loc["S:D215G\n(A22206G)", "P1"] == 1.0 +def test_prepare_variant_heatmap_matrix_normalises_starred_synonymous_label(): + table = pd.DataFrame( + [ + { + "gene": "ORF1ab", + "amino_acid_consequence": "924F", + "nsp_aa_change": "", + "type_of_change": "*synonymous", + "type_of_variant": "snp", + "alt_freq": "0.5 / 0.0", + "samples": "P0 / P1", + "variant": "C3037T", + "start": 3037, + }, + { + "gene": "ORF1ab", + "amino_acid_consequence": "924F", + "nsp_aa_change": "", + "type_of_change": "synonymous", + "type_of_variant": "snp", + "alt_freq": "0.0 / 0.6", + "samples": "P0 / P1", + "variant": "C3037T", + "start": 3037, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix(table, ["P0", "P1"], 0.2, 0.3) + + assert list(matrix.index) == ["nsp3_PLpro:F106=\n(C3037T)"] + assert matrix.loc["nsp3_PLpro:F106=\n(C3037T)", "P0"] == 0.5 + assert matrix.loc["nsp3_PLpro:F106=\n(C3037T)", "P1"] == 0.6 + + +def test_prepare_variant_heatmap_matrix_repairs_stale_stop_gained_nsp_label(): + table = pd.DataFrame( + [ + { + "gene": "ORF1ab", + "amino_acid_consequence": "L889*", + "nsp_aa_change": "nsp3_PLpro:71L", + "type_of_change": "stop_gained", + "type_of_variant": "snp", + "alt_freq": "0.052", + "samples": "P0", + "variant": "T2931A", + "start": 2931, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix(table, ["P0"], 0.0, 0.0) + + assert list(matrix.index) == ["nsp3_PLpro:L71*\n(T2931A)"] + + def test_prepare_variant_heatmap_matrix_excludes_selected_consequence_types(): table = pd.DataFrame( [ @@ -329,7 +386,7 @@ def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): "aa1_weight": "", "aa2_weight": "", "weight_difference": "", - "type_of_change": "joint_joint_frameshift", + "type_of_change": "joint_*frameshift", }, { "start": 101, diff --git a/tests/test_analysis_launcher.py b/tests/test_analysis_launcher.py index 977f832..9bdfb6c 100644 --- a/tests/test_analysis_launcher.py +++ b/tests/test_analysis_launcher.py @@ -50,6 +50,7 @@ def test_validate_lofreq_primer_rescue_on_requires_primer_bed(): min_alt_count=95, min_qual=100, max_ref_count=20, + max_minor_alt_fraction=0.05, ) @@ -63,4 +64,17 @@ def test_validate_lofreq_primer_rescue_rejects_invalid_threshold(): min_alt_count=95, min_qual=100, max_ref_count=20, + max_minor_alt_fraction=0.05, + ) + + with pytest.raises(ValueError, match="between 0 and 1"): + _validate_lofreq_primer_rescue( + "auto", + None, + min_af=0.95, + min_dp=100, + min_alt_count=95, + min_qual=100, + max_ref_count=20, + max_minor_alt_fraction=1.5, ) diff --git a/tests/test_constants.py b/tests/test_constants.py new file mode 100644 index 0000000..f6866a0 --- /dev/null +++ b/tests/test_constants.py @@ -0,0 +1,12 @@ +"""Tests for reference-coordinate formatting helpers.""" + +from __future__ import annotations + +from vartracker.constants import reformat_csq_notation + + +def test_reformat_csq_notation_preserves_orf1ab_stop_gained_nsp_change(): + reformatted, nsp_change = reformat_csq_notation("ORF1ab", "889L>889*") + + assert reformatted == "L889*" + assert nsp_change == "nsp3_PLpro:L71*" diff --git a/tests/test_lofreq_primer_rescue.py b/tests/test_lofreq_primer_rescue.py index 6822785..17476ed 100644 --- a/tests/test_lofreq_primer_rescue.py +++ b/tests/test_lofreq_primer_rescue.py @@ -1,7 +1,10 @@ import subprocess from pathlib import Path -from vartracker.lofreq_primer_rescue import rescue_lofreq_primer_variants +from vartracker.lofreq_primer_rescue import ( + lofreq_filter_with_audit, + rescue_lofreq_primer_variants, +) def test_rescue_lofreq_primer_variants_keeps_pass_and_rescues_overlap( @@ -15,7 +18,10 @@ def test_rescue_lofreq_primer_variants_keeps_pass_and_rescues_overlap( '##INFO=\n' "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" "chr1\t10\t.\tA\tG\t200\t.\tAF=0.98;DP=120;DP4=1,0,118,0\n" - "chr1\t30\t.\tC\tT\t150\t.\tAF=0.40;DP=120;DP4=30,30,30,30\n", + "chr1\t15\t.\tT\tC\t210\t.\tAF=0.98;DP=125;DP4=0,5,114,5\n" + "chr1\t18\t.\tG\tT\t220\t.\tAF=0.98;DP=125;DP4=0,4,115,5\n" + "chr1\t30\t.\tC\tT\t150\t.\tAF=0.40;DP=120;DP4=30,30,30,30\n" + "chr1\t50\t.\tT\tC\t160\t.\tAF=0.96;DP=130;DP4=2,2,60,65\n", encoding="utf-8", ) primers = tmp_path / "primers.bed" @@ -27,11 +33,14 @@ def test_rescue_lofreq_primer_variants_keeps_pass_and_rescues_overlap( '##INFO=\n' "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" "chr1\t10\t.\tA\tG\t200\tstrandbias\tAF=0.98;DP=120;DP4=1,0,118,0\n" + "chr1\t15\t.\tT\tC\t210\tsb_fdr\tAF=0.98;DP=125;DP4=0,5,114,5\n" + "chr1\t18\t.\tG\tT\t220\tlowqual\tAF=0.98;DP=125;DP4=0,4,115,5\n" "chr1\t30\t.\tC\tT\t150\tPASS\tAF=0.40;DP=120;DP4=30,30,30,30\n" + "chr1\t50\t.\tT\tC\t160\tsb_fdr\tAF=0.96;DP=130;DP4=2,2,60,65\n" ) def fake_run(cmd, check): - assert cmd[:4] == ["lofreq", "filter", "-i", str(raw_vcf)] + assert cmd[:5] == ["lofreq", "filter", "--print-all", "-i", str(raw_vcf)] assert check is True output_path = Path(cmd[cmd.index("-o") + 1]) output_path.write_text(default_filtered_vcf, encoding="utf-8") @@ -41,21 +50,26 @@ def fake_run(cmd, check): output_vcf = tmp_path / "final.vcf" rescued_tsv = tmp_path / "rescued.tsv" + filtered_out_tsv = tmp_path / "filtered_out.tsv" result = rescue_lofreq_primer_variants( raw_vcf=raw_vcf, primers_bed=primers, output_vcf=output_vcf, rescued_tsv=rescued_tsv, + filtered_out_tsv=filtered_out_tsv, ) assert result.normal_passed == 1 - assert result.rescued == 1 - assert result.discarded == 0 + assert result.rescued == 2 + assert result.discarded == 2 + assert result.filtered_out == 2 output_text = output_vcf.read_text(encoding="utf-8") assert "##FILTER=\n' + '##INFO=\n' + '##INFO=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" + "chr1\t10\t.\tA\tG\t200\t.\tAF=0.98;DP=120;DP4=1,0,118,0\n" + "chr1\t50\t.\tT\tC\t160\t.\tAF=0.96;DP=130;DP4=2,2,60,65\n", + encoding="utf-8", + ) + default_filtered_vcf = ( + "##fileformat=VCFv4.2\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n" + "chr1\t10\t.\tA\tG\t200\tPASS\tAF=0.98;DP=120;DP4=1,0,118,0\n" + "chr1\t50\t.\tT\tC\t160\tsb_fdr\tAF=0.96;DP=130;DP4=2,2,60,65\n" + ) + + def fake_run(cmd, check): + assert cmd[:5] == ["lofreq", "filter", "--print-all", "-i", str(raw_vcf)] + assert check is True + output_path = Path(cmd[cmd.index("-o") + 1]) + output_path.write_text(default_filtered_vcf, encoding="utf-8") + return subprocess.CompletedProcess(cmd, 0) + + monkeypatch.setattr("vartracker.lofreq_primer_rescue.subprocess.run", fake_run) + + output_vcf = tmp_path / "final.vcf" + filtered_out_tsv = tmp_path / "filtered_out.tsv" + result = lofreq_filter_with_audit( + raw_vcf=raw_vcf, + output_vcf=output_vcf, + filtered_out_tsv=filtered_out_tsv, + ) + + assert result.normal_passed == 1 + assert result.filtered_out == 1 + output_text = output_vcf.read_text(encoding="utf-8") + assert "chr1\t10\t.\tA\tG\t200\tPASS" in output_text + assert "chr1\t50\t.\tT\tC" not in output_text + filtered_out_lines = filtered_out_tsv.read_text(encoding="utf-8").splitlines() + assert filtered_out_lines[0] == "variant\treason_filtered\tmetrics" + assert filtered_out_lines[1].startswith("T50C\tsb_fdr\tAF=0.96") diff --git a/tests/test_main.py b/tests/test_main.py index b2ae61b..3885c39 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -54,12 +54,15 @@ def test_bam_parser_accepts_lofreq_primer_rescue_options(): "off", "--lofreq-rescue-min-af", "0.9", + "--lofreq-rescue-max-minor-alt-fraction", + "0.1", ] ) assert args.primer_bed == "primers.bed" assert args.lofreq_primer_rescue == "off" assert args.lofreq_rescue_min_af == 0.9 + assert args.lofreq_rescue_max_minor_alt_fraction == 0.1 def test_drop_exact_duplicate_result_rows_removes_only_exact_duplicates(capsys): @@ -1072,6 +1075,7 @@ def fake_vcf(args): assert recorded["workflow_kwargs"]["ampliconclip_tolerance"] == 2 assert recorded["workflow_kwargs"]["lofreq_primer_rescue"] == "auto" assert recorded["workflow_kwargs"]["lofreq_rescue_min_af"] == 0.95 + assert recorded["workflow_kwargs"]["lofreq_rescue_max_minor_alt_fraction"] == 0.05 assert recorded["vcf_input"] == str(updated_csv) assert modes_checked == ["e2e"] diff --git a/tests/test_snakemake_workflow.py b/tests/test_snakemake_workflow.py index b100cb9..a29e3b3 100644 --- a/tests/test_snakemake_workflow.py +++ b/tests/test_snakemake_workflow.py @@ -33,7 +33,23 @@ def test_snakemake_rules_write_logs_under_outdir(): assert "--tolerance {params.tolerance}" in snakefile assert "--no-default-filter" in snakefile assert "_variants.raw.vcf.gz" in snakefile + assert "_variants.raw.vcf.gz.tbi" in snakefile + assert ( + 'vcf_raw = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz")' + not in snakefile + ) + assert ( + 'vcf_raw = f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz"' in snakefile + ) + assert ( + 'vcf_raw_tbi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz.tbi"' + in snakefile + ) + assert "_variants.filtered_out.tsv" in snakefile assert "LOFREQ_PRIMER_RESCUE_ENABLED" in snakefile + assert "LOFREQ_RESCUE_MAX_MINOR_ALT_FRACTION" in snakefile + assert "max_minor_alt_fraction=LOFREQ_RESCUE_MAX_MINOR_ALT_FRACTION" in snakefile + assert "lofreq_filter_with_audit" in snakefile assert "rescue_lofreq_primer_variants" in snakefile assert "_variants.rescued.tsv" in snakefile assert "_validate_primer_bed_reference(PRIMER_BED, REF)" in snakefile @@ -41,4 +57,6 @@ def test_snakemake_rules_write_logs_under_outdir(): assert "_iupac_consensus.fasta" in snakefile assert "df['consensus']" in snakefile assert "df['iupac_consensus']" in snakefile + assert "df['raw_vcf']" in snakefile assert "df['lofreq_rescued_tsv']" in snakefile + assert "df['lofreq_filtered_out_tsv']" in snakefile diff --git a/tests/test_vcf_processing.py b/tests/test_vcf_processing.py index d4f589a..6816f40 100644 --- a/tests/test_vcf_processing.py +++ b/tests/test_vcf_processing.py @@ -401,8 +401,8 @@ def test_process_vcf_splits_sample_specific_bcsq_annotations(tmp_path): "start": 5, "amino_acid_consequence": "K2T", "bcsq_aa_notation": "2K>2T", - "presence_absence": "N / Y", - "alt_freq": ". / 0.100", + "presence_absence": "Y / Y", + "alt_freq": "0.100 / 0.100", }, ] ) @@ -410,6 +410,85 @@ def test_process_vcf_splits_sample_specific_bcsq_annotations(tmp_path): pd.testing.assert_frame_equal(observed, expected) +def test_process_vcf_merges_starred_and_unstarred_equivalent_bcsq(tmp_path): + vcf_path = tmp_path / "annotated.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\ts2\n" + "chr1\t4\t.\tA\tG\t.\tPASS\tDP=100;AF=0.1;" + "BCSQ=*synonymous|GENE1|tx|protein_coding|+|2K|4A>G," + "synonymous|GENE1|tx|protein_coding|+|2K|4A>G" + "\tGT:DP:AF:BCSQ\t1:100:0.1:1\t1:100:0.2:4\n", + encoding="utf-8", + ) + + cov1 = tmp_path / "s1.depth.txt" + cov2 = tmp_path / "s2.depth.txt" + _write_depth_file(cov1) + _write_depth_file(cov2) + + table = process_vcf(str(vcf_path), [str(cov1), str(cov2)], 10, ["s1", "s2"]) + + assert len(table) == 1 + row = table.iloc[0] + assert row["variant"] == "A4G" + assert row["type_of_change"] == "synonymous" + assert row["presence_absence"] == "Y / Y" + assert row["alt_freq"] == "0.100 / 0.200" + assert row["variant_status"] == "original" + assert row["persistence_status"] == "original_retained" + + +def test_process_vcf_keeps_single_site_variant_present_when_sample_has_joint_csq( + tmp_path, +): + vcf_path = tmp_path / "annotated.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\ts2\n" + "chr1\t4\t.\tA\tG\t.\tPASS\tDP=100;AF=0.2;" + "BCSQ=missense|GENE1|tx|protein_coding|+|2K>2A|4A>G," + "frameshift|GENE1|tx|protein_coding|+|2KAAAAAAAAAAAA>2K|4A>G+5A>C" + "\tGT:DP:AF:BCSQ\t1:100:0.1:1\t1:100:0.2:4\n", + encoding="utf-8", + ) + + cov1 = tmp_path / "s1.depth.txt" + cov2 = tmp_path / "s2.depth.txt" + _write_depth_file(cov1) + _write_depth_file(cov2) + + table = process_vcf(str(vcf_path), [str(cov1), str(cov2)], 10, ["s1", "s2"]) + simple = table[table["bcsq_nt_notation"].eq("4A>G")].iloc[0] + joint = table[table["bcsq_nt_notation"].str.contains("\\+", regex=True)].iloc[0] + + assert simple["type_of_change"] == "missense" + assert simple["presence_absence"] == "Y / Y" + assert simple["alt_freq"] == "0.100 / 0.200" + assert simple["variant_status"] == "original" + assert simple["persistence_status"] == "original_retained" + + assert joint["type_of_change"] == "frameshift" + assert joint["presence_absence"] == "N / Y" + assert joint["alt_freq"] == ". / 0.200" + + @pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") def test_merge_then_annotate_preserves_joint_annotations_across_samples(tmp_path): ref, gff = _write_minimal_reference_bundle(tmp_path) @@ -525,7 +604,7 @@ def test_merge_then_annotate_preserves_joint_annotations_across_samples(tmp_path "start": 5, "variant": "A5C", "amino_acid_consequence": "K2T", - "presence_absence": "N / Y", + "presence_absence": "Y / Y", "type_of_change": "missense", "joint_variant": False, }, diff --git a/vartracker/Snakefile b/vartracker/Snakefile index 1da3c2e..d5700ed 100644 --- a/vartracker/Snakefile +++ b/vartracker/Snakefile @@ -4,12 +4,14 @@ try: from vartracker.consensus import consensus_genotype_for_variant from vartracker.lofreq_primer_rescue import ( PrimerRescueThresholds, + lofreq_filter_with_audit, rescue_lofreq_primer_variants, ) except ModuleNotFoundError: from consensus import consensus_genotype_for_variant from lofreq_primer_rescue import ( PrimerRescueThresholds, + lofreq_filter_with_audit, rescue_lofreq_primer_variants, ) @@ -33,6 +35,9 @@ LOFREQ_RESCUE_MIN_DP = int(config.get("lofreq_rescue_min_dp", 100)) LOFREQ_RESCUE_MIN_ALT_COUNT = int(config.get("lofreq_rescue_min_alt_count", 95)) LOFREQ_RESCUE_MIN_QUAL = float(config.get("lofreq_rescue_min_qual", 100.0)) LOFREQ_RESCUE_MAX_REF_COUNT = int(config.get("lofreq_rescue_max_ref_count", 20)) +LOFREQ_RESCUE_MAX_MINOR_ALT_FRACTION = float( + config.get("lofreq_rescue_max_minor_alt_fraction", 0.05) +) if LOFREQ_PRIMER_RESCUE not in {"auto", "on", "off"}: raise ValueError("lofreq_primer_rescue must be one of: auto, on, off") @@ -120,6 +125,10 @@ rule all: input: expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz.tbi", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.rescued.tsv", sample=SAMPLES.keys()), + expand(f"{OUTDIR}/{{sample}}/{{sample}}_variants.filtered_out.tsv", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_depth.txt", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_consensus.fasta", sample=SAMPLES.keys()), expand(f"{OUTDIR}/{{sample}}/{{sample}}_iupac_consensus.fasta", sample=SAMPLES.keys()), @@ -279,9 +288,11 @@ rule lofreq_call: bam = f"{OUTDIR}/{{sample}}/{{sample}}_aligned.indelqual.bam", ref = REF output: - vcf_raw = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz"), + vcf_raw = f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz", + vcf_raw_tbi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.raw.vcf.gz.tbi", vcf_filtered = temp(f"{OUTDIR}/{{sample}}/{{sample}}_variants.filtered.vcf"), rescued_tsv = f"{OUTDIR}/{{sample}}/{{sample}}_variants.rescued.tsv", + filtered_out_tsv = f"{OUTDIR}/{{sample}}/{{sample}}_variants.filtered_out.tsv", vcf = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz", csi = f"{OUTDIR}/{{sample}}/{{sample}}_variants.vcf.gz.csi" log: @@ -304,12 +315,14 @@ rule lofreq_call: min_alt_count=LOFREQ_RESCUE_MIN_ALT_COUNT, min_qual=LOFREQ_RESCUE_MIN_QUAL, max_ref_count=LOFREQ_RESCUE_MAX_REF_COUNT, + max_minor_alt_fraction=LOFREQ_RESCUE_MAX_MINOR_ALT_FRACTION, ) result = rescue_lofreq_primer_variants( raw_vcf=str(output.vcf_raw), primers_bed=PRIMER_BED, output_vcf=str(output.vcf_filtered), rescued_tsv=str(output.rescued_tsv), + filtered_out_tsv=str(output.filtered_out_tsv), tmp_dir=str(log_path.parent), thresholds=thresholds, ) @@ -317,17 +330,24 @@ rule lofreq_call: handle.write( "lofreq primer rescue: " f"normal_passed={result.normal_passed} " - f"rescued={result.rescued} discarded={result.discarded}\n" + f"rescued={result.rescued} discarded={result.discarded} " + f"filtered_out={result.filtered_out}\n" ) else: - shell( - "lofreq filter -i {output.vcf_raw} " - "-o {output.vcf_filtered} 2>> {log}" + result = lofreq_filter_with_audit( + raw_vcf=str(output.vcf_raw), + output_vcf=str(output.vcf_filtered), + filtered_out_tsv=str(output.filtered_out_tsv), + tmp_dir=str(log_path.parent), ) with open(output.rescued_tsv, "w", encoding="utf-8") as handle: handle.write("variant\treason_filtered\treason_rescued\tmetrics\n") with log_path.open("a", encoding="utf-8") as handle: - handle.write("lofreq primer rescue: disabled\n") + handle.write( + "lofreq primer rescue: disabled " + f"normal_passed={result.normal_passed} " + f"filtered_out={result.filtered_out}\n" + ) shell("bgzip -c {output.vcf_filtered} > {output.vcf} 2>> {log}") shell("bcftools index -f {output.vcf} 2>> {log}") @@ -597,12 +617,18 @@ rule update_csv: df['vcf'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_variants.vcf.gz") ) + df['raw_vcf'] = df['sample_name'].apply( + lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_variants.raw.vcf.gz") + ) df['coverage'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_depth.txt") ) df['lofreq_rescued_tsv'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_variants.rescued.tsv") ) + df['lofreq_filtered_out_tsv'] = df['sample_name'].apply( + lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_variants.filtered_out.tsv") + ) df['consensus'] = df['sample_name'].apply( lambda x: os.path.abspath(f"{params.outdir}/{x}/{x}_consensus.fasta") ) diff --git a/vartracker/_version.py b/vartracker/_version.py index 4e0c032..11e633a 100644 --- a/vartracker/_version.py +++ b/vartracker/_version.py @@ -9,7 +9,7 @@ # NOTE: When bumping the project version remember to update this fallback value # alongside the version declared in pyproject.toml. -_FALLBACK_VERSION = "2.2.0" +_FALLBACK_VERSION = "2.2.1" try: __version__ = metadata.version("vartracker") diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 9a395ad..e895d93 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -34,7 +34,7 @@ def _ensure_joint_prefix(change_type: object) -> str: text = str(change_type or "").strip() if not text: return "joint" - normalised = re.sub(r"^(joint_)+", "", text) + normalised = re.sub(r"^(joint_)+", "", text).lstrip("*") return f"joint_{normalised}" if normalised else "joint" @@ -493,20 +493,35 @@ def _build_gene_order_map( return gene_order_map, ordered -def _map_orf1ab_position_to_nsp(aa_position: int) -> str: - """Map an ORF1ab amino acid position to its NSP name.""" +def _map_orf1ab_position_to_nsp_info(aa_position: int) -> tuple[str, int]: + """Map an ORF1ab amino acid position to its NSP name and local position.""" nsp_products = cast(Sequence[str], NSPS.get("product", [])) aa_starts = cast(Sequence[int], NSPS.get("aa_start", [])) - for name, start in zip(nsp_products, aa_starts): - length = NSP_LENGTHS.get(str(name)) - if length is None: - continue - end = start + length - 1 + for idx, (name, start) in enumerate(zip(nsp_products, aa_starts)): + end = aa_starts[idx + 1] - 1 if idx + 1 < len(aa_starts) else 999999 if start <= aa_position <= end: - return name + return str(name), aa_position - start + 1 + + return "ORF1ab", aa_position + + +def _map_orf1ab_position_to_nsp(aa_position: int) -> str: + """Map an ORF1ab amino acid position to its NSP name.""" + return _map_orf1ab_position_to_nsp_info(aa_position)[0] + - return "ORF1ab" +def _format_orf1ab_change_as_nsp(amino_change: object) -> tuple[str, str] | None: + """Convert an ORF1ab amino-acid label to NSP-local notation when possible.""" + text = str(amino_change or "").strip() + match = re.fullmatch(r"([A-Za-z*]*)(\d+)([A-Za-z*]*)", text) + if not match: + return None + ref = match.group(1) + aa_position = int(match.group(2)) + alt = match.group(3) + gene_label, local_position = _map_orf1ab_position_to_nsp_info(aa_position) + return gene_label, f"{ref}{local_position}{alt}" def _extract_numeric_position(value: str) -> Optional[int]: @@ -547,7 +562,8 @@ def _resolve_variant_labels(row) -> Tuple[str, str, str]: gene = getattr(row, "gene", "") amino_change = getattr(row, "amino_acid_consequence", "") nsp_change = getattr(row, "nsp_aa_change", "") - change_type = str(getattr(row, "type_of_change", "")) + change_type = str(getattr(row, "type_of_change", "")).lstrip("*") + change_type = re.sub(r"^(joint_)+", "", change_type) gene_label = gene aa_label = amino_change @@ -559,10 +575,14 @@ def _resolve_variant_labels(row) -> Tuple[str, str, str]: gene_label = gene_part if aa_part: aa_label = aa_part + if "*" in str(amino_change) and "*" not in str(aa_label): + mapped_change = _format_orf1ab_change_as_nsp(amino_change) + if mapped_change is not None: + gene_label, aa_label = mapped_change else: - position = _extract_numeric_position(amino_change) - if position is not None: - gene_label = _map_orf1ab_position_to_nsp(position) + mapped_change = _format_orf1ab_change_as_nsp(amino_change) + if mapped_change is not None: + gene_label, aa_label = mapped_change if not aa_label or str(aa_label) in {"", "None"}: if isinstance(nsp_change, str) and ":" in nsp_change: @@ -710,7 +730,7 @@ def _prepare_variant_heatmap_matrix( ) records: List[Dict[str, Union[str, float, int]]] = [] - seen_labels = set() + record_index_by_label: dict[str, int] = {} qc_maps: dict[str, dict[str, str]] = {} for row in table.itertuples(index=False): @@ -729,10 +749,9 @@ def _prepare_variant_heatmap_matrix( gene_label, display_label, base_label = _resolve_variant_labels(row) - if base_label in seen_labels: - continue - - change_type = str(getattr(row, "type_of_change", "")).strip().lower() + change_type = ( + str(getattr(row, "type_of_change", "")).strip().lower().lstrip("*") + ) if not include_joint and change_type.startswith("joint"): continue if included_patterns and not any( @@ -842,8 +861,33 @@ def _prepare_variant_heatmap_matrix( for sample, value in zip(ordered_samples, row_values): record[sample] = value + if display_label in record_index_by_label: + existing = records[record_index_by_label[display_label]] + existing["gene_order"] = min( + int(existing.get("gene_order", record["gene_order"])), + int(record["gene_order"]), + ) + existing["start"] = min( + int(existing.get("start", record["start"])), + int(record["start"]), + ) + existing_qc = qc_maps.setdefault(display_label, {}) + for sample, value in zip(ordered_samples, row_values): + current_value = existing.get(sample, 0.0) + current_frequency = ( + float(current_value) + if isinstance(current_value, (int, float)) + else _coerce_frequency(str(current_value)) + ) + if value > current_frequency: + existing[sample] = value + existing_qc[sample] = row_qc_map.get(sample, "") + elif sample not in existing_qc: + existing_qc[sample] = row_qc_map.get(sample, "") + continue + + record_index_by_label[display_label] = len(records) records.append(record) - seen_labels.add(base_label) qc_maps[display_label] = { sample: row_qc_map.get(sample, "") for sample in ordered_samples } diff --git a/vartracker/analysis_launcher.py b/vartracker/analysis_launcher.py index 791fe68..28d3e15 100644 --- a/vartracker/analysis_launcher.py +++ b/vartracker/analysis_launcher.py @@ -75,6 +75,7 @@ def _validate_lofreq_primer_rescue( min_alt_count: int, min_qual: float, max_ref_count: int, + max_minor_alt_fraction: float, ) -> None: if mode not in {"auto", "on", "off"}: raise ValueError("lofreq_primer_rescue must be one of: auto, on, off") @@ -90,6 +91,8 @@ def _validate_lofreq_primer_rescue( raise ValueError("lofreq_rescue_min_qual must be >= 0") if max_ref_count < 0: raise ValueError("lofreq_rescue_max_ref_count must be >= 0") + if not 0 <= max_minor_alt_fraction <= 1: + raise ValueError("lofreq_rescue_max_minor_alt_fraction must be between 0 and 1") def run_workflow( @@ -114,6 +117,7 @@ def run_workflow( lofreq_rescue_min_alt_count: int = 95, lofreq_rescue_min_qual: float = 100.0, lofreq_rescue_max_ref_count: int = 20, + lofreq_rescue_max_minor_alt_fraction: float = 0.05, ) -> Optional[str]: """Run the lofreq variant calling workflow via the Snakemake API. @@ -146,6 +150,7 @@ def run_workflow( lofreq_rescue_min_alt_count, lofreq_rescue_min_qual, lofreq_rescue_max_ref_count, + lofreq_rescue_max_minor_alt_fraction, ) Path(outdir).mkdir(parents=True, exist_ok=True) @@ -167,6 +172,7 @@ def run_workflow( "lofreq_rescue_min_alt_count": lofreq_rescue_min_alt_count, "lofreq_rescue_min_qual": lofreq_rescue_min_qual, "lofreq_rescue_max_ref_count": lofreq_rescue_max_ref_count, + "lofreq_rescue_max_minor_alt_fraction": (lofreq_rescue_max_minor_alt_fraction), } if primer_bed_path: config_dict["primer_bed"] = primer_bed_path diff --git a/vartracker/constants.py b/vartracker/constants.py index 172314e..3fb663f 100644 --- a/vartracker/constants.py +++ b/vartracker/constants.py @@ -145,24 +145,16 @@ def bcf_orf1ab_to_nsp(mutation): Raises: ValueError: If mutation format is invalid """ - # Parse out important bits from mutation + # Parse out important bits from mutation. Stop codons are represented as "*" + # and must be retained in both ORF1ab and NSP notation. no_gene = re.sub(".*:", "", mutation) - parsed = re.findall(r"([A-Z]+)", no_gene) - - if len(parsed) == 2: - ref = parsed[0] - alt = parsed[1] - elif len(parsed) == 1: - ref = "" - alt = parsed[0] - else: + parsed = re.fullmatch(r"([A-Z*]*)(\d+)([A-Z*]*)", no_gene) + if not parsed: raise ValueError(f"Invalid mutation format: {mutation}") - pos_match = re.search(r"(\d+)", no_gene) - if not pos_match: - raise ValueError(f"No position found in mutation: {mutation}") - - pos = int(pos_match.group()) + ref = parsed.group(1) + pos = int(parsed.group(2)) + alt = parsed.group(3) # Find corresponding nsp idx = bisect_right(NSPS["aa_start"], pos) - 1 @@ -264,9 +256,18 @@ def reformat_csq_notation(gene, string): if splitter == -1: return (string, "") - ref = re.sub("[0-9]", "", string[:splitter]) - pos = re.sub("[A-Za-z]", "", string[:splitter]) - alt = re.sub("[0-9]", "", string[splitter + 1 :]) + left = string[:splitter] + right = string[splitter + 1 :] + left_match = re.fullmatch(r"(\d+)([A-Za-z*]+)", left) + right_match = re.fullmatch(r"(\d+)([A-Za-z*]+)", right) + if left_match and right_match: + ref = left_match.group(2) + pos = left_match.group(1) + alt = right_match.group(2) + else: + ref = re.sub("[0-9]", "", left) + pos = re.sub("[A-Za-z*]", "", left) + alt = re.sub("[0-9]", "", right) reformatted = ref + pos + alt if gene == "ORF1ab": diff --git a/vartracker/lofreq_primer_rescue.py b/vartracker/lofreq_primer_rescue.py index d0a80cb..0cbaaec 100644 --- a/vartracker/lofreq_primer_rescue.py +++ b/vartracker/lofreq_primer_rescue.py @@ -35,6 +35,7 @@ class PrimerRescueThresholds: min_alt_count: int = 95 min_qual: float = 100.0 max_ref_count: int = 20 + max_minor_alt_fraction: float = 0.05 @dataclass(frozen=True) @@ -45,6 +46,17 @@ class PrimerRescueResult: rescued: int discarded: int rescued_tsv: str + filtered_out: int = 0 + filtered_out_tsv: str | None = None + + +@dataclass(frozen=True) +class LofreqFilterAuditResult: + """Summary of default LoFreq filtering with an audit table.""" + + normal_passed: int + filtered_out: int + filtered_out_tsv: str def _open_text(path: str | Path, mode: str = "rt") -> TextIO: @@ -126,6 +138,13 @@ def _filter_is_pass(fields: list[str]) -> bool: return fields[6] in ("PASS", ".") +def _is_strand_bias_filter(reason: str) -> bool: + return any( + item.lower() in {"sb_fdr", "strandbias", "strand_bias"} + for item in reason.split(";") + ) + + def read_default_filter_results( path: str | Path, ) -> tuple[set[tuple[str, str, str, str]], dict[tuple[str, str, str, str], str]]: @@ -150,9 +169,13 @@ def rescue_metrics( fields: list[str], primers: dict[str, list[tuple[int, int]]], thresholds: PrimerRescueThresholds, + reason_filtered: str, ) -> str | None: """Return rescue metrics for a candidate record, or ``None`` if it fails.""" + if not _is_strand_bias_filter(reason_filtered): + return None + chrom = fields[0] pos = int(fields[1]) ref = fields[3] @@ -206,17 +229,52 @@ def rescue_metrics( return None if ref_count > thresholds.max_ref_count: return None - if not (alt_fwd == 0 or alt_rev == 0): + minor_alt_fraction = min(alt_fwd, alt_rev) / alt_count + if minor_alt_fraction > thresholds.max_minor_alt_fraction: return None dp4_af = alt_count / dp4_total return ( f"AF={af:g},DP={dp},QUAL={qual:g}," f"DP4={ref_fwd}/{ref_rev}/{alt_fwd}/{alt_rev}," - f"alt_count={alt_count},ref_count={ref_count},dp4_af={dp4_af:g}" + f"alt_count={alt_count},ref_count={ref_count},dp4_af={dp4_af:g}," + f"minor_alt_fraction={minor_alt_fraction:g}" ) +def variant_metrics(fields: list[str]) -> str: + """Return concise metrics for a raw LoFreq record.""" + if len(fields) < 8: + return "." + + info = parse_info(fields[7]) + metrics = [] + if "AF" in info: + try: + metrics.append(f"AF={_first_float(info['AF']):g}") + except ValueError: + metrics.append(f"AF={info['AF']}") + if "DP" in info: + try: + metrics.append(f"DP={_first_int(info['DP'])}") + except ValueError: + metrics.append(f"DP={info['DP']}") + if fields[5] != ".": + metrics.append(f"QUAL={fields[5]}") + counts = _dp4_counts(info.get("DP4", "")) + if counts is not None: + ref_fwd, ref_rev, alt_fwd, alt_rev = counts + ref_count = ref_fwd + ref_rev + alt_count = alt_fwd + alt_rev + total = ref_count + alt_count + metrics.append(f"DP4={ref_fwd}/{ref_rev}/{alt_fwd}/{alt_rev}") + metrics.append(f"alt_count={alt_count}") + metrics.append(f"ref_count={ref_count}") + if total: + metrics.append(f"dp4_af={alt_count / total:g}") + return ",".join(metrics) if metrics else "." + + def _add_info_flag(info_text: str, flag: str) -> str: if info_text in ("", "."): return flag @@ -284,71 +342,186 @@ def _write_headers(headers: list[str], output: TextIO) -> None: output.write(line) -def write_final_vcf( +def _run_lofreq_filter_print_all( + raw_vcf: str | Path, + tmp_dir: str | Path | None = None, + lofreq: str = "lofreq", +) -> tuple[set[tuple[str, str, str, str]], dict[tuple[str, str, str, str], str]]: + with tempfile.TemporaryDirectory(dir=tmp_dir) as run_tmp_dir: + filtered_vcf = os.path.join(run_tmp_dir, "lofreq.default_filtered.vcf") + subprocess.run( + [lofreq, "filter", "--print-all", "-i", str(raw_vcf), "-o", filtered_vcf], + check=True, + ) + return read_default_filter_results(filtered_vcf) + + +def write_default_filtered_vcf( raw_vcf: str | Path, output_vcf: str | Path, - rescued_tsv_path: str | Path, + filtered_out_tsv_path: str | Path, pass_keys: set[tuple[str, str, str, str]], filter_reasons: dict[tuple[str, str, str, str], str], - primers: dict[str, list[tuple[int, int]]], - thresholds: PrimerRescueThresholds, -) -> PrimerRescueResult: - """Write the default-filtered VCF plus any rescued primer-overlap records.""" +) -> LofreqFilterAuditResult: + """Write normal LoFreq PASS records plus a filtered-out audit table.""" normal_passed = 0 - rescued = 0 - discarded = 0 + filtered_out = 0 headers: list[str] = [] with ( _open_text(raw_vcf) as raw, _open_text(output_vcf, "wt") as output, - Path(rescued_tsv_path).open("w", encoding="utf-8") as rescued_tsv, + Path(filtered_out_tsv_path).open("w", encoding="utf-8") as filtered_out_tsv, ): - rescued_tsv.write("variant\treason_filtered\treason_rescued\tmetrics\n") - + filtered_out_tsv.write("variant\treason_filtered\tmetrics\n") for line in raw: if line.startswith("#"): headers.append(line) continue if headers: - _write_headers(headers, output) + output.writelines(headers) headers = [] fields = line.rstrip("\n").split("\t") if len(fields) < 8: - discarded += 1 continue key = _record_key(fields) if key in pass_keys: output.write("\t".join(_pass_fields(fields)) + "\n") normal_passed += 1 - else: - metrics = rescue_metrics(fields, primers, thresholds) - if metrics is None: + continue + + reason_filtered = filter_reasons.get( + key, "not_passed_default_lofreq_filter" + ) + filtered_out_tsv.write( + f"{_variant_name(fields)}\t{reason_filtered}\t" + f"{variant_metrics(fields)}\n" + ) + filtered_out += 1 + + if headers: + output.writelines(headers) + + return LofreqFilterAuditResult( + normal_passed=normal_passed, + filtered_out=filtered_out, + filtered_out_tsv=str(filtered_out_tsv_path), + ) + + +def write_final_vcf( + raw_vcf: str | Path, + output_vcf: str | Path, + rescued_tsv_path: str | Path, + filtered_out_tsv_path: str | Path | None, + pass_keys: set[tuple[str, str, str, str]], + filter_reasons: dict[tuple[str, str, str, str], str], + primers: dict[str, list[tuple[int, int]]], + thresholds: PrimerRescueThresholds, +) -> PrimerRescueResult: + """Write the default-filtered VCF plus any rescued primer-overlap records.""" + + normal_passed = 0 + rescued = 0 + discarded = 0 + filtered_out = 0 + headers: list[str] = [] + + with ( + _open_text(raw_vcf) as raw, + _open_text(output_vcf, "wt") as output, + Path(rescued_tsv_path).open("w", encoding="utf-8") as rescued_tsv, + ): + rescued_tsv.write("variant\treason_filtered\treason_rescued\tmetrics\n") + filtered_out_tsv = ( + Path(filtered_out_tsv_path).open("w", encoding="utf-8") + if filtered_out_tsv_path is not None + else None + ) + if filtered_out_tsv is not None: + filtered_out_tsv.write("variant\treason_filtered\tmetrics\n") + + try: + for line in raw: + if line.startswith("#"): + headers.append(line) + continue + + if headers: + _write_headers(headers, output) + headers = [] + + fields = line.rstrip("\n").split("\t") + if len(fields) < 8: discarded += 1 continue - output.write("\t".join(_rescued_fields(fields)) + "\n") - reason_filtered = filter_reasons.get( - key, "not_passed_default_lofreq_filter" - ) - rescued_tsv.write( - f"{_variant_name(fields)}\t{reason_filtered}\t" - f"overlap_primer_interval\t{metrics}\n" - ) - rescued += 1 + key = _record_key(fields) + if key in pass_keys: + output.write("\t".join(_pass_fields(fields)) + "\n") + normal_passed += 1 + else: + reason_filtered = filter_reasons.get( + key, "not_passed_default_lofreq_filter" + ) + metrics = rescue_metrics( + fields, primers, thresholds, reason_filtered + ) + if metrics is None: + discarded += 1 + if filtered_out_tsv is not None: + filtered_out_tsv.write( + f"{_variant_name(fields)}\t{reason_filtered}\t" + f"{variant_metrics(fields)}\n" + ) + filtered_out += 1 + continue + + output.write("\t".join(_rescued_fields(fields)) + "\n") + rescued_tsv.write( + f"{_variant_name(fields)}\t{reason_filtered}\t" + f"overlap_primer_interval\t{metrics}\n" + ) + rescued += 1 - if headers: - _write_headers(headers, output) + if headers: + _write_headers(headers, output) + finally: + if filtered_out_tsv is not None: + filtered_out_tsv.close() return PrimerRescueResult( normal_passed=normal_passed, rescued=rescued, discarded=discarded, rescued_tsv=str(rescued_tsv_path), + filtered_out=filtered_out, + filtered_out_tsv=( + str(filtered_out_tsv_path) if filtered_out_tsv_path is not None else None + ), + ) + + +def lofreq_filter_with_audit( + raw_vcf: str | Path, + output_vcf: str | Path, + filtered_out_tsv: str | Path, + tmp_dir: str | Path | None = None, + lofreq: str = "lofreq", +) -> LofreqFilterAuditResult: + """Apply default LoFreq filtering and write filtered-out variant details.""" + + pass_keys, filter_reasons = _run_lofreq_filter_print_all(raw_vcf, tmp_dir, lofreq) + return write_default_filtered_vcf( + raw_vcf, + output_vcf, + filtered_out_tsv, + pass_keys, + filter_reasons, ) @@ -357,6 +530,7 @@ def rescue_lofreq_primer_variants( primers_bed: str | Path, output_vcf: str | Path, rescued_tsv: str | Path | None = None, + filtered_out_tsv: str | Path | None = None, tmp_dir: str | Path | None = None, lofreq: str = "lofreq", thresholds: PrimerRescueThresholds | None = None, @@ -366,21 +540,18 @@ def rescue_lofreq_primer_variants( thresholds = thresholds or PrimerRescueThresholds() if rescued_tsv is None: rescued_tsv = f"{output_vcf}.rescued.tsv" + if filtered_out_tsv is None: + filtered_out_tsv = f"{output_vcf}.filtered_out.tsv" primers = load_primers(primers_bed) - with tempfile.TemporaryDirectory(dir=tmp_dir) as run_tmp_dir: - filtered_vcf = os.path.join(run_tmp_dir, "lofreq.default_filtered.vcf") - subprocess.run( - [lofreq, "filter", "-i", str(raw_vcf), "-o", filtered_vcf], - check=True, - ) - pass_keys, filter_reasons = read_default_filter_results(filtered_vcf) + pass_keys, filter_reasons = _run_lofreq_filter_print_all(raw_vcf, tmp_dir, lofreq) return write_final_vcf( raw_vcf, output_vcf, rescued_tsv, + filtered_out_tsv, pass_keys, filter_reasons, primers, @@ -404,6 +575,10 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "--rescued-tsv", help="TSV of rescued variants; default: OUTPUT_VCF.rescued.tsv", ) + parser.add_argument( + "--filtered-out-tsv", + help="TSV of raw variants retained by LoFreq calling but filtered from final VCF", + ) parser.add_argument("--tmp-dir", help="Optional temporary directory") parser.add_argument("--lofreq", default="lofreq", help="lofreq executable") parser.add_argument( @@ -436,6 +611,15 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: default=PrimerRescueThresholds.max_ref_count, help="Maximum DP4 ref count for rescue candidates only (default: 20)", ) + parser.add_argument( + "--max-minor-alt-fraction", + type=float, + default=PrimerRescueThresholds.max_minor_alt_fraction, + help=( + "Maximum minor ALT strand fraction for rescue candidates only " + "(default: 0.05)" + ), + ) return parser.parse_args(argv) @@ -447,19 +631,21 @@ def main(argv: list[str] | None = None) -> int: min_alt_count=args.min_alt_count, min_qual=args.min_qual, max_ref_count=args.max_ref_count, + max_minor_alt_fraction=args.max_minor_alt_fraction, ) result = rescue_lofreq_primer_variants( raw_vcf=args.raw_vcf, primers_bed=args.primers_bed, output_vcf=args.output_vcf, rescued_tsv=args.rescued_tsv, + filtered_out_tsv=args.filtered_out_tsv, tmp_dir=args.tmp_dir, lofreq=args.lofreq, thresholds=thresholds, ) print( f"normal_passed={result.normal_passed} rescued={result.rescued} " - f"discarded={result.discarded}", + f"discarded={result.discarded} filtered_out={result.filtered_out}", file=sys.stderr, ) return 0 diff --git a/vartracker/main.py b/vartracker/main.py index c9fb2c3..ddbc2b1 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -891,6 +891,15 @@ def _add_lofreq_primer_rescue_arguments(group: argparse._ArgumentGroup) -> None: default=PrimerRescueThresholds.max_ref_count, help="Maximum DP4 reference count for primer rescue candidates only (default: 20)", ) + group.add_argument( + "--lofreq-rescue-max-minor-alt-fraction", + type=float, + default=PrimerRescueThresholds.max_minor_alt_fraction, + help=( + "Maximum minor ALT strand fraction for primer rescue candidates only " + "(default: 0.05)" + ), + ) def _move_action_group_after( @@ -2245,6 +2254,9 @@ def _run_e2e_command(args): lofreq_rescue_min_alt_count=args.lofreq_rescue_min_alt_count, lofreq_rescue_min_qual=args.lofreq_rescue_min_qual, lofreq_rescue_max_ref_count=args.lofreq_rescue_max_ref_count, + lofreq_rescue_max_minor_alt_fraction=( + args.lofreq_rescue_max_minor_alt_fraction + ), ) if rulegraph_path: @@ -2397,6 +2409,9 @@ def _run_bam_command(args): lofreq_rescue_min_alt_count=args.lofreq_rescue_min_alt_count, lofreq_rescue_min_qual=args.lofreq_rescue_min_qual, lofreq_rescue_max_ref_count=args.lofreq_rescue_max_ref_count, + lofreq_rescue_max_minor_alt_fraction=( + args.lofreq_rescue_max_minor_alt_fraction + ), ) if rulegraph_path: diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index bf1c3e0..0c5f257 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -873,18 +873,54 @@ def _decode_sample_bcsq_annotations(v, samples, annotations): return decoded -def _mask_allele_frequencies_for_annotation( - annotation, allele_freqs, samples, sample_bcsq_map +def _normalise_bcsq_change_type(change_type): + """Remove bcftools' leading compound-context marker from a consequence type.""" + return str(change_type or "").lstrip("*") + + +def _normalise_bcsq_annotation(annotation): + """Return a BCSQ annotation key that ignores leading consequence markers.""" + parts = str(annotation).split("|", 1) + if not parts: + return str(annotation) + parts[0] = _normalise_bcsq_change_type(parts[0]) + return "|".join(parts) + + +def _group_equivalent_bcsq_annotations(annotations): + """Group BCSQ annotations that represent the same consequence.""" + groups = {} + ordered_groups = [] + for annotation in annotations: + key = _normalise_bcsq_annotation(annotation) + if key not in groups: + groups[key] = [] + ordered_groups.append(groups[key]) + groups[key].append(annotation) + return ordered_groups + + +def _representative_bcsq_annotation(annotations): + """Choose the least-decorated BCSQ annotation from an equivalent group.""" + for annotation in annotations: + change_type = str(annotation).split("|", 1)[0] + if not change_type.startswith("*"): + return annotation + return annotations[0] + + +def _mask_allele_frequencies_for_annotation_group( + annotations, allele_freqs, samples, sample_bcsq_map ): - """Keep allele frequencies only for samples where the annotation applies.""" + """Keep allele frequencies for samples matching any equivalent annotation.""" + annotation_set = set(annotations) masked = [] for sample, allele_freq in zip(samples, allele_freqs): if allele_freq == ".": masked.append(".") continue - masked.append( - allele_freq if annotation in sample_bcsq_map.get(sample, []) else "." - ) + sample_annotations = set(sample_bcsq_map.get(sample, [])) + masked.append(allele_freq if annotation_set & sample_annotations else ".") return masked @@ -980,6 +1016,29 @@ def _annotation_alt_for_record(v, anno): return None +def _annotation_is_single_record_allele(v, anno, annotation_alt) -> bool: + """Return True when a BCSQ annotation describes only the current allele.""" + if len(anno) <= 6 or (len(anno) == 1 and str(anno[0]).startswith("@")): + return False + + dna_change = str(anno[6] or "").strip() + if not dna_change or "+" in dna_change: + return False + + match = re.fullmatch(r"(\d+)([^>]+)>([^>]+)", dna_change) + if not match: + return False + + pos = int(match.group(1)) + ref = match.group(2) + alt = match.group(3) + if pos != int(v.POS) or ref != str(v.REF): + return False + if annotation_alt is not None and alt != str(annotation_alt): + return False + return alt in {str(value) for value in (v.ALT or [])} + + def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): """ Process VCF file and extract variant information. @@ -1047,11 +1106,13 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): # Process annotations if "BCSQ" in info: annotations = v.INFO["BCSQ"].split(",") + annotation_groups = _group_equivalent_bcsq_annotations(annotations) sample_bcsq_map = _decode_sample_bcsq_annotations(v, samples, annotations) produced_annotation_specific_row = False if sample_bcsq_map: - for annot in annotations: + for annotation_group in annotation_groups: + annot = _representative_bcsq_annotation(annotation_group) anno = annot.split("|") annotation_alt = _annotation_alt_for_record(v, anno) annotation_allele_freqs = ( @@ -1059,9 +1120,17 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): if annotation_alt is not None else allele_freqs ) - masked_allele_freqs = _mask_allele_frequencies_for_annotation( - annot, annotation_allele_freqs, samples, sample_bcsq_map - ) + if _annotation_is_single_record_allele(v, anno, annotation_alt): + masked_allele_freqs = annotation_allele_freqs + else: + masked_allele_freqs = ( + _mask_allele_frequencies_for_annotation_group( + annotation_group, + annotation_allele_freqs, + samples, + sample_bcsq_map, + ) + ) masked_trajectory = _summarise_sample_trajectory( masked_allele_freqs, samples ) @@ -1088,7 +1157,8 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): produced_annotation_specific_row = True if not produced_annotation_specific_row: - for annot in annotations: + for annotation_group in annotation_groups: + annot = _representative_bcsq_annotation(annotation_group) anno = annot.split("|") annotation_alt = _annotation_alt_for_record(v, anno) annotation_allele_freqs = ( @@ -1179,6 +1249,7 @@ def _process_annotation( ): """Process a single annotation from bcftools csq.""" selected_alt = alt_allele or (v.ALT[0] if v.ALT else "") + type_of_change = _normalise_bcsq_change_type(anno[0] if anno else "") if len(anno) == 1 and anno[0].startswith("@"): # Joint variant annotation return { @@ -1194,7 +1265,7 @@ def _process_annotation( "bcsq_nt_notation": anno[0], "bcsq_aa_notation": anno[0], "type_of_variant": v.var_type, - "type_of_change": anno[0], + "type_of_change": type_of_change, "variant_status": variant_status, "persistence_status": persistent_status, "presence_absence": " / ".join(presence_absence), @@ -1240,7 +1311,7 @@ def _process_annotation( "bcsq_nt_notation": anno[6] if len(anno) > 5 else "", "bcsq_aa_notation": anno[5] if len(anno) > 5 else "", "type_of_variant": v.var_type, - "type_of_change": anno[0], + "type_of_change": type_of_change, "variant_status": variant_status, "persistence_status": persistent_status, "presence_absence": " / ".join(presence_absence), From dfcbbac09c3321bbc59645406d6c2804704d481c Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Wed, 29 Jul 2026 10:29:12 +1000 Subject: [PATCH 14/20] Document custom literature-CSV creation for non-SARS-CoV-2 pathogens Adds a README subsection explaining how to build a --literature-csv lookup table for pathogens other than SARS-CoV-2, covering the required gene/mutation columns, gene-name matching against the supplied GFF3 annotation, and pointers to the bundled mock template. --- README.md | 43 +++++++++++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 30a3e45..1c8c3b2 100755 --- a/README.md +++ b/README.md @@ -326,12 +326,12 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test Consequence-calling note: - Vartracker keeps distinct ALT alleles at the same position separate during preprocessing, then rejoins them immediately before `bcftools csq` so codon-level consequences can still be inferred correctly. -- If more than two ALT alleles remain present in a single sample at one genomic position after frequency filtering, vartracker defaults to stopping with an informative error before `bcftools csq`. This is the safest behavior and the default `--multiallelic-overflow error` mode. +- If more than two ALT alleles remain present in a single sample at one genomic position after frequency filtering, vartracker defaults to stopping with an informative error before `bcftools csq`. This is the safest behaviour and the default `--multiallelic-overflow error` mode. - `--multiallelic-overflow drop-lowest-af` continues by removing the lowest-frequency retained ALT allele(s) for the affected sample before `bcftools csq`, and prints a warning describing the site and the dropped allele(s). - `--multiallelic-overflow skip-site` continues by skipping consequence calling for the affected site entirely, leaving those variants in the results as unannotated rows and printing a warning describing the site. Heatmap filtering: -- `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customize heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. +- `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customise heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. - By default, all consequence classes are included except joint variants. Use `--include-joint` to show joint variants. - `--aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. - `--aa-include`: comma-separated `type_of_change` patterns to include. @@ -371,7 +371,7 @@ Genome plot options: - `--gene`: zoom to a single gene region. - `--aa-scale`: with `--gene`, use amino-acid coordinates on the x-axis. - `--cds-scale`: with `--gene`, use CDS-relative nucleotide coordinates on the x-axis. -- `--focus-coords`: highlight nucleotide or amino-acid coordinate ranges, depending on the current x-axis mode. Separate color groups with `;`, ranges within a group with `,`, and optionally prefix a group with `Name:`. +- `--focus-coords`: highlight nucleotide or amino-acid coordinate ranges, depending on the current x-axis mode. Separate colour groups with `;`, ranges within a group with `,`, and optionally prefix a group with `Name:`. - `--focus-region-file`: read named focus region groups from a `.json`, `.csv`, or `.tsv` file for an inset legend. - `--show-intersections`: add a compact `Region | Variant` table below the genome plot for highlighted-region hits. - In the genome plot, undetected samples are rendered at the detection threshold rather than zero; by default this floor is `0.03`, or `--min-af` if supplied, and the dashed guide line follows that same threshold. @@ -430,6 +430,41 @@ vartracker [mode] input_data.csv --literature-csv pokay_database.csv -o results/ Alternatively, pass `--search-pokay` to automatically download and search against the Pokay SARS-CoV-2 literature database. +#### Building a custom literature database for other pathogens + +`--search-pokay` only covers SARS-CoV-2. For other pathogens, supply your own CSV via +`--literature-csv`. Variant lookup during vartracker analysis is based on a CSV-format file that +is either generated automatically (`--search-pokay`) or supplied by the user (`--literature-csv +`). The expected structure of the file is described using the `vartracker schema literature` +command. In brief, after deriving appropriate information from the scientific literature, users +can create their own lookup table by creating a new CSV file whereby each row corresponds to a +variant of interest, with `gene` and `mutation` required and `category`, `information`, and +`reference` recommended: + +- **`gene`**: must exactly match (case-sensitive) the gene/product name assigned to that variant + by `bcftools csq` using the GFF3/GenBank annotation supplied via `--gff3`. For non-SARS-CoV-2 + pathogens this is simply the gene name as it appears in your annotation file — vartracker's + SARS-CoV-2-specific remapping of `ORF1ab` into individual `nsp1`–`nsp16` names does not apply + outside SARS-CoV-2, so for other pathogens use the gene names exactly as they appear in your + GFF3. +- **`mutation`**: the amino acid consequence in short-hand notation *without* a gene prefix (e.g. + `D614G`, not `S:D614G`). Each row describes a single mutation; if you have information on + several mutations in the same gene, add one row per mutation. Note that matching is done via + substring containment on this column, so avoid overly short or ambiguous notations that could + unintentionally match unrelated variants (e.g. a bare position number). +- **`category`**: a free-text label used to group/colour variants in output tables and the + heatmap. There's no fixed vocabulary — choose categories meaningful for your pathogen (e.g. + "resistance", "immune_escape", "homoplasy"). +- **`information`**: free-text description of the mutation's putative effect, drawn from the + literature. +- **`reference`**: one or more supporting DOIs or URLs, semicolon-delimited if there are multiple. + +A minimal template with this exact structure is provided at +`test_data/mock_literature/mock_literature.csv`; the SARS-CoV-2-specific `pokay_database.csv` +generated by `--search-pokay` follows the same schema and can also be used as a real-world +formatting reference, bearing in mind its `ORF1ab`/`nsp` gene naming is SARS-CoV-2-specific and +shouldn't be copied for other pathogens. + ### Command Line Reference ``` @@ -552,7 +587,7 @@ vartracker schema literature The pipeline performs the following analysis: -1. **VCF Standardization**: Normalizes and standardizes input VCF files, preserving distinct ALT alleles at the same genomic position +1. **VCF Standardisation**: Normalises and standardises input VCF files, preserving distinct ALT alleles at the same genomic position 2. **Variant Merging**: Combines all longitudinal samples 3. **Annotation**: Adds amino acid consequences using `bcftools csq` on the merged VCF so sample-specific joint consequences are inferred from each sample's surviving ALT combination 4. **Comprehensive Analysis**: For each variant, determines: From 85202f84abfb361336a85960875f07a48725ce28 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Wed, 29 Jul 2026 10:30:46 +1000 Subject: [PATCH 15/20] Fix silent merging of same-named genes across contigs/replicons Bacterial references commonly reuse gene names across replicons (e.g. repA on both a chromosome and a plasmid). gene_lengths_from_gff3 and generate_gene_table previously aggregated CDS lengths and mutation counts by gene name alone, so two unrelated genes sharing a name on different contigs had their stats silently summed into one row. - gene_lengths_from_gff3 now tracks CDS lengths per (contig, gene) and only disambiguates the output label, e.g. "repA (chrom1)" vs "repA (plasmid1)", when a name actually collides across contigs. Names confined to a single contig stay unqualified. - generate_gene_table splits colliding gene names by chrom using the results table's own chrom column. - Added ambiguous_gene_names() to flag collisions directly from the annotation, closing an edge case where a colliding gene only has variants on one contig in a given run and would otherwise be dropped when merged against the gene-length scaffold. - No change in output for single-contig references or multi-segment references with uniquely-named segments (e.g. 8-segment influenza). Adds tests/test_annotation_processing.py and three new tests in tests/test_analysis.py covering the single-contig baseline, the colliding-name split, and the ambiguous-genes edge case. Starts CHANGELOG.md to track revision-cycle changes going forward. --- CHANGELOG.md | 40 +++++++++ tests/test_analysis.py | 84 +++++++++++++++++++ tests/test_annotation_processing.py | 53 ++++++++++++ vartracker/analysis.py | 31 ++++++- vartracker/annotation_processing.py | 124 ++++++++++++++++++++++------ vartracker/main.py | 6 +- 6 files changed, 308 insertions(+), 30 deletions(-) create mode 100644 CHANGELOG.md create mode 100644 tests/test_annotation_processing.py diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..b29a55f --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,40 @@ +# Changelog + +This file tracks changes made during the NARGAB manuscript revision cycle, in +addition to (not replacing) the project's versioned GitHub releases. + +Format loosely follows [Keep a Changelog](https://keepachangelog.com/). + +## [Unreleased] + +### Fixed + +- **Multi-contig gene-name collisions no longer silently merged.** Bacterial + reference genomes commonly reuse gene names across replicons (e.g. `repA` + on both a chromosome and a plasmid). Previously, `gene_lengths_from_gff3` + and `generate_gene_table` aggregated CDS lengths and mutation statistics by + gene name alone, so two unrelated genes sharing a name on different contigs + had their lengths and variant counts silently summed into one row. + - `gene_lengths_from_gff3` (`annotation_processing.py`) now tracks CDS + lengths per `(contig, gene)` internally and only disambiguates the + output label (e.g. `repA (chrom1)` / `repA (plasmid1)`) when a gene name + actually collides across contigs. Gene names confined to a single contig + are returned unqualified, so single-contig (viral) references and + existing multi-segment references with uniquely-named segments (e.g. + 8-segment influenza) are unaffected. + - `generate_gene_table` (`analysis.py`) now splits colliding gene names by + `chrom` in the same way, using the results table's own `chrom` column. + - Added `ambiguous_gene_names()` (`annotation_processing.py`), which + identifies gene names that collide across contigs directly from the + reference annotation. This closes an edge case where a colliding gene + only has variants on one of its contigs in a given dataset: without it, + that gene's real data could be silently dropped when merged against the + gene-length scaffold, because the scaffold would only contain the + contig-qualified labels while the variants table only produced the bare, + unqualified label. + - Added test coverage: `tests/test_annotation_processing.py` (new file) + and three new tests in `tests/test_analysis.py` covering the + single-contig baseline, the colliding-gene-name split, and the + ambiguous-genes data-loss edge case. + - No changes to output schema, column names, or behaviour for any + single-contig or uniquely-named multi-segment reference. diff --git a/tests/test_analysis.py b/tests/test_analysis.py index a97c007..62def45 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -11,9 +11,93 @@ _prepare_variant_heatmap_matrix, process_joint_variants, generate_variant_heatmap, + generate_gene_table, ) +def _gene_table_rows(gene_table, gene): + rows = gene_table[gene_table["gene"] == gene] + return dict(zip(rows["type"], rows["number"])) + + +def test_generate_gene_table_keeps_single_contig_genes_unqualified(): + """Baseline: single-contig data (e.g. viral references) is unaffected.""" + table = pd.DataFrame( + { + "gene": ["S", "S", "N"], + "chrom": ["NC_045512.2", "NC_045512.2", "NC_045512.2"], + "type_of_change": ["missense", "missense", "synonymous"], + "presence_absence": ["NY", "NY", "YY"], + } + ) + + gene_table = generate_gene_table(table, gene_lengths={"S": 3822, "N": 1260}) + + assert set(gene_table["gene"]) == {"S", "N"} + assert _gene_table_rows(gene_table, "S")["total"] == 2 + + +def test_generate_gene_table_splits_colliding_gene_names_across_contigs(): + """Same gene name on chromosome + plasmid must not be summed together.""" + table = pd.DataFrame( + { + "gene": ["repA", "repA", "dnaA"], + "chrom": ["chrom1", "plasmid1", "chrom1"], + "type_of_change": ["missense", "missense", "missense"], + "presence_absence": ["NY", "NY", "NY"], + } + ) + + gene_table = generate_gene_table( + table, + gene_lengths={ + "dnaA": 300, + "repA (chrom1)": 201, + "repA (plasmid1)": 150, + }, + ) + + genes_present = set(gene_table["gene"]) + assert "repA (chrom1)" in genes_present + assert "repA (plasmid1)" in genes_present + assert "repA" not in genes_present + + # Each contig's copy of repA should retain its own count, not a merged total. + assert _gene_table_rows(gene_table, "repA (chrom1)")["total"] == 1 + assert _gene_table_rows(gene_table, "repA (plasmid1)")["total"] == 1 + + +def test_generate_gene_table_ambiguous_genes_param_prevents_data_loss(): + """A gene flagged ambiguous by the annotation, but with variants on only + one contig in this particular table, must still surface under its + qualified label so it matches the gene-length scaffold and isn't dropped + by the scaffold merge.""" + table = pd.DataFrame( + { + "gene": ["repA"], + "chrom": ["chrom1"], + "type_of_change": ["missense"], + "presence_absence": ["NY"], + } + ) + + gene_table = generate_gene_table( + table, + gene_lengths={ + "repA (chrom1)": 201, + "repA (plasmid1)": 150, + }, + ambiguous_genes={"repA"}, + ) + + genes_present = set(gene_table["gene"]) + assert "repA (chrom1)" in genes_present + assert "repA" not in genes_present + assert _gene_table_rows(gene_table, "repA (chrom1)")["total"] == 1 + # The scaffold still surfaces the other contig's copy with a zero count. + assert _gene_table_rows(gene_table, "repA (plasmid1)")["total"] == 0 + + def test_search_literature_handles_nullable_boolean_masks(tmp_path): table = pd.DataFrame( { diff --git a/tests/test_annotation_processing.py b/tests/test_annotation_processing.py new file mode 100644 index 0000000..f2bdd74 --- /dev/null +++ b/tests/test_annotation_processing.py @@ -0,0 +1,53 @@ +"""Tests for annotation_processing helpers, especially multi-contig handling.""" + +from __future__ import annotations + +from vartracker.annotation_processing import ( + ambiguous_gene_names, + gene_lengths_from_gff3, +) + +SINGLE_CONTIG_GFF3 = """\ +##gff-version 3 +chrom1\tcustom\tCDS\t1\t300\t.\t+\t0\tID=cds-dnaA;gene=dnaA +chrom1\tcustom\tCDS\t400\t600\t.\t+\t0\tID=cds-repA;gene=repA +""" + +MULTI_CONTIG_COLLISION_GFF3 = """\ +##gff-version 3 +chrom1\tcustom\tCDS\t1\t300\t.\t+\t0\tID=cds-dnaA;gene=dnaA +chrom1\tcustom\tCDS\t400\t600\t.\t+\t0\tID=cds-repA-chrom;gene=repA +plasmid1\tcustom\tCDS\t1\t150\t.\t+\t0\tID=cds-repA-plasmid;gene=repA +plasmid1\tcustom\tCDS\t200\t350\t.\t+\t0\tID=cds-mobA;gene=mobA +""" + + +def test_gene_lengths_single_contig_unqualified(tmp_path): + """Single-contig references (e.g. viral genomes) keep plain gene names.""" + gff_path = tmp_path / "single.gff3" + gff_path.write_text(SINGLE_CONTIG_GFF3) + + lengths = gene_lengths_from_gff3(gff_path) + + assert lengths["dnaA"] == 300 + assert lengths["repA"] == 201 + assert ambiguous_gene_names(gff_path) == set() + + +def test_gene_lengths_disambiguates_colliding_names_across_contigs(tmp_path): + """Identically-named genes on different contigs/replicons must not be summed.""" + gff_path = tmp_path / "multi.gff3" + gff_path.write_text(MULTI_CONTIG_COLLISION_GFF3) + + lengths = gene_lengths_from_gff3(gff_path) + + # Colliding gene name is split per-contig, each keeping its own length. + assert lengths["repA (chrom1)"] == 201 + assert lengths["repA (plasmid1)"] == 150 + assert "repA" not in lengths + + # Non-colliding gene names on either contig stay unqualified. + assert lengths["dnaA"] == 300 + assert lengths["mobA"] == 151 + + assert ambiguous_gene_names(gff_path) == {"repA"} diff --git a/vartracker/analysis.py b/vartracker/analysis.py index e895d93..e6f409d 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -282,13 +282,23 @@ def generate_cumulative_lineplot(table, pname, sample_number_list, outname): def generate_gene_table( - table: pd.DataFrame, gene_lengths: Dict[str, int] | None = None + table: pd.DataFrame, + gene_lengths: Dict[str, int] | None = None, + ambiguous_genes: set | None = None, ): """ Generate gene-wise mutation statistics table. Args: table (pd.DataFrame): Variants table + gene_lengths (dict, optional): Per-gene CDS lengths, e.g. from + `gene_lengths_from_gff3`. + ambiguous_genes (set, optional): Gene names known (from the reference + annotation, e.g. via `ambiguous_gene_names`) to occur on more than + one contig/replicon. These are always split per-contig even if a + given results table only contains variants for one of the + colliding copies, so real data for the other copy is never + silently dropped when the gene-length scaffold is merged in. Returns: pd.DataFrame: Gene statistics table @@ -353,15 +363,28 @@ def make_gene_rows(gene, df): return pd.DataFrame(gene_result) - # Process original genes + # Process original genes. Genes sharing a name across more than one contig + # (e.g. bacterial plasmids reusing names like 'repA') are disambiguated by + # contig so their variant counts are never silently summed together; genes + # confined to a single contig keep their plain name, matching prior + # behaviour for single-contig/viral references and multi-segment genomes + # with unique per-segment gene names. result = [] genes_in_table = table["gene"].unique() + gene_chrom_counts = ( + table.groupby("gene")["chrom"].nunique() if "chrom" in table.columns else {} + ) + known_ambiguous = ambiguous_genes or set() for gene in genes_in_table: if gene == "INTERGENIC": continue - df = table[table["gene"] == gene] - result.append(make_gene_rows(gene, df)) + gene_df = table[table["gene"] == gene] + if gene in known_ambiguous or gene_chrom_counts.get(gene, 1) > 1: + for chrom, chrom_df in gene_df.groupby("chrom"): + result.append(make_gene_rows(f"{gene} ({chrom})", chrom_df)) + else: + result.append(make_gene_rows(gene, gene_df)) if include_nsps: table = table.copy() diff --git a/vartracker/annotation_processing.py b/vartracker/annotation_processing.py index 7a40543..cef7f5b 100644 --- a/vartracker/annotation_processing.py +++ b/vartracker/annotation_processing.py @@ -55,15 +55,27 @@ def _parse_attrs(attr: str) -> dict: return out -def gene_lengths_from_gff3(gff_path: str | Path) -> Dict[str, int]: +def _parse_gene_cds_by_contig( + gff_path: str | Path, +) -> Tuple[ + Dict[Tuple[str, str], int], + Dict[Tuple[str, str], int], + Tuple[int, int] | None, + int | None, + int | None, +]: """ - Parse a GFF3 and return {gene_name: CDS_length}, plus 5'/3' UTR and INTERGENIC=1. + Parse a GFF3 and return per-(contig, gene) CDS lengths and first-start + positions, plus whole-genome bounds for UTR computation. Robust to feature order (e.g., CDS before mRNA). Prefers: 1) CDS attribute 'gene' 2) Parent=gene: → gene 'Name' 3) Parent=transcript: → mRNA 'Name' or its parent gene 'Name' 4) As last resort, protein_id/tx id (rare). + + Shared by `gene_lengths_from_gff3` and `ambiguous_gene_names` so both draw + on the same single parse of the annotation. """ gff_path = Path(gff_path) @@ -71,10 +83,11 @@ def gene_lengths_from_gff3(gff_path: str | Path) -> Dict[str, int]: gene_id_to_name: Dict[str, str] = {} # gene-id -> gene symbol/name tx_id_to_gene_name: Dict[str, str] = {} # transcript-id -> gene name - # Accumulators - gene_cds_bp: Dict[str, int] = defaultdict(int) - gene_first_start: Dict[str, int] = {} - pending_by_tx: Dict[str, List[tuple[int, int]]] = defaultdict(list) + # Accumulators, keyed by (seqid, gene_name) so identically-named genes on + # different contigs are never conflated. + gene_cds_bp: Dict[Tuple[str, str], int] = defaultdict(int) + gene_first_start: Dict[Tuple[str, str], int] = {} + pending_by_tx: Dict[str, List[Tuple[int, int, str]]] = defaultdict(list) # For UTR computation (whole-genome) seq_range: Tuple[int, int] | None = None @@ -149,34 +162,94 @@ def gene_lengths_from_gff3(gff_path: str | Path) -> Dict[str, int]: gname = tx_id_to_gene_name.get(tid) if not gname: # Defer until after we see the mRNA - pending_by_tx[tid].append((cds_len, start_i)) + pending_by_tx[tid].append((cds_len, start_i, seqid)) continue # will add later once tx->gene is known if not gname: # Skip unresolvable CDS entries; they will be handled later continue - gene_cds_bp[str(gname)] += cds_len - if gname is not None: - current = gene_first_start.get(str(gname)) - gene_first_start[str(gname)] = ( - start_i if current is None else min(current, start_i) - ) + key = (seqid, str(gname)) + gene_cds_bp[key] += cds_len + current = gene_first_start.get(key) + gene_first_start[key] = ( + start_i if current is None else min(current, start_i) + ) # --- Resolve any pending CDS whose transcripts appeared after - for tid, lengths in pending_by_tx.items(): + for tid, entries in pending_by_tx.items(): gname = tx_id_to_gene_name.get(tid) if not gname: # Last resort: use the transcript id as a stand-in (or skip) gname = tid if not gname: continue - gene_cds_bp[str(gname)] += sum(length for length, _ in lengths) - if gname is not None: - current = gene_first_start.get(str(gname)) - first_start = min(start for _, start in lengths) - gene_first_start[str(gname)] = ( - first_start if current is None else min(current, first_start) + # All entries for a given transcript id share the same contig. + seqid = entries[0][2] + key = (seqid, str(gname)) + gene_cds_bp[key] += sum(length for length, _, _ in entries) + current = gene_first_start.get(key) + first_start = min(start for _, start, _ in entries) + gene_first_start[key] = ( + first_start if current is None else min(current, first_start) + ) + + return gene_cds_bp, gene_first_start, seq_range, first_cds_start, last_cds_end + + +def ambiguous_gene_names(gff_path: str | Path) -> set[str]: + """ + Return gene names that occur on more than one contig/replicon in a GFF3. + + Useful for other per-gene aggregations (e.g. results-table gene summaries) + that need to disambiguate the same gene names vartracker's own + `gene_lengths_from_gff3` disambiguates, even for genes with no CDS-length + contribution in a particular downstream table. + """ + gene_cds_bp, _, _, _, _ = _parse_gene_cds_by_contig(gff_path) + gname_to_seqids: Dict[str, set] = defaultdict(set) + for seqid, gname in gene_cds_bp: + gname_to_seqids[gname].add(seqid) + return {gname for gname, seqids in gname_to_seqids.items() if len(seqids) > 1} + + +def gene_lengths_from_gff3(gff_path: str | Path) -> Dict[str, int]: + """ + Parse a GFF3 and return {gene_name: CDS_length}, plus 5'/3' UTR and INTERGENIC=1. + + Gene names are keyed per-contig internally. If the same gene name occurs on + more than one contig/replicon (common for bacterial plasmids reusing names + like 'repA'), the returned keys are disambiguated as " ()" so + unrelated genes are never silently summed together. Gene names that occur + on only one contig are returned unqualified, matching prior behaviour. + """ + ( + gene_cds_bp, + gene_first_start, + seq_range, + first_cds_start, + last_cds_end, + ) = _parse_gene_cds_by_contig(gff_path) + + # --- Collapse (seqid, gene) keys into final gene labels. Gene names that + # only ever appear on a single contig keep their plain name (unchanged + # behaviour for single-contig/viral references and multi-segment genomes + # with unique per-segment gene names). Names shared across contigs are + # disambiguated so their lengths are never summed together. + gname_to_seqids: Dict[str, set] = defaultdict(set) + for seqid, gname in gene_cds_bp: + gname_to_seqids[gname].add(seqid) + + labelled_cds_bp: Dict[str, int] = defaultdict(int) + labelled_first_start: Dict[str, int] = {} + for (seqid, gname), length in gene_cds_bp.items(): + label = gname if len(gname_to_seqids[gname]) == 1 else f"{gname} ({seqid})" + labelled_cds_bp[label] += length + start = gene_first_start.get((seqid, gname)) + if start is not None: + current = labelled_first_start.get(label) + labelled_first_start[label] = ( + start if current is None else min(current, start) ) # --- Add UTRs if we have genome bounds @@ -184,16 +257,17 @@ def gene_lengths_from_gff3(gff_path: str | Path) -> Dict[str, int]: seq_start, seq_end = seq_range five_utr = max(0, first_cds_start - seq_start) # 1..first_cds_start-1 three_utr = max(0, seq_end - last_cds_end) # last_cds_end+1..end - gene_cds_bp["5' UTR"] = five_utr - gene_cds_bp["3' UTR"] = three_utr + labelled_cds_bp["5' UTR"] = five_utr + labelled_cds_bp["3' UTR"] = three_utr # Conventional placeholder - gene_cds_bp["INTERGENIC"] = 1 + labelled_cds_bp["INTERGENIC"] = 1 ordered = { - gene: gene_cds_bp[gene] + gene: labelled_cds_bp[gene] for gene in sorted( - gene_cds_bp.keys(), key=lambda g: gene_first_start.get(g, float("inf")) + labelled_cds_bp.keys(), + key=lambda g: labelled_first_start.get(g, float("inf")), ) } diff --git a/vartracker/main.py b/vartracker/main.py index ddbc2b1..4ded196 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -75,6 +75,7 @@ ) from .data import parse_pokay as parse_pokay_module from .annotation_processing import ( + ambiguous_gene_names, gene_lengths_from_gff3, validate_reference_and_annotation, ) @@ -2782,7 +2783,10 @@ def _process_files( os.path.join(args.outdir, "cumulative_mutations.pdf"), ) - gene_table = generate_gene_table(table, gene_lengths) + ambiguous_genes = None + if gene_lengths is not None and getattr(args, "gff3", None): + ambiguous_genes = ambiguous_gene_names(args.gff3) + gene_table = generate_gene_table(table, gene_lengths, ambiguous_genes) plot_gene_table(gene_table, pname, args.outdir) plot_variant_turnover( *apply_shared_plot_filters( From adb346ea7dd10848a8bfa298bf90069bca152980 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Wed, 29 Jul 2026 11:23:46 +1000 Subject: [PATCH 16/20] Cap gene-wise plot to top genes by newly emerged variants for large genomes --- README.md | 28 +++++ tests/test_analysis.py | 225 +++++++++++++++++++++++++++++++++++++++++ vartracker/analysis.py | 163 +++++++++++++++++++++++++++-- vartracker/main.py | 38 ++++++- 4 files changed, 447 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 1c8c3b2..18bc14a 100755 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ A bioinformatics pipeline to summarise variants called against a reference in a - [Quick Start](#quick-start) - [Output](#output) - [What does vartracker do?](#what-does-vartracker-do) +- [Limitations](#limitations) - [Citation](#citation) - [License](#license) - [Contributing](#contributing) @@ -311,6 +312,8 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test - `vartracker vcf` – accepts core analysis options such as `--min-snv-freq`, `--min-indel-freq`, `--allele-frequency-tag`, `--multiallelic-overflow`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. + `--max-plot-genes` and `--plot-genes` control the gene-wise summary figure only (see + [Limitations](#limitations)); the tabular/TSV output always includes every annotated gene. - `vartracker bam` – everything from `vcf`, plus Snakemake options: `--snakemake-outdir`, `--cores`, `--snakemake-dryrun`, `--verbose`, `--redo`, `--rulegraph`, `--primer-bed`, `--lofreq-primer-rescue`, `--consensus-snp-min-af`, @@ -601,6 +604,31 @@ The pipeline performs the following analysis: 5. **Visualization**: Generates plots for mutation accumulation and gene-wise statistics 6. **Functional Annotation**: (optional) Searches against literature databases for known functional impacts +## Limitations + +vartracker was designed for viral pathogens with small, compact genomes (SARS-CoV-2: ~30 kb, +12 genes). The underlying analysis - VCF standardisation, merging, annotation, and +original/new/persistent/transient classification - scales to larger genomes without modification. +The practical constraint on larger genomes (e.g. bacterial pathogens, which can carry thousands of +annotated genes) is **visualisation**, not computation: + +- The gene-wise summary figure (`mutations_per_gene.pdf`) plots one bar per gene per panel. On a + genome with thousands of annotated genes this becomes unreadable regardless of how many variants + are actually present, because the plot iterates over every annotated gene, not just genes that + carry a variant. +- By default, the figure is capped to the top 30 genes, ranked by number of newly emerged variants + (ties broken by total variant count), via `--max-plot-genes`. Use `--plot-genes` to instead name + an explicit set of genes to plot. **This cap applies to the figure only** - the tabular/TSV output + always contains every annotated gene, so no data is discarded by this option. +- When the figure is truncated, this is stated directly on the figure itself (e.g. "top 30 of 412 + genes with variants"); if nothing was truncated, no such note is shown. + +Separately, the bundled `pokay` functional-annotation database +(see [Using Literature Database](#using-literature-database)) is specific to SARS-CoV-2 mutations +and is not applied to, or meaningful for, other pathogens. A custom literature CSV following the +same schema can be supplied via `--literature-csv` for other organisms; see +[Building a custom literature database for other pathogens](#building-a-custom-literature-database-for-other-pathogens). + ## Citation When using vartracker, please cite the software release you used. Citation metadata is provided diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 62def45..867c66c 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -2,8 +2,13 @@ from __future__ import annotations +import os +import string + +import matplotlib.pyplot as plt import pandas as pd import pytest +import seaborn as sns from vartracker.analysis import ( _heatmap_figure_size, @@ -12,6 +17,9 @@ process_joint_variants, generate_variant_heatmap, generate_gene_table, + plot_gene_table, + parse_plot_genes_arg, + select_genes_for_plot, ) @@ -98,6 +106,223 @@ def test_generate_gene_table_ambiguous_genes_param_prevents_data_loss(): assert _gene_table_rows(gene_table, "repA (plasmid1)")["total"] == 0 +def _make_variant_row(gene, type_of_change, presence_absence, chrom="chrom1"): + return { + "gene": gene, + "chrom": chrom, + "type_of_change": type_of_change, + "presence_absence": presence_absence, + } + + +def test_select_genes_for_plot_keeps_everything_when_annotation_under_cap(): + """When the whole annotated gene set already fits under the cap (e.g. a + viral reference), nothing is filtered - this is what keeps SARS-CoV-2 + figures byte-identical regardless of the new default cap.""" + table = pd.DataFrame( + [ + _make_variant_row("S", "missense", "NY"), + _make_variant_row("N", "synonymous", "YY"), + ] + ) + gene_table = generate_gene_table( + table, gene_lengths={"S": 3822, "N": 1260, "E": 228} + ) + + plot_table, subtitle = select_genes_for_plot(gene_table, max_plot_genes=30) + + assert subtitle is None + pd.testing.assert_frame_equal( + plot_table.reset_index(drop=True), gene_table.reset_index(drop=True) + ) + + +def test_select_genes_for_plot_ranks_by_new_mutations_with_total_tiebreak(): + """Ranking must use newly emerged variants first, falling back to total + variants only to break ties - not the other way around.""" + table = pd.DataFrame( + [ + # Gene A: 3 new mutations, 5 total. + _make_variant_row("A", "missense", "NY"), + _make_variant_row("A", "missense", "NY"), + _make_variant_row("A", "missense", "NY"), + _make_variant_row("A", "synonymous", "YY"), + _make_variant_row("A", "synonymous", "YY"), + # Gene B: 3 new mutations (tied with A), only 2 total. + _make_variant_row("B", "missense", "NY"), + _make_variant_row("B", "missense", "NY"), + _make_variant_row("B", "missense", "NY"), + # Gene C: fewer new mutations than A/B, but more total variants. + _make_variant_row("C", "missense", "NY"), + *[_make_variant_row("C", "synonymous", "YY") for _ in range(9)], + # Gene D: no new mutations at all. + _make_variant_row("D", "synonymous", "YY"), + ] + ) + gene_table = generate_gene_table( + table, gene_lengths={"A": 100, "B": 100, "C": 100, "D": 100, "E": 100} + ) + + plot_table, subtitle = select_genes_for_plot(gene_table, max_plot_genes=2) + + assert subtitle == "top 2 of 4 genes with variants" + assert set(plot_table["gene"]) == {"A", "B"} + + +def test_select_genes_for_plot_omits_subtitle_when_variant_genes_fit_cap(): + """Even if the annotated genome is large, if the genes that actually + carry variants already fit under the cap, nothing was truncated and the + subtitle must be omitted (e.g. never print "top 30 of 12").""" + table = pd.DataFrame( + [ + _make_variant_row("A", "missense", "NY"), + _make_variant_row("B", "missense", "NY"), + ] + ) + gene_lengths = {"A": 100, "B": 100} + gene_lengths.update({f"unused{i}": 100 for i in range(20)}) + gene_table = generate_gene_table(table, gene_lengths=gene_lengths) + + plot_table, subtitle = select_genes_for_plot(gene_table, max_plot_genes=10) + + assert subtitle is None + assert set(plot_table["gene"]) == {"A", "B"} + + +def test_select_genes_for_plot_cap_zero_or_negative_raises(): + table = pd.DataFrame([_make_variant_row("A", "missense", "NY")]) + gene_table = generate_gene_table(table, gene_lengths={"A": 100}) + + with pytest.raises(ValueError): + select_genes_for_plot(gene_table, max_plot_genes=0) + + with pytest.raises(ValueError): + select_genes_for_plot(gene_table, max_plot_genes=-5) + + +def test_select_genes_for_plot_explicit_list_warns_and_drops_invalid(capsys): + table = pd.DataFrame( + [ + _make_variant_row("X", "missense", "NY"), + _make_variant_row("Y", "missense", "NY"), + ] + ) + gene_lengths = {"X": 100, "Y": 100, "Z": 100} # Z is annotated but has no variants + gene_table = generate_gene_table(table, gene_lengths=gene_lengths) + + plot_table, subtitle = select_genes_for_plot( + gene_table, plot_genes=["X", "Z", "FAKE"] + ) + + assert subtitle is None + assert set(plot_table["gene"]) == {"X"} + captured = capsys.readouterr() + assert "Z" in captured.out and "no variants" in captured.out + assert "FAKE" in captured.out and "reference annotation" in captured.out + + +def test_select_genes_for_plot_explicit_list_all_invalid_raises(): + table = pd.DataFrame([_make_variant_row("X", "missense", "NY")]) + gene_table = generate_gene_table(table, gene_lengths={"X": 100}) + + with pytest.raises(ValueError): + select_genes_for_plot(gene_table, plot_genes=["FAKE1", "FAKE2"]) + + +def test_select_genes_for_plot_explicit_list_overrides_max_plot_genes(): + table = pd.DataFrame( + [ + _make_variant_row("A", "missense", "NY"), + _make_variant_row("B", "missense", "NY"), + _make_variant_row("C", "missense", "NY"), + ] + ) + gene_table = generate_gene_table(table, gene_lengths={"A": 100, "B": 100, "C": 100}) + + plot_table, subtitle = select_genes_for_plot( + gene_table, max_plot_genes=1, plot_genes=["A", "B"] + ) + + assert subtitle is None + assert set(plot_table["gene"]) == {"A", "B"} + + +def test_parse_plot_genes_arg_comma_separated_deduplicates_and_preserves_order(): + genes = parse_plot_genes_arg("geneA, geneB,geneA, geneC") + assert genes == ["geneA", "geneB", "geneC"] + + +def test_parse_plot_genes_arg_reads_from_file(tmp_path): + gene_file = tmp_path / "genes.txt" + gene_file.write_text("geneA\n\n# a comment\ngeneB\ngeneA\n") + + genes = parse_plot_genes_arg(str(gene_file)) + + assert genes == ["geneA", "geneB"] + + +def _golden_plot_gene_table(gene_table, pname, outdir): + """Frozen copy of `plot_gene_table` as it existed before the + `--max-plot-genes`/`--plot-genes` options were added. Used only as a + ground truth for the byte-identical regression test below.""" + g = sns.catplot( + x="gene", + y="number", + col="type", + col_wrap=3, + data=gene_table, + kind="bar", + height=4, + aspect=1.2, + ) + g.set_axis_labels("", "Number of Mutations") + g.fig.subplots_adjust(top=0.9) + g.fig.suptitle(f"{pname}", weight="bold") + + for ax in g.axes.flat: + for label in ax.get_xticklabels(): + label.set_rotation(90) + + for ax, title in zip(g.fig.axes, list(gene_table["type"].unique())): + if pd.isna(title): + title_str = "None" + else: + title_str = str(title) if not isinstance(title, str) else title + ax.set_title(string.capwords(title_str.replace("_", " "))) + + plt.savefig( + os.path.join(outdir, "mutations_per_gene.pdf"), dpi=300, bbox_inches="tight" + ) + plt.close() + + +def test_plot_gene_table_default_matches_pre_cap_baseline(tmp_path, monkeypatch): + """With the default cap (30) and a SARS-CoV-2-sized (12 gene) reference, + output must be byte-identical to the pre-existing plotting behaviour, so + published manuscript figures don't change.""" + monkeypatch.setenv("SOURCE_DATE_EPOCH", "0") + + gene_names = [f"gene{i}" for i in range(12)] + rows = [] + for i, gene in enumerate(gene_names[:10]): + rows.append(_make_variant_row(gene, "missense", "NY" if i % 2 == 0 else "YY")) + table = pd.DataFrame(rows) + gene_lengths = {gene: 1000 for gene in gene_names} + gene_table = generate_gene_table(table, gene_lengths=gene_lengths) + + golden_dir = tmp_path / "golden" + new_dir = tmp_path / "new" + golden_dir.mkdir() + new_dir.mkdir() + + _golden_plot_gene_table(gene_table.copy(deep=True), "Baseline", str(golden_dir)) + plot_gene_table(gene_table.copy(deep=True), "Baseline", str(new_dir)) + + golden_bytes = (golden_dir / "mutations_per_gene.pdf").read_bytes() + new_bytes = (new_dir / "mutations_per_gene.pdf").read_bytes() + assert golden_bytes == new_bytes + + def test_search_literature_handles_nullable_boolean_masks(tmp_path): table = pd.DataFrame( { diff --git a/vartracker/analysis.py b/vartracker/analysis.py index e6f409d..3dbfd16 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -423,29 +423,180 @@ def assign_nsp(row): return gene_table -def plot_gene_table(gene_table, pname, outdir): +def parse_plot_genes_arg(value: str) -> List[str]: + """ + Parse a `--plot-genes` CLI value into a gene name list. + + `value` is either a comma-separated list of gene names, or a path to a + file with one gene name per line (blank lines and lines starting with + `#` are ignored). Order is preserved and duplicates are dropped. + """ + if os.path.isfile(value): + with open(value, encoding="utf-8") as handle: + raw_values = [ + line.strip() + for line in handle + if line.strip() and not line.strip().startswith("#") + ] + else: + raw_values = [token.strip() for token in value.split(",") if token.strip()] + + seen: set = set() + genes: List[str] = [] + for gene in raw_values: + if gene not in seen: + seen.add(gene) + genes.append(gene) + return genes + + +def select_genes_for_plot( + gene_table: pd.DataFrame, + max_plot_genes: Optional[int] = 30, + plot_genes: Optional[List[str]] = None, +) -> Tuple[pd.DataFrame, Optional[str]]: + """ + Select which genes appear on the gene-wise summary FIGURE. + + This is presentation-layer only: it never modifies `gene_table` itself + (the tabular/TSV output built by `generate_gene_table` always contains + every annotated gene), it only decides which subset is handed to + `plot_gene_table` for rendering. + + When `plot_genes` is not given, and the number of annotated genes + already fits within `max_plot_genes`, every gene is kept (this is what + makes SARS-CoV-2-scale runs, with their handful of genes, identical to + plotting with no cap at all). Otherwise the figure is limited to genes + that carry at least one variant, ranked by number of newly emerged + variants (total variants as tiebreak) - ranking by total variants alone + would surface long genes or reference/sample divergence rather than + anything that changed over the longitudinal series. + + Args: + gene_table: Output of `generate_gene_table`. + max_plot_genes: Cap on number of genes shown, or None for no cap. + plot_genes: Explicit gene names to show; overrides `max_plot_genes`. + Names absent from the reference annotation or with no variants + in this dataset are dropped with a warning. + + Returns: + (plot_table, subtitle): `plot_table` is the subset of `gene_table` + to plot (original gene order preserved); `subtitle` describes the + truncation applied (e.g. "top 30 of 412 genes with variants"), or + None if nothing was truncated. + """ + if gene_table.empty: + return gene_table, None + + genes_in_order = list(dict.fromkeys(gene_table["gene"])) + totals = gene_table.loc[gene_table["type"] == "total"].set_index("gene")["number"] + variant_genes = [g for g in genes_in_order if totals.get(g, 0) > 0] + + if plot_genes: + annotation_genes = set(genes_in_order) + variant_gene_set = set(variant_genes) + valid: List[str] = [] + for gene in plot_genes: + if gene not in annotation_genes: + print( + f"Warning: --plot-genes gene '{gene}' was not found in the " + "reference annotation; skipping." + ) + continue + if gene not in variant_gene_set: + print( + f"Warning: --plot-genes gene '{gene}' has no variants in " + "this dataset; skipping." + ) + continue + valid.append(gene) + + if not valid: + raise ValueError( + "None of the genes named in --plot-genes are valid: each gene " + "must be present in the reference annotation and carry at " + "least one variant in this dataset." + ) + + selected = set(valid) + plot_table = gene_table[gene_table["gene"].isin(selected)] + return plot_table, None + + if max_plot_genes is None: + return gene_table, None + if max_plot_genes <= 0: + raise ValueError( + f"--max-plot-genes must be a positive integer, got {max_plot_genes}." + ) + + if len(genes_in_order) <= max_plot_genes: + return gene_table, None + + n_with_variants = len(variant_genes) + if n_with_variants <= max_plot_genes: + plot_table = gene_table[gene_table["gene"].isin(variant_genes)] + return plot_table, None + + new_counts = gene_table.loc[gene_table["type"] == "new_mutations"].set_index("gene")[ + "number" + ] + ranked = sorted( + variant_genes, key=lambda g: (-new_counts.get(g, 0), -totals.get(g, 0)) + ) + selected = set(ranked[:max_plot_genes]) + plot_table = gene_table[gene_table["gene"].isin(selected)] + subtitle = f"top {max_plot_genes} of {n_with_variants} genes with variants" + return plot_table, subtitle + + +def plot_gene_table( + gene_table, + pname, + outdir, + max_plot_genes: Optional[int] = 30, + plot_genes: Optional[List[str]] = None, +): """ Plot gene-wise mutation statistics. Args: - gene_table (pd.DataFrame): Gene statistics table + gene_table (pd.DataFrame): Gene statistics table (all annotated + genes; the tabular/TSV output is unaffected by this function). pname (str): Project name for plot title outdir (str): Output directory + max_plot_genes: Cap on number of genes shown in the FIGURE, ranked + by newly emerged variants (default: 30). Does not affect any + other output. + plot_genes: Explicit gene names to show in the FIGURE; overrides + `max_plot_genes`. """ try: + plot_table, subtitle = select_genes_for_plot( + gene_table, max_plot_genes=max_plot_genes, plot_genes=plot_genes + ) + g = sns.catplot( x="gene", y="number", col="type", col_wrap=3, - data=gene_table, + data=plot_table, kind="bar", height=4, aspect=1.2, ) - g.set_axis_labels("", "Number of Mutations") + x_label = ( + "Gene (top genes ranked by newly emerged variants; " + "total variants as tiebreak)" + if subtitle + else "" + ) + g.set_axis_labels(x_label, "Number of Mutations") g.fig.subplots_adjust(top=0.9) - g.fig.suptitle(f"{pname}", weight="bold") + title = f"{pname}" + if subtitle: + title = f"{title}\n{subtitle}" if title else subtitle + g.fig.suptitle(title, weight="bold") # Rotate labels by 90 degrees for ax in g.axes.flat: @@ -453,7 +604,7 @@ def plot_gene_table(gene_table, pname, outdir): label.set_rotation(90) # Make subplot titles nicer - for ax, title in zip(g.fig.axes, list(gene_table["type"].unique())): + for ax, title in zip(g.fig.axes, list(plot_table["type"].unique())): # Convert title to string to handle numeric values and NaN if pd.isna(title): title_str = "None" diff --git a/vartracker/main.py b/vartracker/main.py index 4ded196..37d9c9d 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -38,6 +38,7 @@ generate_cumulative_lineplot, generate_gene_table, plot_gene_table, + parse_plot_genes_arg, generate_variant_heatmap, search_literature, ) @@ -784,6 +785,31 @@ def _configure_vcf_parser( help="Only analyse samples with sample_number less than or equal to this value", default=None, ) + analysis_group.add_argument( + "--max-plot-genes", + action="store", + type=int, + default=30, + help=( + "Cap the gene-wise summary FIGURE to the top N genes, ranked by " + "number of newly emerged variants (total variants used as a " + "tiebreak). Does not affect the tabular/TSV output, which always " + "includes every annotated gene. Overridden by --plot-genes. " + "(default: 30)" + ), + ) + analysis_group.add_argument( + "--plot-genes", + action="store", + default=None, + help=( + "Comma-separated gene names, or a path to a file with one gene " + "name per line, to show on the gene-wise summary FIGURE. " + "Overrides --max-plot-genes. Genes absent from the reference " + "annotation or with no variants in this dataset are skipped " + "with a warning." + ), + ) analysis_group.add_argument( "--literature-csv", action="store", @@ -2786,8 +2812,18 @@ def _process_files( ambiguous_genes = None if gene_lengths is not None and getattr(args, "gff3", None): ambiguous_genes = ambiguous_gene_names(args.gff3) + plot_genes = None + if getattr(args, "plot_genes", None): + plot_genes = parse_plot_genes_arg(args.plot_genes) + gene_table = generate_gene_table(table, gene_lengths, ambiguous_genes) - plot_gene_table(gene_table, pname, args.outdir) + plot_gene_table( + gene_table, + pname, + args.outdir, + max_plot_genes=getattr(args, "max_plot_genes", 30), + plot_genes=plot_genes, + ) plot_variant_turnover( *apply_shared_plot_filters( *prepare_plot_inputs(table)[:2], From e28e74249c1d820aa2cb6b4d9bb2680c4f95954e Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Wed, 29 Jul 2026 12:22:44 +1000 Subject: [PATCH 17/20] Address reviewer comments: QC threshold caveats, intermittent persistence labels, primer-bed warning --- README.md | 46 ++++++++++++++++++++-- docs/OUTPUT_SCHEMA.md | 4 +- tests/test_analysis.py | 55 ++++++++++++++++++++++++++ tests/test_main.py | 73 +++++++++++++++++++++++++++++++++++ tests/test_plotting.py | 75 ++++++++++++++++++++++++++++++++++++ tests/test_vcf_processing.py | 66 +++++++++++++++++++++++++++++++ vartracker/analysis.py | 3 +- vartracker/main.py | 54 +++++++++++++++++++++----- vartracker/plotting.py | 9 +++-- vartracker/schemas.py | 21 ++++++++-- vartracker/vcf_processing.py | 32 +++++++++++++-- 11 files changed, 413 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 18bc14a..3eefbac 100755 --- a/README.md +++ b/README.md @@ -327,6 +327,17 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test - `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. - `vartracker plot lifespan` – plot first-to-last detection spans for a selected or auto-ranked subset of variants. +QC threshold note: +- `--min-snv-freq`, `--min-indel-freq`, and `--min-depth` are configurable allele-frequency and + read-depth thresholds applied when summarising and visualising longitudinal variant calls. Their + defaults reflect our own genomic surveillance and longitudinal sequencing workflows and should be + treated as starting points, not universally applicable QC recommendations. The appropriate + thresholds for a given study depend on its objective and on the sequencing protocol, depth, + variant caller, and empirically established error profile of the upstream workflow: use more + stringent thresholds when specificity is prioritised or the input data have higher error rates, + and only lower thresholds for low-frequency variant analysis when this is supported by a suitably + validated upstream workflow. + Consequence-calling note: - Vartracker keeps distinct ALT alleles at the same position separate during preprocessing, then rejoins them immediately before `bcftools csq` so codon-level consequences can still be inferred correctly. - If more than two ALT alleles remain present in a single sample at one genomic position after frequency filtering, vartracker defaults to stopping with an informative error before `bcftools csq`. This is the safest behaviour and the default `--multiallelic-overflow error` mode. @@ -338,7 +349,8 @@ Heatmap filtering: - By default, all consequence classes are included except joint variants. Use `--include-joint` to show joint variants. - `--aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. - `--aa-include`: comma-separated `type_of_change` patterns to include. -- `--only-persistent`: only include `new_persistent` variants. +- `--only-persistent`: only include new variants present at the final timepoint (`new_persistent` or + `new_intermittent`; see [Persistence labels](#persistence-labels)). - `--only-new`: only include variants with `variant_status == new`. - `--gene-include` and `--gene-exclude`: comma-separated gene patterns. - `--variant-type`: comma-separated variant-type patterns such as `snp` or `indel`. @@ -359,7 +371,9 @@ Standalone plot filtering: - `--gene`, `--effect`, `--min-af`, `--max-af`: restrict the plotted result set before ranking/selection. - `--variants` or `--variant-file`: explicitly choose variants and preserve that order. - `--sample-min`, `--sample-max`: restrict the passage/sample-number window. -- `--persistent-only` and `--new-only`: keep only persistent new variants or only variants with `variant_status == new`. +- `--persistent-only` and `--new-only`: keep only new variants present at the final timepoint + (`new_persistent` or `new_intermittent`; see [Persistence labels](#persistence-labels)) or only + variants with `variant_status == new`. - `trajectory` and `lifespan` auto-select a limited subset by default (`--top-n`) to stay readable. - `turnover` uses all filtered variants by default and is also written automatically during the main `vcf`/`bam`/`end-to-end` workflows as `variant_turnover_plot.pdf`. - `genome` uses SNPs only by default, keeps all observed allele-frequency values for each plotted variant, and writes `variant_genome_plot.pdf` during the main workflows. @@ -554,7 +568,8 @@ vartracker produces several output files: - **`_variants.rescued.tsv`** (`bam`/`end-to-end`): LoFreq primer-overlap rescue audit table, empty when rescue is disabled or no variants are rescued - **`_variants.filtered_out.tsv`** (`bam`/`end-to-end`): Raw LoFreq calls excluded from the final VCF, including filter reason and call metrics - **new_mutations.csv**: Mutations not present in the first sample -- **persistent_new_mutations.csv**: New mutations that persist to the final sample +- **persistent_new_mutations.csv**: New mutations present at the final sample (`new_persistent` or + `new_intermittent`; see [Persistence labels](#persistence-labels)) - **cumulative_mutations.pdf**: Plot showing mutation accumulation over time - **mutations_per_gene.pdf**: Gene-wise mutation statistics - **variant_allele_frequency_heatmap.html**: Interactive heatmap with optional literature annotations @@ -565,6 +580,31 @@ vartracker produces several output files: By default the manifest is lightweight. Use `--manifest-level deep` to checksum all referenced input files (FASTQ/BAM/VCF/coverage) and include file sizes. +### Persistence labels + +The `persistence_status` column classifies each variant from `variant_status` +(`original`: present in the first sample; `new`: absent in the first sample) plus its presence +pattern across the rest of the samples: + +- `original_retained`: an `original` variant continuously present through the final sample. +- `original_intermittent`: an `original` variant present in the final sample, but absent from at + least one sample in between (i.e. lost and regained). +- `original_lost`: an `original` variant absent by the final sample. +- `new_persistent`: a `new` variant continuously present from its first appearance through the + final sample. +- `new_intermittent`: a `new` variant present in the final sample, but absent from at least one + sample between its first appearance and the final sample (i.e. it appeared, disappeared in a + later sample, then reappeared). +- `new_transient`: a `new` variant absent by the final sample. + +These labels are driven by presence/absence, not allele frequency, and depend only on the first, +last, and intervening samples - they say nothing on their own about whether an intervening absence +reflects genuine loss or a QC dropout (see the `per_sample_variant_qc` column in +[Output schema](#output-schema)). `--only-persistent` / `--persistent-only` filters (heatmap and +standalone plots) and `persistent_new_mutations.csv` include both `new_persistent` and +`new_intermittent` variants, since both reached the final timepoint; the label only distinguishes +the path taken to get there. + ### Output schema The results table schema is documented in `docs/OUTPUT_SCHEMA.md`. You can also print it from the CLI: diff --git a/docs/OUTPUT_SCHEMA.md b/docs/OUTPUT_SCHEMA.md index 578c9d1..11ce759 100644 --- a/docs/OUTPUT_SCHEMA.md +++ b/docs/OUTPUT_SCHEMA.md @@ -24,13 +24,13 @@ Columns that encode per-sample values are slash-separated and ordered by the inp | type_of_variant | string | Variant type derived from the VCF entry. | | snp, indel | | type_of_change | string | Functional change classification from bcftools csq. | | synonymous, missense, frameshift, ... | | variant_status | string | Whether the variant is present in the first sample. | | original, new | -| persistence_status | string | Persistence class based on first and last sample presence. | | original_retained, original_lost, new_persistent, new_transient | +| persistence_status | string | Persistence class based on presence at the first and last sample, and whether presence was continuous in between. The '_intermittent' classes mean the variant was present at both the first/last (original) or reappeared by the last sample (new), but was absent from at least one sample in between. | | original_retained, original_intermittent, original_lost, new_persistent, new_intermittent, new_transient | | presence_absence | string (slash-separated) | Per-sample presence (Y) or absence (N), ordered by input. | | Y/N | | first_appearance | string | Sample name where the variant first appears. | | | | last_appearance | string | Sample name where the variant last appears. | | | | all_samples_pass_qc | boolean | True if every sample passes per-sample variant QC. | | true, false | | proportion_samples_passing_qc | number | Proportion of samples passing per-sample variant QC. | fraction | 0-1 | -| per_sample_variant_qc | string (slash-separated) | Per-sample QC flags (P/F) ordered by input. | | P, F | +| per_sample_variant_qc | string (slash-separated) | Per-sample QC flags (P/F) ordered by input. A sample fails ('F') when there is no variant-supporting read and site coverage is below --min-depth, i.e. when genuine absence of the variant cannot be distinguished from dropout/non-detection; 'P' means absence (or presence) was confidently called at that sample. | | P, F | | aa1_total_properties | string | Physicochemical properties for the reference amino acid. | | semicolon-separated properties | | aa2_total_properties | string | Physicochemical properties for the alternate amino acid. | | semicolon-separated properties | | aa1_unique_properties | string | Properties unique to the reference amino acid. | | semicolon-separated properties | diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 867c66c..c324e7a 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -677,6 +677,61 @@ def test_prepare_variant_heatmap_matrix_applies_extended_filters(): assert list(matrix.index) == ["S:D215G\n(A22206G)"] +def test_prepare_variant_heatmap_matrix_only_persistent_includes_new_intermittent(): + """--only-persistent must keep new_intermittent alongside new_persistent, + since both reached the final timepoint - only the path there differs.""" + table = pd.DataFrame( + [ + { + "gene": "S", + "amino_acid_consequence": "D215G", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "persistence_status": "new_persistent", + "alt_freq": "0.0 / 0.6 / 0.7", + "samples": "P0 / P1 / P2", + "variant": "A22206G", + "start": 22206, + }, + { + "gene": "S", + "amino_acid_consequence": "E484K", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "persistence_status": "new_intermittent", + "alt_freq": "0.0 / 0.6 / 0.7", + "samples": "P0 / P1 / P2", + "variant": "G23012A", + "start": 23012, + }, + { + "gene": "S", + "amino_acid_consequence": "N501Y", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "persistence_status": "new_transient", + "alt_freq": "0.0 / 0.6 / 0.0", + "samples": "P0 / P1 / P2", + "variant": "A23063T", + "start": 23063, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix( + table, + ["P0", "P1", "P2"], + 0.2, + 0.2, + only_persistent=True, + ) + + assert set(matrix.index) == {"S:D215G\n(A22206G)", "S:E484K\n(G23012A)"} + + def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): csv_path = tmp_path / "results.csv" pd.DataFrame( diff --git a/tests/test_main.py b/tests/test_main.py index 3885c39..5d65799 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1080,6 +1080,79 @@ def fake_vcf(args): assert modes_checked == ["e2e"] +def _prepare_e2e_run(monkeypatch, tmp_path): + updated_csv = tmp_path / "samples_updated.csv" + vcf_out = tmp_path / "vcf.gz" + cov_out = tmp_path / "coverage.txt" + vcf_out.write_text("", encoding="utf-8") + cov_out.write_text("", encoding="utf-8") + updated_csv.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + f"Sample1,0,,,,{vcf_out},{cov_out}\n", + encoding="utf-8", + ) + + monkeypatch.setattr(main_module, "validate_dependencies", lambda mode="vcf": None) + monkeypatch.setattr( + main_module, "run_e2e_workflow", lambda **kwargs: str(updated_csv) + ) + monkeypatch.setattr(main_module, "_run_vcf_command", lambda args: 0) + + reads1 = tmp_path / "reads1.fastq" + reads2 = tmp_path / "reads2.fastq" + reads1.write_text("", encoding="utf-8") + reads2.write_text("", encoding="utf-8") + samples_csv = tmp_path / "reads.csv" + samples_csv.write_text( + "sample_name,sample_number,reads1,reads2,bam,vcf,coverage\n" + f"Sample1,0,{reads1},{reads2},,,\n", + encoding="utf-8", + ) + return samples_csv + + +def test_e2e_warns_when_primer_bed_missing(monkeypatch, tmp_path, capsys): + samples_csv = _prepare_e2e_run(monkeypatch, tmp_path) + + exit_code = main_module.main( + [ + "end-to-end", + str(samples_csv), + "--reference", + "ref.fasta", + "--outdir", + str(tmp_path / "results"), + ] + ) + + assert exit_code == 0 + captured = capsys.readouterr() + assert "no --primer-bed supplied" in captured.out + + +def test_e2e_no_warning_when_primer_bed_supplied(monkeypatch, tmp_path, capsys): + samples_csv = _prepare_e2e_run(monkeypatch, tmp_path) + primer_bed = tmp_path / "primers.bed" + primer_bed.write_text("", encoding="utf-8") + + exit_code = main_module.main( + [ + "end-to-end", + str(samples_csv), + "--reference", + "ref.fasta", + "--primer-bed", + str(primer_bed), + "--outdir", + str(tmp_path / "results"), + ] + ) + + assert exit_code == 0 + captured = capsys.readouterr() + assert "no --primer-bed supplied" not in captured.out + + def test_e2e_dryrun_skips_vcf(monkeypatch, tmp_path_factory): reads1 = tmp_path_factory.mktemp("reads") / "reads1.fastq" reads2 = reads1.parent / "reads2.fastq" diff --git a/tests/test_plotting.py b/tests/test_plotting.py index fb3b200..4f974ca 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -142,6 +142,81 @@ def test_prepare_plot_inputs_and_filters(): assert filtered_long["sample_number"].max() == 3 +def test_apply_shared_plot_filters_persistent_only_includes_new_intermittent(): + """persistent_only must keep new_intermittent alongside new_persistent - + both are new variants present at the final timepoint, just differing in + whether presence was continuous along the way.""" + table = pd.DataFrame( + [ + { + "chrom": "segA", + "start": 23403, + "end": 23403, + "gene": "S", + "variant": "A23403G", + "amino_acid_consequence": "S:D614G", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_persistent", + "presence_absence": "N / Y / Y / Y", + "alt_freq": "0.0 / 0.20 / 0.45 / 0.70", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + { + "chrom": "segA", + "start": 23012, + "end": 23012, + "gene": "S", + "variant": "G23012A", + "amino_acid_consequence": "S:E484K", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_intermittent", + "presence_absence": "N / Y / N / Y", + "alt_freq": "0.0 / 0.20 / 0.0 / 0.45", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + { + "chrom": "segA", + "start": 23063, + "end": 23063, + "gene": "S", + "variant": "A23063T", + "amino_acid_consequence": "S:N501Y", + "nsp_aa_change": "", + "type_of_variant": "snp", + "type_of_change": "missense", + "variant_status": "new", + "persistence_status": "new_transient", + "presence_absence": "N / Y / N / N", + "alt_freq": "0.0 / 0.20 / 0.0 / 0.0", + "samples": "P0 / P1 / P2 / P3", + "sample_number": "0 / 1 / 2 / 3", + "per_sample_variant_qc": "P / P / P / P", + "reference": "", + }, + ] + ) + + summary, long_df, _, _ = prepare_plot_inputs(table) + filtered_summary, _ = apply_shared_plot_filters( + summary, long_df, persistent_only=True + ) + + labels = {label.split(" (", 1)[0] for label in filtered_summary["variant_label"]} + assert labels == {"S:D614G", "S:E484K"} + + def test_auto_select_variants_prefers_literature_and_persistent(): summary, _, _, _ = prepare_plot_inputs(_results_table()) diff --git a/tests/test_vcf_processing.py b/tests/test_vcf_processing.py index 6816f40..3df76e7 100644 --- a/tests/test_vcf_processing.py +++ b/tests/test_vcf_processing.py @@ -14,6 +14,7 @@ from vartracker.analysis import process_joint_variants from vartracker.vcf_processing import ( _derive_vcf_output_paths, + _summarise_sample_trajectory, annotate_vcf, format_vcf, merge_consequences, @@ -32,6 +33,71 @@ def test_derive_vcf_output_paths_handles_gz(tmp_path): assert os.path.dirname(out) == str(tmp_path) +def test_summarise_sample_trajectory_original_retained_no_gap(): + result = _summarise_sample_trajectory(["0.1", "0.1", "0.1"], ["s0", "s1", "s2"]) + assert result["variant_status"] == "original" + assert result["persistent_status"] == "original_retained" + assert result["presence_absence"] == ["Y", "Y", "Y"] + + +def test_summarise_sample_trajectory_original_intermittent_has_gap(): + """Present at first and last timepoints, but absent in between - this is + the "original" side of the reviewer's ambiguity, distinct from + continuous presence.""" + result = _summarise_sample_trajectory(["0.1", ".", "0.1"], ["s0", "s1", "s2"]) + assert result["variant_status"] == "original" + assert result["persistent_status"] == "original_intermittent" + assert result["presence_absence"] == ["Y", "N", "Y"] + assert result["first_appearance"] == "s0" + assert result["last_appearance"] == "s2" + + +def test_summarise_sample_trajectory_original_lost(): + result = _summarise_sample_trajectory(["0.1", "0.1", "."], ["s0", "s1", "s2"]) + assert result["persistent_status"] == "original_lost" + + +def test_summarise_sample_trajectory_new_persistent_no_gap(): + result = _summarise_sample_trajectory([".", "0.1", "0.1"], ["s0", "s1", "s2"]) + assert result["variant_status"] == "new" + assert result["persistent_status"] == "new_persistent" + + +def test_summarise_sample_trajectory_new_intermittent_has_gap(): + """The exact scenario the reviewer flagged: a new mutation that appears, + disappears in a subsequent passage, then reappears by the final + timepoint. Previously indistinguishable from continuous persistence.""" + result = _summarise_sample_trajectory( + [".", "0.1", ".", "0.1"], ["s0", "s1", "s2", "s3"] + ) + assert result["variant_status"] == "new" + assert result["persistent_status"] == "new_intermittent" + assert result["first_appearance"] == "s1" + assert result["last_appearance"] == "s3" + + +def test_summarise_sample_trajectory_new_transient(): + result = _summarise_sample_trajectory([".", "0.1", "."], ["s0", "s1", "s2"]) + assert result["persistent_status"] == "new_transient" + + +def test_summarise_sample_trajectory_absent_throughout_is_new_transient(): + result = _summarise_sample_trajectory([".", ".", "."], ["s0", "s1", "s2"]) + assert result["persistent_status"] == "new_transient" + assert result["first_appearance"] == "None" + assert result["last_appearance"] == "None" + + +def test_summarise_sample_trajectory_leading_absence_is_not_a_gap(): + """A run of absence *before* the first appearance of a new variant is + expected (that's just what "new" means) and must not itself be treated + as a gap.""" + result = _summarise_sample_trajectory( + [".", ".", "0.1", "0.1"], ["s0", "s1", "s2", "s3"] + ) + assert result["persistent_status"] == "new_persistent" + + def test_derive_vcf_output_paths_handles_plain_vcf(tmp_path): out, csq, log = _derive_vcf_output_paths("/data/alpha.vcf", tmp_path, "alpha") diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 3dbfd16..8ad335f 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -19,6 +19,7 @@ import seaborn as sns from .constants import REF_GENE_LENGTHS, NSP_LENGTHS, NSPS from .core import get_logo +from .vcf_processing import NEW_ENDS_PRESENT_STATUSES # Global configuration for plotting plt.rcdefaults() @@ -938,7 +939,7 @@ def _prepare_variant_heatmap_matrix( if ( only_persistent and str(getattr(row, "persistence_status", "")).strip().lower() - != "new_persistent" + not in NEW_ENDS_PRESENT_STATUSES ): continue if ( diff --git a/vartracker/main.py b/vartracker/main.py index 37d9c9d..7a459bf 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -32,7 +32,13 @@ ProcessingError, FILE_COLUMNS, ) -from .vcf_processing import annotate_vcf, format_vcf, merge_consequences, process_vcf +from .vcf_processing import ( + NEW_ENDS_PRESENT_STATUSES, + annotate_vcf, + format_vcf, + merge_consequences, + process_vcf, +) from .analysis import ( process_joint_variants, generate_cumulative_lineplot, @@ -191,7 +197,10 @@ def add_legacy(name: str, dest: str, **kwargs) -> None: action="store_true", default=False, dest="heatmap_only_persistent", - help="Only include variants with persistence_status == new_persistent", + help=( + "Only include new variants present at the final timepoint " + "(persistence_status new_persistent or new_intermittent)" + ), ) add_legacy("only-persistent", "heatmap_only_persistent", action="store_true") group.add_argument( @@ -387,7 +396,10 @@ def _add_shared_plot_filter_arguments(group: argparse._ArgumentGroup) -> None: "--persistent-only", action="store_true", default=False, - help="Only include variants with persistence_status == new_persistent", + help=( + "Only include new variants present at the final timepoint " + "(persistence_status new_persistent or new_intermittent)" + ), ) group.add_argument( "--new-only", @@ -676,6 +688,7 @@ def _configure_vcf_parser( include_input_csv: bool, input_csv_required: bool = False, include_consensus_options: bool = False, + consensus_group: argparse._ArgumentGroup | None = None, ) -> None: if include_input_csv: if input_csv_required: @@ -745,10 +758,17 @@ def _configure_vcf_parser( required=False, type=int, default=10, - help="Minimum depth threshold for variant QC (default: 10)", + help=( + "Minimum depth threshold for variant QC (default: 10). Below " + "this depth at a sample, genuine absence of a variant cannot be " + "distinguished from dropout/non-detection; that sample is " + "flagged 'F' in the per_sample_variant_qc output column rather " + "than treated as confident absence." + ), ) if include_consensus_options: - analysis_group.add_argument( + group = consensus_group or analysis_group + group.add_argument( "--consensus-snp-min-af", action="store", required=False, @@ -759,7 +779,7 @@ def _configure_vcf_parser( "considered for consensus (default: 0.25)" ), ) - analysis_group.add_argument( + group.add_argument( "--consensus-snp-thresh", action="store", required=False, @@ -770,7 +790,7 @@ def _configure_vcf_parser( "consensus base (default: 0.75)" ), ) - analysis_group.add_argument( + group.add_argument( "--consensus-indel-thresh", action="store", required=False, @@ -1061,6 +1081,7 @@ def _add_bam_subparser(subparsers): include_input_csv=True, input_csv_required=False, include_consensus_options=True, + consensus_group=snk_group, ) _move_action_group_after( bam_parser, "Snakemake options", "Vartracker Analysis Options" @@ -1672,7 +1693,10 @@ def _add_plot_genome_subparser(subparsers): "--persistent-only", action="store_true", default=False, - help="Only include variants with persistence_status == new_persistent", + help=( + "Only include new variants present at the final timepoint " + "(persistence_status new_persistent or new_intermittent)" + ), ) filter_group.add_argument( "--new-only", @@ -2175,6 +2199,7 @@ def _add_e2e_subparser(subparsers): e2e_parser, include_input_csv=False, include_consensus_options=True, + consensus_group=snk_group, ) _move_action_group_after( e2e_parser, "Snakemake options", "Vartracker Analysis Options" @@ -2249,6 +2274,15 @@ def _run_e2e_command(args): optional_empty={"bam", "vcf", "coverage", "reads2"}, ) + if not args.primer_bed: + print( + "\033[93mWarning:\033[0m no --primer-bed supplied. If this data was " + "generated with amplicon sequencing, primer-binding sites will not " + "be clipped and may produce spurious variant calls near " + "primer-overlap regions. If your library prep is not amplicon-based, " + "this does not apply and can be ignored." + ) + print(get_logo()) rulegraph_path = _normalise_rulegraph_path(args.rulegraph) @@ -2862,7 +2896,9 @@ def _process_files( ].reset_index(drop=True) new_mutations.to_csv(os.path.join(args.outdir, "new_mutations.csv"), index=None) - persistent_mutations = table[table.persistence_status == "new_persistent"][ + persistent_mutations = table[ + table.persistence_status.isin(NEW_ENDS_PRESENT_STATUSES) + ][ ["gene", "variant", "amino_acid_consequence", "nsp_aa_change"] ].reset_index(drop=True) persistent_mutations.to_csv( diff --git a/vartracker/plotting.py b/vartracker/plotting.py index 0ae53b0..5d00f25 100644 --- a/vartracker/plotting.py +++ b/vartracker/plotting.py @@ -22,6 +22,7 @@ _resolve_variant_labels, ) from .core import InputValidationError, ProcessingError +from .vcf_processing import NEW_ENDS_PRESENT_STATUSES DEFAULT_TRAJECTORY_TOP_N = 12 DEFAULT_LIFESPAN_TOP_N = 20 @@ -249,7 +250,9 @@ def prepare_plot_inputs( summary["duration"] = summary["last_seen"].fillna(0).astype(float) - summary[ "first_seen" ].fillna(0).astype(float) - summary["is_persistent_new"] = summary["persistence_status"].eq("new_persistent") + summary["is_persistent_new"] = summary["persistence_status"].isin( + NEW_ENDS_PRESENT_STATUSES + ) summary["is_nonsynonymous"] = ( ~summary["type_of_change"].str.lower().str.contains("synonymous", na=False) ) @@ -308,7 +311,7 @@ def apply_shared_plot_filters( if persistent_only: filtered_summary = filtered_summary[ - filtered_summary["persistence_status"].eq("new_persistent") + filtered_summary["persistence_status"].isin(NEW_ENDS_PRESENT_STATUSES) ] if new_only: @@ -1242,7 +1245,7 @@ def _subset_collapsed_variants( if max_af is not None: subset = subset[subset["summary_af"] <= max_af] if persistent_only: - subset = subset[subset["persistence_status"].eq("new_persistent")] + subset = subset[subset["persistence_status"].isin(NEW_ENDS_PRESENT_STATUSES)] if new_only: subset = subset[subset["variant_status"].eq("new")] if subset.empty: diff --git a/vartracker/schemas.py b/vartracker/schemas.py index 8ea37b2..b69e510 100644 --- a/vartracker/schemas.py +++ b/vartracker/schemas.py @@ -113,9 +113,18 @@ { "name": "persistence_status", "type": "string", - "description": "Persistence class based on first and last sample presence.", + "description": ( + "Persistence class based on presence at the first and last " + "sample, and whether presence was continuous in between. The " + "'_intermittent' classes mean the variant was present at both " + "the first/last (original) or reappeared by the last sample " + "(new), but was absent from at least one sample in between." + ), "units": "", - "values": "original_retained, original_lost, new_persistent, new_transient", + "values": ( + "original_retained, original_intermittent, original_lost, " + "new_persistent, new_intermittent, new_transient" + ), }, { "name": "presence_absence", @@ -155,7 +164,13 @@ { "name": "per_sample_variant_qc", "type": "string (slash-separated)", - "description": "Per-sample QC flags (P/F) ordered by input.", + "description": ( + "Per-sample QC flags (P/F) ordered by input. A sample fails ('F') " + "when there is no variant-supporting read and site coverage is " + "below --min-depth, i.e. when genuine absence of the variant " + "cannot be distinguished from dropout/non-detection; 'P' means " + "absence (or presence) was confidently called at that sample." + ), "units": "", "values": "P, F", }, diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index 0c5f257..31241d4 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -799,27 +799,51 @@ def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): return result +# A "new" variant present at the final timepoint is new_persistent when its +# presence was continuous from first appearance onward, or new_intermittent +# when it disappeared and reappeared along the way. Both are new variants +# that reached the final timepoint, just with a different path there - +# consumers that filter on "did this new variant persist to the end" should +# match this whole set, not new_persistent alone. +NEW_ENDS_PRESENT_STATUSES = frozenset({"new_persistent", "new_intermittent"}) + + def _summarise_sample_trajectory(allele_freqs, samples): """Derive trajectory metadata from per-sample allele frequencies.""" presence_absence = ["N" if x == "." else "Y" for x in allele_freqs] variant_status = "new" if allele_freqs[0] == "." else "original" + if "Y" in presence_absence: + first_present_idx = presence_absence.index("Y") + last_present_idx = rindex(presence_absence, "Y") + # A gap is an absence sandwiched between the first and last presence, + # i.e. the variant disappeared and later reappeared rather than + # persisting continuously - see original_intermittent/new_intermittent + # below. + has_gap = "N" in presence_absence[first_present_idx : last_present_idx + 1] + else: + first_present_idx = None + last_present_idx = None + has_gap = False + if allele_freqs[0] != "." and allele_freqs[-1] == ".": persistent_status = "original_lost" elif allele_freqs[0] != "." and allele_freqs[-1] != ".": - persistent_status = "original_retained" + persistent_status = ( + "original_intermittent" if has_gap else "original_retained" + ) elif allele_freqs[0] == "." and allele_freqs[-1] != ".": - persistent_status = "new_persistent" + persistent_status = "new_intermittent" if has_gap else "new_persistent" elif allele_freqs[0] == "." and allele_freqs[-1] == ".": persistent_status = "new_transient" else: persistent_status = "unknown" first_appearance = ( - samples[presence_absence.index("Y")] if "Y" in presence_absence else "None" + samples[first_present_idx] if first_present_idx is not None else "None" ) last_appearance = ( - samples[rindex(presence_absence, "Y")] if "Y" in presence_absence else "None" + samples[last_present_idx] if last_present_idx is not None else "None" ) return { From 15359aa9e16ab48d721aa57da8b478a5064d6ee9 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Mon, 3 Aug 2026 12:23:59 +1000 Subject: [PATCH 18/20] Bump to v2.3.0: fix bacterial-genome joint-variant/canonical-row handling and duplicate-row plotting, add PAO1 validation scripts, correct README/CHANGELOG --- .gitignore | 3 + CHANGELOG.md | 93 ++++-- CITATION.cff | 8 +- README.md | 100 +++++- docs/OUTPUT_SCHEMA.md | 4 +- pyproject.toml | 2 +- scripts/validation/pao1/README.md | 73 +++++ scripts/validation/pao1/run_validation.sh | 61 ++++ scripts/validation/pao1/simulate_pao1.py | 370 ++++++++++++++++++++++ tests/test_analysis.py | 261 +++++++++++++++ tests/test_cli.py | 3 +- tests/test_plotting.py | 228 ++++++++++++- tests/test_vcf_processing.py | 96 ++++++ vartracker/_version.py | 30 +- vartracker/analysis.py | 269 +++++++++++++++- vartracker/main.py | 62 +++- vartracker/plotting.py | 125 ++++++-- vartracker/schemas.py | 15 +- vartracker/vcf_processing.py | 202 +++++++++--- 19 files changed, 1861 insertions(+), 144 deletions(-) create mode 100644 scripts/validation/pao1/README.md create mode 100644 scripts/validation/pao1/run_validation.sh create mode 100644 scripts/validation/pao1/simulate_pao1.py diff --git a/.gitignore b/.gitignore index d548845..9f32e87 100755 --- a/.gitignore +++ b/.gitignore @@ -119,3 +119,6 @@ vartracker/.snakemake/ !vartracker/test_data/ !vartracker/test_data/** vartracker/test_data/**/.DS_Store + +# Claude Code local config +.claude/ diff --git a/CHANGELOG.md b/CHANGELOG.md index b29a55f..6938a42 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,40 +1,69 @@ # Changelog -This file tracks changes made during the NARGAB manuscript revision cycle, in -addition to (not replacing) the project's versioned GitHub releases. +All notable changes to vartracker are documented in this file, in addition to +(not replacing) the project's versioned GitHub releases. Format loosely follows [Keep a Changelog](https://keepachangelog.com/). -## [Unreleased] +## [2.3.0] - 2026-07-31 + +### Added + +- `--local-csq` option (`vartracker vcf`/`bam`/`end-to-end`) to switch `bcftools csq` to + independent, per-variant consequence calling instead of the default joint/compound calling - + useful for gene-dense, high-variant-density data where unphased, sub-consensus variants can + cluster in the same gene without genotype evidence they co-occur. +- `--out`/`--outdir` for `vartracker plot heatmap`, to control the output filename and location. +- `--multiallelic-overflow` (`error`/`drop-lowest-af`/`skip-site`) to control how sites with more + than two surviving ALT alleles in a single sample are handled before `bcftools csq`. +- LoFreq primer-overlap rescue for amplicon data (`--primer-bed`, `--lofreq-primer-rescue`, + `--lofreq-rescue-*` thresholds), with raw, rescued, and filtered-out LoFreq calls written to + per-sample audit TSVs. +- Consensus genome generation in the `bam`/`end-to-end` workflows (`--consensus-snp-min-af`, + `--consensus-snp-thresh`, `--consensus-indel-thresh`), writing both a simple and an + IUPAC-aware consensus FASTA per sample. +- Gene-name disambiguation across contigs/replicons: `ambiguous_gene_names()` and updates to + `gene_lengths_from_gff3`/`generate_gene_table` so genes that legitimately reuse a name on + different contigs (e.g. `repA` on a chromosome and a plasmid) are tracked and reported separately. +- Gene-wise summary plot (`mutations_per_gene.pdf`) is now capped to the top genes by newly + emerged variants on large (e.g. bacterial) genomes, via `--max-plot-genes`/`--plot-genes`; the + tabular output is unaffected. +- New standalone plotting subcommands (`plot heatmap`, `plot genome`, `plot trajectory`, + `plot turnover`, `plot lifespan`) and heatmap filtering options (`--aa-exclude`, `--aa-include`, + `--only-persistent`, `--only-new`, `--gene-include`/`--gene-exclude`, `--variant-type`, `--qc`, + `--min-prop-passing-qc`, `--min-persistence`, `--min-max-af`, `--min-sample-af`, + `--sample-subset`, `--hide-singletons`, `--min-depth`, `--x-labels`, `--title`, + `--literature-csv`). +- Structured logging for Snakemake pipeline steps. +- Documentation: a guide for building a custom `--literature-csv` for non-SARS-CoV-2 pathogens, a + QC-column interpretation guide, and a bacterial-genome validation/limitations section. + +### Changed + +- `vartracker plot heatmap` CLI flags were de-prefixed for standalone use (e.g. + `--heatmap-aa-exclude` -> `--aa-exclude`); the legacy prefixed forms are kept as aliases. +- The default heatmap now always shows a variant's canonical row, even if that row is + joint/compound, instead of unconditionally hiding all joint rows; `--include-joint` reveals + any additional joint/compound annotation-group rows a variant has. +- Per-sample variant QC is now reported directly (`per_sample_variant_qc`, `P`/`F`) rather than + only as an overall pass/fail, and the default heatmap visually flags failing cells. ### Fixed -- **Multi-contig gene-name collisions no longer silently merged.** Bacterial - reference genomes commonly reuse gene names across replicons (e.g. `repA` - on both a chromosome and a plasmid). Previously, `gene_lengths_from_gff3` - and `generate_gene_table` aggregated CDS lengths and mutation statistics by - gene name alone, so two unrelated genes sharing a name on different contigs - had their lengths and variant counts silently summed into one row. - - `gene_lengths_from_gff3` (`annotation_processing.py`) now tracks CDS - lengths per `(contig, gene)` internally and only disambiguates the - output label (e.g. `repA (chrom1)` / `repA (plasmid1)`) when a gene name - actually collides across contigs. Gene names confined to a single contig - are returned unqualified, so single-contig (viral) references and - existing multi-segment references with uniquely-named segments (e.g. - 8-segment influenza) are unaffected. - - `generate_gene_table` (`analysis.py`) now splits colliding gene names by - `chrom` in the same way, using the results table's own `chrom` column. - - Added `ambiguous_gene_names()` (`annotation_processing.py`), which - identifies gene names that collide across contigs directly from the - reference annotation. This closes an edge case where a colliding gene - only has variants on one of its contigs in a given dataset: without it, - that gene's real data could be silently dropped when merged against the - gene-length scaffold, because the scaffold would only contain the - contig-qualified labels while the variants table only produced the bare, - unqualified label. - - Added test coverage: `tests/test_annotation_processing.py` (new file) - and three new tests in `tests/test_analysis.py` covering the - single-contig baseline, the colliding-gene-name split, and the - ambiguous-genes data-loss edge case. - - No changes to output schema, column names, or behaviour for any - single-contig or uniquely-named multi-segment reference. +- Duplicate `results.csv` rows for the same variant (one per `bcftools csq` annotation group) were + previously plotted as independent data points, inflating turnover counts; plotting now collapses + to one row per (variant, sample), taking the union of presence and the maximum allele frequency. +- `joint_variant` could miss rows with a genuinely compound `bcsq_nt_notation` when a variant had + several sibling rows at the same position; it now flags every such row. +- Heatmap label truncation could make two distinct variants render an identical label, silently + dropping one of them; colliding labels are now disambiguated with a positional suffix. +- Gene names that collide across contigs/replicons were previously merged into a single row in + gene-length and gene-wise summary tables; they are now tracked and reported per-contig. +- Percent-encoded gene names read from `bcftools`' `BCSQ` field (e.g. `ercS%27`) were not + URL-decoded on the variant-calling side, so they never matched the (correctly decoded) + annotation side and were silently dropped from gene-wise summaries. +- Corrected several low-frequency multiallelic-site handling bugs, so distinct ALT alleles at the + same position are kept separate through preprocessing and correctly rejoined before + `bcftools csq`. +- Fixed potential results-row duplication for some samples during heatmap generation. +- Fixed longitudinal variant-rescue edge cases and LoFreq strand-bias rescue thresholds. diff --git a/CITATION.cff b/CITATION.cff index 55d2fd2..1381fd3 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -1,8 +1,8 @@ cff-version: 1.2.0 message: "If you use this software, please cite it as below." title: "vartracker" -version: "2.2.1" -date-released: "2026-05-06" +version: "2.3.0" +date-released: "2026-07-31" license: MIT repository-code: "https://github.com/charlesfoster/vartracker" url: "https://github.com/charlesfoster/vartracker" @@ -22,7 +22,7 @@ preferred-citation: - family-names: Foster given-names: Charles title: "vartracker" - version: "2.2.1" + version: "2.3.0" doi: "10.5281/zenodo.18452274" url: "https://github.com/charlesfoster/vartracker" - date-released: "2026-05-06" + date-released: "2026-07-31" diff --git a/README.md b/README.md index 3eefbac..bf41af9 100755 --- a/README.md +++ b/README.md @@ -168,7 +168,7 @@ Docker is a self-contained reproducible option. If you publish the image, record set it when running to include it in the run manifest: ```bash -export VARTRACKER_CONTAINER_IMAGE=ghcr.io/your-org/vartracker:2.2.1 +export VARTRACKER_CONTAINER_IMAGE=ghcr.io/your-org/vartracker:2.3.0 export VARTRACKER_CONTAINER_DIGEST=sha256:... ``` @@ -310,8 +310,9 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test ### Mode-specific options - `vartracker vcf` – accepts core analysis options such as `--min-snv-freq`, `--min-indel-freq`, - `--allele-frequency-tag`, `--multiallelic-overflow`, `--name`, `--outdir`, `--sample-cap`, `--manifest-level`, and literature controls - (`--search-pokay`, `--literature-csv`). Use `--test` to run the bundled smoke test. + `--allele-frequency-tag`, `--multiallelic-overflow`, `--local-csq`, `--name`, `--outdir`, + `--sample-cap`, `--manifest-level`, and literature controls (`--search-pokay`, + `--literature-csv`). Use `--test` to run the bundled smoke test. `--max-plot-genes` and `--plot-genes` control the gene-wise summary figure only (see [Limitations](#limitations)); the tabular/TSV output always includes every annotated gene. - `vartracker bam` – everything from `vcf`, plus Snakemake options: @@ -321,7 +322,7 @@ for both `.depth.txt` and `_depth.txt` patterns when preparing its internal test - `vartracker end-to-end` – similar to `bam`, with optional amplicon clipping controls: `--primer-bed` and `--ampliconclip-tolerance` (default: `1`). Supplying `--primer-bed` also enables LoFreq primer-overlap rescue by default. -- `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV, including all heatmap customization filters. +- `vartracker plot heatmap` (`hm`) – regenerate the heatmap from an existing vartracker results CSV, including all heatmap customisation filters. - `vartracker plot genome` – plot SNP positions along the genome or a selected gene region using all observed allele-frequency values for each variant. - `vartracker plot trajectory` – plot allele-frequency trajectories for a selected or auto-ranked subset of variants, optionally in takeover mode using threshold lines and threshold-based filtering. - `vartracker plot turnover` – plot new-versus-lost longitudinal turnover from the filtered result set. @@ -346,7 +347,9 @@ Consequence-calling note: Heatmap filtering: - `vcf`, `bam`, and `end-to-end` always write the default heatmap. To customise heatmap content after a run, use `vartracker plot heatmap results.csv [options]`. -- By default, all consequence classes are included except joint variants. Use `--include-joint` to show joint variants. +- By default, each variant is shown once, using its canonical row (whether that row is joint or + not - see [Limitations](#limitations)). Use `--include-joint` to additionally reveal extra + joint/compound annotation-group rows for variants that have more than one. - `--aa-exclude`: comma-separated `type_of_change` patterns to exclude. Wildcards are supported. - `--aa-include`: comma-separated `type_of_change` patterns to include. - `--only-persistent`: only include new variants present at the final timepoint (`new_persistent` or @@ -365,6 +368,10 @@ Heatmap filtering: - `--x-labels sample-number`: label heatmap x-axis columns by `sample_number` instead of sample name. - `--title`: set the heatmap plot title. The default is `Variant allele frequencies`. - `--literature-csv`: include literature links in the interactive HTML heatmap using a literature hits CSV. +- `--out` (`vartracker plot heatmap` only): write the heatmap using this path as the base name, + e.g. `--out plots/myheatmap` writes `plots/myheatmap.pdf` and `plots/myheatmap.html`. +- `--outdir` (`vartracker plot heatmap` only): output directory for heatmap files (default: + beside `results.csv`). - Example: `--aa-exclude "synonymous,*frameshift*,stop_gained"` Standalone plot filtering: @@ -605,6 +612,39 @@ standalone plots) and `persistent_new_mutations.csv` include both `new_persisten `new_intermittent` variants, since both reached the final timepoint; the label only distinguishes the path taken to get there. +### Interpreting the QC columns + +`results.csv` records presence/absence per sample (`presence_absence`, `Y`/`N`), but an `N` does +not always mean the variant was confidently confirmed absent. At low sequencing depth, a variant +can go undetected simply because there was insufficient coverage to call it either way - this is +indistinguishable, from the VCF alone, from genuine absence. The QC columns exist to flag this: + +- `per_sample_variant_qc`: a per-sample `P`/`F` flag. `F` means that sample had no + variant-supporting read *and* site coverage below `--min-depth` (default: 10) - i.e. absence + could not be confidently distinguished from dropout/non-detection at that sample. `P` means the + call (presence or absence) was made with confidence. +- `all_samples_pass_qc`: `true` only if every sample is `P`. +- `proportion_samples_passing_qc`: the fraction of samples that are `P`. + +**Practical guidance:** if `all_samples_pass_qc` is `false` for a variant, inspect +`per_sample_variant_qc` to see exactly which sample(s) it failed at - e.g. `P / P / F / P / P / P` +identifies the third sample as the QC failure. Before treating an `N` in `presence_absence` as +evidence a variant was truly lost or never present, check the corresponding position in +`per_sample_variant_qc`: an `N` paired with `F` should be read as "not detected", not "confirmed +absent" - this is especially relevant for low-frequency variants near the allele-frequency or depth +thresholds (`--min-snv-freq`, `--min-indel-freq`, `--min-depth`), where dropout is more likely than +at high-confidence, high-depth sites. This ambiguity also propagates into `persistence_status` (see +[Persistence labels](#persistence-labels)): an apparent loss-then-reappearance (`*_intermittent`) +may reflect genuine intermittent presence, or simply a low-coverage sample in between. + +**QC in the heatmap.** The default heatmap marks `F` cells visually rather than just via colour: the +static PDF draws an unfilled black-bordered rectangle over any cell whose sample failed QC for that +variant; the interactive HTML version uses a dark inset ring plus a hover tooltip reading +`QC=FAIL`. To exclude variants that don't pass QC from a plot entirely (rather than just flagging +the cells), use `--qc` and `--min-prop-passing-qc` (see +[Mode-specific options](#mode-specific-options)), or inspect `per_sample_variant_qc` directly for +the samples of interest. + ### Output schema The results table schema is documented in `docs/OUTPUT_SCHEMA.md`. You can also print it from the CLI: @@ -653,7 +693,7 @@ The practical constraint on larger genomes (e.g. bacterial pathogens, which can annotated genes) is **visualisation**, not computation: - The gene-wise summary figure (`mutations_per_gene.pdf`) plots one bar per gene per panel. On a - genome with thousands of annotated genes this becomes unreadable regardless of how many variants + genome with thousands of annotated genes this becomes unreadable as a static image regardless of how many variants are actually present, because the plot iterates over every annotated gene, not just genes that carry a variant. - By default, the figure is capped to the top 30 genes, ranked by number of newly emerged variants @@ -663,6 +703,52 @@ annotated genes) is **visualisation**, not computation: - When the figure is truncated, this is stated directly on the figure itself (e.g. "top 30 of 412 genes with variants"); if nothing was truncated, no such note is shown. +### Bacterial genomes + +vartracker works well and efficiently at bacterial genome scale. It has been validated in `vcf` +mode (i.e. from pre-called VCFs and coverage files, not the `bam`/`end-to-end` read-mapping +workflow) against simulated *Pseudomonas aeruginosa* PAO1 data (NC_002516.2, 6.26 Mb, 5,573 CDS +features) across two scenarios - 80 and 1,000 simulated variants, each across 6 timepoints. The +smaller, 80-variant scenario completed in approximately 11 seconds of wall-clock time with +approximately 0.9 GiB peak memory; the larger, 1,000-variant scenario completed in approximately +22 seconds with approximately 2.8 GiB peak memory. At this scale, the practical caveat is not +runtime or memory but the **interpretability of joint/compound amino-acid consequences in +gene-dense hotspots**, discussed below. + +**Joint vs local `bcftools csq` calling.** By default, vartracker calls consequences jointly (the +`bcftools csq` default), so that variants close enough together to plausibly affect the same +codon(s) are described together as a single, compound amino-acid change. This is the correct +behaviour for genuinely linked variants, but on gene-dense, high-variant-density data - common in +bacterial within-host or experimental-evolution datasets, and rare in vartracker's original viral +use case - many unphased, sub-consensus variants can cluster in the same gene without genotype +evidence that they actually co-occur on the same haplotype. Joint calling then produces long, +compound descriptions that are technically correct but hard to read, and can fragment a single +variant's presence/absence trajectory across samples. The `--local-csq` option (see +`vartracker --help`) switches to independent, SnpEff-like per-variant consequence calling, at the +cost of no longer detecting genuinely combined effects between physically linked variants. Whichever +mode is used, rows describing a joint/compound consequence are flagged in the `joint_variant` +column of `results.csv`, which can be used to identify or filter these rows after the fact. This +column is now fully reliable: on the PAO1 validation dataset, 100% of genuinely compound +`bcftools csq` rows are correctly flagged. + +**Heatmap `--include-joint` semantics.** By default, the heatmap shows one row per variant - its +canonical row, whether that row happens to be joint or not. `--include-joint` additionally reveals +extra joint/compound annotation-group rows for variants that have more than one. This is a different +kind of control from the gene-wise figure's `--max-plot-genes` cap: there is no row cap on the +heatmap. + +**Heatmap legibility at bacterial scale.** Unlike the gene-wise figure, the heatmap has no built-in +row cap, and is effectively illegible as a static image. However, you can still open it and scroll to read the rows. Alternatively, for large numbers of variants, narrow the heatmap using the +"Heatmap filtering" options described under [Mode-specific options](#mode-specific-options) - for +example `--gene-include`, `--hide-singletons`, `--min-max-af`, and `--only-persistent` - or generate +multiple heatmaps over subsets of genes/samples rather than relying on a single, unfiltered plot. + +**Coverage-file disk and memory footprint.** Disk and memory usage for coverage/depth files scale +with genome length multiplied by timepoint count. For reference, the PAO1 validation used 6 depth +files at approximately 141 MB each (846 MB total) for one 6.3 Mb genome across 6 timepoints. Users +planning many-timepoint experimental-evolution designs (often dozens of timepoints) on genomes +larger than PAO1 should budget disk and memory accordingly. + Separately, the bundled `pokay` functional-annotation database (see [Using Literature Database](#using-literature-database)) is specific to SARS-CoV-2 mutations and is not applied to, or meaningful for, other pathogens. A custom literature CSV following the @@ -674,7 +760,7 @@ same schema can be supplied via `--literature-csv` for other organisms; see When using vartracker, please cite the software release you used. Citation metadata is provided in `CITATION.cff`, and GitHub releases are archived on Zenodo. -- Foster, C. (2026). *vartracker* (Version 2.2.1). Zenodo. https://doi.org/10.5281/zenodo.18452274 +- Foster, C. (2026). *vartracker* (Version 2.3.0). Zenodo. https://doi.org/10.5281/zenodo.18452274 Note: the DOI above is the Zenodo concept DOI for all versions; a version-specific DOI is minted by Zenodo after each GitHub release. diff --git a/docs/OUTPUT_SCHEMA.md b/docs/OUTPUT_SCHEMA.md index 11ce759..ff824f1 100644 --- a/docs/OUTPUT_SCHEMA.md +++ b/docs/OUTPUT_SCHEMA.md @@ -28,9 +28,9 @@ Columns that encode per-sample values are slash-separated and ordered by the inp | presence_absence | string (slash-separated) | Per-sample presence (Y) or absence (N), ordered by input. | | Y/N | | first_appearance | string | Sample name where the variant first appears. | | | | last_appearance | string | Sample name where the variant last appears. | | | -| all_samples_pass_qc | boolean | True if every sample passes per-sample variant QC. | | true, false | +| all_samples_pass_qc | boolean | True only if every sample is 'P' in per_sample_variant_qc. False means at least one sample's presence/absence call could not be confidently distinguished from dropout/non-detection - inspect per_sample_variant_qc to see which sample(s) are affected before treating this variant's presence/absence pattern as reliable. | | true, false | | proportion_samples_passing_qc | number | Proportion of samples passing per-sample variant QC. | fraction | 0-1 | -| per_sample_variant_qc | string (slash-separated) | Per-sample QC flags (P/F) ordered by input. A sample fails ('F') when there is no variant-supporting read and site coverage is below --min-depth, i.e. when genuine absence of the variant cannot be distinguished from dropout/non-detection; 'P' means absence (or presence) was confidently called at that sample. | | P, F | +| per_sample_variant_qc | string (slash-separated) | Per-sample QC flags (P/F) ordered by input. A sample fails ('F') when there is no variant-supporting read and site coverage is below --min-depth, i.e. when genuine absence of the variant cannot be distinguished from dropout/non-detection; 'P' means absence (or presence) was confidently called at that sample. E.g. 'P / P / F / P / P / P' identifies the third sample as the one where QC failed. The default heatmap marks 'F' cells visually: an unfilled black-bordered rectangle in the static PDF, and a dark inset ring plus a 'QC=FAIL' hover tooltip in the interactive HTML version. | | P, F | | aa1_total_properties | string | Physicochemical properties for the reference amino acid. | | semicolon-separated properties | | aa2_total_properties | string | Physicochemical properties for the alternate amino acid. | | semicolon-separated properties | | aa1_unique_properties | string | Properties unique to the reference amino acid. | | semicolon-separated properties | diff --git a/pyproject.toml b/pyproject.toml index 790e3cc..99696a7 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "vartracker" -version = "2.2.1" +version = "2.3.0" authors = [ {name = "Dr Charles Foster"}, ] diff --git a/scripts/validation/pao1/README.md b/scripts/validation/pao1/README.md new file mode 100644 index 0000000..ef3415e --- /dev/null +++ b/scripts/validation/pao1/README.md @@ -0,0 +1,73 @@ +# PAO1 bacterial-genome validation + +Scripts to reproduce the *Pseudomonas aeruginosa* PAO1 (NC_002516.2) bacterial-scale validation +referenced in the main README's [Bacterial genomes](../../../README.md#bacterial-genomes) section +and in the accompanying manuscript. This regenerates two synthetic longitudinal datasets from +scratch, then runs `vartracker vcf` on each: + +- **scenario_a**: 80 variants, uniform gene weighting, 6 timepoints. +- **scenario_b**: 1,000 variants, Zipf-biased gene weighting (concentrating variants in a smaller + set of genes, to stress-test gene-dense hotspots), 6 timepoints. + +## What's here + +- `simulate_pao1.py`: builds both scenarios against a supplied PAO1 reference FASTA/GFF3. For each + scenario it assigns simulated variants a `persistence_status` pattern (`original_retained`, + `original_lost`, `new_persistent`, `new_transient`, `original_intermittent`, + `new_intermittent`), synthesises per-timepoint VCFs and matching depth/coverage files, and writes + a `_ground_truth.csv` (intended pattern, presence/absence, allele frequency per + timepoint) alongside the `_input.csv` manifest `vartracker vcf` expects. All randomness + is seeded (`scenario_a`: seed 42; `scenario_b`: seed 43; coverage: seed 123), so re-running + produces identical output. +- `run_validation.sh`: end-to-end driver - builds the reference bundle, runs the simulation, then + runs `vartracker vcf` on both scenarios (benchmarked with GNU time if available). + +## Requirements + +- `vartracker` installed and on `PATH` (see the main [README](../../../README.md#installation)). +- `bcftools` and `bgzip` (external dependencies of `simulate_pao1.py` and of vartracker itself). +- `python3` with `numpy` and `pandas`. +- Optional, to reproduce reported timing/peak-memory figures: GNU time - `gtime` on macOS + (`brew install gnu-time`) or `/usr/bin/time -v` on Linux. Without it the scenarios still run, just + unbenchmarked. + +## Usage + +```bash +scripts/validation/pao1/run_validation.sh [outdir] # outdir defaults to ./pao1_validation +``` + +This is equivalent to running the steps individually: + +```bash +vartracker prepare reference --accessions NC_002516.2 \ + --outdir pao1_validation/refs/pao1 --prefix pao1_ref + +python3 scripts/validation/pao1/simulate_pao1.py \ + --fasta pao1_validation/refs/pao1/pao1_ref.fa \ + --gff3 pao1_validation/refs/pao1/pao1_ref.gff3 \ + --base-outdir pao1_validation + +vartracker vcf pao1_validation/scenario_a/scenario_a_input.csv \ + --reference pao1_validation/refs/pao1/pao1_ref.fa \ + --gff3 pao1_validation/refs/pao1/pao1_ref.gff3 \ + --outdir pao1_validation/results/pao1_80_variants + +vartracker vcf pao1_validation/scenario_b/scenario_b_input.csv \ + --reference pao1_validation/refs/pao1/pao1_ref.fa \ + --gff3 pao1_validation/refs/pao1/pao1_ref.gff3 \ + --outdir pao1_validation/results/pao1_1000_variants +``` + +## Checking the results + +- Compare `results/pao1_*_variants/results.csv` (specifically `variant`, `variant_status`, and + `persistence_status`) against the corresponding `scenario_{a,b}_ground_truth.csv` to confirm every + simulated variant and its intended temporal pattern was recovered. +- Timing and peak memory are hardware-dependent, so exact figures will vary by machine; the README + and manuscript report approximately 11 s / 0.9 GiB peak memory for scenario_a and approximately + 22 s / 2.8 GiB peak memory for scenario_b, measured via `gtime -v` on the reference machine used + for validation. +- Neither the generated coverage files (large - tens to hundreds of MB per timepoint) nor the + reference bundle are committed to this repository; both are deterministically regenerated by the + steps above. diff --git a/scripts/validation/pao1/run_validation.sh b/scripts/validation/pao1/run_validation.sh new file mode 100644 index 0000000..c649fda --- /dev/null +++ b/scripts/validation/pao1/run_validation.sh @@ -0,0 +1,61 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Reproduces the PAO1 bacterial-genome validation described in the vartracker +# README ("Bacterial genomes" section under Limitations) and manuscript: a +# synthetic 80-variant scenario and a synthetic 1000-variant scenario, each +# spanning 6 longitudinal timepoints, simulated against Pseudomonas aeruginosa +# PAO1 (NC_002516.2) and analysed with `vartracker vcf`. +# +# Usage: +# scripts/validation/pao1/run_validation.sh [outdir] +# +# Requires on PATH: vartracker, bcftools, bgzip, python3 (with numpy and +# pandas). Optional, to reproduce the reported timing/peak-memory figures: +# GNU time - `gtime` on macOS (brew install gnu-time) or `/usr/bin/time -v` +# on Linux. Without either, the analysis still runs, just unbenchmarked. + +OUTDIR="${1:-pao1_validation}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +mkdir -p "$OUTDIR" + +echo "== Step 1: build the PAO1 reference bundle (NC_002516.2) ==" +vartracker prepare reference \ + --accessions NC_002516.2 \ + --outdir "$OUTDIR/refs/pao1" \ + --prefix pao1_ref + +echo "== Step 2: simulate scenario_a (80 variants, uniform gene weighting) and scenario_b (1000 variants, Zipf-biased gene weighting) ==" +python3 "$SCRIPT_DIR/simulate_pao1.py" \ + --fasta "$OUTDIR/refs/pao1/pao1_ref.fa" \ + --gff3 "$OUTDIR/refs/pao1/pao1_ref.gff3" \ + --base-outdir "$OUTDIR" + +if command -v gtime >/dev/null 2>&1; then + TIME_CMD="gtime -v" +elif /usr/bin/time -v true >/dev/null 2>&1; then + TIME_CMD="/usr/bin/time -v" +else + echo "No GNU time found (gtime / time -v); continuing without wall-clock/memory benchmarking." >&2 + TIME_CMD="" +fi + +echo "== Step 3: run vartracker vcf on scenario_a ==" +$TIME_CMD vartracker vcf "$OUTDIR/scenario_a/scenario_a_input.csv" \ + --reference "$OUTDIR/refs/pao1/pao1_ref.fa" \ + --gff3 "$OUTDIR/refs/pao1/pao1_ref.gff3" \ + --outdir "$OUTDIR/results/pao1_80_variants" + +echo "== Step 4: run vartracker vcf on scenario_b ==" +$TIME_CMD vartracker vcf "$OUTDIR/scenario_b/scenario_b_input.csv" \ + --reference "$OUTDIR/refs/pao1/pao1_ref.fa" \ + --gff3 "$OUTDIR/refs/pao1/pao1_ref.gff3" \ + --outdir "$OUTDIR/results/pao1_1000_variants" + +echo +echo "Done." +echo "Cross-check $OUTDIR/results/pao1_*_variants/results.csv against" +echo "$OUTDIR/scenario_{a,b}/scenario_{a,b}_ground_truth.csv, and (if GNU time" +echo "ran above) compare Elapsed / Maximum resident set size against the" +echo "figures quoted in the README's Bacterial genomes section." diff --git a/scripts/validation/pao1/simulate_pao1.py b/scripts/validation/pao1/simulate_pao1.py new file mode 100644 index 0000000..cef7053 --- /dev/null +++ b/scripts/validation/pao1/simulate_pao1.py @@ -0,0 +1,370 @@ +#!/usr/bin/env python3 +"""Simulate longitudinal VCF + coverage data against the PAO1 reference bundle +for vartracker bacterial-scale validation (Scenario A ~80 variants, Scenario B ~1000). +""" +import argparse +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +N_TIMEPOINTS = 6 +BASES = ["A", "C", "G", "T"] + +PATTERNS = [ + "original_retained", + "original_lost", + "new_persistent", + "new_transient", + "original_intermittent", + "new_intermittent", +] +PATTERN_WEIGHTS = [0.30, 0.15, 0.25, 0.15, 0.08, 0.07] +MIN_INTERMITTENT = 3 + + +def parse_fasta(path): + header = None + seq_chunks = [] + with open(path) as fh: + for line in fh: + line = line.rstrip("\n") + if line.startswith(">"): + header = line[1:].split()[0] + else: + seq_chunks.append(line) + return header, "".join(seq_chunks).upper() + + +def parse_cds(gff3_path, contig): + cds = [] + with open(gff3_path) as fh: + for line in fh: + if line.startswith("#"): + continue + parts = line.rstrip("\n").split("\t") + if len(parts) < 9 or parts[0] != contig or parts[2] != "CDS": + continue + start, end, strand, attrs = int(parts[3]), int(parts[4]), parts[6], parts[8] + gene = None + for kv in attrs.split(";"): + if kv.startswith("gene="): + gene = kv.split("=", 1)[1] + break + if gene is None or end - start < 20: + continue + cds.append({"gene": gene, "start": start, "end": end, "strand": strand}) + return cds + + +def zipf_weights(n, rng, exponent=1.1): + order = rng.permutation(n) + ranks = np.empty(n, dtype=float) + ranks[order] = np.arange(1, n + 1) + w = 1.0 / (ranks ** exponent) + return w / w.sum() + + +def choose_positions(cds_list, n_variants, ref_seq, rng, indel_fraction=0.08, biased=False): + """Pick n_variants distinct (pos, ref, alt, gene, var_type) tuples from CDS regions.""" + genes = np.array([c["gene"] for c in cds_list]) + unique_genes, gene_index = np.unique(genes, return_inverse=True) + # map each unique gene name -> list of cds_list indices (handles duplicate-name genes e.g. speA) + gene_to_cds_idx = {g: [] for g in unique_genes} + for i, g in enumerate(genes): + gene_to_cds_idx[g].append(i) + + if biased: + weights = zipf_weights(len(unique_genes), rng) + else: + weights = np.full(len(unique_genes), 1.0 / len(unique_genes)) + + used_positions = set() + variants = [] + attempts = 0 + max_attempts = n_variants * 50 + + while len(variants) < n_variants and attempts < max_attempts: + attempts += 1 + gene_choice = rng.choice(unique_genes, p=weights) + cds_idx = rng.choice(gene_to_cds_idx[gene_choice]) + cds = cds_list[cds_idx] + is_indel = rng.random() < indel_fraction + span = rng.integers(1, 3) if is_indel else 0 + lo, hi = cds["start"] + 3, cds["end"] - 3 - span + if hi <= lo: + continue + pos = int(rng.integers(lo, hi)) + span_positions = set(range(pos, pos + span + 2)) + if span_positions & used_positions: + continue + + ref_base = ref_seq[pos - 1] + if ref_base not in BASES: + continue + + if not is_indel: + alt_base = rng.choice([b for b in BASES if b != ref_base]) + ref, alt, var_type = ref_base, alt_base, "snp" + else: + if rng.random() < 0.5: + # insertion + inserted = "".join(rng.choice(BASES, size=span)) + ref, alt, var_type = ref_base, ref_base + inserted, "indel" + else: + # deletion: anchor base + (span) deleted ref bases + del_seq = ref_seq[pos - 1: pos - 1 + span + 1] + if len(del_seq) < span + 1 or any(b not in BASES for b in del_seq): + continue + ref, alt, var_type = del_seq, del_seq[0], "indel" + + used_positions |= span_positions + variants.append( + { + "chrom": None, # filled by caller + "pos": pos, + "ref": ref, + "alt": alt, + "gene": gene_choice, + "var_type": var_type, + } + ) + + if len(variants) < n_variants: + print( + f"WARNING: only placed {len(variants)}/{n_variants} variants after {attempts} attempts", + file=sys.stderr, + ) + return variants + + +def gen_presence_af(pattern, var_type, rng): + af_floor = 0.12 if var_type == "indel" else 0.05 + af_ceiling = 0.95 + present = [False] * N_TIMEPOINTS + + if pattern == "original_retained": + present = [True] * N_TIMEPOINTS + base = rng.uniform(0.2, 0.5) + af_vals = [np.clip(base + rng.uniform(-0.05, 0.05), af_floor, af_ceiling) for _ in range(N_TIMEPOINTS)] + + elif pattern == "original_lost": + last_present = int(rng.integers(1, 5)) # 1..4 inclusive -> absent by tp5 + present = [i <= last_present for i in range(N_TIMEPOINTS)] + n_present = last_present + 1 + af_vals = list(np.clip(np.linspace(0.55, 0.15, n_present) + rng.uniform(-0.03, 0.03, n_present), af_floor, af_ceiling)) + + elif pattern == "new_persistent": + first_present = int(rng.integers(1, 5)) # 1..4 + present = [i >= first_present for i in range(N_TIMEPOINTS)] + n_present = N_TIMEPOINTS - first_present + af_vals = list(np.clip(np.linspace(0.12, 0.85, n_present) + rng.uniform(-0.03, 0.03, n_present), af_floor, af_ceiling)) + + elif pattern == "new_transient": + first_present = int(rng.integers(1, 4)) # 1..3 + last_present = int(rng.integers(first_present, min(first_present + 3, 5))) # < 5 + present = [first_present <= i <= last_present for i in range(N_TIMEPOINTS)] + n_present = last_present - first_present + 1 + base = rng.uniform(0.15, 0.4) + af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + + elif pattern == "original_intermittent": + gap_start = int(rng.integers(1, 4)) # 1..3 + gap_len = int(rng.integers(1, 5 - gap_start)) # keep gap inside 1..4, tp5 stays present + present = [True] * N_TIMEPOINTS + for i in range(gap_start, min(gap_start + gap_len, 5)): + present[i] = False + present[0] = True + present[5] = True + n_present = sum(present) + base = rng.uniform(0.2, 0.45) + af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + + elif pattern == "new_intermittent": + first_present = int(rng.integers(1, 3)) # 1..2 + gap_start = int(rng.integers(first_present + 1, 5)) # inside 1..4 + gap_len = int(rng.integers(1, max(2, 5 - gap_start))) + present = [False] * N_TIMEPOINTS + for i in range(first_present, N_TIMEPOINTS): + present[i] = True + for i in range(gap_start, min(gap_start + gap_len, 5)): + present[i] = False + present[0] = False + present[5] = True + n_present = sum(present) + base = rng.uniform(0.15, 0.4) + af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + + else: + raise ValueError(pattern) + + # sanity checks on invariants + if pattern.startswith("original"): + assert present[0] is True + else: + assert present[0] is False + if pattern in ("original_retained", "new_persistent", "original_intermittent", "new_intermittent"): + assert present[5] is True + else: + assert present[5] is False + + af_iter = iter(af_vals) + af_full = [round(float(next(af_iter)), 4) if p else None for p in present] + dp_full = [int(rng.integers(80, 200)) if p else None for p in present] + return present, af_full, dp_full + + +def assign_patterns(n, rng): + patterns = list(rng.choice(PATTERNS, size=n, p=PATTERN_WEIGHTS)) + for target in ("original_intermittent", "new_intermittent"): + count = patterns.count(target) + if count < MIN_INTERMITTENT: + donor_pool = [i for i, p in enumerate(patterns) if p not in ("original_intermittent", "new_intermittent")] + n_needed = MIN_INTERMITTENT - count + idx_to_convert = rng.choice(donor_pool, size=n_needed, replace=False) + for idx in idx_to_convert: + patterns[idx] = target + return patterns + + +def simulate_scenario(name, n_variants, cds_list, contig, ref_seq, seed, biased, outdir: Path): + rng = np.random.default_rng(seed) + variants = choose_positions(cds_list, n_variants, ref_seq, rng, biased=biased) + for v in variants: + v["chrom"] = contig + + patterns = assign_patterns(len(variants), rng) + + records = [] + for v, pattern in zip(variants, patterns): + present, af, dp = gen_presence_af(pattern, v["var_type"], rng) + variant_status = "original" if pattern.startswith("original") else "new" + records.append( + { + **v, + "pattern": pattern, + "variant_status": variant_status, + "present": present, + "af": af, + "dp": dp, + } + ) + + records.sort(key=lambda r: r["pos"]) + + # sanity: no duplicate positions + positions = [r["pos"] for r in records] + assert len(positions) == len(set(positions)), "duplicate positions generated" + + outdir.mkdir(parents=True, exist_ok=True) + vcf_dir = outdir / "vcfs" + vcf_dir.mkdir(exist_ok=True) + + for tp in range(N_TIMEPOINTS): + rows = [r for r in records if r["present"][tp]] + rows.sort(key=lambda r: r["pos"]) + raw_path = vcf_dir / f"pao1_tp{tp}.raw.vcf" + with open(raw_path, "w") as fh: + fh.write("##fileformat=VCFv4.2\n") + fh.write(f"##contig=\n") + fh.write('##INFO=\n') + fh.write('##INFO=\n') + fh.write("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n") + for r in rows: + af_val = r["af"][tp] + dp_val = r["dp"][tp] + qual = 60 + fh.write(f"{r['chrom']}\t{r['pos']}\t.\t{r['ref']}\t{r['alt']}\t{qual}\tPASS\tDP={dp_val};AF={af_val}\n") + + final_gz = vcf_dir / f"pao1_tp{tp}.vcf.gz" + sort_cmd = f"bcftools sort -Ov {raw_path} 2>/dev/null | bgzip > {final_gz}" + subprocess.run(sort_cmd, shell=True, check=True) + subprocess.run(["bcftools", "index", "-f", str(final_gz)], check=True) + raw_path.unlink() + + # ground truth table for cross-checking step 4 + gt_rows = [] + for r in records: + gt_rows.append( + { + "chrom": r["chrom"], + "pos": r["pos"], + "ref": r["ref"], + "alt": r["alt"], + "gene": r["gene"], + "var_type": r["var_type"], + "variant_status": r["variant_status"], + "intended_pattern": r["pattern"], + "presence": "/".join("Y" if p else "N" for p in r["present"]), + "af": "/".join("" if a is None else str(a) for a in r["af"]), + } + ) + gt_df = pd.DataFrame(gt_rows).sort_values(["gene", "pos"]) + gt_df.to_csv(outdir / f"{name}_ground_truth.csv", index=False) + + # input CSV for vartracker vcf mode + csv_rows = [] + for tp in range(N_TIMEPOINTS): + csv_rows.append( + { + "sample_name": f"pao1_tp{tp}", + "sample_number": tp, + "reads1": "", + "reads2": "", + "bam": "", + "vcf": str((vcf_dir / f"pao1_tp{tp}.vcf.gz").resolve()), + "coverage": str((outdir.parent / "coverage" / f"pao1_tp{tp}_depth.txt").resolve()), + } + ) + pd.DataFrame(csv_rows).to_csv(outdir / f"{name}_input.csv", index=False) + + # summary counts + pattern_counts = pd.Series([r["pattern"] for r in records]).value_counts() + gene_counts = pd.Series([r["gene"] for r in records]).value_counts() + print(f"[{name}] {len(records)} distinct variants placed across {gene_counts.shape[0]} genes") + print(f"[{name}] pattern mix:\n{pattern_counts}") + print(f"[{name}] top genes by variant count:\n{gene_counts.head(10)}") + return records + + +def generate_coverage(contig, genome_length, coverage_dir: Path, seed=123): + coverage_dir.mkdir(parents=True, exist_ok=True) + pos = np.arange(1, genome_length + 1, dtype=np.int64) + for tp in range(N_TIMEPOINTS): + out_path = coverage_dir / f"pao1_tp{tp}_depth.txt" + if out_path.exists(): + continue + rng = np.random.default_rng(seed + tp) + depth = rng.integers(80, 151, size=genome_length, dtype=np.int32) + df = pd.DataFrame({"chrom": contig, "pos": pos, "depth": depth}) + df.to_csv(out_path, sep="\t", header=False, index=False) + print(f"wrote coverage file {out_path} ({genome_length} lines)") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--fasta", required=True) + ap.add_argument("--gff3", required=True) + ap.add_argument("--base-outdir", required=True) + args = ap.parse_args() + + base = Path(args.base_outdir) + contig, ref_seq = parse_fasta(args.fasta) + print(f"Loaded reference contig {contig}, length {len(ref_seq)}") + cds_list = parse_cds(args.gff3, contig) + print(f"Parsed {len(cds_list)} CDS features on {contig}") + + generate_coverage(contig, len(ref_seq), base / "coverage") + + simulate_scenario( + "scenario_a", 80, cds_list, contig, ref_seq, seed=42, biased=False, outdir=base / "scenario_a" + ) + simulate_scenario( + "scenario_b", 1000, cds_list, contig, ref_seq, seed=43, biased=True, outdir=base / "scenario_b" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_analysis.py b/tests/test_analysis.py index c324e7a..83edacf 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -11,7 +11,10 @@ import seaborn as sns from vartracker.analysis import ( + _build_variant_key, _heatmap_figure_size, + disambiguate_base_label, + find_colliding_base_labels, search_literature, _prepare_variant_heatmap_matrix, process_joint_variants, @@ -19,6 +22,7 @@ generate_gene_table, plot_gene_table, parse_plot_genes_arg, + select_canonical_row_positions, select_genes_for_plot, ) @@ -600,6 +604,74 @@ def test_prepare_variant_heatmap_matrix_excludes_wildcard_consequence_types(): assert list(matrix.index) == ["S:N501Y\n(A23063T)"] +def test_prepare_variant_heatmap_matrix_disambiguates_colliding_base_labels(): + """Two genuinely different variants that render to the same (truncated) + base label - a real bcftools-csq artefact on joint-called bacterial data + - must not collapse into a single heatmap row. A third, non-colliding + variant must come out byte-identical to plain `_resolve_variant_labels` + output, i.e. no unwanted disambiguation suffix.""" + long_aa = "CLLLDEFGHIJKLMNOPQRPA" # 21 chars: long enough to truncate + table = pd.DataFrame( + [ + { + "gene": "PA3025", + "amino_acid_consequence": long_aa, + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "alt_freq": "0.6", + "samples": "P0", + "variant": "C100T", + "start": 100, + }, + { + "gene": "PA3025", + "amino_acid_consequence": long_aa, + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "alt_freq": "0.7", + "samples": "P0", + "variant": "C200A", + "start": 200, + }, + { + "gene": "PA1000", + "amino_acid_consequence": "D50G", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "alt_freq": "0.8", + "samples": "P0", + "variant": "A300G", + "start": 300, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix(table, ["P0"], 0.0, 0.0) + + head = long_aa[:10] + tail = long_aa[-10:] + middle = len(long_aa) - 20 + truncated = f"{head}+{middle}{tail}" if middle > 0 else long_aa + + labels = list(matrix.index) + colliding_labels = [label for label in labels if f"PA3025:{truncated}" in label] + + # Both distinct variants survive (previously the second was silently + # dropped by `drop_duplicates(subset=["base_label"])`), and their + # displayed labels differ from one another. + assert len(colliding_labels) == 2 + assert colliding_labels[0] != colliding_labels[1] + assert any("@100" in label for label in colliding_labels) + assert any("@200" in label for label in colliding_labels) + + # The non-colliding variant is untouched: identical to what + # `_resolve_variant_labels` alone would produce. + assert "PA1000:D50G\n(A300G)" in labels + + def test_prepare_variant_heatmap_matrix_applies_extended_filters(): table = pd.DataFrame( [ @@ -732,6 +804,71 @@ def test_prepare_variant_heatmap_matrix_only_persistent_includes_new_intermitten assert set(matrix.index) == {"S:D215G\n(A22206G)", "S:E484K\n(G23012A)"} +def test_prepare_variant_heatmap_matrix_default_includes_canonical_joint_only_variant(): + """When every row for a variant is joint, the default (`include_joint=False`) + heatmap must not silently drop it entirely - its (only, hence canonical) + row is exempted from the implicit joint exclusion.""" + table = pd.DataFrame( + [ + { + "gene": "PA1199", + "amino_acid_consequence": "D50G", + "nsp_aa_change": "", + "type_of_change": "joint_missense", + "type_of_variant": "snp", + "alt_freq": "0.6", + "samples": "P0", + "variant": "A100G", + "start": 100, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix(table, ["P0"], 0.0, 0.0) + + assert list(matrix.index) == ["PA1199:D50G\n(A100G)"] + + +def test_prepare_variant_heatmap_matrix_default_hides_noncanonical_joint_fragment(): + """A variant with both a canonical (non-joint) row and a joint fragment + row must still hide the joint fragment by default - the canonical-row + exemption only rescues variants whose ONLY row is joint, it must not + let every joint row through.""" + table = pd.DataFrame( + [ + { + "gene": "S", + "amino_acid_consequence": "D50G", + "nsp_aa_change": "", + "type_of_change": "missense", + "type_of_variant": "snp", + "presence_absence": "Y / Y / Y", + "alt_freq": "0.5 / 0.6 / 0.7", + "samples": "P0 / P1 / P2", + "variant": "A100G", + "start": 100, + }, + { + "gene": "S", + "amino_acid_consequence": "D50G+X60Y", + "nsp_aa_change": "", + "type_of_change": "joint_missense", + "type_of_variant": "snp", + "presence_absence": "N / N / Y", + "alt_freq": "0.0 / 0.0 / 0.9", + "samples": "P0 / P1 / P2", + "variant": "A100G", + "start": 100, + }, + ] + ) + + matrix = _prepare_variant_heatmap_matrix(table, ["P0", "P1", "P2"], 0.0, 0.0) + + assert list(matrix.index) == ["S:D50G\n(A100G)"] + assert matrix.loc["S:D50G\n(A100G)", "P2"] == 0.7 + + def test_process_joint_variants_only_adds_single_joint_prefix(tmp_path): csv_path = tmp_path / "results.csv" pd.DataFrame( @@ -841,6 +978,40 @@ def test_process_joint_variants_matches_main_row_by_presence_pattern(tmp_path): assert result.loc[2, "type_of_change"] == "joint_missense" +def test_process_joint_variants_flags_compound_row_without_at_pointer(tmp_path): + """A row whose own `bcsq_nt_notation` is genuinely compound (multiple + '+'-joined `pos ref>alt` terms) must be flagged joint even when no + `@`-pointer row in the table points at it - the one-to-one `@`-pointer + matching only ever touches one sibling row per position, so a compound + sibling that nothing points to would otherwise be missed entirely.""" + csv_path = tmp_path / "results.csv" + pd.DataFrame( + [ + { + "start": 100, + "gene": "PA3025", + "amino_acid_consequence": "CLLL", + "nsp_aa_change": "", + "bcsq_nt_notation": "100A>T+150C>G", + "bcsq_aa_notation": "p.CLLL", + "aa1_total_properties": "", + "aa2_total_properties": "", + "aa1_unique_properties": "", + "aa2_unique_properties": "", + "aa1_weight": "", + "aa2_weight": "", + "weight_difference": "", + "type_of_change": "missense", + }, + ] + ).to_csv(csv_path, index=False) + + result = process_joint_variants(str(csv_path)) + + assert bool(result.loc[0, "joint_variant"]) is True + assert result.loc[0, "type_of_change"] == "joint_missense" + + def test_process_joint_variants_is_order_invariant_for_overlapping_gene_rows(tmp_path): shared_rows = [ { @@ -1013,3 +1184,93 @@ def test_heatmap_figure_size_enforces_minimum_row_height(): assert width == 4.0 assert height == pytest.approx(5.2) + + +def test_select_canonical_row_positions_prefers_union_over_non_joint(): + """When every row for a variant is joint, the row carrying the true + (union) trajectory must win even though every candidate is joint - this + is the PA1199-shaped case where "prefer non-joint" alone would pick a + fragment row with the wrong trajectory.""" + table = pd.DataFrame( + [ + { + "chrom": "NC_002516.2", + "start": 1300712, + "ref": "AG", + "alt": "A", + "presence_absence": "N / N / Y / N / N / N", + "joint_variant": True, + "amino_acid_consequence": "short", + "type_of_change": "joint_stop_gained&frameshift", + }, + { + "chrom": "NC_002516.2", + "start": 1300712, + "ref": "AG", + "alt": "A", + "presence_absence": "N / N / N / Y / N / N", + "joint_variant": True, + "amino_acid_consequence": "short", + "type_of_change": "joint_missense&inframe_altering", + }, + { + "chrom": "NC_002516.2", + "start": 1300712, + "ref": "AG", + "alt": "A", + "presence_absence": "N / N / Y / Y / Y / Y", + "joint_variant": True, + "amino_acid_consequence": "short", + "type_of_change": "joint_stop_gained&frameshift", + }, + ] + ) + keys = [ + _build_variant_key(row, "fallback") for row in table.itertuples(index=False) + ] + + canonical = select_canonical_row_positions(table, keys) + + assert len(canonical) == 1 + (position,) = canonical.values() + assert table.iloc[position]["presence_absence"] == "N / N / Y / Y / Y / Y" + + +def test_select_canonical_row_positions_prefers_standalone_when_it_covers_the_union(): + """The pre-existing viral case: a standalone (non-joint) row already + carries the full trajectory alongside a joint row masked to a subset of + samples - the standalone row must still win.""" + table = pd.DataFrame( + [ + { + "chrom": "segA", + "start": 100, + "ref": "A", + "alt": "T", + "presence_absence": "Y / Y / Y", + "joint_variant": False, + "amino_acid_consequence": "D614G", + "type_of_change": "missense", + }, + { + "chrom": "segA", + "start": 100, + "ref": "A", + "alt": "T", + "presence_absence": "N / N / Y", + "joint_variant": True, + "amino_acid_consequence": "D614G+N501Y", + "type_of_change": "joint_missense", + }, + ] + ) + keys = [ + _build_variant_key(row, "fallback") for row in table.itertuples(index=False) + ] + + canonical = select_canonical_row_positions(table, keys) + + assert len(canonical) == 1 + (position,) = canonical.values() + assert table.iloc[position]["presence_absence"] == "Y / Y / Y" + assert bool(table.iloc[position]["joint_variant"]) is False diff --git a/tests/test_cli.py b/tests/test_cli.py index 3f29f0c..40db029 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -34,7 +34,8 @@ def test_plot_heatmap_help_uses_standalone_option_names(): assert "--x-labels" in heatmap_options assert "--heatmap-aa-exclude" not in out assert "--heatmap-include-joint" not in out - assert "--outdir" not in out + assert "--out" in out + assert "--outdir" in out assert "--name" not in out assert "--min-snv-freq" not in out assert "--min-indel-freq" not in out diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 4f974ca..8b9190f 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -365,17 +365,241 @@ def test_prepare_plot_inputs_uses_unique_internal_variant_keys(): summary, long_df, _, _ = prepare_plot_inputs(table) assert summary["variant_id"].nunique() == 2 + # Both variants would otherwise render to the same "INTERGENIC:None" + # display label. variant_id (the internal key) was always unique for + # this case; now the *displayed* label is also disambiguated (via + # find_colliding_base_labels/disambiguate_base_label) with a positional + # + ref>alt suffix, so the two variants are distinguishable on plots too. assert {label.split(" (", 1)[0] for label in summary["variant_label"]} == { - "INTERGENIC:None" + "INTERGENIC:None @100A>G", + "INTERGENIC:None @200C>T", } assert ( long_df.loc[ - long_df["variant_label"].str.startswith("INTERGENIC:None"), "variant_id" + long_df["variant_label"].str.startswith("INTERGENIC:None @"), + "variant_id", ].nunique() == 2 ) +def test_prepare_plot_inputs_collapses_duplicate_rows_for_trajectory(): + # Models the real PA1199 case from PAO1 validation + # (NC_002516.2:1300712 AG>A): one physical variant described by + # bcftools-csq under three distinct joint annotation groups, each row + # masked to a different single sample, collectively present in all 3 + # samples even though no single row shows that on its own. + base = { + "chrom": "NC_002516.2", + "start": 1300712, + "end": 1300713, + "ref": "AG", + "alt": "A", + "gene": "PA1199", + "variant": "1300712_AGdelG", + "nsp_aa_change": "", + "type_of_variant": "indel", + "variant_status": "new", + "persistence_status": "new_persistent", + "samples": "P1 / P2 / P3", + "sample_number": "1 / 2 / 3", + "per_sample_variant_qc": "P / P / P", + "joint_variant": True, + } + table = pd.DataFrame( + [ + { + **base, + "amino_acid_consequence": "PA1199:frameshift_group_A", + "type_of_change": "joint_frameshift", + "presence_absence": "Y / N / N", + "alt_freq": "0.50 / 0.03 / 0.02", + }, + { + **base, + "amino_acid_consequence": "PA1199:frameshift_group_B", + "type_of_change": "joint_frameshift", + "presence_absence": "N / Y / N", + "alt_freq": "0.03 / 0.60 / 0.04", + }, + { + **base, + "amino_acid_consequence": "PA1199:frameshift_group_C", + "type_of_change": "joint_frameshift", + "presence_absence": "N / N / Y", + "alt_freq": "0.02 / 0.05 / 0.70", + }, + ] + ) + + summary, long_df, _, _ = prepare_plot_inputs(table) + + assert long_df["variant_id"].nunique() == 1 + variant_id = long_df["variant_id"].iloc[0] + variant_rows = long_df[long_df["variant_id"] == variant_id].sort_values( + "sample_number" + ) + + assert len(variant_rows) == 3 + assert list(variant_rows["sample_number"]) == [1, 2, 3] + assert variant_rows["present"].all() + # Each sample's AF must be the max across the variant's 3 source rows for + # that sample, not zero and not an arbitrary single row's value. + assert list(variant_rows["allele_frequency"].round(2)) == [0.50, 0.60, 0.70] + + +def test_prepare_plot_inputs_collapses_duplicate_rows_for_turnover_counts(): + # Same fragmented-row scenario as above, but shaped for a clear + # new/lost/new transition sequence across 4 samples so turnover counting + # can be checked against a known-correct answer. + base = { + "chrom": "NC_002516.2", + "start": 1300712, + "end": 1300713, + "ref": "AG", + "alt": "A", + "gene": "PA1199", + "variant": "1300712_AGdelG", + "amino_acid_consequence": "PA1199:frameshift", + "nsp_aa_change": "", + "type_of_variant": "indel", + "type_of_change": "joint_frameshift", + "variant_status": "new", + "persistence_status": "new_transient", + "samples": "P1 / P2 / P3 / P4", + "sample_number": "1 / 2 / 3 / 4", + "per_sample_variant_qc": "P / P / P / P", + "joint_variant": True, + } + table = pd.DataFrame( + [ + { + **base, + "presence_absence": "Y / N / N / N", + "alt_freq": "0.40 / 0.0 / 0.0 / 0.0", + }, + { + **base, + "presence_absence": "N / N / Y / Y", + "alt_freq": "0.0 / 0.0 / 0.30 / 0.50", + }, + ] + ) + + summary, long_df, _, _ = prepare_plot_inputs(table) + + # No duplicate (variant_id, sample_number) pairs: exactly one row per + # sample for this variant, as plot_variant_turnover assumes. + assert not long_df.duplicated(subset=["variant_id", "sample_number"]).any() + + variant_id = long_df["variant_id"].iloc[0] + trajectory = long_df[long_df["variant_id"] == variant_id].sort_values( + "sample_number" + ) + presence_sequence = list(trajectory["present"]) + # Union of the two rows' presence vectors: present, absent, present, present. + assert presence_sequence == [True, False, True, True] + + # Naive new/lost transition scan mirroring what plot_variant_turnover + # does over a collapsed long_df (one row per variant per sample). + transitions = 0 + previous_present = False + for present in presence_sequence: + if present != previous_present: + transitions += 1 + previous_present = present + # new at sample 1, lost at sample 2, new (re-emergence) at sample 3. + assert transitions == 3 + + +def test_prepare_plot_inputs_summary_uses_canonical_row_metadata(): + variant_a_base = { + "chrom": "NC_002516.2", + "start": 2000000, + "end": 2000000, + "ref": "C", + "alt": "T", + "gene": "PA2000", + "variant": "2000000C>T", + "nsp_aa_change": "", + "type_of_variant": "snp", + "variant_status": "new", + "persistence_status": "new_persistent", + "samples": "P1 / P2 / P3", + "sample_number": "1 / 2 / 3", + "per_sample_variant_qc": "P / P / P", + } + variant_b_base = { + "chrom": "NC_002516.2", + "start": 3000000, + "end": 3000000, + "ref": "G", + "alt": "A", + "gene": "PA3000", + "variant": "3000000G>A", + "nsp_aa_change": "", + "type_of_variant": "snp", + "variant_status": "new", + "persistence_status": "new_persistent", + "samples": "P1 / P2 / P3", + "sample_number": "1 / 2 / 3", + "per_sample_variant_qc": "P / P / P", + } + + table = pd.DataFrame( + [ + # Variant A: one joint fragment row (masked, NOT the union) and + # one standalone row that IS the union -> standalone wins + # because it's both the union match AND non-joint. + { + **variant_a_base, + "amino_acid_consequence": "PA2000:joint_frag", + "type_of_change": "joint_missense", + "presence_absence": "N / N / Y", + "alt_freq": "0.0 / 0.0 / 0.30", + "joint_variant": True, + }, + { + **variant_a_base, + "amino_acid_consequence": "PA2000:D10E", + "type_of_change": "missense", + "presence_absence": "Y / Y / Y", + "alt_freq": "0.20 / 0.25 / 0.30", + "joint_variant": False, + }, + # Variant B: both rows are joint; one is a masked fragment, the + # other is the union -> the union row wins even though it is + # also joint (no non-joint candidate exists to prefer instead). + { + **variant_b_base, + "amino_acid_consequence": "PA3000:joint_frag_x", + "type_of_change": "joint_missense_x", + "presence_absence": "N / N / Y", + "alt_freq": "0.0 / 0.0 / 0.40", + "joint_variant": True, + }, + { + **variant_b_base, + "amino_acid_consequence": "PA3000:joint_frag_y", + "type_of_change": "joint_missense_y", + "presence_absence": "Y / Y / Y", + "alt_freq": "0.10 / 0.20 / 0.40", + "joint_variant": True, + }, + ] + ) + + summary, long_df, _, _ = prepare_plot_inputs(table) + + row_a = summary[summary["gene"] == "PA2000"] + assert len(row_a) == 1 + assert row_a["type_of_change"].iloc[0] == "missense" + + row_b = summary[summary["gene"] == "PA3000"] + assert len(row_b) == 1 + assert row_b["type_of_change"].iloc[0] == "joint_missense_y" + + def test_genome_plot_defaults_to_snps_only(capsys): table = pd.DataFrame( [ diff --git a/tests/test_vcf_processing.py b/tests/test_vcf_processing.py index 3df76e7..d6468ed 100644 --- a/tests/test_vcf_processing.py +++ b/tests/test_vcf_processing.py @@ -513,6 +513,34 @@ def test_process_vcf_merges_starred_and_unstarred_equivalent_bcsq(tmp_path): assert row["persistence_status"] == "original_retained" +def test_process_vcf_decodes_percent_encoded_gene_name_in_bcsq(tmp_path): + vcf_path = tmp_path / "annotated.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\n" + "chr1\t4\t.\tA\tG\t.\tPASS\tDP=100;AF=0.1;" + "BCSQ=missense|ercS%27|tx|protein_coding|+|10K>10A|100A>G" + "\tGT:DP:AF:BCSQ\t1:100:0.1:1\n", + encoding="utf-8", + ) + + cov1 = tmp_path / "s1.depth.txt" + _write_depth_file(cov1) + + table = process_vcf(str(vcf_path), [str(cov1)], 10, ["s1"]) + + assert len(table) == 1 + assert table.iloc[0]["gene"] == "ercS'" + + def test_process_vcf_keeps_single_site_variant_present_when_sample_has_joint_csq( tmp_path, ): @@ -555,6 +583,74 @@ def test_process_vcf_keeps_single_site_variant_present_when_sample_has_joint_csq assert joint["alt_freq"] == ". / 0.200" +def test_process_vcf_recovers_presence_when_no_sample_gets_a_standalone_bcsq( + tmp_path, +): + """When a variant sits in a gene dense enough that every sample's own + joint-consequence call pairs it with a *different* co-occurring + neighbour, bcftools never reports it standalone for any sample, so each + annotation-group row is masked to a different, non-overlapping subset of + samples - fragmenting a variant that is genuinely present in all three + samples across two partial rows, with neither showing its true, full + trajectory. This is the bacterial-scale hotspot-gene scenario (see the + PAO1 validation): the variant's real presence must still surface in one + row so persistence_status isn't computed from a fragment. The two joint + rows themselves are kept unconditionally (see the module-level comment + in process_vcf on why row count is never reduced by deleting them) - + a single-sample joint call can be genuine, novel biology, not just an + artefact, and there is no reliable way to tell the two apart from + presence data alone.""" + vcf_path = tmp_path / "annotated.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "##contig=\n" + '##INFO=\n' + '##INFO=\n' + '##INFO=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + '##FORMAT=\n' + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\ts2\ts3\n" + "chr1\t4\t.\tA\tG\t.\tPASS\tDP=100;AF=0.2;" + "BCSQ=missense|GENE1|tx|protein_coding|+|2K>2A|4A>G+5A>C," + "frameshift|GENE1|tx|protein_coding|+|2KAAAAAAAAAAAA>2K|4A>G+6A>T" + "\tGT:DP:AF:BCSQ\t1:100:0.1:1\t1:100:0.2:4\t1:100:0.3:0\n", + encoding="utf-8", + ) + + cov1 = tmp_path / "s1.depth.txt" + cov2 = tmp_path / "s2.depth.txt" + cov3 = tmp_path / "s3.depth.txt" + _write_depth_file(cov1) + _write_depth_file(cov2) + _write_depth_file(cov3) + + table = process_vcf( + str(vcf_path), [str(cov1), str(cov2), str(cov3)], 10, ["s1", "s2", "s3"] + ) + + # Both joint-detail rows survive, each masked to the one sample whose + # own bitmask matched that specific compound description. + joint_with_5 = table[table["bcsq_nt_notation"].eq("4A>G+5A>C")].iloc[0] + assert joint_with_5["presence_absence"] == "Y / N / N" + + joint_with_6 = table[table["bcsq_nt_notation"].eq("4A>G+6A>T")].iloc[0] + assert joint_with_6["presence_absence"] == "N / Y / N" + + # No group's bitmask matched s3, yet s3 genuinely carries the ALT allele + # (GT=1, AF=0.3) - the fix must add one extra row from the true, + # unmasked per-sample presence so this isn't silently dropped. + recovered = table[table["presence_absence"].eq("Y / Y / Y")] + assert len(recovered) == 1 + recovered_row = recovered.iloc[0] + assert recovered_row["alt_freq"] == "0.100 / 0.200 / 0.300" + assert recovered_row["variant_status"] == "original" + assert recovered_row["persistence_status"] == "original_retained" + + assert len(table) == 3 + + @pytest.mark.skipif(shutil.which("bcftools") is None, reason="bcftools not available") def test_merge_then_annotate_preserves_joint_annotations_across_samples(tmp_path): ref, gff = _write_minimal_reference_bundle(tmp_path) diff --git a/vartracker/_version.py b/vartracker/_version.py index 11e633a..e463cb7 100644 --- a/vartracker/_version.py +++ b/vartracker/_version.py @@ -2,18 +2,40 @@ from __future__ import annotations +from pathlib import Path + try: # Python 3.8+ from importlib import metadata except ImportError: # pragma: no cover -- fallback for very old Pythons import importlib_metadata as metadata # type: ignore -# NOTE: When bumping the project version remember to update this fallback value -# alongside the version declared in pyproject.toml. -_FALLBACK_VERSION = "2.2.1" +_LAST_RESORT_VERSION = "0.0.0+unknown" + + +def _read_pyproject_version() -> str | None: + """Read `project.version` straight from pyproject.toml. + + Only reachable when vartracker isn't pip-installed (no dist-info to read + metadata from), e.g. a git checkout run in place - so pyproject.toml is + still sitting one directory above this file. Keeps the version fallback + in sync with pyproject.toml automatically, rather than a second + hand-maintained copy of the version string. + """ + pyproject_path = Path(__file__).resolve().parent.parent / "pyproject.toml" + if not pyproject_path.exists(): + return None + try: + import tomllib + except ImportError: # pragma: no cover -- Python <3.11 + return None + with pyproject_path.open("rb") as handle: + data = tomllib.load(handle) + return data.get("project", {}).get("version") + try: __version__ = metadata.version("vartracker") except metadata.PackageNotFoundError: # pragma: no cover - during local dev - __version__ = _FALLBACK_VERSION + __version__ = _read_pyproject_version() or _LAST_RESORT_VERSION __all__ = ["__version__"] diff --git a/vartracker/analysis.py b/vartracker/analysis.py index 8ad335f..5a32cac 100644 --- a/vartracker/analysis.py +++ b/vartracker/analysis.py @@ -101,6 +101,130 @@ def _presence_vector(value: object) -> tuple[str, ...]: return tuple(_parse_slash_separated_tokens(value)) +def _build_variant_key(row: object, fallback_label: str) -> str: + chrom = str(getattr(row, "chrom", "")).strip() + start = str(getattr(row, "start", "")).strip() + ref = str(getattr(row, "ref", "")).strip() + alt = str(getattr(row, "alt", "")).strip() + variant_name = str(getattr(row, "variant", "")).strip() + + if chrom and start and ref and alt: + return f"{chrom}:{start}:{ref}>{alt}" + if chrom and start and variant_name: + return f"{chrom}:{start}:{variant_name}" + if variant_name: + return variant_name + return fallback_label + + +def _row_is_joint(row: object) -> bool: + """Best-effort joint/compound flag for a results-table row. + + Prefers the explicit `joint_variant` column (set by + `process_joint_variants`); falls back to the `type_of_change` prefix for + tables that predate that column (e.g. synthetic test fixtures). + """ + joint_variant = getattr(row, "joint_variant", None) + if joint_variant is not None and not ( + isinstance(joint_variant, float) and pd.isna(joint_variant) + ): + return bool(joint_variant) + change_type = str(getattr(row, "type_of_change", "")).strip().lstrip("*") + return re.sub(r"^(joint_)+", "", change_type) != change_type + + +def select_canonical_row_positions( + table: pd.DataFrame, variant_keys: Sequence[str] +) -> dict[str, int]: + """Map each variant key to the positional index of its canonical row. + + `variant_keys[i]` is the caller's key for the i-th row of `table` (as + produced by `_build_variant_key`), so callers that already build keys in + their own pass cannot drift from this one. + + A single physical variant can have more than one row when bcftools csq + describes it under several distinct joint/compound annotation groups + (see the row-canonicalisation logic in `vcf_processing.py::process_vcf`). + The canonical row is chosen, in order: + + 1. Among rows whose `presence_absence` equals the element-wise union of + all of the variant's rows' presence vectors (i.e. the row - or rows - + that alone already carry the variant's true, full trajectory). If no + row matches the union exactly (not expected, but cheap to guard), + fall back to the row(s) with the most "Y" tokens. + 2. Among those, prefer a non-joint row (see `_row_is_joint`) - this + reproduces today's behaviour exactly for the standalone-vs-joint-detail + case that already existed before row-canonicalisation. + 3. Tie-break deterministically: shortest `amino_acid_consequence` (least + compound description), then lowest positional index. + + Returns a dict of variant_key -> positional index into `table` (i.e. an + index usable with `table.iloc[...]`, not `table.loc[...]`). + """ + positions_by_key: dict[str, list[int]] = {} + for position, key in enumerate(variant_keys): + positions_by_key.setdefault(key, []).append(position) + + canonical: dict[str, int] = {} + for key, positions in positions_by_key.items(): + if len(positions) == 1: + canonical[key] = positions[0] + continue + + presence_vectors = { + position: _presence_vector( + getattr(table.iloc[position], "presence_absence", "") + ) + for position in positions + } + length = max((len(vec) for vec in presence_vectors.values()), default=0) + union = tuple( + ( + "Y" + if any( + vec[idx] == "Y" + for vec in presence_vectors.values() + if idx < len(vec) + ) + else "N" + ) + for idx in range(length) + ) + + candidates = [ + position for position in positions if presence_vectors[position] == union + ] + if not candidates: + max_present = max( + sum(token == "Y" for token in vec) for vec in presence_vectors.values() + ) + candidates = [ + position + for position in positions + if sum(token == "Y" for token in presence_vectors[position]) + == max_present + ] + + rows_by_position = {position: table.iloc[position] for position in candidates} + non_joint_candidates = [ + position + for position in candidates + if not _row_is_joint(rows_by_position[position]) + ] + if non_joint_candidates: + candidates = non_joint_candidates + + def _tie_break_key(position: int) -> tuple[int, int]: + aa_consequence = str( + getattr(rows_by_position[position], "amino_acid_consequence", "") + ) + return (len(aa_consequence), position) + + canonical[key] = min(candidates, key=_tie_break_key) + + return canonical + + def _find_joint_main_index(tab: pd.DataFrame, joint_index: int, main_pos: int) -> int: candidates = tab[ (tab["start"] == main_pos) @@ -164,10 +288,6 @@ def process_joint_variants(path): # Get indices where bcsq_aa_notation starts with "@" idx = tab[tab["bcsq_aa_notation"].str.startswith("@", na=False)].index - # If no joint variants, return the table as is - if idx.empty: - return tab - for i in idx: try: # Find the main pos and main index number @@ -207,6 +327,21 @@ def process_joint_variants(path): print(f"Warning: Could not process joint variant at index {i}: {str(e)}") continue + # Second, additive pass: some rows carry a genuinely compound + # `bcsq_nt_notation` (e.g. "190263TCC>T+190295T>G+190354G>A") - more than + # one "pos ref>alt" term joined by "+" - but were never flagged joint by + # the `@`-pointer matching above. This happens when a variant has several + # sibling rows at the same position (routine now that rows are + # canonicalised), so the one-to-one `@`-pointer match only touches one + # sibling even though every sibling's own notation is plainly compound. + # Flag any such row directly from its own notation, so `joint_variant` + # and the `joint_` `type_of_change` prefix never disagree with what the + # row itself encodes. + compound_nt_mask = tab["bcsq_nt_notation"].astype(str).str.count(">") > 1 + for i in tab.index[compound_nt_mask]: + tab.at[i, "joint_variant"] = True + tab.at[i, "type_of_change"] = _ensure_joint_prefix(tab.at[i, "type_of_change"]) + tab.to_csv(path, index=None) return tab @@ -538,9 +673,9 @@ def select_genes_for_plot( plot_table = gene_table[gene_table["gene"].isin(variant_genes)] return plot_table, None - new_counts = gene_table.loc[gene_table["type"] == "new_mutations"].set_index("gene")[ - "number" - ] + new_counts = gene_table.loc[gene_table["type"] == "new_mutations"].set_index( + "gene" + )["number"] ranked = sorted( variant_genes, key=lambda g: (-new_counts.get(g, 0), -totals.get(g, 0)) ) @@ -804,6 +939,99 @@ def _resolve_variant_labels(row) -> Tuple[str, str, str]: return gene_label, display_label, base_label +def find_colliding_base_labels( + base_labels: Sequence[str], variant_keys: Sequence[str] +) -> set[str]: + """Base labels that map to more than one distinct variant key. + + `base_labels` and `variant_keys` are parallel sequences, one entry per + results-table row (see `_build_variant_key` for how a variant key is + derived). The same physical variant can legitimately appear on more than + one row (e.g. once standalone and once as a joint/compound fragment) and + still share a `base_label` - that is not a collision. A collision is a + `base_label` shared by rows whose variant keys differ, i.e. genuinely + different variants that `_resolve_variant_labels` happens to render + identically. This includes, but is not limited to, truncation of long + compound amino-acid-consequence strings to the same head+N+tail form - + on real bacterial (joint-csq) data, bcftools csq can also give two + different nucleotide variants an identical, untruncated amino-acid + description. + """ + keys_by_label: dict[str, set[str]] = {} + for label, key in zip(base_labels, variant_keys): + keys_by_label.setdefault(label, set()).add(key) + return {label for label, keys in keys_by_label.items() if len(keys) > 1} + + +def disambiguate_base_label(base_label: str, row: object) -> str: + """Append a positional suffix to a base label already known to collide. + + Only call this for labels already returned by `find_colliding_base_labels` + - i.e. labels genuinely shared by more than one distinct variant. The + row's `start` is usually enough to disambiguate two otherwise + identically-rendered variants (e.g. `"PA3025:CLLL+255RPA @3388987"`); + `ref`/`alt` are folded in too whenever available, so that two distinct + variants at the exact same position (a multiallelic site) still end up + with different labels rather than requiring a second "is this still + colliding?" pass. + + Note: a disambiguated label no longer matches the literature-anchor + `_canonical_label` lookup. That's acceptable - collisions only arise for + long compound joint labels, which never matched the literature database + anyway (see `_canonical_label`). + """ + start = str(getattr(row, "start", "")).strip() + candidate = f"{base_label} @{start}" + ref = str(getattr(row, "ref", "")).strip() + alt = str(getattr(row, "alt", "")).strip() + if ref or alt: + candidate = f"{candidate}{ref}>{alt}" + return candidate + + +def compute_variant_labeling( + table: pd.DataFrame, +) -> tuple[list[str], set[str], dict[str, int]]: + """Table-wide bookkeeping shared by every plot/heatmap caller: a variant + key per row, the set of base labels that collide across distinct + variants, and each variant's canonical row position. + + Returns (variant_keys, colliding_base_labels, canonical_positions), + where variant_keys[i] is the key for table's i-th row (see + `_build_variant_key`). + """ + base_labels: list[str] = [] + variant_keys: list[str] = [] + for row in table.itertuples(index=False): + _, _, row_base_label = _resolve_variant_labels(row) + base_labels.append(row_base_label) + variant_keys.append(_build_variant_key(row, row_base_label)) + + colliding_base_labels = find_colliding_base_labels(base_labels, variant_keys) + canonical_positions = select_canonical_row_positions(table, variant_keys) + return variant_keys, colliding_base_labels, canonical_positions + + +def disambiguate_labels_if_colliding( + base_label: str, + display_label: str, + row: object, + colliding_base_labels: set[str], +) -> tuple[str, str]: + """Apply `disambiguate_base_label` to (base_label, display_label) when + base_label is a known collision; otherwise return them unchanged. + """ + if base_label not in colliding_base_labels: + return base_label, display_label + base_label = disambiguate_base_label(base_label, row) + nuc_change = str(getattr(row, "variant", "")).strip() + if nuc_change and nuc_change not in {"", "None"}: + display_label = f"{base_label}\n({nuc_change})" + else: + display_label = base_label + return base_label, display_label + + def _prepare_variant_heatmap_matrix( table: pd.DataFrame, sample_names: Sequence[str], @@ -904,11 +1132,20 @@ def _prepare_variant_heatmap_matrix( ) ) + # Computed over the whole, unfiltered table (not just the rows that will + # survive filtering) so label-collision detection and canonical-row + # selection see the complete picture. Positional indices line up 1:1 + # with the `enumerate(table.itertuples(...))` loop that follows, since + # both walk the same unfiltered `table` in the same order. + variant_keys, colliding_base_labels, canonical_positions = compute_variant_labeling( + table + ) + records: List[Dict[str, Union[str, float, int]]] = [] record_index_by_label: dict[str, int] = {} qc_maps: dict[str, dict[str, str]] = {} - for row in table.itertuples(index=False): + for pos, row in enumerate(table.itertuples(index=False)): gene_value = getattr(row, "gene", "") if str(gene_value) in {"5' UTR", "3' UTR", "INTERGENIC"}: continue @@ -923,11 +1160,18 @@ def _prepare_variant_heatmap_matrix( continue gene_label, display_label, base_label = _resolve_variant_labels(row) + base_label, display_label = disambiguate_labels_if_colliding( + base_label, display_label, row, colliding_base_labels + ) change_type = ( str(getattr(row, "type_of_change", "")).strip().lower().lstrip("*") ) - if not include_joint and change_type.startswith("joint"): + if ( + not include_joint + and change_type.startswith("joint") + and pos != canonical_positions.get(variant_keys[pos]) + ): continue if included_patterns and not any( fnmatch.fnmatch(change_type, pattern) for pattern in included_patterns @@ -1205,6 +1449,7 @@ def _write_interactive_heatmap_html( cli_command: Optional[str], sample_labels: Sequence[str] | None = None, plot_title: str | None = None, + filename_stem: str = "variant_allele_frequency_heatmap", ) -> None: if matrix.empty: return @@ -1565,7 +1810,7 @@ def _frequency_to_color(value: float) -> tuple[str, str]: """ - output_path = os.path.join(outdir, "variant_allele_frequency_heatmap.html") + output_path = os.path.join(outdir, f"{filename_stem}.html") with open(output_path, "w", encoding="utf-8") as handle: handle.write(html_doc) @@ -1600,6 +1845,7 @@ def generate_variant_heatmap( cli_command: Optional[str] = None, x_tick_labels: Sequence[str] | None = None, plot_title: str | None = None, + filename_stem: str = "variant_allele_frequency_heatmap", ): """Generate a heatmap of variant allele frequencies across passages.""" @@ -1631,7 +1877,7 @@ def generate_variant_heatmap( print("No variant data available for heatmap; skipping plot.") return - heatmap_path = os.path.join(outdir, "variant_allele_frequency_heatmap.pdf") + heatmap_path = os.path.join(outdir, f"{filename_stem}.pdf") fig_width, fig_height = _heatmap_figure_size( len(heatmap_data.index), len(heatmap_data.columns) @@ -1708,6 +1954,7 @@ def generate_variant_heatmap( cli_command, sample_labels=tick_labels, plot_title=heatmap_title, + filename_stem=filename_stem, ) except Exception as html_exc: # pragma: no cover - best-effort UX print(f"Warning: failed to generate interactive heatmap: {html_exc}") diff --git a/vartracker/main.py b/vartracker/main.py index 7a459bf..6da2720 100644 --- a/vartracker/main.py +++ b/vartracker/main.py @@ -189,7 +189,11 @@ def add_legacy(name: str, dest: str, **kwargs) -> None: action="store_true", default=False, dest="heatmap_include_joint", - help="Include joint variants in heatmaps (default: exclude them)", + help=( + "Also show non-canonical joint/compound annotation-group rows in " + "heatmaps (a variant's canonical row is always shown by default, " + "even if joint)" + ), ) add_legacy("include-joint", "heatmap_include_joint", action="store_true") group.add_argument( @@ -751,6 +755,24 @@ def _configure_vcf_parser( "one sample after filtering (default: error)" ), ) + analysis_group.add_argument( + "--local-csq", + action="store_true", + default=False, + help=( + "Pass bcftools csq the -l/--local-csq flag, so each variant is " + "annotated on its own rather than jointly with nearby co-occurring " + "variants ('joint_*' consequence types). Off by default, since " + "joint calling is the correct behaviour for genuinely linked " + "variants (e.g. adjacent-codon changes on the same haplotype). " + "Consider enabling this for datasets with many unphased, " + "sub-consensus variants clustered in the same gene, where " + "bcftools has no genotype evidence that co-occurring variants are " + "actually on the same haplotype and joint calling can otherwise " + "fragment a single variant's presence/absence trajectory across " + "samples." + ), + ) analysis_group.add_argument( "-d", "--min-depth", @@ -1511,6 +1533,21 @@ def _add_plot_heatmap_subparser(subparsers): _add_heatmap_option_arguments( heatmap_group, prefix="", add_legacy_prefixed_aliases=True ) + output_group = parser.add_argument_group("Output") + output_group.add_argument( + "--out", + default=None, + help=( + "Write the heatmap using this path as the base name (extension, if " + "any, is ignored) - e.g. --out plots/myheatmap writes " + "plots/myheatmap.pdf and plots/myheatmap.html" + ), + ) + output_group.add_argument( + "--outdir", + default=None, + help="Output directory for heatmap files (default: beside results.csv)", + ) parser.set_defaults(handler=_run_plot_heatmap_command) @@ -2534,7 +2571,15 @@ def _run_plot_heatmap_command(args): if not results_csv.exists(): raise InputValidationError(f"Results CSV not found: {results_csv}") - outdir = results_csv.parent + filename_stem = "variant_allele_frequency_heatmap" + if args.out: + out_path = Path(args.out).expanduser().resolve() + outdir = out_path.parent + filename_stem = out_path.stem or filename_stem + elif args.outdir: + outdir = Path(args.outdir).expanduser().resolve() + else: + outdir = results_csv.parent outdir.mkdir(parents=True, exist_ok=True) table = pd.read_csv(results_csv, keep_default_na=False) @@ -2602,9 +2647,13 @@ def _run_plot_heatmap_command(args): plot_title=args.title, literature_hits=literature_df, literature_table_path=literature_path, + filename_stem=filename_stem, **_collect_heatmap_kwargs(args), ) - print(f"\nFinished: find results in {outdir}\n") + print( + f"\nFinished: wrote {outdir / f'{filename_stem}.pdf'} and " + f"{outdir / f'{filename_stem}.html'}\n" + ) return 0 except (InputValidationError, ProcessingError) as exc: print(f"\nERROR: {exc}\n") @@ -2784,6 +2833,7 @@ def _process_files( args.gff3, args.debug, args.multiallelic_overflow, + local_csq=args.local_csq, ) # Process VCF and extract variants @@ -2898,9 +2948,9 @@ def _process_files( persistent_mutations = table[ table.persistence_status.isin(NEW_ENDS_PRESENT_STATUSES) - ][ - ["gene", "variant", "amino_acid_consequence", "nsp_aa_change"] - ].reset_index(drop=True) + ][["gene", "variant", "amino_acid_consequence", "nsp_aa_change"]].reset_index( + drop=True + ) persistent_mutations.to_csv( os.path.join(args.outdir, "persistent_new_mutations.csv"), index=None ) diff --git a/vartracker/plotting.py b/vartracker/plotting.py index 5d00f25..07a9991 100644 --- a/vartracker/plotting.py +++ b/vartracker/plotting.py @@ -20,6 +20,8 @@ _coerce_frequency, _extract_numeric_position, _resolve_variant_labels, + compute_variant_labeling, + disambiguate_labels_if_colliding, ) from .core import InputValidationError, ProcessingError from .vcf_processing import NEW_ENDS_PRESENT_STATUSES @@ -114,22 +116,6 @@ def _project_name_from_results(table: pd.DataFrame, explicit_name: str | None) - return names[0] if len(names) == 1 else "" -def _build_variant_key(row: object, fallback_label: str) -> str: - chrom = str(getattr(row, "chrom", "")).strip() - start = str(getattr(row, "start", "")).strip() - ref = str(getattr(row, "ref", "")).strip() - alt = str(getattr(row, "alt", "")).strip() - variant_name = str(getattr(row, "variant", "")).strip() - - if chrom and start and ref and alt: - return f"{chrom}:{start}:{ref}>{alt}" - if chrom and start and variant_name: - return f"{chrom}:{start}:{variant_name}" - if variant_name: - return variant_name - return fallback_label - - def _display_label_prefix(label: object) -> str: return str(label).split(" (", 1)[0].strip() @@ -147,11 +133,21 @@ def prepare_plot_inputs( table: pd.DataFrame, ) -> tuple[pd.DataFrame, pd.DataFrame, list[str], list[int]]: sample_names, sample_numbers = _get_sample_axis(table) + dedup_table = table.drop_duplicates().reset_index(drop=True) + + all_variant_keys, colliding_base_labels, canonical_positions = ( + compute_variant_labeling(dedup_table) + ) + long_records: list[dict[str, object]] = [] - for row in table.drop_duplicates().itertuples(index=False): - _, display_label, base_label = _resolve_variant_labels(row) - variant_key = _build_variant_key(row, base_label) + for position, row in enumerate(dedup_table.itertuples(index=False)): + variant_key = all_variant_keys[position] + canonical_row = dedup_table.iloc[canonical_positions[variant_key]] + _, display_label, base_label = _resolve_variant_labels(canonical_row) + base_label, display_label = disambiguate_labels_if_colliding( + base_label, display_label, canonical_row, colliding_base_labels + ) names = _parse_slash_tokens(getattr(row, "samples", "")) numbers = _parse_slash_tokens(getattr(row, "sample_number", "")) freqs = [ @@ -185,13 +181,19 @@ def prepare_plot_inputs( { "variant_id": variant_key, "variant_label": display_label.replace("\n", " "), - "variant_name": str(getattr(row, "variant", "")).strip(), - "gene": str(getattr(row, "gene", "")).strip(), - "type_of_change": str(getattr(row, "type_of_change", "")).strip(), - "type_of_variant": str(getattr(row, "type_of_variant", "")).strip(), - "variant_status": str(getattr(row, "variant_status", "")).strip(), + "variant_name": str(getattr(canonical_row, "variant", "")).strip(), + "gene": str(getattr(canonical_row, "gene", "")).strip(), + "type_of_change": str( + getattr(canonical_row, "type_of_change", "") + ).strip(), + "type_of_variant": str( + getattr(canonical_row, "type_of_variant", "") + ).strip(), + "variant_status": str( + getattr(canonical_row, "variant_status", "") + ).strip(), "persistence_status": str( - getattr(row, "persistence_status", "") + getattr(canonical_row, "persistence_status", "") ).strip(), "sample_name": sample_name, "sample_number": int(float(sample_number)), @@ -222,6 +224,36 @@ def prepare_plot_inputs( ["sample_number", "variant_id", "sample_name"] ).reset_index(drop=True) + # Collapse duplicate (variant_id, sample_number) rows created when a + # variant has multiple bcftools-csq annotation-group rows (see + # _build_variant_key / process_vcf's row-canonicalisation): a row's + # presence/AF is masked to whichever samples that specific annotation + # group matched, so the same sample can appear as absent on one row and + # present on another for the same variant. Take the union of presence + # and the max AF across a variant's rows for each sample, so plots see + # one point per (variant, sample) instead of drawing/counting duplicates. + # Metadata columns are already uniform per variant_id at this point + # (sourced from each variant's canonical row above), so "first" here is + # well-defined rather than arbitrary. + long_df = long_df.groupby(["variant_id", "sample_number"], as_index=False).agg( + variant_label=("variant_label", "first"), + variant_name=("variant_name", "first"), + gene=("gene", "first"), + type_of_change=("type_of_change", "first"), + type_of_variant=("type_of_variant", "first"), + variant_status=("variant_status", "first"), + persistence_status=("persistence_status", "first"), + sample_name=("sample_name", "first"), + allele_frequency=("allele_frequency", "max"), + present=("present", "max"), + sample_qc=("sample_qc", "first"), + sample_pass_qc=("sample_pass_qc", "first"), + has_literature=("has_literature", "max"), + ) + long_df = long_df.sort_values( + ["sample_number", "variant_id", "sample_name"] + ).reset_index(drop=True) + summary = long_df.groupby("variant_id", as_index=False).agg( variant_label=("variant_label", "first"), variant_name=("variant_name", "first"), @@ -1146,10 +1178,16 @@ def load_reference_feature_metadata(results_csv: str | Path) -> dict[str, object def _collapse_variants_for_genome_plot( table: pd.DataFrame, ) -> pd.DataFrame: + dedup_table = table.drop_duplicates().reset_index(drop=True) + + all_variant_keys, colliding_base_labels, canonical_positions = ( + compute_variant_labeling(dedup_table) + ) + records: list[dict[str, object]] = [] - for row in table.drop_duplicates().itertuples(index=False): + for position, row in enumerate(dedup_table.itertuples(index=False)): gene_label, display_label, base_label = _resolve_variant_labels(row) - variant_key = _build_variant_key(row, base_label) + variant_key = all_variant_keys[position] af_values = [ _coerce_frequency(token) for token in _parse_slash_tokens(getattr(row, "alt_freq", "")) @@ -1196,6 +1234,39 @@ def _collapse_variants_for_genome_plot( ) continue first = group.iloc[0].to_dict() + canonical_row = dedup_table.iloc[canonical_positions[variant_id]] + canonical_gene_label, canonical_display_label, canonical_base_label = ( + _resolve_variant_labels(canonical_row) + ) + canonical_base_label, canonical_display_label = ( + disambiguate_labels_if_colliding( + canonical_base_label, + canonical_display_label, + canonical_row, + colliding_base_labels, + ) + ) + first["variant_label"] = canonical_display_label.replace("\n", " ") + first["variant_name"] = str(getattr(canonical_row, "variant", "")).strip() + first["plot_gene"] = str(canonical_gene_label).strip() + first["raw_gene"] = str(getattr(canonical_row, "gene", "")).strip() + first["type_of_change"] = str( + getattr(canonical_row, "type_of_change", "") + ).strip() + first["type_of_variant"] = str( + getattr(canonical_row, "type_of_variant", "") + ).strip() + first["variant_status"] = str( + getattr(canonical_row, "variant_status", "") + ).strip() + first["persistence_status"] = str( + getattr(canonical_row, "persistence_status", "") + ).strip() + first["aa_position"] = _extract_numeric_position( + canonical_base_label.split(":", 1)[1] + if ":" in canonical_base_label + else canonical_base_label + ) merged_af_values: list[float] = [] for value in group["af_values"].tolist(): merged_af_values.extend([float(item) for item in cast(list[float], value)]) diff --git a/vartracker/schemas.py b/vartracker/schemas.py index b69e510..0c513fe 100644 --- a/vartracker/schemas.py +++ b/vartracker/schemas.py @@ -150,7 +150,13 @@ { "name": "all_samples_pass_qc", "type": "boolean", - "description": "True if every sample passes per-sample variant QC.", + "description": ( + "True only if every sample is 'P' in per_sample_variant_qc. False " + "means at least one sample's presence/absence call could not be " + "confidently distinguished from dropout/non-detection - inspect " + "per_sample_variant_qc to see which sample(s) are affected before " + "treating this variant's presence/absence pattern as reliable." + ), "units": "", "values": "true, false", }, @@ -169,7 +175,12 @@ "when there is no variant-supporting read and site coverage is " "below --min-depth, i.e. when genuine absence of the variant " "cannot be distinguished from dropout/non-detection; 'P' means " - "absence (or presence) was confidently called at that sample." + "absence (or presence) was confidently called at that sample. " + "E.g. 'P / P / F / P / P / P' identifies the third sample as the " + "one where QC failed. The default heatmap marks 'F' cells " + "visually: an unfilled black-bordered rectangle in the static " + "PDF, and a dark inset ring plus a 'QC=FAIL' hover tooltip in the " + "interactive HTML version." ), "units": "", "values": "P, F", diff --git a/vartracker/vcf_processing.py b/vartracker/vcf_processing.py index 31241d4..64d2370 100644 --- a/vartracker/vcf_processing.py +++ b/vartracker/vcf_processing.py @@ -11,6 +11,7 @@ import shutil from pathlib import Path from typing import TypedDict +from urllib.parse import unquote import pandas as pd import numpy as np @@ -563,8 +564,18 @@ def annotate_vcf( annotation, debug, multiallelic_overflow="error", + local_csq=False, ): - """Annotate a merged multi-sample VCF with bcftools csq.""" + """Annotate a merged multi-sample VCF with bcftools csq. + + local_csq: pass bcftools csq -l/--local-csq, so each variant is + considered on its own rather than jointly with nearby co-occurring + variants. Off by default - joint calling is correct for genuinely + linked variants, but can fragment a single variant's presence across + samples when co-occurrence is just an artefact of unphased, sub-consensus + calls (see analysis.py::process_joint_variants and the row-splitting + this can otherwise cause in process_vcf below). + """ if multiallelic_overflow not in {"error", "drop-lowest-af", "skip-site"}: raise RuntimeError( "multiallelic_overflow must be one of: error, drop-lowest-af, skip-site" @@ -690,8 +701,9 @@ def annotate_vcf( shutil.copyfile(skipped_vcf, output_file) return + local_flag = " -l" if local_csq else "" cmd = ( - f"bcftools csq -p R -f {reference} -g {annotation} --force " + f"bcftools csq -p R{local_flag} -f {reference} -g {annotation} --force " f"-Ov -o {annotated_only_vcf} {prepared_vcf}" ) @@ -709,12 +721,32 @@ def annotate_vcf( shutil.copyfile(annotated_only_vcf, output_file) -def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): +def _mean_depth_in_range(arr, lo: int, hi: int): + """Mean depth over the 1-based inclusive position range [lo, hi]. + + Positions outside the array (beyond the coverage file's span) or marked + missing (-1, the fill value for positions absent from the coverage file) + are excluded rather than treated as zero coverage. Returns None if no + position in range has recorded depth. + """ + lo = max(lo, 1) + hi = min(hi, arr.shape[0] - 1) + if hi < lo: + return None + window = arr[lo : hi + 1] + valid = window[window >= 0] + if valid.size == 0: + return None + return valid.mean() + + +def calculate_variant_site_depths(depth_arrays, v, samples, min_depth: int): """ Calculate depth metrics for variant sites. Args: - cov_df (pd.DataFrame): Coverage data + depth_arrays (dict[str, np.ndarray]): Per-sample coverage depth, + indexed directly by 1-based genome position (index 0 unused). v: Variant object from cyvcf2 samples (list): List of sample names @@ -729,33 +761,23 @@ def calculate_variant_site_depths(cov_df, v, samples, min_depth: int): end = v.end if v.var_type == "snp": - site_depths = list(cov_df.loc[cov_df["pos"] == start]["depth"]) + site_depths = [] + for sample in samples: + arr = depth_arrays[sample] + if 0 <= start < arr.shape[0] and arr[start] >= 0: + site_depths.append(int(arr[start])) else: - site_depths = ( - cov_df[cov_df["pos"].between(start, end, inclusive="both")] - .groupby("sample")["depth"] - .mean() - .reset_index() - ) - site_depths["order"] = pd.Categorical( - site_depths["sample"], categories=samples, ordered=True - ) - site_depths = [ - int(round(x, 0)) for x in list(site_depths.sort_values("order")["depth"]) - ] - - window_depth = ( - cov_df[cov_df["pos"].between(start - 10, end + 10, inclusive="neither")] - .groupby("sample")["depth"] - .mean() - .reset_index() - ) - window_depth["order"] = pd.Categorical( - window_depth["sample"], categories=samples, ordered=True - ) - window_depth = [ - int(round(x, 0)) for x in list(window_depth.sort_values("order")["depth"]) - ] + site_depths = [] + for sample in samples: + mean_depth = _mean_depth_in_range(depth_arrays[sample], start, end) + if mean_depth is not None: + site_depths.append(int(round(mean_depth, 0))) + + window_depth = [] + for sample in samples: + mean_depth = _mean_depth_in_range(depth_arrays[sample], start - 9, end + 9) + if mean_depth is not None: + window_depth.append(int(round(mean_depth, 0))) try: dp_values = v.format("DP").tolist() @@ -829,9 +851,7 @@ def _summarise_sample_trajectory(allele_freqs, samples): if allele_freqs[0] != "." and allele_freqs[-1] == ".": persistent_status = "original_lost" elif allele_freqs[0] != "." and allele_freqs[-1] != ".": - persistent_status = ( - "original_intermittent" if has_gap else "original_retained" - ) + persistent_status = "original_intermittent" if has_gap else "original_retained" elif allele_freqs[0] == "." and allele_freqs[-1] != ".": persistent_status = "new_intermittent" if has_gap else "new_persistent" elif allele_freqs[0] == "." and allele_freqs[-1] == ".": @@ -1093,22 +1113,36 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): "Number of coverage files does not match number of VCF samples." ) - # Load coverage data - cov_list = [] + # Load coverage data. Coverage files follow the `samtools depth -aa` + # convention (one line per genome position, contiguous from position 1), + # so per-sample depth is stored as a plain numpy array indexed directly + # by 1-based position rather than as rows in a concatenated DataFrame - + # for a bacterial-scale genome the DataFrame form holds tens of millions + # of rows (with per-row "ref"/"sample" strings) in memory simultaneously, + # while the array form needs only a few bytes per position per sample. + depth_arrays: dict[str, "np.ndarray"] = {} total_cov_list = [] - for i, cov_file in enumerate(covs): + for sample, cov_file in zip(samples, covs): try: - cov = pd.read_csv(cov_file, sep="\t", names=["ref", "pos", "depth"]) + cov = pd.read_csv( + cov_file, + sep="\t", + names=["ref", "pos", "depth"], + usecols=["pos", "depth"], + dtype={"pos": "int64", "depth": "int32"}, + ) total_cov = "{:.2f}".format(((sum(cov.depth > 10) / cov.shape[0]) * 100)) total_cov_list.append(total_cov) - cov["sample"] = samples[i] - cov_list.append(cov) + + positions = cov["pos"].to_numpy() + depths = cov["depth"].to_numpy() + arr = np.full(int(positions.max()) + 1, -1, dtype=np.int32) + arr[positions] = depths + depth_arrays[sample] = arr except Exception as e: raise RuntimeError(f"Error reading coverage file {cov_file}: {str(e)}") - cov_df = pd.concat(cov_list).set_index("sample") - # Process variants for v in vcf: info = dict(list(v.INFO)) @@ -1118,7 +1152,7 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): allele_freqs = ["."] * max(1, len(samples)) trajectory = _summarise_sample_trajectory(allele_freqs, samples) - depths_qc = calculate_variant_site_depths(cov_df, v, samples, min_depth) + depths_qc = calculate_variant_site_depths(depth_arrays, v, samples, min_depth) all_samples_pass_qc = "F" not in depths_qc["variant_qc"] proportion_samples_passing_qc = ( sum(flag == "P" for flag in depths_qc["variant_qc"]) @@ -1135,6 +1169,8 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): produced_annotation_specific_row = False if sample_bcsq_map: + trajectories_by_alt: dict = {} + groups_by_alt: dict = {} for annotation_group in annotation_groups: annot = _representative_bcsq_annotation(annotation_group) anno = annot.split("|") @@ -1158,9 +1194,29 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): masked_trajectory = _summarise_sample_trajectory( masked_allele_freqs, samples ) + + if annotation_alt is not None: + groups_by_alt.setdefault(annotation_alt, []).append( + annotation_group + ) + trajectories_by_alt.setdefault(annotation_alt, []).append( + tuple(masked_trajectory["presence_absence"]) + ) + if "Y" not in masked_trajectory["presence_absence"]: continue + # Every joint/compound annotation-group row is kept, not + # just standalone ones: a compound consequence that only + # ever shows up in one sample can still be real biology + # (e.g. a newly emerging linked pair), and recurrence + # across samples can only ever confirm something after + # the fact - it cannot distinguish a genuine first + # occurrence from an artefact, so filtering on it would + # silently discard exactly the most novel findings along + # with the noise. `joint_variant` (set downstream by + # analysis.py::process_joint_variants) flags which rows + # these are, so nothing here is deleted, only labelled. result = _process_annotation( v, anno, @@ -1180,6 +1236,61 @@ def process_vcf(vcf_file, covs, min_depth, sample_names_override=None): results.append(result) produced_annotation_specific_row = True + # A gene with many closely-spaced variants can lead bcftools + # to describe every sample's occurrence of this same allele + # as a *different* joint/compound consequence (whichever + # neighbouring variants happened to co-occur in that + # particular sample), with no sample ever getting a + # standalone description. Each annotation-group row above is + # then masked to only the sample(s) matching its own exact + # compound description - e.g. six samples can each land in a + # *different* one-sample-wide row, so the union of all rows + # covers every sample yet no single row ever shows the + # variant's true, continuous trajectory (and persistence + # classification, computed per-row, is wrong for every one + # of them). Whenever no already-emitted row's trajectory for + # this allele exactly matches the true, unmasked per-sample + # presence, synthesise one extra row carrying that true + # trajectory (the same genotype/AF-derived data used for the + # single-record-allele case above) alongside the existing + # joint-detail rows, so the variant's real presence and + # persistence_status is available from at least one row. + for alt, true_allele_freqs in alt_frequency_map.items(): + true_trajectory = _summarise_sample_trajectory( + true_allele_freqs, samples + ) + true_tuple = tuple(true_trajectory["presence_absence"]) + if "Y" not in true_tuple: + continue + if true_tuple in trajectories_by_alt.get(alt, []): + continue + + candidate_groups = groups_by_alt.get(alt) or annotation_groups + representative_group = min( + candidate_groups, + key=lambda g: len(_representative_bcsq_annotation(g)), + ) + annot = _representative_bcsq_annotation(representative_group) + anno = annot.split("|") + result = _process_annotation( + v, + anno, + true_trajectory["variant_status"], + true_trajectory["persistent_status"], + true_trajectory["presence_absence"], + true_trajectory["first_appearance"], + true_trajectory["last_appearance"], + all_samples_pass_qc, + proportion_samples_passing_qc, + depths_qc, + true_allele_freqs, + samples, + total_cov_list, + alt, + ) + results.append(result) + produced_annotation_specific_row = True + if not produced_annotation_specific_row: for annotation_group in annotation_groups: annot = _representative_bcsq_annotation(annotation_group) @@ -1314,10 +1425,11 @@ def _process_annotation( } else: # Regular annotation + gene = unquote(anno[1]) reformatted_aa = ( - reformat_csq_notation(anno[1], anno[5]) + reformat_csq_notation(gene, anno[5]) if len(anno) > 5 - else [anno[1] + ":" + anno[0], ""] + else [gene + ":" + anno[0], ""] ) aa_exploration = AminoAcidChange(reformatted_aa[0].replace("*", "")) @@ -1326,7 +1438,7 @@ def _process_annotation( "chrom": v.CHROM, "start": v.start + 1, "end": v.end, - "gene": anno[1], + "gene": gene, "ref": v.REF, "alt": selected_alt, "variant": v.REF + str(v.POS) + selected_alt, From 36dbbb52bf72b7f92ac89c5756b8bee9608ef825 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Mon, 3 Aug 2026 12:52:17 +1000 Subject: [PATCH 19/20] Fix lint: black-format simulate_pao1.py, drop unused test imports --- scripts/validation/pao1/simulate_pao1.py | 115 +++++++++++++++++++---- tests/test_analysis.py | 2 - 2 files changed, 95 insertions(+), 22 deletions(-) diff --git a/scripts/validation/pao1/simulate_pao1.py b/scripts/validation/pao1/simulate_pao1.py index cef7053..54c8c27 100644 --- a/scripts/validation/pao1/simulate_pao1.py +++ b/scripts/validation/pao1/simulate_pao1.py @@ -2,6 +2,7 @@ """Simulate longitudinal VCF + coverage data against the PAO1 reference bundle for vartracker bacterial-scale validation (Scenario A ~80 variants, Scenario B ~1000). """ + import argparse import subprocess import sys @@ -63,11 +64,13 @@ def zipf_weights(n, rng, exponent=1.1): order = rng.permutation(n) ranks = np.empty(n, dtype=float) ranks[order] = np.arange(1, n + 1) - w = 1.0 / (ranks ** exponent) + w = 1.0 / (ranks**exponent) return w / w.sum() -def choose_positions(cds_list, n_variants, ref_seq, rng, indel_fraction=0.08, biased=False): +def choose_positions( + cds_list, n_variants, ref_seq, rng, indel_fraction=0.08, biased=False +): """Pick n_variants distinct (pos, ref, alt, gene, var_type) tuples from CDS regions.""" genes = np.array([c["gene"] for c in cds_list]) unique_genes, gene_index = np.unique(genes, return_inverse=True) @@ -115,7 +118,7 @@ def choose_positions(cds_list, n_variants, ref_seq, rng, indel_fraction=0.08, bi ref, alt, var_type = ref_base, ref_base + inserted, "indel" else: # deletion: anchor base + (span) deleted ref bases - del_seq = ref_seq[pos - 1: pos - 1 + span + 1] + del_seq = ref_seq[pos - 1 : pos - 1 + span + 1] if len(del_seq) < span + 1 or any(b not in BASES for b in del_seq): continue ref, alt, var_type = del_seq, del_seq[0], "indel" @@ -148,31 +151,58 @@ def gen_presence_af(pattern, var_type, rng): if pattern == "original_retained": present = [True] * N_TIMEPOINTS base = rng.uniform(0.2, 0.5) - af_vals = [np.clip(base + rng.uniform(-0.05, 0.05), af_floor, af_ceiling) for _ in range(N_TIMEPOINTS)] + af_vals = [ + np.clip(base + rng.uniform(-0.05, 0.05), af_floor, af_ceiling) + for _ in range(N_TIMEPOINTS) + ] elif pattern == "original_lost": last_present = int(rng.integers(1, 5)) # 1..4 inclusive -> absent by tp5 present = [i <= last_present for i in range(N_TIMEPOINTS)] n_present = last_present + 1 - af_vals = list(np.clip(np.linspace(0.55, 0.15, n_present) + rng.uniform(-0.03, 0.03, n_present), af_floor, af_ceiling)) + af_vals = list( + np.clip( + np.linspace(0.55, 0.15, n_present) + + rng.uniform(-0.03, 0.03, n_present), + af_floor, + af_ceiling, + ) + ) elif pattern == "new_persistent": first_present = int(rng.integers(1, 5)) # 1..4 present = [i >= first_present for i in range(N_TIMEPOINTS)] n_present = N_TIMEPOINTS - first_present - af_vals = list(np.clip(np.linspace(0.12, 0.85, n_present) + rng.uniform(-0.03, 0.03, n_present), af_floor, af_ceiling)) + af_vals = list( + np.clip( + np.linspace(0.12, 0.85, n_present) + + rng.uniform(-0.03, 0.03, n_present), + af_floor, + af_ceiling, + ) + ) elif pattern == "new_transient": first_present = int(rng.integers(1, 4)) # 1..3 - last_present = int(rng.integers(first_present, min(first_present + 3, 5))) # < 5 + last_present = int( + rng.integers(first_present, min(first_present + 3, 5)) + ) # < 5 present = [first_present <= i <= last_present for i in range(N_TIMEPOINTS)] n_present = last_present - first_present + 1 base = rng.uniform(0.15, 0.4) - af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + af_vals = list( + np.clip( + [base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], + af_floor, + af_ceiling, + ) + ) elif pattern == "original_intermittent": gap_start = int(rng.integers(1, 4)) # 1..3 - gap_len = int(rng.integers(1, 5 - gap_start)) # keep gap inside 1..4, tp5 stays present + gap_len = int( + rng.integers(1, 5 - gap_start) + ) # keep gap inside 1..4, tp5 stays present present = [True] * N_TIMEPOINTS for i in range(gap_start, min(gap_start + gap_len, 5)): present[i] = False @@ -180,7 +210,13 @@ def gen_presence_af(pattern, var_type, rng): present[5] = True n_present = sum(present) base = rng.uniform(0.2, 0.45) - af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + af_vals = list( + np.clip( + [base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], + af_floor, + af_ceiling, + ) + ) elif pattern == "new_intermittent": first_present = int(rng.integers(1, 3)) # 1..2 @@ -195,7 +231,13 @@ def gen_presence_af(pattern, var_type, rng): present[5] = True n_present = sum(present) base = rng.uniform(0.15, 0.4) - af_vals = list(np.clip([base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], af_floor, af_ceiling)) + af_vals = list( + np.clip( + [base + rng.uniform(-0.05, 0.05) for _ in range(n_present)], + af_floor, + af_ceiling, + ) + ) else: raise ValueError(pattern) @@ -205,7 +247,12 @@ def gen_presence_af(pattern, var_type, rng): assert present[0] is True else: assert present[0] is False - if pattern in ("original_retained", "new_persistent", "original_intermittent", "new_intermittent"): + if pattern in ( + "original_retained", + "new_persistent", + "original_intermittent", + "new_intermittent", + ): assert present[5] is True else: assert present[5] is False @@ -221,7 +268,11 @@ def assign_patterns(n, rng): for target in ("original_intermittent", "new_intermittent"): count = patterns.count(target) if count < MIN_INTERMITTENT: - donor_pool = [i for i, p in enumerate(patterns) if p not in ("original_intermittent", "new_intermittent")] + donor_pool = [ + i + for i, p in enumerate(patterns) + if p not in ("original_intermittent", "new_intermittent") + ] n_needed = MIN_INTERMITTENT - count idx_to_convert = rng.choice(donor_pool, size=n_needed, replace=False) for idx in idx_to_convert: @@ -229,7 +280,9 @@ def assign_patterns(n, rng): return patterns -def simulate_scenario(name, n_variants, cds_list, contig, ref_seq, seed, biased, outdir: Path): +def simulate_scenario( + name, n_variants, cds_list, contig, ref_seq, seed, biased, outdir: Path +): rng = np.random.default_rng(seed) variants = choose_positions(cds_list, n_variants, ref_seq, rng, biased=biased) for v in variants: @@ -270,13 +323,17 @@ def simulate_scenario(name, n_variants, cds_list, contig, ref_seq, seed, biased, fh.write("##fileformat=VCFv4.2\n") fh.write(f"##contig=\n") fh.write('##INFO=\n') - fh.write('##INFO=\n') + fh.write( + '##INFO=\n' + ) fh.write("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n") for r in rows: af_val = r["af"][tp] dp_val = r["dp"][tp] qual = 60 - fh.write(f"{r['chrom']}\t{r['pos']}\t.\t{r['ref']}\t{r['alt']}\t{qual}\tPASS\tDP={dp_val};AF={af_val}\n") + fh.write( + f"{r['chrom']}\t{r['pos']}\t.\t{r['ref']}\t{r['alt']}\t{qual}\tPASS\tDP={dp_val};AF={af_val}\n" + ) final_gz = vcf_dir / f"pao1_tp{tp}.vcf.gz" sort_cmd = f"bcftools sort -Ov {raw_path} 2>/dev/null | bgzip > {final_gz}" @@ -315,7 +372,9 @@ def simulate_scenario(name, n_variants, cds_list, contig, ref_seq, seed, biased, "reads2": "", "bam": "", "vcf": str((vcf_dir / f"pao1_tp{tp}.vcf.gz").resolve()), - "coverage": str((outdir.parent / "coverage" / f"pao1_tp{tp}_depth.txt").resolve()), + "coverage": str( + (outdir.parent / "coverage" / f"pao1_tp{tp}_depth.txt").resolve() + ), } ) pd.DataFrame(csv_rows).to_csv(outdir / f"{name}_input.csv", index=False) @@ -323,7 +382,9 @@ def simulate_scenario(name, n_variants, cds_list, contig, ref_seq, seed, biased, # summary counts pattern_counts = pd.Series([r["pattern"] for r in records]).value_counts() gene_counts = pd.Series([r["gene"] for r in records]).value_counts() - print(f"[{name}] {len(records)} distinct variants placed across {gene_counts.shape[0]} genes") + print( + f"[{name}] {len(records)} distinct variants placed across {gene_counts.shape[0]} genes" + ) print(f"[{name}] pattern mix:\n{pattern_counts}") print(f"[{name}] top genes by variant count:\n{gene_counts.head(10)}") return records @@ -359,10 +420,24 @@ def main(): generate_coverage(contig, len(ref_seq), base / "coverage") simulate_scenario( - "scenario_a", 80, cds_list, contig, ref_seq, seed=42, biased=False, outdir=base / "scenario_a" + "scenario_a", + 80, + cds_list, + contig, + ref_seq, + seed=42, + biased=False, + outdir=base / "scenario_a", ) simulate_scenario( - "scenario_b", 1000, cds_list, contig, ref_seq, seed=43, biased=True, outdir=base / "scenario_b" + "scenario_b", + 1000, + cds_list, + contig, + ref_seq, + seed=43, + biased=True, + outdir=base / "scenario_b", ) diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 83edacf..3030c57 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -13,8 +13,6 @@ from vartracker.analysis import ( _build_variant_key, _heatmap_figure_size, - disambiguate_base_label, - find_colliding_base_labels, search_literature, _prepare_variant_heatmap_matrix, process_joint_variants, From 5f9df912c321f10428f17a2d5bca63eae036a098 Mon Sep 17 00:00:00 2001 From: charlesfoster Date: Mon, 3 Aug 2026 13:10:28 +1000 Subject: [PATCH 20/20] Fix CI lint job: run under Python 3.11 to match mypy's target and avoid numpy 2.5's Python-3.12-only PEP 695 stub syntax --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 666a65a..de5a510 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: - python-version: '3.12' + python-version: '3.11' - name: Install dependencies run: | python -m pip install --upgrade pip