Add more systematic base info for nodes in recombinant report

Author: jeromekelleher

sc2ts/info.py

@@ -451,11 +451,12 @@ def __init__(
         quick=False,
         show_progress=True,
         pango_source="Viridian_pangolin",
+        scorpio_source="Viridian_scorpio",
         sample_group_id_prefix_len=10,
     ):
         self.ts = ts
         self.pango_source = pango_source
-        self.scorpio_source = "Viridian_scorpio"
+        self.scorpio_source = scorpio_source
         self.strain_map = {}
         self.recombinants = np.where(ts.nodes_flags == core.NODE_IS_RECOMBINANT)[0]
 
@@ -970,6 +971,19 @@ def recombinants_summary(
     ):
         if parent_pango_source is None:
             parent_pango_source = self.pango_source
+
+        def node_info(node, label):
+            datum = {label: node}
+            datum[f"{label}_pango"] = self.nodes_metadata[node].get(
+                self.pango_source, "Unknown"
+            )
+            datum[f"{label}_scorpio"] = self.nodes_metadata[node].get(
+                self.scorpio_source, "Unknown"
+            )
+            datum[f"{label}_time"] = self.ts.nodes_time[node]
+            datum[f"{label}_date"] = self.nodes_date[node]
+            return datum
+
         data = []
         for u in self.recombinants:
             md = dict(self.nodes_metadata[u]["sc2ts"])
@@ -1011,30 +1025,33 @@ def recombinants_summary(
             interval = breakpoint_intervals[0]
             parent_left = hmm_match["path"][0]["parent"]
             parent_right = hmm_match["path"][1]["parent"]
-            data.append(
-                {
-                    "recombinant": u,
-                    "descendants": self.nodes_max_descendant_samples[u],
-                    "sample": v,
-                    "sample_pango": causal_lineages[v],
-                    "num_samples": len(samples),
-                    "distinct_sample_pango": len(set(causal_lineages.values())),
-                    "interval_left": interval[0][0],
-                    "interval_right": interval[0][1],
-                    "parent_left": parent_left,
-                    "parent_right": parent_right,
-                    "parent_left_pango": self.nodes_metadata[parent_left].get(
-                        parent_pango_source,
-                        "Unknown",
-                    ),
-                    "parent_right_pango": self.nodes_metadata[parent_right].get(
-                        parent_pango_source,
-                        "Unknown",
-                    ),
-                    "num_mutations": len(hmm_match["mutations"]),
-                    **md,
-                }
-            )
+
+            datum = {
+                "num_descendant_samples": self.nodes_max_descendant_samples[u],
+                "num_samples": len(samples),
+                "distinct_sample_pango": len(set(causal_lineages.values())),
+                "interval_left": interval[0][0],
+                "interval_right": interval[0][1],
+                "num_mutations": len(hmm_match["mutations"]),
+                "Viridian_amplicon_scheme": self.nodes_metadata[v].get(
+                    "Viridian_amplicon_scheme", "Unknown"
+                ),
+                "Artic_primer_version": self.nodes_metadata[v].get(
+                    "Artic_primer_version", "Unknown"
+                ),
+                **md,
+            }
+
+            for node, label in [
+                (u, "recombinant"),
+                (v, "sample"),
+                (parent_left, "parent_left"),
+                (parent_right, "parent_right"),
+            ]:
+                datum = {**datum, **node_info(node, label)}
+
+            data.append(datum)
+
         # Compute the MRCAs by iterating along trees in order of
         # breakpoint. We use the right interval
         df = pd.DataFrame(data).sort_values("interval_right")
@@ -1051,10 +1068,11 @@ def recombinants_summary(
             left_path = jit.get_root_path(tree, row.parent_left)
             assert tree.parent(row.recombinant) == row.parent_left
             mrca = jit.get_path_mrca(left_path, right_path, self.ts.nodes_time)
-            mrca_data.append(mrca)
-        mrca_data = np.array(mrca_data)
-        df["mrca"] = mrca_data
-        df["t_mrca"] = self.ts.nodes_time[mrca_data]
+            mrca_data.append(node_info(mrca, "parent_mrca"))
+
+        mrca_df = pd.DataFrame(mrca_data)
+        for col in mrca_df:
+            df[col] = mrca_df[col]
 
         if characterise_copying:
             # Slow - don't do this unless we really want to.

tests/test_info.py

@@ -176,7 +176,7 @@ def test_recombinants_summary_example_1(self, fx_ti_recombinant_example_1):
         df = fx_ti_recombinant_example_1.recombinants_summary()
         assert df.shape[0] == 1
         row = df.iloc[0]
-        assert row.descendants == 2
+        assert row.num_descendant_samples == 2
         assert row["sample"] == 53
         assert row.num_samples == 2
         assert row.group_size == 3
@@ -189,8 +189,8 @@ def test_recombinants_summary_example_1(self, fx_ti_recombinant_example_1):
         assert row.parent_right == 46
         assert row.parent_right_pango == "Unknown"
         assert row.num_mutations == 0
-        assert row.mrca == 1
-        assert row.t_mrca == 51
+        assert row.parent_mrca == 1
+        assert row.parent_mrca_time == 51
         assert "diffs" not in df
 
         df2 = fx_ti_recombinant_example_1.recombinants_summary(
@@ -206,17 +206,19 @@ def test_recombinants_summary_example_2(self, fx_recombinant_example_2):
         df = ti.recombinants_summary(characterise_copying=True, show_progress=False)
         assert df.shape[0] == 1
         row = df.iloc[0]
-        assert row.descendants == 1
+        assert row.num_descendant_samples == 1
         assert row["sample"] == 55
         assert row["distinct_sample_pango"] == 1
         assert row["recombinant"] == 56
+        assert row["recombinant_pango"] == "Unknown"
+        assert row["recombinant_time"] == 0.000001
         assert row["sample_pango"] == "Unknown"
         assert row["num_mutations"] == 0
         assert row["parent_left"] == 53
         assert row["parent_left_pango"] == "Unknown"
         assert row["parent_right"] == 54
         assert row["parent_right_pango"] == "Unknown"
-        assert row["mrca"] == 48
+        assert row["parent_mrca"] == 48
         assert row["group_size"] == 2
         assert row["diffs"] == 6
         assert row["max_run_length"] == 2
@@ -243,8 +245,6 @@ def test_example_node(self, fx_ts_min_2020_02_15, fx_ti_2020_02_15):
         nt.assert_array_equal(
             ti.nodes_max_descendant_samples, df["max_descendant_samples"]
         )
-        print(ti.nodes_date.dtype)
-        print(df["date"].dtype)
         nt.assert_array_equal(ti.nodes_date, df["date"])
         assert list(np.where(df["is_recombinant"])[0]) == list(ti.recombinants)
         assert list(np.where(df["is_sample"])[0]) == list(ts.samples())