Add keep_all option for additional_inputs and sequence_source column

.gitignore

@@ -33,3 +33,6 @@ Thumbs.db
 
 # Jupyter notebook cache files
 .ipynb_checkpoints
+
+# Claude generated files
+.claude

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

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"

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

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} \

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:

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