Merge pull request #103 from emilydolson/packaging-test

Implement collapse unifurcations feature

systematics_bindings.cpp

@@ -352,6 +352,18 @@ PYBIND11_MODULE(systematics, m) {
             val : bool 
                 Value representing whether a synchronous population is being tracked.
         )mydelimiter")
+        .def("set_collapse_unifurcations", static_cast<void (sys_t::*) (bool)>(&sys_t::SetCollapseUnifurcations), R"mydelimiter(
+            A setter method to configure whether unifurcations (sequences of extinct parents with exactly one offspring taxon)
+            should be collapsed (i.e. extinct taxon with one offspring will be removed and the offspring's parent will be set to be 
+            what was formerly its grandparent). This setting will save space and be more consistent with a lot of biological representations
+            but will sacrifice information.
+            This option defaults to False.
+
+            Parameters
+            ----------
+            val : bool 
+                Value representing whether to collapse unifurcations
+        )mydelimiter")        
         .def("set_update", static_cast<void (sys_t::*) (size_t)>(&sys_t::SetUpdate), R"mydelimiter(
             A setter method to modify the current time step. This should be used if you want PhylotrackPy to track when events occur.
 
@@ -388,6 +400,10 @@ PYBIND11_MODULE(systematics, m) {
             It is recommended to verify this setting before any tracking.
             Can be set using the `set_track_synchronous()` method.
         )mydelimiter")
+        .def("get_collapse_unifurcations", static_cast<bool (sys_t::*) () const>(&sys_t::GetCollapseUnifurcations), R"mydelimiter(
+            Whether the Systematics Manager is configured to collapse unifurcations.
+            Can be set using the `set_collapse_unifurcations()` method.
+        )mydelimiter")
         .def("get_update", static_cast<size_t (sys_t::*) () const>(&sys_t::GetUpdate), R"mydelimiter(
             Returns the current timestep of the simulation.
             This time step can be overriden using the `set_update()` method.
@@ -586,7 +602,7 @@ PYBIND11_MODULE(systematics, m) {
             Set a custom function that is triggered every time a taxon is pruned from the tree. This occurs when a taxon and all its descendants go extinct.
             The function must take a single argument: a `taxon_t` object representing the taxon getting pruned.
             The custom function will be triggered at the beginning of the taxon pruning process.
-            This allows the user to customize the way objects are represented interlally by the systematics manager, or to implement extra bookkeeping functionality.
+            This allows the user to customize the way objects are represented internally by the systematics manager, or to implement extra bookkeeping functionality.
 
             Parameters
             ----------

test/test_systematics.py

@@ -38,9 +38,24 @@ def test_systematics_by_position():
     sys.add_org_by_position(org, org_pos)
     child_pos = systematics.WorldPosition(2, 0)
     child_org = ExampleOrg("hello2")
+    assert sys.is_taxon_at(child_pos) is False
     sys.add_org_by_position(child_org, child_pos, org_pos)
+    assert sys.is_taxon_at(child_pos) is True
+    assert sys.get_taxon_at(org_pos).get_info() == "hello"
+    sys.swap_positions(child_pos, org_pos)
+    assert sys.get_taxon_at(org_pos).get_info() == "hello2"
+    sys.swap_positions(child_pos, org_pos)
     sys.remove_org_by_position(org_pos)
-    # sys.remove_org_by_position((2,0))
+    assert sys.is_taxon_at(org_pos) is False
+    sys.remove_org_by_position_after_repro(child_pos)
+    assert sys.is_taxon_at(child_pos) is True
+    sys.add_org_by_position(org, org_pos, child_pos)
+    assert sys.is_taxon_at(child_pos) is False
+    sys.set_next_parent_by_position(org_pos)
+    assert sys.get_next_parent() == sys.get_taxon_at(org_pos)
+    sys.add_org_by_position(ExampleOrg("test"), child_pos)
+    assert sys.get_taxon_at(child_pos).get_parent() == sys.get_taxon_at(org_pos)
+    assert sys.get_next_parent() is None
 
 
 def test_construct_systematics():
@@ -468,3 +483,48 @@ class Organism():
     org0 = Organism()
     org0.genotype
     syst.add_org(org0)
+
+
+def test_collapse_unifurcations():
+    org_info_1, org_info_2 = [1, 2]
+    sys = systematics.Systematics(taxon_info_fun, True, True, False, False)
+    assert sys.get_collapse_unifurcations() is False
+    sys.set_collapse_unifurcations(True)
+    assert sys.get_collapse_unifurcations() is True
+    org = ExampleOrg(org_info_1)
+    org2 = ExampleOrg(org_info_2)
+    org_tax = sys.add_org(org)
+    org2_tax = sys.add_org(org2, org_tax)
+    org3_tax = sys.add_org(org2, org_tax)
+    sys.remove_org(org_tax)
+    org4_tax = sys.add_org(org, org2_tax)
+    org5_tax = sys.add_org(org, org4_tax)
+    assert sys.get_num_ancestors() == 1
+    sys.remove_org(org2_tax)
+    assert sys.get_num_ancestors() == 1
+    sys.remove_org(org4_tax)
+    assert sys.get_num_ancestors() == 1
+    sys.remove_org(org5_tax)
+    assert sys.get_num_ancestors() == 0
+
+
+tax_sum = 0
+
+
+def test_on_prune():
+    sys = systematics.Systematics(lambda x: x, True, True, False, False)
+
+    def prune_fun(x):
+        global tax_sum
+        tax_sum += x.get_info()
+
+    sys.on_prune(prune_fun)
+    tax1 = sys.add_org(1)
+    tax2 = sys.add_org(2, tax1)
+    tax3 = sys.add_org(3, tax2)
+    sys.remove_org(tax2)
+    assert tax_sum == 0
+    sys.remove_org(tax3)
+    assert tax_sum == 5
+    sys.remove_org(tax1)
+    assert tax_sum == 6