Stuff for nexus

Author: jeromekelleher

python/tests/test_fasta.py

@@ -316,7 +316,49 @@ def test_unwrapped(self):
 
     def test_missing_data(self):
         ts = missing_data_example()
-        ref = "-" * int(ts.sequence_length)
-        text = ts.as_fasta(reference_sequence=ref, wrap_width=0)
+        ref = "A" * int(ts.sequence_length)
+        text = ts.as_fasta(reference_sequence=ref)
         alignment_map = self.parse(text)
         assert get_alignment_map(ts, ref) == alignment_map
+
+
+class TestDendropyMissingData:
+    """
+    Test that we detect missing data correctly in dendropy under
+    various combinations of options.
+    """
+
+    # 2.00┊   4     ┊
+    #     ┊ ┏━┻┓    ┊
+    # 1.00┊ ┃  3    ┊
+    #     ┊ ┃ ┏┻┓   ┊
+    # 0.00┊ 0 1 2 5 ┊
+    #     0        10
+    #      |      |
+    #  pos 2      9
+    #  anc A      T
+
+    def ts(self):
+        ts = tskit.Tree.generate_balanced(3, span=10).tree_sequence
+        tables = ts.dump_tables()
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.sites.add_row(2, ancestral_state="A")
+        tables.sites.add_row(9, ancestral_state="T")
+        tables.mutations.add_row(site=0, node=0, derived_state="G")
+        tables.mutations.add_row(site=1, node=3, derived_state="C")
+        return tables.tree_sequence()
+
+    def test_defaults_A_ref(self):
+        ts = self.ts()
+        ref = "A" * int(ts.sequence_length)
+        text = ts.as_fasta(reference_sequence=ref, missing_data_character="N")
+        # print(text)
+        d = dendropy.DnaCharacterMatrix.get(data=text, schema="fasta")
+        assert d.sequence_size == 10
+        # for k, v in d.items():
+        #     print(k, v)
+        # FINISH ME
+        # a0 = d["n0"]
+        # a5 = d["n5"]
+        # print(a0.character_type_at(0))
+        # print(a5)

python/tests/test_genotypes.py

@@ -1158,6 +1158,31 @@ def test_fasta_reference_sequence(self):
         )
         assert expected == self.ts().as_fasta(reference_sequence=ref)
 
+    def test_nexus_reference_sequence(self):
+        ref = "0123456789"
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS n0 n1 n2;
+            END;
+            BEGIN DATA;
+              DIMENSIONS NCHAR=10;
+              FORMAT datatype=dna missing=-;
+              MATRIX
+                n0 01G345678T
+                n1 01A345678C
+                n2 01A345678C
+              ;
+            END;
+            BEGIN TREES;
+              TREE t0^10 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == self.ts().as_nexus(reference_sequence=ref)
+
 
 class TestMissingDataExample:
     # 2.00┊   4     ┊
@@ -1303,6 +1328,85 @@ def test_fasta_impute_missing(self):
             reference_sequence=ref, isolated_as_missing=False
         )
 
+    # Note: the nexus tree output isn't compatible with our representation of
+    # missing data as trees with isolated roots (newick parsers won't accept
+    # this as valid input), so we set include_trees=False for thes examples.
+    def test_nexus_reference_sequence(self):
+        ref = "0123456789"
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=4;
+              TAXLABELS n0 n1 n2 n5;
+            END;
+            BEGIN DATA;
+              DIMENSIONS NCHAR=10;
+              FORMAT datatype=dna missing=-;
+              MATRIX
+                n0 01G345678T
+                n1 01A345678C
+                n2 01A345678C
+                n5 01-345678-
+            END;
+            """
+        )
+        assert expected == self.ts().as_nexus(
+            reference_sequence=ref, include_trees=False
+        )
+
+    def test_nexus_reference_sequence_missing_data_char(self):
+        ref = "0123456789"
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=4;
+              TAXLABELS n0 n1 n2 n5;
+            END;
+            BEGIN DATA;
+              DIMENSIONS NCHAR=10;
+              FORMAT datatype=dna missing=Q;
+              MATRIX
+                n0 01G345678T
+                n1 01A345678C
+                n2 01A345678C
+                n5 01Q345678Q
+            END;
+            """
+        )
+        assert expected == self.ts().as_nexus(
+            reference_sequence=ref,
+            missing_data_character="Q",
+            include_trees=False,
+        )
+
+    def test_nexus_impute_missing(self):
+        ref = "0123456789"
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=4;
+              TAXLABELS n0 n1 n2 n5;
+            END;
+            BEGIN DATA;
+              DIMENSIONS NCHAR=10;
+              FORMAT datatype=dna missing=-;
+              MATRIX
+                n0 01G345678T
+                n1 01A345678C
+                n2 01A345678C
+                n5 01A345678T
+            END;
+            """
+        )
+        assert expected == self.ts().as_nexus(
+            reference_sequence=ref,
+            isolated_as_missing=False,
+            include_trees=False,
+        )
+
 
 class TestAlignmentsErrors:
     @tests.cached_example