From bf9273f878e54ffae450f92e54161e090ce2af7b Mon Sep 17 00:00:00 2001 From: Charles Foster Date: Wed, 25 Mar 2026 14:12:58 +1100 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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