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} + """