Merge additional data during phylogenetic (instead of ingest) workflow #68

j23414 · web-flow · commit 4e8387132c58 · 2025-03-03T15:34:36.000-08:00
diff --git a/.github/workflows/phylogenetic.yaml b/.github/workflows/phylogenetic.yaml
@@ -124,7 +124,7 @@ jobs:
         nextstrain build \
           phylogenetic \
             deploy_all \
-            --configfile build-configs/nextstrain-automation/config.yaml build-configs/washington-state/config.yaml build-configs/washington-state/config-public.yaml \
+            --configfile build-configs/nextstrain-automation/config.yaml build-configs/washington-state/config.yaml \
             $CONFIG_OVERRIDES
       # Specifying artifact name to differentiate ingest build outputs from
       # the phylogenetic build outputs
diff --git a/ingest/defaults/annotations.tsv b/ingest/defaults/annotations.tsv
@@ -1,21 +1,3 @@
-MH507691	location	San Diego
-OQ721209	location	Yolo
-OQ721236	location	Sacramento
-KX547286	location	Chautauqua
-KX547169	location	Oswego
-OQ721825	location	Larimer
-MH507691	latitude	32.715736
-OQ721209	latitude	38.732967
-OQ721236	latitude	38.575764
-KX547286	latitude	42.209869
-KX547169	latitude	43.455345
-OQ721825	latitude	40.66641
-MH507691	longitude	-117.161087
-OQ721209	longitude	-121.807281
-OQ721236	longitude	-121.478851
-KX547286	longitude	-79.470428
-KX547169	longitude	-76.510498
-OQ721825	longitude	-105.46116
 AF380138	country	Democratic Republic of the Congo
 AY741551	country	Sierra Leone
 DQ011153	country	USA
diff --git a/ingest/defaults/config.yaml b/ingest/defaults/config.yaml
@@ -99,11 +99,6 @@ curate:
   output_sequence_field: "sequence"
   # The field in the NDJSON record that contains the actual GenBank accession
   genbank_accession: 'accession'
