Merge pull request #16 from alife-data-standards/support-forest

Author: emilydolson

ALifeStdDev/phylogeny/loader.py

@@ -5,8 +5,12 @@
 
 def load_phylogeny_to_pandas_df(filename):
     data = pd.read_csv(filename)
+    data.columns = data.columns.str.replace(' ', '')
+
     # The ancestor_list column contains a string of lists of strings
     # Let's turn it into an actual list of strings
+    data.loc[:, "ancestor_list"] = \
+        data.loc[:, "ancestor_list"].apply(lambda x: x.strip())
     data.loc[:, "ancestor_list"] = \
         data.loc[:, "ancestor_list"].apply(ast.literal_eval)
     data.set_index("id", inplace=True)
@@ -69,7 +73,11 @@ def networkx_to_pandas_df(g, bonus_cols=None):
         bonus_cols = {}
     df = pd.DataFrame()
     df["id"] = g.nodes.keys()
-    df["ancestor_list"] = [list(g.predecessors(i)) for i in df["id"]]
+    df["ancestor_list"] = [list(g.predecessors(i)) if len(list(g.predecessors(i))) > 0 else [None] for i in df["id"]]
+
+    # roots = df[df["ancestor_list"].apply(lambda x: len(x) == 0)].index
+    # df.loc[roots, "ancestor_list"] = "[None]"
+
     for name1, name2 in bonus_cols.items():
         df[name1] = [g.nodes[i][name2] for i in df["id"]]
 

ALifeStdDev/phylogeny/utils.py

@@ -517,32 +517,36 @@ def abstract_asexual_phylogeny(phylogeny, attribute_list,
 
     # Start with the root nodes
     # TODO: Make work for forests (multiple roots)
-    root_id = get_root_ids(phylogeny)[0]
+    to_process = []
+    root_ids = get_root_ids(phylogeny)
     state_id = 0
+    next_id = 1
     # Add the first phylogeny state to the abstract phylogeny
-    abstract_phylogeny.add_node(state_id)
-    abstract_phylogeny.nodes[state_id]["state_id"] = state_id
-    abstract_phylogeny.nodes[state_id]["node_state"] = \
-        [phylogeny.nodes[root_id][attr] for attr in attribute_list]
-    # Add attributes to state
-    for attr in attribute_list:
-        abstract_phylogeny.nodes[state_id][attr] = phylogeny.nodes[root_id][attr]
-    if track_origin:
-        abstract_phylogeny.nodes[state_id]["origin_time"] = \
-            phylogeny.nodes[root_id][origin_time_attr]
-    if track_destruction:  # (this might get updated as we go)
-        dest_time = phylogeny.nodes[root_id][destruction_time_attr]
-        if dest_time == "none" or float(dest_time) == float("inf"):
-            dest_time = -1
-        dest_time = float(dest_time)
-        abstract_phylogeny.nodes[state_id]["destruction_time"] = dest_time
-
-    # Add first member
-    abstract_phylogeny.nodes[state_id]["members"] = {root_id:
-                                                   phylogeny.nodes[root_id]}
-
-    next_id = state_id + 1
-    to_process = [(root_id, state_id)]
+    for root_id in root_ids:
+        abstract_phylogeny.add_node(state_id)
+        abstract_phylogeny.nodes[state_id]["state_id"] = state_id
+        abstract_phylogeny.nodes[state_id]["node_state"] = \
+            [phylogeny.nodes[root_id][attr] for attr in attribute_list]
+        # Add attributes to state
+        for attr in attribute_list:
+            abstract_phylogeny.nodes[state_id][attr] = phylogeny.nodes[root_id][attr]
+        if track_origin:
+            abstract_phylogeny.nodes[state_id]["origin_time"] = \
+                phylogeny.nodes[root_id][origin_time_attr]
+        if track_destruction:  # (this might get updated as we go)
+            dest_time = phylogeny.nodes[root_id][destruction_time_attr]
+            if dest_time == "none" or float(dest_time) == float("inf"):
+                dest_time = -1
+            dest_time = float(dest_time)
+            abstract_phylogeny.nodes[state_id]["destruction_time"] = dest_time
+
+        # Add first member
+        abstract_phylogeny.nodes[state_id]["members"] = {root_id:
+                                                    phylogeny.nodes[root_id]}
+
+        to_process.append((root_id, state_id))
+        state_id = next_id
+        next_id += 1
 
     while to_process:
         phylogeny_id, state_id = to_process.pop()
@@ -592,7 +596,7 @@ def abstract_asexual_phylogeny(phylogeny, attribute_list,
                         phylogeny.nodes[id][attr]
                 # Add first member
                 abstract_phylogeny.nodes[next_state_id]["members"] = \
-                    {phylogeny_id: phylogeny.nodes[id]}
+                    {id: phylogeny.nodes[id]}
             to_process.append((id, next_state_id))
 
     return abstract_phylogeny