fix relabel

historydag/dag.py

@@ -3532,6 +3532,51 @@ def _unrooted_from_tree(treeroot):
     return HistoryDag(dagroot)
 
 
+def ascii_compare_histories(
+    history1,
+    history2,
+    name_func,
+    name_func2 = None,
+    show_internal=False,
+    sort=False,
+    compact=False,
+):
+    """A convenience function to print two histories as ascii art trees side-by-side.
+
+    Args:
+        history1: The first history to compare. Will appear on the left
+        history2: The second history to compare. Will appear on the right
+        name_func: A function mapping each HistoryDagNode to a node name string.
+        name_func2: A different name_func to be used for history2. If not provided,
+            ``name_func`` will be used.
+        show_internal: whether to show internal node names
+        sort: Whether to sort trees by node names
+        compact: Whether to represent trees in a more compact way"""
+    if name_func2 is None:
+        name_func2 = name_func
+    t1 = history1.to_ete(name_func=name_func)
+    t2 = history2.to_ete(name_func=name_func2)
+    trees = [t1, t2]
+    if sort:
+        for tree in trees:
+            for node in tree.traverse(strategy="postorder"):
+                node.children.sort(key=lambda n: n.name)
+
+    a1, a2 = [t.get_ascii(show_internal=show_internal, compact=compact).split('\n') for t in trees]
+
+    # There are no tabs in the ascii lines
+    offset = max(len(l) for l in a1) + 2
+    # expand all lines in a1 so they have len offset
+    a1padded = []
+    for line in a1:
+        add = offset - len(line)
+        a1padded.append(line + (" " * add))
+
+    # print
+    for l1, l2 in zip(a1padded, a2):
+        print(l1 + l2)
+
+
 def history_dag_from_trees(
     treelist: List[ete3.TreeNode],
     label_features: List[str],

tests/test_mutation_annotated_dag.py

@@ -1,5 +1,6 @@
 from historydag.mutation_annotated_dag import (
     load_MAD_protobuf_file,
+    load_MAD_protobuf,
     load_json_file,
 )
 from historydag.sequence_dag import SequenceHistoryDag
@@ -21,14 +22,12 @@ def test_load_protobuf():
 
     for dag in [cg_nid_dag, cg_dag]:
         assert dag.optimal_weight_annotate() == 428
-        assert dag.optimal_weight_annotate() == 1008
+        assert dag.optimal_weight_annotate(optimal_func=max) == 1008
 
-    test_filename = "_test_write_pb.pb"
-    cg_nid_dag.to_protobuf_file(test_filename)
-    ndag = load_MAD_protobuf_file(test_filename)
+    ndag = load_MAD_protobuf(cg_nid_dag.to_protobuf())
+    ndag._check_valid()
+    ndag = load_MAD_protobuf(cg_dag.to_protobuf())
     ndag._check_valid()
-    ndag.convert_to_collapsed()
-    assert cg_nid_dag.weight_count() == ndag.weight_count()
 
 
 def test_load_json():