-  # Added metadata columns
-  added_columns: [
-    'longitude',
-    'latitude'
-  ]
 
   # The list of metadata columns to keep in the final output of the curation pipeline.
   metadata_columns: [
@@ -128,8 +123,6 @@ curate:
     #'full_authors',
     #'reverse',
     'authors',
-    'latitude',
-    'longitude',
     #'institution',
     #'title',
     #'journal',
diff --git a/ingest/rules/curate.smk b/ingest/rules/curate.smk
@@ -68,7 +68,6 @@ rule curate:
         authors_default_value=config["curate"]["authors_default_value"],
         abbr_authors_field=config["curate"]["abbr_authors_field"],
         annotations_id=config["curate"]["annotations_id"],
-        added_columns=config["curate"]["added_columns"],
         id_field=config["curate"]["output_id_field"],
         sequence_field=config["curate"]["output_sequence_field"],
     shell:
@@ -97,8 +96,6 @@ rule curate:
                 --geolocation-rules {input.all_geolocation_rules} \
             | ./scripts/transform-state-names \
             | ./scripts/post_process_metadata.py \
-            | ./scripts/add-field-names \
-                --metadata-columns {params.added_columns} \
             | ./scripts/transform-new-fields \
                 --map-tsv {input.manual_mapping} \
                 --map-id host \
diff --git a/ingest/scripts/add-field-names b/ingest/scripts/add-field-names
diff --git a/phylogenetic/Snakefile b/phylogenetic/Snakefile
@@ -34,12 +34,11 @@ rule all:
 # custom_rules imported below to ensure that the core workflow is not complicated
 # by build specific rules.
 
+include: "rules/merge_additional_inputs.smk"
+include: "rules/prepare_sequences.smk"
 include: "rules/subsampling_manual.smk"
-
 if config.get("subsampling", False):
     include: "rules/subsampling_configurable.smk"
-
-include: "rules/prepare_sequences.smk"
 include: "rules/construct_phylogeny.smk"
 include: "rules/annotate_phylogeny.smk"
 include: "rules/export.smk"
diff --git a/phylogenetic/build-configs/ci/config.yaml b/phylogenetic/build-configs/ci/config.yaml
@@ -2,8 +2,10 @@
 # for the CI workflow to run with the example data.
 
 # Pull in metadata and sequences from the example_data directory
-input_metadata: "example_data/metadata.tsv"
-input_sequences: "example_data/sequences.fasta"
+inputs:
+  - name: example
+    metadata: "example_data/metadata.tsv"
+    sequences: "example_data/sequences.fasta"
 
 ## Custom rules to run as part of the CI automated workflow
 ## The paths should be relative to the phylogenetic directory.
diff --git a/phylogenetic/build-configs/ci/copy_example_data.smk b/phylogenetic/build-configs/ci/copy_example_data.smk
diff --git a/phylogenetic/build-configs/washington-state/config-public.yaml b/phylogenetic/build-configs/washington-state/config-public.yaml
diff --git a/phylogenetic/build-configs/washington-state/config.yaml b/phylogenetic/build-configs/washington-state/config.yaml
@@ -6,9 +6,17 @@ reference: "defaults/wa/reference.gb"
 # Use 'IS88' as the root strain on the phylogenetic tree to place samples within the global context
 root: "AF481864"
 
-# Pull in metadata and sequences from the ingest directory after it has been annotated with washington-state specific metadata
-input_metadata: "../ingest/results/metadata.tsv"
-input_sequences: "../ingest/results/sequences.fasta"
+# Sequences must be FASTA and metadata must be TSV
+# Both files must be zstd compressed
+inputs:
+  - name: ncbi
+    metadata: "s3://nextstrain-data/files/workflows/WNV/metadata.tsv.zst"
+    sequences: "s3://nextstrain-data/files/workflows/WNV/sequences.fasta.zst"
+
+additional_inputs:
+  - name: private-lat-longs
+    metadata: "data-private/metadata.tsv"
+    sequences: "data-private/sequences.fasta"
 
 builds: ['wa']
 
diff --git a/phylogenetic/data-private/README.md b/phylogenetic/data-private/README.md
@@ -0,0 +1,4 @@
+# Example Private Data
+
+This is some example private data that can be merged with the existing dataset.
+The use-case is if we have more fine-grained latitude and longitude locations for mosquito traps.
diff --git a/phylogenetic/data-private/metadata.tsv b/phylogenetic/data-private/metadata.tsv
@@ -0,0 +1,7 @@
+accession	latitude	longitude
+MH507691	32.715736	-117.161087
+OQ721209	38.732967	-121.807281
+OQ721236	38.575764	-121.478851
+KX547286	42.209869	-79.470428
+KX547169	43.455345	-76.510498
+OQ721825	40.66641	-105.46116
diff --git a/phylogenetic/data-private/sequences.fasta b/phylogenetic/data-private/sequences.fasta
diff --git a/phylogenetic/defaults/config.yaml b/phylogenetic/defaults/config.yaml
@@ -6,12 +6,10 @@ root: "mid_point"
 
 # Sequences must be FASTA and metadata must be TSV
 # Both files must be zstd compressed
-sequences_url: "https://data.nextstrain.org/files/workflows/WNV/sequences.fasta.zst"
-metadata_url: "https://data.nextstrain.org/files/workflows/WNV/metadata.tsv.zst"
-
-# Pull in metadata and sequences from the ingest workflow
-input_metadata: "data/metadata.tsv"
-input_sequences: "data/sequences.fasta"
+inputs:
+  - name: ncbi
+    metadata: "s3://nextstrain-data/files/workflows/WNV/metadata.tsv.zst"
+    sequences: "s3://nextstrain-data/files/workflows/WNV/sequences.fasta.zst"
 
 builds: ['global']
 
diff --git a/phylogenetic/rules/merge_additional_inputs.smk b/phylogenetic/rules/merge_additional_inputs.smk
@@ -0,0 +1,152 @@
+"""
+This part of the workflow merges additional metadata and sequences.
+
+REQUIRED INPUTS:
+
+    config = "defaults/config.yaml" with an `input` dictionary of metadata and sequence pairs
+
+OUTPUTS:
+
+    merged_sequences = results/sequences_merged.fasta
+    merged_metadata = results/metadata_merged.tsv
+
+This part of the workflow usually includes the following steps:
+
+    - augur merge
+See Augur's usage docs for these commands for more details.
+"""
+
+# ------------- helper functions to collect, merge & download input files ------------------- #
+NEXTSTRAIN_PUBLIC_BUCKET = "s3://nextstrain-data/"
+
+def _parse_config_input(input):
+    """
+    Parses information from an individual config-defined input, i.e. an element within `config.inputs` or `config.additional_inputs`
+    and returns information snakemake rules can use to obtain the underlying data.
+
+    The structure of `input` is a dictionary with keys:
+    - name:string (required)
+    - metadata:string (optional) - a s3 URI or a local file path
+    - sequences:string|dict[string,string] (optional) - either a s3 URI or a local file path, in which case
+      it must include a '{segment}' wildcard substring, or a dict of segment → s3 URI or local file path,
+      in which case it must not include the wildcard substring.
+
+    Returns a dictionary with optional keys:
+    - metadata:string - the relative path to the metadata file. If the original data was remote then this represents
+      the output of a rule which downloads the file
+    - metadata_location:string - the URI for the remote file if applicable else `None`
+    - sequences:function. Takes in wildcards and returns the relative path to the sequences FASTA for the provided
+      segment wildcard, or returns `None` if this input doesn't define sequences for the provided segment.
+    - sequences_location:function. Takes in wildcards and returns the URI for the remote file, or `None`, where applicable.
+
+    Raises InvalidConfigError
+    """
+    name = input['name']
+    lambda_none = lambda w: None
+
+    info = {'metadata': None, 'metadata_location': None, 'sequences': lambda_none, 'sequences_location': lambda_none}
+
+    def _source(uri, *,  s3, local):
+        if uri.startswith('s3://'):
+            return s3
+        elif uri.lower().startswith('http://') or uri.lower().startswith('https://'):
+            raise InvalidConfigError("Workflow cannot yet handle HTTP[S] inputs")
+        return local
+
+    if location:=input.get('metadata', False):
+        info['metadata'] = _source(location,  s3=f"data/{name}/metadata.tsv", local=location)
+        info['metadata_location'] = _source(location,  s3=location, local=None)
+
+    if location:=input.get('sequences', False):
+        info['sequences'] = _source(location,  s3=f"data/{name}/sequences.fasta", local=location)
+        info['sequences_location'] = _source(location,  s3=location, local=None)
+
+    return info
+
+
+def _gather_inputs():
+    all_inputs = [*config['inputs'], *config.get('additional_inputs', [])]
+
+    if len(all_inputs)==0:
+        raise InvalidConfigError("Config must define at least one element in config.inputs or config.additional_inputs lists")
+    if not all([isinstance(i, dict) for i in all_inputs]):
+        raise InvalidConfigError("All of the elements in config.inputs and config.additional_inputs lists must be dictionaries"
+            "If you've used a command line '--config' double check your quoting.")
+    if len({i['name'] for i in all_inputs})!=len(all_inputs):
+        raise InvalidConfigError("Names of inputs (config.inputs and config.additional_inputs) must be unique")
+    if not all(['name' in i and ('sequences' in i or 'metadata' in i) for i in all_inputs]):
+        raise InvalidConfigError("Each input (config.inputs and config.additional_inputs) must have a 'name' and 'metadata' and/or 'sequences'")
+
+    return {i['name']: _parse_config_input(i) for i in all_inputs}
+
+input_sources = _gather_inputs()
+
+def input_metadata(wildcards):
+    inputs = [info['metadata'] for info in input_sources.values() if info.get('metadata', None)]
+    return inputs[0] if len(inputs)==1 else "results/metadata_merged.tsv"
+
+def input_sequences(wildcards):
+    inputs = [info['sequences'] for info in input_sources.values() if info.get('sequences', None)]
+    return inputs[0] if len(inputs)==1 else "results/sequences_merged.fasta"
+
+rule download_s3_sequences:
+    output:
+        sequences = "data/{input_name}/sequences.fasta",
+    params:
+        address = lambda w: input_sources[w.input_name]['sequences_location'],
+        no_sign_request=lambda w: "--no-sign-request" \
+            if input_sources[w.input_name]['sequences_location'].startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
+            else "",
+    shell:
+        """
+        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.sequences}
+        """
+
+rule download_s3_metadata:
+    output:
+        metadata = "data/{input_name}/metadata.tsv",
+    params:
+        address = lambda w: input_sources[w.input_name]['metadata_location'],
+        no_sign_request=lambda w: "--no-sign-request" \
+            if input_sources[w.input_name]['metadata_location'].startswith(NEXTSTRAIN_PUBLIC_BUCKET) \
+            else "",
+    shell:
+        """
+        aws s3 cp {params.no_sign_request:q} {params.address:q} - | zstd -d > {output.metadata}
+        """
+
+rule merge_metadata:
+    """
+    This rule should only be invoked if there are multiple defined metadata inputs
+    (config.inputs + config.additional_inputs)
+    """
+    input:
+        **{name: info['metadata'] for name,info in input_sources.items() if info.get('metadata', None)}
+    params:
+        metadata = lambda w, input: list(map("=".join, input.items())),
+        id_field = config['strain_id_field'],
+    output:
+        metadata = "results/metadata_merged.tsv"
+    shell:
+        r"""
+        augur merge \
+            --metadata {params.metadata:q} \
+            --metadata-id-columns {params.id_field} \
+            --output-metadata {output.metadata}
+        """
+
+rule merge_sequences:
+    """
+    This rule should only be invoked if there are multiple defined sequences inputs
+    (config.inputs + config.additional_inputs) for this particular segment
+    """
+    input:
+        **{name: info['sequences'] for name,info in input_sources.items() if info.get('sequences', None)}
+    output:
+        sequences = "results/sequences_merged.fasta"
+    shell:
+        r"""
+        seqkit rmdup {input:q} > {output.sequences:q}
+        """
+
+# -------------------------------------------------------------------------------------------- #
diff --git a/phylogenetic/rules/prepare_sequences.smk b/phylogenetic/rules/prepare_sequences.smk
@@ -21,36 +21,6 @@ This part of the workflow usually includes the following steps:
 See Augur's usage docs for these commands for more details.
 """
 
-
-rule download:
-    """Downloading sequences and metadata from data.nextstrain.org"""
-    output:
-        sequences = "data/sequences.fasta.zst",
-        metadata = "data/metadata.tsv.zst"
-    params:
-        sequences_url = config["sequences_url"],
-        metadata_url = config["metadata_url"],
-    shell:
-        """
-        curl -fsSL --compressed {params.sequences_url:q} --output {output.sequences}
-        curl -fsSL --compressed {params.metadata_url:q} --output {output.metadata}
-        """
-
-rule decompress:
-    """Decompressing sequences and metadata"""
-    input:
-        sequences = "data/sequences.fasta.zst",
-        metadata = "data/metadata.tsv.zst"
-    output:
-        sequences = "data/sequences.fasta",
-        metadata = "data/metadata.tsv"
-    shell:
-        """
-        zstd -d -c {input.sequences} > {output.sequences}
-        zstd -d -c {input.metadata} > {output.metadata}
-        """
-
-
 rule align:
     input:
         sequences = "results/{build}/sequences_filtered.fasta",
diff --git a/phylogenetic/rules/subsampling_configurable.smk b/phylogenetic/rules/subsampling_configurable.smk
diff --git a/phylogenetic/rules/subsampling_manual.smk b/phylogenetic/rules/subsampling_manual.smk
