From 8fff3fbcd596aeb16bd65746bc1254c5de4fde51 Mon Sep 17 00:00:00 2001 From: jbloom Date: Mon, 9 Mar 2026 07:37:26 -0700 Subject: [PATCH] Add keep_all option for additional_inputs and sequence_source column Additional inputs specified under `additional_inputs` in the config are now subsampled by default (preserving previous behavior). Setting `keep_all: True` on an additional input force-includes all its sequences, bypassing subsampling and quality filters. A `sequence_source` metadata column is added whenever any additional inputs are present, labeling sequences by their input name or "Pathoplexus" for background sequences. This enables coloring and filtering by origin in Auspice. Addresses [this issue](https://github.com/nextstrain/rsv/issues/122). Co-Authored-By: Claude Opus 4.6 --- .gitignore | 3 + README.md | 47 ++++- Snakefile | 1 + config/auspice_config.json | 8 +- workflow/snakemake_rules/core.smk | 2 + workflow/snakemake_rules/merge_inputs.smk | 22 +- workflow/snakemake_rules/sequence_source.smk | 203 +++++++++++++++++++ 7 files changed, 275 insertions(+), 11 deletions(-) create mode 100644 workflow/snakemake_rules/sequence_source.smk diff --git a/.gitignore b/.gitignore index ec1121d..28f45ea 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,6 @@ Thumbs.db # Jupyter notebook cache files .ipynb_checkpoints + +# Claude generated files +.claude diff --git a/README.md b/README.md index fe171cc..d787cb8 100644 --- a/README.md +++ b/README.md @@ -84,11 +84,52 @@ additional_inputs: sequences: example_data/{a_or_b}/sequences.fasta ``` -Note that the additional inputs also require the `{a_or_b}` expandable field. -If you only have data for a single subtype, then you can do so with +By default, additional inputs are subsampled alongside the default inputs. +To force-include all sequences from an additional input (bypassing subsampling), +add `keep_all: True`. If `keep_all` is missing or `False`, then the additional +inputs are subsampled alongside other inputs: ```yaml -serotypes: ["a"] +additional_inputs: + - name: example-data + metadata: example_data/{a_or_b}/metadata.tsv + sequences: example_data/{a_or_b}/sequences.fasta + keep_all: True +``` + +Note that `keep_all` sequences bypass quality filters (coverage, missing data, +etc.) in addition to subsampling, so they do not need Nextclade QC columns. +Subsampled additional inputs must pass the same quality filters as default +inputs and therefore need the same metadata columns. + +#### Metadata requirements + +Additional input metadata TSV files must include these columns: + +| Column | Description | +|--------|-------------| +| `accession` | Must match the FASTA header exactly | +| `strain` | Display name for the sequence in the tree | +| `date` | Collection date (`YYYY-MM-DD`, `YYYY-MM`, or `YYYY`; use `XXXX-XX-XX` if unknown) | + +Recommended columns for better tree visualization: `country` (eg, "USA"), `region` (eg, "North America"), +`division` (eg, "New York"). Any column set to `?` or left empty is treated as unknown. + +#### `sequence_source` column + +When `additional_inputs` are present, a `sequence_source` metadata column is +added: additional sequences are labeled with their input `name`, and all other +sequences are labeled "Pathoplexus". This enables coloring and filtering by +origin in Auspice. Do not use "Pathoplexus" as an additional input name. + +#### Paths and subtypes + +The additional inputs require the `{a_or_b}` expandable field in paths. +If you only have data for a single subtype, restrict the entire build with +`subtypes`: + +```yaml +subtypes: ["a"] additional_inputs: - name: private metadata: private/a/metadata.tsv diff --git a/Snakefile b/Snakefile index 33e77ff..6364743 100644 --- a/Snakefile +++ b/Snakefile @@ -40,6 +40,7 @@ rule all: # remote_files.smk must be before merge_inputs.smk include: "shared/vendored/snakemake/remote_files.smk" include: "workflow/snakemake_rules/merge_inputs.smk" +include: "workflow/snakemake_rules/sequence_source.smk" include: "workflow/snakemake_rules/core.smk" include: "workflow/snakemake_rules/export.smk" diff --git a/config/auspice_config.json b/config/auspice_config.json index b641b53..ffdaab8 100644 --- a/config/auspice_config.json +++ b/config/auspice_config.json @@ -42,6 +42,11 @@ "title": "Region", "type": "categorical" }, + { + "key": "sequence_source", + "title": "Sequence source", + "type": "categorical" + }, { "key": "glyc", "title": "Glycosylation", @@ -118,6 +123,7 @@ "division", "country", "region", - "clade_membership" + "clade_membership", + "sequence_source" ] } diff --git a/workflow/snakemake_rules/core.smk b/workflow/snakemake_rules/core.smk index 6e6a2d5..12d6910 100644 --- a/workflow/snakemake_rules/core.smk +++ b/workflow/snakemake_rules/core.smk @@ -65,6 +65,7 @@ rule filter_recent: metadata="results/{a_or_b}/metadata.tsv", sequence_index=rules.index_sequences.output, exclude=config["exclude"], + additional_include="results/{a_or_b}/additional_include.txt", output: sequences=build_dir + "/{a_or_b}/{build_name}/{resolution}/filtered_recent.fasta", @@ -93,6 +94,7 @@ rule filter_recent: --metadata {input.metadata} \ --metadata-id-columns {params.strain_id} \ --exclude {input.exclude} \ + --include {input.additional_include} \ --exclude-where {params.exclude_where:q} \ --min-date {params.min_date} \ --min-length {params.min_length} \ diff --git a/workflow/snakemake_rules/merge_inputs.smk b/workflow/snakemake_rules/merge_inputs.smk index 7d66df0..4e88129 100644 --- a/workflow/snakemake_rules/merge_inputs.smk +++ b/workflow/snakemake_rules/merge_inputs.smk @@ -3,7 +3,7 @@ This part of the workflow merges inputs based on what is defined in the config. OUTPUTS: - metadata = results/{a_or_b}/metadata.tsv + metadata = results/{a_or_b}/metadata_merged.tsv sequences = results/{a_or_b}/sequences.fasta The config dict is expected to have a top-level `inputs` list that defines the @@ -46,11 +46,19 @@ def _gather_inputs(): if not any (['sequences' in i for i in all_inputs]): raise InvalidConfigError("At least one input must have 'sequences'") - available_keys = set(['name', 'metadata', 'sequences']) - if any([len(set(el.keys())-available_keys)>0 for el in all_inputs]): - raise InvalidConfigError(f"Each input (config.inputs and config.additional_inputs) can only include keys of {', '.join(available_keys)}") + available_keys_inputs = {'name', 'metadata', 'sequences'} + available_keys_additional = {'name', 'metadata', 'sequences', 'keep_all'} + for el in config.get('inputs', []): + extra = set(el.keys()) - available_keys_inputs + if extra: + raise InvalidConfigError(f"Each input can only include keys of {', '.join(sorted(available_keys_inputs))}. Got extra: {', '.join(sorted(extra))}") + for el in config.get('additional_inputs', []): + extra = set(el.keys()) - available_keys_additional + if extra: + raise InvalidConfigError(f"Each additional_input can only include keys of {', '.join(sorted(available_keys_additional))}. Got extra: {', '.join(sorted(extra))}") - return {el['name']: {k:(v if k=='name' else path_or_url(v)) for k,v in el.items()} for el in all_inputs} + path_keys = {'metadata', 'sequences'} + return {el['name']: {k:(path_or_url(v) if k in path_keys else v) for k,v in el.items()} for el in all_inputs} input_sources = _gather_inputs() _input_metadata = [info['metadata'] for info in input_sources.values() if info.get('metadata', None)] @@ -68,7 +76,7 @@ if len(_input_metadata) == 1: input: metadata = _input_metadata[0], output: - metadata = "results/{a_or_b}/metadata.tsv", + metadata = "results/{a_or_b}/metadata_merged.tsv", log: "logs/decompress_metadata_{a_or_b}.txt" benchmark: @@ -93,7 +101,7 @@ else: metadata = lambda w, input: list(map("=".join, input.items())), id_field = config['strain_id_field'], output: - metadata = "results/{a_or_b}/metadata.tsv" + metadata = "results/{a_or_b}/metadata_merged.tsv" log: "logs/merge_metadata_{a_or_b}.txt" benchmark: diff --git a/workflow/snakemake_rules/sequence_source.smk b/workflow/snakemake_rules/sequence_source.smk new file mode 100644 index 0000000..774c541 --- /dev/null +++ b/workflow/snakemake_rules/sequence_source.smk @@ -0,0 +1,203 @@ +""" +This part of the workflow adds a `sequence_source` column to the metadata to +distinguish sequences from different input sources. + +If `additional_inputs` is defined in the config, sequences from those inputs +are labeled with their input name, and all other sequences are labeled as +"Pathoplexus". If no additional inputs are defined, the metadata is passed +through unchanged. + +Additional inputs with `keep_all: True` bypass subsampling (force-included +via augur filter --include). All other additional inputs are subsampled +normally alongside the default inputs. + +INPUTS: + metadata_merged = results/{a_or_b}/metadata_merged.tsv + +OUTPUTS: + metadata = results/{a_or_b}/metadata.tsv + additional_include = results/{a_or_b}/additional_include.txt +""" + +_additional_inputs = config.get("additional_inputs", []) + +# Validate that no additional input uses the reserved name "Pathoplexus" +for _ai in _additional_inputs: + if _ai["name"].lower() == "pathoplexus": + raise ValueError( + f"Additional input name '{_ai['name']}' conflicts with the reserved " + "name used for background sequences. Please choose a different name." + ) + +# Validate keep_all is boolean if present +for _ai in _additional_inputs: + if "keep_all" in _ai and not isinstance(_ai["keep_all"], bool): + raise ValueError( + f"Additional input '{_ai['name']}' has keep_all={_ai['keep_all']!r} " + "but it must be True or False (a boolean, not a string)." + ) + + +if _additional_inputs: + + def _get_additional_sequence_files(wildcards): + """Get all sequence files from additional_inputs for this subtype.""" + files = [] + for ai in _additional_inputs: + if "sequences" in ai: + seq_path = ai["sequences"].replace("{a_or_b}", wildcards.a_or_b) + files.append(seq_path) + return files + + def _get_keep_all_sequence_files(wildcards): + """Get sequence files from additional_inputs with keep_all: True.""" + files = [] + for ai in _additional_inputs: + if ai.get("keep_all", False) and "sequences" in ai: + seq_path = ai["sequences"].replace("{a_or_b}", wildcards.a_or_b) + files.append(seq_path) + return files + + rule add_sequence_source: + """ + Add sequence_source column to metadata based on additional_inputs. + Sequences found in additional_inputs FASTAs are labeled with their + input name; all others are labeled 'Pathoplexus'. + """ + input: + metadata="results/{a_or_b}/metadata_merged.tsv", + additional_sequences=_get_additional_sequence_files, + output: + metadata="results/{a_or_b}/metadata.tsv", + log: + "logs/add_sequence_source_{a_or_b}.txt", + benchmark: + "benchmarks/add_sequence_source_{a_or_b}.txt" + run: + import csv + from Bio import SeqIO + + with open(log[0], "w") as log_file: + # Build mapping: accession -> source name from additional FASTAs + accession_to_source = {} + for ai in _additional_inputs: + if "sequences" not in ai: + continue + seq_path = ai["sequences"].replace("{a_or_b}", wildcards.a_or_b) + count = 0 + for record in SeqIO.parse(seq_path, "fasta"): + accession_to_source[record.id] = ai["name"] + count += 1 + print( + f"Found {count} sequences from additional input '{ai['name']}' " + f"in {seq_path}", + file=log_file, + ) + + print( + f"Total additional accessions: {len(accession_to_source)}", + file=log_file, + ) + + # Read merged metadata, add sequence_source column, write out + with open(input.metadata) as f_in, \ + open(output.metadata, "w", newline="") as f_out: + reader = csv.DictReader(f_in, delimiter="\t") + fieldnames = list(reader.fieldnames) + if "sequence_source" not in fieldnames: + fieldnames.append("sequence_source") + writer = csv.DictWriter( + f_out, + fieldnames=fieldnames, + delimiter="\t", + lineterminator="\n", + ) + writer.writeheader() + n_additional = 0 + n_pathoplexus = 0 + for row in reader: + accession = row.get("accession", "") + source = accession_to_source.get(accession, "Pathoplexus") + row["sequence_source"] = source + if source != "Pathoplexus": + n_additional += 1 + else: + n_pathoplexus += 1 + writer.writerow(row) + + print( + f"Labeled {n_additional} additional and {n_pathoplexus} Pathoplexus sequences", + file=log_file, + ) + + + rule generate_additional_include_list: + """ + Generate a list of accessions from additional_inputs with keep_all: True + to force-include in subsampling filters. + """ + input: + keep_all_sequences=_get_keep_all_sequence_files, + output: + include_list="results/{a_or_b}/additional_include.txt", + log: + "logs/generate_additional_include_list_{a_or_b}.txt", + benchmark: + "benchmarks/generate_additional_include_list_{a_or_b}.txt" + run: + from Bio import SeqIO + + accessions = [] + for seq_file in input.keep_all_sequences: + for record in SeqIO.parse(seq_file, "fasta"): + accessions.append(record.id) + + with open(output.include_list, "w") as f: + for acc in accessions: + f.write(acc + "\n") + + with open(log[0], "w") as log_file: + print( + f"Wrote {len(accessions)} keep_all accessions to {output.include_list}", + file=log_file, + ) + + +else: + + rule passthrough_metadata: + """ + No additional inputs defined; pass metadata through unchanged. + """ + input: + metadata="results/{a_or_b}/metadata_merged.tsv", + output: + metadata="results/{a_or_b}/metadata.tsv", + log: + "logs/passthrough_metadata_{a_or_b}.txt", + benchmark: + "benchmarks/passthrough_metadata_{a_or_b}.txt" + shell: + r""" + exec &> >(tee {log:q}) + + cp {input.metadata} {output.metadata} + """ + + + rule generate_empty_additional_include_list: + """ + No additional inputs; create empty include list (no-op for augur filter). + """ + output: + include_list="results/{a_or_b}/additional_include.txt", + log: + "logs/generate_empty_additional_include_list_{a_or_b}.txt", + benchmark: + "benchmarks/generate_empty_additional_include_list_{a_or_b}.txt" + shell: + r""" + exec &> >(tee {log:q}) + + touch {output.include_list} + """