Add reference as a sample, and include reference sequence

Closes #152
Closes  #22

Author: jeromekelleher

sc2ts/core.py

@@ -82,8 +82,8 @@ def get_problematic_sites():
 def get_reference_sequence():
     global __cached_reference
     if __cached_reference is None:
-        reader = FastaReader(data_path / "reference.fasta")
-        __cached_reference = reader[REFERENCE_GENBANK]
+        reader = pyfaidx.Fasta(str(data_path / "reference.fasta"))
+        __cached_reference = "X" + str(reader[REFERENCE_GENBANK])
     return __cached_reference
 
 

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,