Merge pull request #244 from jeromekelleher/add_ref_as_sample

Add reference as sample, and include reference sequence in ts

Author: jeromekelleher

sc2ts/core.py

@@ -79,12 +79,16 @@ def get_problematic_sites():
 __cached_reference = None
 
 
-def get_reference_sequence():
+def get_reference_sequence(as_array=False):
     global __cached_reference
     if __cached_reference is None:
-        reader = FastaReader(data_path / "reference.fasta")
+        reader = pyfaidx.Fasta(str(data_path / "reference.fasta"))
         __cached_reference = reader[REFERENCE_GENBANK]
-    return __cached_reference
+    if as_array:
+        h = np.array(__cached_reference).astype(str)
+        return np.append(["X"], h)
+    else:
+        return "X" + str(__cached_reference)
 
 
 __cached_genes = None

sc2ts/inference.py

@@ -4,6 +4,7 @@
 import datetime
 import dataclasses
 import collections
+import json
 import pickle
 import os
 import sqlite3
@@ -220,6 +221,12 @@ def initial_ts():
     tables = tskit.TableCollection(L)
     tables.time_units = core.TIME_UNITS
     base_schema = tskit.MetadataSchema.permissive_json().schema
+    tables.reference_sequence.metadata_schema = tskit.MetadataSchema(base_schema)
+    tables.reference_sequence.metadata = {
+        "genbank_id": core.REFERENCE_GENBANK,
+        "notes": "X prepended to alignment to map from 1-based to 0-based coordinates"
+    }
+    tables.reference_sequence.data = reference
 
     tables.metadata_schema = tskit.MetadataSchema(base_schema)
 
@@ -235,10 +242,10 @@ def initial_ts():
             tables.sites.add_row(pos, reference[pos], metadata={"masked_samples": 0})
     # TODO should probably make the ultimate ancestor time something less
     # plausible or at least configurable. However, this will be removed
-    # in later versions when we remove the dependence on tskit.
+    # in later versions when we remove the dependence on tsinfer.
     tables.nodes.add_row(time=1, metadata={"strain": "Vestigial_ignore"})
     tables.nodes.add_row(
-        time=0, metadata={"strain": core.REFERENCE_STRAIN, "date": core.REFERENCE_DATE}
+        flags=tskit.NODE_IS_SAMPLE, time=0, metadata={"strain": core.REFERENCE_STRAIN, "date": core.REFERENCE_DATE}
     )
     tables.edges.add_row(0, L, 0, 1)
     return tables.tree_sequence()

tests/test_inference.py

@@ -8,6 +8,22 @@
 import util
 
 
+class TestInitialTs:
+    def test_reference_sequence(self):
+        ts = sc2ts.initial_ts()
+        assert ts.reference_sequence.metadata["genbank_id"] == "MN908947"
+        assert ts.reference_sequence.data == sc2ts.core.get_reference_sequence()
+
+    def test_reference_sample(self):
+        ts = sc2ts.initial_ts()
+        assert ts.num_samples == 1
+        node = ts.node(ts.samples()[0])
+        assert node.time == 0
+        assert node.metadata == {"date": "2019-12-26", "strain": "Wuhan/Hu-1/2019"}
+        alignment = next(ts.alignments())
+        assert alignment == sc2ts.core.get_reference_sequence()
+
+
 class TestAddMatchingResults:
     def add_matching_results(
         self,
@@ -132,7 +148,7 @@ def test_one_sample_one_mutation(self, tmp_path):
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
         assert ts2.num_mutations == 1
         var = next(ts2.variants())
-        assert var.alleles[var.genotypes[0]] == "X"
+        assert var.alleles[var.genotypes[1]] == "X"
 
     def test_one_sample_one_mutation_filtered(self, tmp_path):
         ts = sc2ts.initial_ts()
@@ -171,7 +187,7 @@ def test_two_samples_one_mutation_one_filtered(self, tmp_path):
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
         assert ts2.num_mutations == 1
         var = next(ts2.variants())
-        assert var.alleles[var.genotypes[0]] == "X"
+        assert var.alleles[var.genotypes[1]] == "X"
 
 
 class TestMatchTsinfer:
@@ -187,7 +203,7 @@ def test_match_reference(self, mirror):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         h = ma.alignment[ts.sites_position.astype(int)]
         samples[0].alignment = h
@@ -204,7 +220,7 @@ def test_match_reference_one_mutation(self, mirror, site_id):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         h = ma.alignment[ts.sites_position.astype(int)]
         # Mutate to gap
@@ -230,7 +246,7 @@ def test_match_reference_all_same(self, mirror, allele):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         ref = ma.alignment[ts.sites_position.astype(int)]
         h = np.zeros_like(ref) + allele
@@ -251,7 +267,7 @@ def match_path_ts(self, samples, ts):
         # FIXME this API is terrible
         ts2 = sc2ts.match_path_ts(samples, ts, samples[0].path, [])
         assert ts2.num_samples == len(samples)
-        for u, sample in zip(ts.samples(), samples):
+        for u, sample in zip(ts.samples()[1:], samples):
             node = ts.node(u)
             assert node.time == 0
             assert node.metadata == sample.metadata

tests/test_utils.py

@@ -9,7 +9,7 @@
 class TestPadSites:
     def check_site_padding(self, ts):
         ts = sc2ts.utils.pad_sites(ts)
-        ref = sc2ts.core.get_reference_sequence()
+        ref = sc2ts.core.get_reference_sequence(as_array=True)
         assert ts.num_sites == len(ref) - 1
         ancestral_state = ts.tables.sites.ancestral_state.view("S1").astype(str)
         assert np.all(ancestral_state == ref[1:])
@@ -67,7 +67,7 @@ def test_no_recombinants(self, ts, tmp_path):
 
     def test_one_sample_recombinant(self, tmp_path):
         ts = self.make_recombinant_tree(tmp_path / "match.db")
-        assert ts.num_samples == 3
+        assert ts.num_samples == 4
         re_nodes = [
             node.id for node in ts.nodes() if node.flags & sc2ts.NODE_IS_RECOMBINANT
         ]
@@ -93,6 +93,6 @@ def test_one_sample_recombinant(self, tmp_path):
     def test_two_sample_recombinant(self, tmp_path):
         """Test that we don't detach anything if the recombinant node is not a singleton"""
         ts = self.make_recombinant_tree(num_samples=2, db_path=tmp_path / "match.db")
-        assert ts.num_samples == 4
+        assert ts.num_samples == 5
         ts2 = utils.detach_singleton_recombinants(ts)
         ts.tables.assert_equals(ts2.tables, ignore_provenance=True)