Add migrations argument to load_text and dump_text, and add parse_migrations.

Author: szhan

docs/python-api.md

@@ -1041,6 +1041,7 @@ Move some or all of these examples into a suitable alternative chapter.
   parse_mutations
   parse_individuals
   parse_populations
+  parse_migrations
   pack_strings
   unpack_strings
   pack_bytes
@@ -1143,6 +1144,7 @@ The following constants are used throughout the `tskit` API.
 .. autofunction:: parse_mutations
 .. autofunction:: parse_nodes
 .. autofunction:: parse_populations
+.. autofunction:: parse_migrations
 .. autofunction:: parse_sites
 .. autofunction:: random_nucleotides
 .. autofunction:: register_metadata_codec

python/tests/test_highlevel.py

@@ -3093,6 +3093,38 @@ def convert(v):
             else:
                 assert repr(population.metadata) == splits[1]
 
+    def verify_migrations_format(self, ts, migrations_file, precision, base64_metadata):
+        """
+        Verifies that the migrations we output have the correct form.
+        """
+
+        def convert(v):
+            return "{:.{}f}".format(v, precision)
+
+        output_migrations = migrations_file.read().splitlines()
+        assert len(output_migrations) - 1 == ts.num_migrations
+        assert list(output_migrations[0].split()) == [
+            "left",
+            "right",
+            "node",
+            "source",
+            "dest",
+            "time",
+            "metadata",
+        ]
+        for migration, line in zip(ts.migrations(), output_migrations[1:]):
+            splits = line.split("\t")
+            assert str(migration.left) == splits[0]
+            assert str(migration.right) == splits[1]
+            assert str(migration.node) == splits[2]
+            assert str(migration.source) == splits[3]
+            assert str(migration.dest) == splits[4]
+            assert str(migration.time) == splits[5]
+            if isinstance(migration.metadata, bytes) and base64_metadata:
+                assert tests.base64_encode(migration.metadata) == splits[6]
+            else:
+                assert repr(migration.metadata) == splits[6]
+
     @pytest.mark.parametrize(("precision", "base64_metadata"), [(2, True), (7, False)])
     def test_output_format(self, precision, base64_metadata):
         for ts in get_example_tree_sequences():
@@ -3102,6 +3134,7 @@ def test_output_format(self, precision, base64_metadata):
             mutations_file = io.StringIO()
             individuals_file = io.StringIO()
             populations_file = io.StringIO()
