speed up writing hdf5 by a huge factor

Author: alpae

edgehog/hdf5.py

@@ -16,6 +16,11 @@
     sys.exit(2)
 
 
+def fast_max_except_minus1(x):
+    arr = x.values
+    mask = arr != -1
+    return arr[mask].max() if mask.any() else -1
+
 class HDF5Writer:
     def __init__(self, fname, oma_db_fn, date_edges=False, orient_edges=False):
         self.fname = fname
@@ -95,8 +100,8 @@ def add_graph_at_level(self, taxid, tree_node):
             "Weight": "max",
             "Evidence": "min",
             # -1 indicates 'n/a', all other values should be consistent.
-            "LCA_taxid": lambda x: x[x != -1].max() if not x[x != -1].empty else -1,
-            "Orientation": lambda x: x[x != -1].max() if not x[x != -1].empty else -1,
+            "LCA_taxid": fast_max_except_minus1,
+            "Orientation": fast_max_except_minus1,
             "OrientationScore": "max",
 
         })
@@ -167,8 +172,8 @@ def add_extant_graph(self, taxid, tree_node):
             "Weight": "max",
             "Evidence": "min",
             # -1 indicates 'n/a', all other values should be consistent.
-            "LCA_taxid": lambda x: x[x != -1].max() if not x[x != -1].empty else -1,
-            "Orientation": lambda x: x[x != -1].max() if not x[x != -1].empty else -1,
+            "LCA_taxid": fast_max_except_minus1,
+            "Orientation": fast_max_except_minus1,
             "OrientationScore": "max",
         })
         as_array = sumdf.astype({col: dtype.fields[col][0] for col in dtype.names}).to_records(index=False)

edgehog/infer_ancestral_synteny_graphs.py

@@ -138,10 +138,6 @@ def build_transitive_graph(graph, orient_edges, max_gap_length = 2):
                                             children = list(set(graph[left_gene][gene]["children"] + graph[gene][right_gene]["children"])),
                                             extant_descendants = list(set(graph[left_gene][gene]["extant_descendants"] + graph[gene][right_gene]["extant_descendants"])))
                              if orient_edges:
-                                 try:
-                                     print(left_gene.hog_id, right_gene.hog_id)
-                                 except:
-                                     print(left_gene.prot_id, right_gene.prot_id)
                                  wA, uA, cA, dA = graph[left_gene][gene]["weight"], graph[left_gene][gene]["unidirectional"], graph[left_gene][gene]["convergent"], graph[left_gene][gene]["divergent"]
                                  wB, uB, cB, dB = graph[gene][right_gene]["weight"], graph[gene][right_gene]["unidirectional"], graph[gene][right_gene]["convergent"], graph[gene][right_gene]["divergent"]
                                  u, c, d = orient_reconnected_edges(wA, uA, cA, dA, wB, uB, cB, dB)