Added the write_ms function to write out ms-style output from a tree sequence

python/tests/test_ms.py

@@ -0,0 +1,344 @@
+# MIT License
+#
+# Copyright (c) 2018-2020 Tskit Developers
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+Test cases for ms output in tskit. All of these tests have separate versions
+for the cases of single replicate and multiple replicates. This is because
+msprime.simulate generates a tree_sequence object if the num_replicates argument
+is not used but an iterator over tree_sequences if the num_replicates argument
+is used.
+"""
+import collections
+import itertools
+import os
+import tempfile
+import unittest
+
+import msprime
+
+import tskit as ts
+
+length = 1e2
+mutation_rate = 1e-2
+num_replicates = 3
+
+
+def get_ms_file_quantity(ms_file, quantity):
+    quantities = {}
+    num_replicates = 0
+    num_sites = []
+    num_positions = []
+    num_haplotypes = []
+    genotypes = []
+    positions = []
+    gens = []
+    for line in ms_file:
+        if len(line.split()) > 0:
+            if line.split()[0] == "//":
+                num_replicates = num_replicates + 1
+                num_haplotypes.append(0)
+                if len(gens) > 0:
+                    genotypes.append(gens)
+                    gens = []
+            if line.split()[0] == "segsites:":
+                num_sites.append(int(line.split()[1]))
+            if line.split()[0] == "positions:":
+                num_positions.append(len(line.split()) - 1)
+                positions.append(line[11:].rstrip())
+            if (
+                line[0:2] == "00"
+                or line[0:2] == "01"
+                or line[0:2] == "10"
+                or line[0:2] == "11"
+            ):
+                num_haplotypes[-1] = num_haplotypes[-1] + 1
+                gens.append(line.rstrip())
+    genotypes.append(gens)
+    quantities["num_replicates"] = num_replicates
+    quantities["num_sites"] = num_sites
+    quantities["num_positions"] = num_positions
+    quantities["num_haplotypes"] = num_haplotypes
+    quantities["genotypes"] = genotypes
+    quantities["positions"] = positions
+
+    return quantities[quantity]
+
+
+class TestNumReplicates(unittest.TestCase):
+    """
+    Tests that the number of replicates written out is the same as
+    the number of replicates simulated
+    """
+
+    def verify_num_replicates(self, tree_seq, num_replicates):
+        if isinstance(tree_seq, collections.Iterable):
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(
+                        tree_seq,
+                        f,
+                        num_replicates=num_replicates,
+                    )
+                with open(ms_file_path) as handle:
+                    num_replicates_file = get_ms_file_quantity(handle, "num_replicates")
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(tree_seq, f)
+                with open(ms_file_path) as handle:
+                    num_replicates_file = get_ms_file_quantity(handle, "num_replicates")
+        self.assertEqual(num_replicates, num_replicates_file)
+
+    def test_num_replicates(self):
+        tree_seq = msprime.simulate(
+            25, length=length, mutation_rate=mutation_rate, random_seed=123
+        )
+        self.verify_num_replicates(tree_seq, 1)
+
+    def test_num_replicates_multiple(self):
+        tree_seq = msprime.simulate(
+            25,
+            length=length,
+            mutation_rate=mutation_rate,
+            random_seed=123,
+            num_replicates=num_replicates,
+        )
+        self.verify_num_replicates(tree_seq, num_replicates)
+
+
+class TestNumHaplotypes(unittest.TestCase):
+    """
+    Tests that the number of haplotypes output are the same as the
+    number of individuals simulated.
+    """
+
+    def verify_num_haplotypes(self, tree_seq, tree_seq2, num_replicates):
+        if isinstance(tree_seq, collections.Iterable):
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(
+                        tree_seq,
+                        f,
+                        num_replicates=num_replicates,
+                    )
+                with open(ms_file_path) as handle:
+                    num_haplotypes = get_ms_file_quantity(handle, "num_haplotypes")
+            j = 0
+            for ts_indv in tree_seq2:
+                self.assertEqual(ts_indv.num_samples, num_haplotypes[j])
+                j = j + 1
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(tree_seq, f)
+                with open(ms_file_path) as handle:
+                    num_haplotypes = get_ms_file_quantity(handle, "num_haplotypes")
+            self.assertEqual(tree_seq.num_samples, num_haplotypes[0])
+
+    def test_num_haplotypes(self):
+        tree_seq = msprime.simulate(
+            25, length=length, mutation_rate=mutation_rate, random_seed=123
+        )
+        self.verify_num_haplotypes(tree_seq, tree_seq, 1)
+
+    def test_num_haplotypes_replicates(self):
+        tree_seq = msprime.simulate(
+            25,
+            length=length,
+            mutation_rate=mutation_rate,
+            random_seed=123,
+            num_replicates=num_replicates,
+        )
+        tree_seq, tree_seq2 = itertools.tee(tree_seq)
+        self.verify_num_haplotypes(tree_seq, tree_seq2, num_replicates)
+
+
+class TestNumSites(unittest.TestCase):
+    """
+    Tests that the number of sites written out as well as the length
+    of the positions list match the number of variants in the tree sequence
+    """
+
+    def verify_num_sites(self, tree_seq, tree_seq2, num_replicates):
+        if isinstance(tree_seq, collections.Iterable):
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(
+                        tree_seq,
+                        f,
+                        num_replicates=num_replicates,
+                    )
+                with open(ms_file_path) as handle:
+                    num_sites = get_ms_file_quantity(handle, "num_sites")
+                with open(ms_file_path) as handle:
+                    num_positions = get_ms_file_quantity(handle, "num_positions")
+            j = 0
+            for ts_indv in tree_seq2:
+                self.assertEqual(ts_indv.num_sites, num_sites[j])
+                self.assertEqual(ts_indv.num_sites, num_positions[j])
+                j = j + 1
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(tree_seq, f)
+                with open(ms_file_path) as handle:
+                    num_sites = get_ms_file_quantity(handle, "num_sites")
+                with open(ms_file_path) as handle:
+                    num_positions = get_ms_file_quantity(handle, "num_positions")
+            self.assertEqual(tree_seq.num_sites, num_sites[0])
+            self.assertEqual(tree_seq.num_sites, num_positions[0])
+
+    def test_num_sites(self):
+        tree_seq = msprime.simulate(
+            25, length=length, mutation_rate=mutation_rate, random_seed=123
+        )
+        self.verify_num_sites(tree_seq, tree_seq, 1)
+
+    def test_num_sites_replicates(self):
+        tree_seq = msprime.simulate(
+            25,
+            length=length,
+            mutation_rate=mutation_rate,
+            random_seed=123,
+            num_replicates=num_replicates,
+        )
+        tree_seq, tree_seq2 = itertools.tee(tree_seq)
+        self.verify_num_sites(tree_seq, tree_seq2, num_replicates)
+
+
+class TestGenotypes(unittest.TestCase):
+    """
+    Tests that the haplotypes written out are the same as the haplotypes generated.
+    """
+
+    def get_genotypes(self, tree_seq):
+        genotypes = tree_seq.genotype_matrix()
+        gens_array = []
+        for k in range(tree_seq.num_samples):
+            tmp_str = "".join(map(str, genotypes[:, k]))
+            gens_array.append(tmp_str)
+        return gens_array
+
+    def verify_genotypes(self, tree_seq, tree_seq2, num_replicates):
+        if isinstance(tree_seq, collections.Iterable):
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(
+                        tree_seq,
+                        f,
+                        num_replicates=num_replicates,
+                    )
+                with open(ms_file_path) as handle:
+                    genotypes = get_ms_file_quantity(handle, "genotypes")
+            j = 0
+            for ts_indv in tree_seq2:
+                self.assertEqual(self.get_genotypes(ts_indv), genotypes[j])
+                j = j + 1
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(tree_seq, f)
+                with open(ms_file_path) as handle:
+                    genotypes = get_ms_file_quantity(handle, "genotypes")
+                self.assertEqual(self.get_genotypes(tree_seq), genotypes[0])
+
+    def test_genotypes(self):
+        tree_seq = msprime.simulate(
+            25, length=length, mutation_rate=mutation_rate, random_seed=123
+        )
+        self.verify_genotypes(tree_seq, tree_seq, 1)
+
+    def test_genotypes_replicates(self):
+        tree_seq = msprime.simulate(
+            25,
+            length=length,
+            mutation_rate=mutation_rate,
+            random_seed=123,
+            num_replicates=num_replicates,
+        )
+        tree_seq, tree_seq2 = itertools.tee(tree_seq)
+        self.verify_genotypes(tree_seq, tree_seq2, num_replicates)
+
+
+class TestPositions(unittest.TestCase):
+    """
+    Tests that the positions for the mutations written out are the same as the
+    positions generated.
+    """
+
+    def get_positions(self, tree_seq):
+        positions = []
+        for i in range(tree_seq.num_sites):
+            positions.append(
+                "{:.4f}".format(tree_seq.site(i).position / tree_seq.sequence_length)
+            )
+        positions = " ".join(positions)
+        return positions
+
+    def verify_positions(self, tree_seq, tree_seq2, num_replicates):
+        if isinstance(tree_seq, collections.Iterable):
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(
+                        tree_seq,
+                        f,
+                        num_replicates=num_replicates,
+                    )
+                with open(ms_file_path) as handle:
+                    positions = get_ms_file_quantity(handle, "positions")
+            j = 0
+            for ts_indv in tree_seq2:
+                self.assertEqual(self.get_positions(ts_indv), positions[j])
+                j = j + 1
+        else:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
+                with open(ms_file_path, "w") as f:
+                    ts.write_ms(tree_seq, f)
+                with open(ms_file_path) as handle:
+                    positions = get_ms_file_quantity(handle, "positions")
+            self.assertEqual(self.get_positions(tree_seq), positions[0])
+
+    def test_positions(self):
+        tree_seq = msprime.simulate(
+            25, length=length, mutation_rate=mutation_rate, random_seed=123
+        )
+        self.verify_positions(tree_seq, tree_seq, 1)
+
+    def test_positions_replicates(self):
+        tree_seq = msprime.simulate(
+            25,
+            length=length,
+            mutation_rate=mutation_rate,
+            random_seed=123,
+            num_replicates=num_replicates,
+        )
+        tree_seq, tree_seq2 = itertools.tee(tree_seq)
+        self.verify_positions(tree_seq, tree_seq2, num_replicates)