cleanup

Author: nspope

tests/test_mutation_rate_map.py

@@ -23,6 +23,8 @@
 Test mutation rate map support
 """
 
+from string import ascii_lowercase, ascii_uppercase
+
 import numpy as np
 import pytest
 import tskit
@@ -82,21 +84,31 @@ def example_transform_pair():
         left=[0, 0, 25, 25, 50, 50, 0, 0, 25, 50, 0, 0, 75, 75, 50, 50],
         right=[25, 25, 50, 50, 100, 100, 100, 25, 50, 100, 50, 50, 100, 100, 75, 75],
     )
-    ancestral_state = [str(x).encode("ascii") for x in range(11)]
+    site_id = list(range(11))
+    ancestral_state = [str(x).encode("ascii") for x in site_id]
     ancestral_state, ancestral_state_offset = tskit.pack_bytes(ancestral_state)
+    site_metadata = [ascii_uppercase[i].encode("ascii") for i in site_id]
+    site_metadata, site_metadata_offset = tskit.pack_bytes(site_metadata)
     tab.sites.set_columns(
         position=[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 99],
         ancestral_state=ancestral_state,
         ancestral_state_offset=ancestral_state_offset,
+        metadata=site_metadata,
+        metadata_offset=site_metadata_offset,
     )
-    derived_state = [str(x).encode("ascii") for x in range(16)]
+    mut_id = list(range(16))
+    derived_state = [str(x).encode("ascii") for x in mut_id]
     derived_state, derived_state_offset = tskit.pack_bytes(derived_state)
+    mut_metadata = [ascii_lowercase[i].encode("ascii") for i in mut_id]
+    mut_metadata, mut_metadata_offset = tskit.pack_bytes(mut_metadata)
     tab.mutations.set_columns(
         site=[0, 1, 1, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8, 9, 9, 10],
         parent=[-1, -1, 1, -1, -1, 4, -1, -1, 7, -1, -1, 10, -1, -1, 13, -1],
         node=[7] * 16,
         derived_state=derived_state,
         derived_state_offset=derived_state_offset,
+        metadata=mut_metadata,
+        metadata_offset=mut_metadata_offset,
     )
     tab.sequence_length = 100
     ts = tab.tree_sequence()
@@ -114,21 +126,31 @@ def example_transform_pair():
         left=[0, 0, 2, 2, 0, 0, 2, 0, 0, 2, 2],
         right=[2, 2, 52, 52, 52, 2, 52, 2, 2, 52, 52],
     )
-    ancestral_state = [str(x).encode("ascii") for x in [0, 1, 8, 9, 10]]
+    site_subset = [0, 1, 8, 9, 10]
+    ancestral_state = [str(x).encode("ascii") for x in site_subset]
     ancestral_state, ancestral_state_offset = tskit.pack_bytes(ancestral_state)
+    site_metadata = [ascii_uppercase[i].encode("ascii") for i in site_subset]
+    site_metadata, site_metadata_offset = tskit.pack_bytes(site_metadata)
     trans_tab.sites.set_columns(
         position=[0, 1, 12, 32, 50],
         ancestral_state=ancestral_state,
         ancestral_state_offset=ancestral_state_offset,
+        metadata=site_metadata,
+        metadata_offset=site_metadata_offset,
     )
-    derived_state = [str(x).encode("ascii") for x in [0, 1, 2, 12, 13, 14, 15]]
+    mut_subset = [0, 1, 2, 12, 13, 14, 15]
+    derived_state = [str(x).encode("ascii") for x in mut_subset]
     derived_state, derived_state_offset = tskit.pack_bytes(derived_state)
+    mut_metadata = [ascii_lowercase[i].encode("ascii") for i in mut_subset]
+    mut_metadata, mut_metadata_offset = tskit.pack_bytes(mut_metadata)
     trans_tab.mutations.set_columns(
         site=[0, 1, 1, 2, 3, 3, 4],
         parent=[-1, -1, 1, -1, -1, 4, -1],
         node=[7] * 7,
         derived_state=derived_state,
         derived_state_offset=derived_state_offset,
+        metadata=mut_metadata,
+        metadata_offset=mut_metadata_offset,
     )
     trans_tab.sequence_length = 52
     trans_ts = trans_tab.tree_sequence()
@@ -142,6 +164,11 @@ def test_transform_coordinates_by_ratemap():
     ts, trans_ts, ratemap = example_transform_pair()
     trans_ts_ck = tsdate.util.transform_coordinates_by_ratemap(ts, ratemap)
     assert trans_ts_ck == trans_ts
+    # manual metadata checking, FIXME: is this done by ts equality operator?
+    for m, m_ck in zip(trans_ts.mutations(), trans_ts_ck.mutations()):
+        assert m.metadata == m_ck.metadata
+    for s, s_ck in zip(trans_ts.sites(), trans_ts_ck.sites()):
+        assert s.metadata == s_ck.metadata
 
 
 def test_transform_coordinates_identity():

tsdate/util.py

@@ -771,10 +771,10 @@ def transform_coordinates_by_ratemap(
 ) -> tskit.TreeSequence:
     """
     Return a copy of the tree sequence in the coordinate system created by `y =
-    ratemap.get_cumulative_mass(x)`. Zero length edges in the new coordinate system
-    are removed, as are any sites and mutations that fall within zero-rate
-    intervals. All nodes are retained, even if they are disconnected in the transformed
-    tree sequence.
+    ratemap.get_cumulative_mass(x)`. Zero-length edges in the new coordinate
+    system are removed, as are any sites and mutations that fall within
+    intervals with zero or NaN rates. All nodes are retained, even if they are
+    disconnected in the transformed tree sequence.
     """
     assert ratemap.sequence_length == ts.sequence_length, "Ratemap has the wrong length"
 
@@ -787,10 +787,6 @@ def transform_coordinates_by_ratemap(
     site_map = tab.sites.keep_rows(ratemap.get_rate(tab.sites.position) > 0.0)
     tab.sites.position = ratemap.get_cumulative_mass(tab.sites.position)
     tab.mutations.site = site_map[tab.mutations.site]
-    tab.mutations.keep_rows(tab.mutations.site != tskit.NULL)
-
-    if tab.sites.num_rows != ts.num_sites:
-        tab.build_index()
-        tab.compute_mutation_parents()
+    tab.mutations.keep_rows(tab.mutations.site != tskit.NULL)  # updates parent column
 
     return tab.tree_sequence()