From bdc17751fbbb6871bd71b74031ce76aec95b8643 Mon Sep 17 00:00:00 2001
From: John Huddleston <huddlej@gmail.com>
Date: Thu, 4 Aug 2022 15:32:04 -0700
Subject: [PATCH] Support annotating metadata with "name" column

Replaces "merge" between metadata and sequence index with a "join" that
uses the index of the input data frames to merge their content. This
change allows us to support metadata where strains are indexed by a "name"
column instead of a "strain" column. The sequence index will always have
a "strain" column (by design), so we load the sequence index and tell
pandas the name of its column to index by. This additional change allows
the "join" command to work as expected without ever needing to specify
the name of the strain column in the metadata.

This commit adds a functional test for this expected behavior.
---
 scripts/annotate_metadata_with_index.py    |  8 +++-----
 tests/annotate_metadata_with_index.t       | 13 +++++++++++++
 tests/numeric_strains/metadata_by_name.tsv |  3 +++
 3 files changed, 19 insertions(+), 5 deletions(-)
 create mode 100644 tests/numeric_strains/metadata_by_name.tsv

diff --git a/scripts/annotate_metadata_with_index.py b/scripts/annotate_metadata_with_index.py
index 79d173d5d..fd643dc87 100644
--- a/scripts/annotate_metadata_with_index.py
+++ b/scripts/annotate_metadata_with_index.py
@@ -18,7 +18,8 @@
     index = pd.read_csv(
         args.sequence_index,
         sep="\t",
-        dtype={"strain": "string"}
+        dtype={"strain": "string"},
+        index_col="strain",
     ).drop(
         columns=["length"],
     )
@@ -26,15 +27,12 @@
     new_columns = {
         column: f"_{column}"
         for column in index.columns
-        if column != "strain"
     }
     index = index.rename(columns=new_columns)
 
-    metadata.merge(
+    metadata.join(
         index,
-        on="strain",
     ).to_csv(
         args.output,
         sep="\t",
-        index=False,
     )
diff --git a/tests/annotate_metadata_with_index.t b/tests/annotate_metadata_with_index.t
index 73f9337e1..dbbc95dc5 100644
--- a/tests/annotate_metadata_with_index.t
+++ b/tests/annotate_metadata_with_index.t
@@ -12,4 +12,17 @@ Annotate metadata with an index where strains have numeric names.
   \s*3 .* (re)
 
   $ rm -f "$TMP/metadata_with_index.tsv"
+
+Annotate metadata with an index where metadata uses the "name" column for strain names.
+This should produce the same result as above, except with a "name" column in the output.
+
+  $ python3 ../scripts/annotate_metadata_with_index.py \
+  >  --metadata numeric_strains/metadata_by_name.tsv \
+  >  --sequence-index numeric_strains/sequence_index.tsv \
+  >  --output "$TMP/metadata_with_index.tsv"
+  $ wc -l "$TMP/metadata_with_index.tsv"
+  \s*3 .* (re)
+
+  $ rm -f "$TMP/metadata_with_index.tsv"
+
   $ popd > /dev/null
diff --git a/tests/numeric_strains/metadata_by_name.tsv b/tests/numeric_strains/metadata_by_name.tsv
new file mode 100644
index 000000000..f6c721c83
--- /dev/null
+++ b/tests/numeric_strains/metadata_by_name.tsv
@@ -0,0 +1,3 @@
+name	virus	gisaid_epi_isl	genbank_accession	date	region	country	division	location	region_exposure	country_exposure	division_exposure	segment	length	host	age	sex	originating_lab	submitting_lab	authors	url	title	date_submitted
+1	ncov	?	MT450922	2020-03-05	Oceania	Australia	Victoria		Oceania	Australia	Victoria	genome	29812	Homo sapiens	?	?	?	?	Seemann et al	https://www.ncbi.nlm.nih.gov/nuccore/MT450922	Severe acute respiratory syndrome coronavirus 2 isolate SARS-CoV-2/human/AUS/VIC05/2020 ORF1ab polyprotein (ORF1ab), ORF1a polyprotein (ORF1ab), surface glycoprotein (S), ORF3a protein (ORF3a), envelope protein (E), membrane glycoprotein (M), ORF6 protein (ORF6), ORF7a protein (ORF7a), ORF7b (ORF7b), ORF8 protein (ORF8), nucleocapsid phosphoprotein (N), and ORF10 protein (ORF10) genes, complete cds	2020-05-11
+2	ncov	?	MT451684	2020-03-31	Oceania	Australia	Victoria		Oceania	Australia	Victoria	genome	29806	Homo sapiens	?	?	?	?	Seemann et al	https://www.ncbi.nlm.nih.gov/nuccore/MT451684	Severe acute respiratory syndrome coronavirus 2 isolate SARS-CoV-2/human/AUS/VIC1000/2020, complete genome	2020-05-11