+            migrations_file = io.StringIO()
             provenances_file = io.StringIO()
             ts.dump_text(
                 nodes=nodes_file,
@@ -3110,6 +3143,7 @@ def test_output_format(self, precision, base64_metadata):
                 mutations=mutations_file,
                 individuals=individuals_file,
                 populations=populations_file,
+                migrations=migrations_file,
                 provenances=provenances_file,
                 precision=precision,
                 base64_metadata=base64_metadata,
@@ -3120,6 +3154,7 @@ def test_output_format(self, precision, base64_metadata):
             mutations_file.seek(0)
             individuals_file.seek(0)
             populations_file.seek(0)
+            migrations_file.seek(0)
             self.verify_nodes_format(ts, nodes_file, precision, base64_metadata)
             self.verify_edges_format(ts, edges_file, precision, base64_metadata)
             self.verify_sites_format(ts, sites_file, precision, base64_metadata)
@@ -3130,6 +3165,9 @@ def test_output_format(self, precision, base64_metadata):
             self.verify_populations_format(
                 ts, populations_file, precision, base64_metadata
             )
+            self.verify_migrations_format(
+                ts, migrations_file, precision, base64_metadata
+            )
 
     def verify_approximate_equality(self, ts1, ts2):
         """
@@ -3143,6 +3181,7 @@ def verify_approximate_equality(self, ts1, ts2):
         assert ts1.num_sites == ts2.num_sites
         assert ts1.num_mutations == ts2.num_mutations
         assert ts1.num_populations == ts2.num_populations
+        assert ts1.num_migrations == ts2.num_migrations
 
         checked = 0
         for n1, n2 in zip(ts1.nodes(), ts2.nodes()):
@@ -3182,6 +3221,18 @@ def verify_approximate_equality(self, ts1, ts2):
             assert s1.metadata == s2.metadata
         assert ts1.num_mutations == checked
 
+        checked = 0
+        for s1, s2 in zip(ts1.migrations(), ts2.migrations()):
+            checked += 1
+            assert s1.left == s2.left
+            assert s1.right == s2.right
+            assert s1.node == s2.node
+            assert s1.source == s2.source
+            assert s1.dest == s2.dest
+            assert s1.time == s2.time
+            assert s1.metadata == s2.metadata
+        assert ts1.num_migrations == checked
+
         # Check the trees
         check = 0
         for t1, t2 in zip(ts1.trees(), ts2.trees()):
@@ -3199,13 +3250,15 @@ def test_text_record_round_trip(self):
                 mutations_file = io.StringIO()
                 individuals_file = io.StringIO()
                 populations_file = io.StringIO()
+                migrations_file = io.StringIO()
                 ts1.dump_text(
                     nodes=nodes_file,
                     edges=edges_file,
                     sites=sites_file,
                     mutations=mutations_file,
                     individuals=individuals_file,
                     populations=populations_file,
+                    migrations=migrations_file,
                     precision=16,
                 )
                 nodes_file.seek(0)
@@ -3214,13 +3267,15 @@ def test_text_record_round_trip(self):
                 mutations_file.seek(0)
                 individuals_file.seek(0)
                 populations_file.seek(0)
+                migrations_file.seek(0)
                 ts2 = tskit.load_text(
                     nodes=nodes_file,
                     edges=edges_file,
                     sites=sites_file,
                     mutations=mutations_file,
                     individuals=individuals_file,
                     populations=populations_file,
+                    migrations=migrations_file,
                     sequence_length=ts1.sequence_length,
                     strict=True,
                 )
@@ -3231,31 +3286,36 @@ def test_empty_files(self):
         edges_file = io.StringIO("left\tright\tparent\tchild\n")
         sites_file = io.StringIO("position\tancestral_state\n")
         mutations_file = io.StringIO("site\tnode\tderived_state\n")
+        migrations_file = io.StringIO("left\tright\tnode\tsource\tdest\ttime\n")
         with pytest.raises(_tskit.LibraryError):
             tskit.load_text(
                 nodes=nodes_file,
                 edges=edges_file,
                 sites=sites_file,
                 mutations=mutations_file,
+                migrations=migrations_file,
             )
 
     def test_empty_files_sequence_length(self):
         nodes_file = io.StringIO("is_sample\ttime\n")
         edges_file = io.StringIO("left\tright\tparent\tchild\n")
         sites_file = io.StringIO("position\tancestral_state\n")
         mutations_file = io.StringIO("site\tnode\tderived_state\n")
+        migrations_file = io.StringIO("left\tright\tnode\tsource\tdest\ttime\n")
         ts = tskit.load_text(
             nodes=nodes_file,
             edges=edges_file,
             sites=sites_file,
             mutations=mutations_file,
+            migrations=migrations_file,
             sequence_length=100,
         )
         assert ts.sequence_length == 100
         assert ts.num_nodes == 0
         assert ts.num_edges == 0
         assert ts.num_sites == 0
         assert ts.num_edges == 0
+        assert ts.num_migrations == 0
 
     def test_load_text_no_populations(self):
         nodes_file = io.StringIO("is_sample\ttime\tpopulation\n1\t0\t2\n")

python/tests/test_metadata.py

@@ -264,6 +264,23 @@ def test_populations(self):
         for a, b in zip(expected, p):
             assert a.encode("utf8") == b.metadata
 
+    @pytest.mark.parametrize(
+        "base64_metadata,expected", [(True, ["pop", "gen"]), (False, ["cG9w", "Z2Vu"])]
+    )
+    def test_migrations(self, base64_metadata, expected):
+        migrations = io.StringIO(
+            """\
+        left    right    node    source    dest    time    metadata
+        10    100    0    3    4    123.0    cG9w
+        150    360    1    1    2    307.0    Z2Vu
+        """
+        )
+        m = tskit.parse_migrations(
+            migrations, strict=False, encoding="utf8", base64_metadata=base64_metadata
+        )
+        for a, b in zip(expected, m):
+            assert a.encode("utf8") == b.metadata
+
 
 class TestMetadataModule:
     """

python/tskit/text_formats.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2021 Tskit Developers
+# Copyright (c) 2021-2022 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
@@ -261,6 +261,7 @@ def dump_text(
     mutations,
     individuals,
     populations,
+    migrations,
     provenances,
     precision,
     encoding,
@@ -396,6 +397,39 @@ def dump_text(
             row = ("{id}\t" "{metadata}").format(id=population.id, metadata=metadata)
             print(row, file=populations)
 
+    if migrations is not None:
+        print(
+            "left",
+            "right",
+            "node",
+            "source",
+            "dest",
+            "time",
+            "metadata",
+            sep="\t",
+            file=migrations,
+        )
+        for migration in ts.migrations():
+            metadata = text_metadata(base64_metadata, encoding, migration)
+            row = (
+                "{left}\t"
+                "{right}\t"
+                "{node}\t"
+                "{source}\t"
+                "{dest}\t"
+                "{time}\t"
+                "{metadata}\t"
+            ).format(
+                left=migration.left,
+                right=migration.right,
+                node=migration.node,
+                source=migration.source,
+                dest=migration.dest,
+                time=migration.time,
+                metadata=metadata,
+            )
+            print(row, file=migrations)
+
     if provenances is not None:
         print("id", "timestamp", "record", sep="\t", file=provenances)
         for provenance in ts.provenances():