LineageTree 1.6.1__py3-none-any.whl → 1.8.0__py3-none-any.whl

LineageTree/lineageTree.py

@@ -20,7 +20,8 @@ try:
     from edist import uted
 except ImportError:
     warnings.warn(
-        "No edist installed therefore you will not be able to compute the tree edit distance."
+        "No edist installed therefore you will not be able to compute the tree edit distance.",
+        stacklevel=2,
     )
 import matplotlib.pyplot as plt
 import numpy as np
@@ -69,7 +70,12 @@ class lineageTree(lineageTreeLoaders):
             sub = set(self.get_sub_tree(node))
             specific_leaves = sub.intersection(self.leaves)
             for leaf in specific_leaves:
-                self.add_branch(leaf, self.t_e - self.time[leaf], reverse=True)
+                self.add_branch(
+                    leaf,
+                    (self.t_e - self.time[leaf]),
+                    reverse=True,
+                    move_timepoints=True,
+                )
 
     ###TODO pos can be callable and stay motionless (copy the position of the succ node, use something like optical flow)
     def add_branch(
@@ -87,8 +93,8 @@ class lineageTree(lineageTreeLoaders):
             pred (int): Id of the successor (predecessor if reverse is False)
             length (int): The length of the new branch.
             pos (np.ndarray, optional): The new position of the branch. Defaults to None.
-            move_timepoints (bool): Moves the ti Important only if reverse= True
-            reverese (bool): If reverse will add a successor branch instead of a predecessor branch
+            move_timepoints (bool): Moves the time, important only if reverse= True
+            reverese (bool): If True will create a branch that goes forwards in time otherwise backwards.
         Returns:
             (int): Id of the first node of the sublineage.
         """
@@ -129,19 +135,20 @@ class lineageTree(lineageTreeLoaders):
                     )
                     pred = _next
         else:
-            for t in range(length):
+            for _ in range(length):
                 _next = self.add_node(
-                    time + t, succ=pred, pos=self.pos[original], reverse=False
+                    self.time[pred] + 1,
+                    succ=pred,
+                    pos=self.pos[original],
+                    reverse=False,
                 )
                 pred = _next
+            self.successor[self.get_cycle(pred)[-1]] = []
             self.labels[pred] = "New branch"
         if self.time[pred] == self.t_b:
-            self.roots.add(pred)
             self.labels[pred] = "New branch"
         if original in self.roots and reverse is True:
-            self.roots.add(pred)
             self.labels[pred] = "New branch"
-            self.roots.remove(original)
             self.labels.pop(original, -1)
         self.t_e = max(self.time_nodes)
         return pred
@@ -163,9 +170,7 @@ class lineageTree(lineageTreeLoaders):
             self.predecessor.pop(new_lT)
             label_of_root = self.labels.get(cycle[0], cycle[0])
             self.labels[cycle[0]] = f"L-Split {label_of_root}"
-            new_tr = self.add_branch(
-                new_lT, len(cycle) + 1, move_timepoints=False
-            )
+            new_tr = self.add_branch(new_lT, len(cycle), move_timepoints=False)
             self.roots.add(new_tr)
             self.labels[new_tr] = f"R-Split {label_of_root}"
             return new_tr
@@ -205,7 +210,10 @@ class lineageTree(lineageTreeLoaders):
         self.remove_nodes(new_root1)
         self.successor[new_root2].append(next_root1)
         self.predecessor[next_root1] = [new_root2]
-        new_branch = self.add_branch(new_root2, length)
+        new_branch = self.add_branch(
+            new_root2,
+            length - 1,
+        )
         self.labels[new_branch] = f"Fusion of {new_root1} and {new_root2}"
         return new_branch
 
@@ -317,7 +325,7 @@ class lineageTree(lineageTreeLoaders):
             new_length (int): The new length of the tree.
         """
         if new_length <= 1:
-            warnings.warn("New length should be more than 1")
+            warnings.warn("New length should be more than 1", stacklevel=2)
             return None
         cycle = self.get_cycle(node)
         length = len(cycle)
@@ -325,7 +333,10 @@ class lineageTree(lineageTreeLoaders):
         if length == 1 and new_length != 1:
             pred = self.predecessor.pop(node, None)
             new_node = self.add_branch(
-                node, length=new_length, move_timepoints=True, reverse=False
+                node,
+                length=new_length - 1,
+                move_timepoints=True,
+                reverse=False,
             )
             if pred:
                 self.successor[pred[0]].remove(node)
@@ -365,11 +376,40 @@ class lineageTree(lineageTreeLoaders):
         else:
             return None
 
+    @property
+    def time_resolution(self):
+        if not hasattr(self, "_time_resolution"):
+            self.time_resolution = 1
+        return self._time_resolution / 10
+
+    @time_resolution.setter
+    def time_resolution(self, time_resolution):
+        if time_resolution is not None:
+            self._time_resolution = int(time_resolution * 10)
+        else:
+            warnings.warn("Time resolution set to default 0", stacklevel=2)
+            self._time_resolution = 10
+
+    @property
+    def depth(self):
+        if not hasattr(self, "_depth"):
+            self._depth = {}
+            for leaf in self.leaves:
+                self._depth[leaf] = 1
+                while leaf in self.predecessor:
+                    parent = self.predecessor[leaf][0]
+                    current_depth = self._depth.get(parent, 0)
+                    self._depth[parent] = max(
+                        self._depth[leaf] + 1, current_depth
+                    )
+                    leaf = parent
+            for root in self.roots - set(self._depth):
+                self._depth[root] = 1
+        return self._depth
+
     @property
     def roots(self):
-        if not hasattr(self, "_roots"):
-            self._roots = set(self.nodes).difference(self.predecessor)
-        return self._roots
+        return set(self.nodes).difference(self.predecessor)
 
     @property
     def edges(self):
@@ -377,7 +417,7 @@ class lineageTree(lineageTreeLoaders):
 
     @property
     def leaves(self):
-        return set(self.predecessor).difference(self.successor)
+        return {p for p, s in self.successor.items() if s == []}
 
     @property
     def labels(self):
@@ -388,10 +428,10 @@ class lineageTree(lineageTreeLoaders):
                 }
             else:
                 self._labels = {
-                    i: "Unlabeled"
-                    for i in self.roots
-                    for l in self.find_leaves(i)
-                    if abs(self.time[l] - self.time[i])
+                    root: "Unlabeled"
+                    for root in self.roots
+                    for leaf in self.find_leaves(root)
+                    if abs(self.time[leaf] - self.time[root])
                     >= abs(self.t_e - self.t_b) / 4
                 }
         return self._labels
@@ -483,7 +523,7 @@ class lineageTree(lineageTreeLoaders):
 
                 f.write("@5\n")
                 for C in self.to_take_time[t]:
-                    f.write("%f\n" % (manual_labels.get(C, default_label)))
+                    f.write(f"{manual_labels.get(C, default_label):f}\n")
                     f.write(f"{0:f}\n")
 
                 f.write("@6\n")
@@ -502,8 +542,7 @@ class lineageTree(lineageTreeLoaders):
                 f.write("@7\n")
                 for C in self.to_take_time[t]:
                     f.write(
-                        "%f\n"
-                        % (np.linalg.norm(points_v[C][0] - points_v[C][-1]))
+                        f"{np.linalg.norm(points_v[C][0] - points_v[C][-1]):f}\n"
                     )
 
                 f.write("@8\n")
@@ -932,10 +971,10 @@ class lineageTree(lineageTreeLoaders):
             if node_properties:
                 for p_name, (p_dict, default) in node_properties.items():
                     if isinstance(list(p_dict.values())[0], str):
-                        f.write('(property 0 string "%s"\n' % p_name)
+                        f.write(f'(property 0 string "{p_name}"\n')
                         f.write(f"\t(default {default} {default})\n")
                     elif isinstance(list(p_dict.values())[0], Number):
-                        f.write('(property 0 double "%s"\n' % p_name)
+                        f.write(f'(property 0 double "{p_name}"\n')
                         f.write('\t(default "0" "0")\n')
                     for n in nodes_to_use:
                         f.write(
@@ -1009,148 +1048,6 @@ class lineageTree(lineageTreeLoaders):
 
             f.close()
 
-    def read_from_binary(self, fname: str):
-        """
-        Reads a binary lineageTree file name.
-        Format description: see self.to_binary
-
-        Args:
-            fname: string, path to the binary file
-            reverse_time: bool, not used
-        """
-        q_size = struct.calcsize("q")
-        H_size = struct.calcsize("H")
-        d_size = struct.calcsize("d")
-
-        with open(fname, "rb") as f:
-            len_tree = struct.unpack("q", f.read(q_size))[0]
-            len_time = struct.unpack("q", f.read(q_size))[0]
-            len_pos = struct.unpack("q", f.read(q_size))[0]
-            number_sequence = list(
-                struct.unpack("q" * len_tree, f.read(q_size * len_tree))
-            )
-            time_sequence = list(
-                struct.unpack("H" * len_time, f.read(H_size * len_time))
-            )
-            pos_sequence = np.array(
-                struct.unpack("d" * len_pos, f.read(d_size * len_pos))
-            )
-
-            f.close()
-
-        successor = {}
-        predecessor = {}
-        time = {}
-        time_nodes = {}
-        time_edges = {}
-        pos = {}
-        is_root = {}
-        nodes = []
-        edges = []
-        waiting_list = []
-        print(number_sequence[0])
-        i = 0
-        done = False
-        if max(number_sequence[::2]) == -1:
-            tmp = number_sequence[1::2]
-            if len(tmp) * 3 == len(pos_sequence) == len(time_sequence) * 3:
-                time = dict(list(zip(tmp, time_sequence)))
-                for c, t in time.items():
-                    time_nodes.setdefault(t, set()).add(c)
-                pos = dict(
-                    list(zip(tmp, np.reshape(pos_sequence, (len_time, 3))))
-                )
-                is_root = {c: True for c in tmp}
-                nodes = tmp
-                done = True
-        while (
-            i < len(number_sequence) and not done
-        ):  # , c in enumerate(number_sequence[:-1]):
-            c = number_sequence[i]
-            if c == -1:
-                if waiting_list != []:
-                    prev_mother = waiting_list.pop()
-                    successor[prev_mother].insert(0, number_sequence[i + 1])
-                    edges.append((prev_mother, number_sequence[i + 1]))
-                    time_edges.setdefault(t, set()).add(
-                        (prev_mother, number_sequence[i + 1])
-                    )
-                    is_root[number_sequence[i + 1]] = False
-                    t = time[prev_mother] + 1
-                else:
-                    t = time_sequence.pop(0)
-                    is_root[number_sequence[i + 1]] = True
-
-            elif c == -2:
-                successor[waiting_list[-1]] = [number_sequence[i + 1]]
-                edges.append((waiting_list[-1], number_sequence[i + 1]))
-                time_edges.setdefault(t, set()).add(
-                    (waiting_list[-1], number_sequence[i + 1])
-                )
-                is_root[number_sequence[i + 1]] = False
-                pos[waiting_list[-1]] = pos_sequence[:3]
-                pos_sequence = pos_sequence[3:]
-                nodes.append(waiting_list[-1])
-                time[waiting_list[-1]] = t
-                time_nodes.setdefault(t, set()).add(waiting_list[-1])
-                t += 1
-
-            elif number_sequence[i + 1] >= 0:
-                successor[c] = [number_sequence[i + 1]]
-                edges.append((c, number_sequence[i + 1]))
-                time_edges.setdefault(t, set()).add(
-                    (c, number_sequence[i + 1])
-                )
-                is_root[number_sequence[i + 1]] = False
-                pos[c] = pos_sequence[:3]
-                pos_sequence = pos_sequence[3:]
-                nodes.append(c)
-                time[c] = t
-                time_nodes.setdefault(t, set()).add(c)
-                t += 1
-
-            elif number_sequence[i + 1] == -2:
-                waiting_list += [c]
-
-            elif number_sequence[i + 1] == -1:
-                pos[c] = pos_sequence[:3]
-                pos_sequence = pos_sequence[3:]
-                nodes.append(c)
-                time[c] = t
-                time_nodes.setdefault(t, set()).add(c)
-                t += 1
-                i += 1
-                if waiting_list != []:
-                    prev_mother = waiting_list.pop()
-                    successor[prev_mother].insert(0, number_sequence[i + 1])
-                    edges.append((prev_mother, number_sequence[i + 1]))
-                    time_edges.setdefault(t, set()).add(
-                        (prev_mother, number_sequence[i + 1])
-                    )
-                    if i + 1 < len(number_sequence):
-                        is_root[number_sequence[i + 1]] = False
-                    t = time[prev_mother] + 1
-                else:
-                    if len(time_sequence) > 0:
-                        t = time_sequence.pop(0)
-                    if i + 1 < len(number_sequence):
-                        is_root[number_sequence[i + 1]] = True
-            i += 1
-
-        predecessor = {vi: [k] for k, v in successor.items() for vi in v}
-
-        self.successor = successor
-        self.predecessor = predecessor
-        self.time = time
-        self.time_nodes = time_nodes
-        self.time_edges = time_edges
-        self.pos = pos
-        self.nodes = set(nodes)
-        self.t_b = min(time_nodes.keys())
-        self.t_e = max(time_nodes.keys())
-        self.is_root = is_root
-        self.max_id = max(self.nodes)
-
     def write(self, fname: str):
         """
         Write a lineage tree on disk as an .lT file.
@@ -1165,7 +1062,7 @@ class lineageTree(lineageTreeLoaders):
             f.close()
 
     @classmethod
-    def load(clf, fname: str, rm_empty_lists=True):
+    def load(clf, fname: str, rm_empty_lists=False):
         """
         Loading a lineage tree from a ".lT" file.
 
@@ -1178,18 +1075,8 @@ class lineageTree(lineageTreeLoaders):
         with open(fname, "br") as f:
             lT = pkl.load(f)
             f.close()
-        if rm_empty_lists:
-            if [] in lT.successor.values():
-                for node, succ in lT.successor.items():
-                    if succ == []:
-                        lT.successor.pop(node)
-            if [] in lT.predecessor.values():
-                for node, succ in lT.predecessor.items():
-                    if succ == []:
-                        lT.predecessor.pop(node)
-            lT.t_e = max(lT.time_nodes)
-            lT.t_b = min(lT.time_nodes)
-            warnings.warn("Empty lists have been removed")
+        if not hasattr(lT, "time_resolution"):
+            lT.time_resolution = None
         return lT
 
     def get_idx3d(self, t: int) -> tuple:
@@ -1392,8 +1279,10 @@ class lineageTree(lineageTreeLoaders):
         while to_do:
             current = to_do.pop()
             track = self.get_successors(current, end_time=end_time)
-            branches += [track]
-            to_do += self[track[-1]]
+            # if len(track) != 1 or self.time[current] <= end_time:
+            if self.time[track[-1]] <= end_time:
+                branches += [track]
+                to_do += self[track[-1]]
         return branches
 
     def get_all_tracks(self, force_recompute: bool = False) -> list:
@@ -1449,11 +1338,10 @@ class lineageTree(lineageTreeLoaders):
         leaves = set()
         while to_do:
             curr = to_do.pop()
-            succ = self.successor.get(curr)
-            if succ is not None:
+            succ = self.successor.get(curr, [])
+            if not succ:
                 leaves.add(curr)
-            else:
-                to_do += succ
+            to_do += succ
         return leaves
 
     def get_sub_tree(
@@ -1629,6 +1517,8 @@ class lineageTree(lineageTreeLoaders):
         at a given time `time`.
 
         If there is no ancestor, returns `-1`
+        If time is None return the root of the sub tree that spawns
+        the node n.
 
         Args:
             n (int): node for which to look the ancestor
@@ -1645,7 +1535,9 @@ class lineageTree(lineageTreeLoaders):
         if time is None:
             time = self.t_b
         ancestor = n
-        while time < self.time.get(ancestor, -1):
+        while (
+            time < self.time.get(ancestor, -1) and ancestor in self.predecessor
+        ):
             ancestor = self.predecessor.get(ancestor, [-1])[0]
         return ancestor
 
@@ -1674,10 +1566,11 @@ class lineageTree(lineageTreeLoaders):
     def unordered_tree_edit_distances_at_time_t(
         self,
         t: int,
-        delta: callable = None,
-        norm: callable = None,
-        recompute: bool = False,
         end_time: int = None,
+        style="simple",
+        downsample: int = 2,
+        normalize: bool = True,
+        recompute: bool = False,
     ) -> dict:
         """
         Compute all the pairwise unordered tree edit distances from Zhang 996 between the trees spawned at time `t`
@@ -1704,7 +1597,12 @@ class lineageTree(lineageTreeLoaders):
         for n1, n2 in combinations(roots, 2):
             key = tuple(sorted((n1, n2)))
             self.uted[t][key] = self.unordered_tree_edit_distance(
-                n1, n2, end_time=end_time
+                n1,
+                n2,
+                end_time=end_time,
+                style=style,
+                downsample=downsample,
+                normalize=normalize,
             )
         return self.uted[t]
 
@@ -1713,11 +1611,11 @@ class lineageTree(lineageTreeLoaders):
         n1: int,
         n2: int,
         end_time: int = None,
-        style="fragmented",
-        node_lengths: tuple = (1, 5, 7),
+        norm: Union["max", "sum", None] = "max",
+        style="simple",
+        downsample: int = 2,
     ) -> float:
         """
-        TODO: Add option for choosing which tree aproximation should be used (Full, simple, comp)
         Compute the unordered tree edit distance from Zhang 1996 between the trees spawned
         by two nodes `n1` and `n2`. The topology of the trees are compared and the matching
         cost is given by the function delta (see edist doc for more information).
@@ -1736,10 +1634,18 @@ class lineageTree(lineageTreeLoaders):
 
         tree = tree_style[style].value
         tree1 = tree(
-            lT=self, node_length=node_lengths, end_time=end_time, root=n1
+            lT=self,
+            downsample=downsample,
+            end_time=end_time,
+            root=n1,
+            time_scale=1,
         )
         tree2 = tree(
-            lT=self, node_length=node_lengths, end_time=end_time, root=n2
+            lT=self,
+            downsample=downsample,
+            end_time=end_time,
+            root=n2,
+            time_scale=1,
         )
         delta = tree1.delta
         _, times1 = tree1.tree
@@ -1763,71 +1669,140 @@ class lineageTree(lineageTreeLoaders):
             times1=times1,
             times2=times2,
         )
+        norm1 = tree1.get_norm()
+        norm2 = tree2.get_norm()
+        norm_dict = {"max": max, "sum": sum, "None": lambda x: 1}
+        if norm is None:
+            norm = "None"
+        if norm not in norm_dict:
+            raise Warning(
+                "Select a viable normalization method (max, sum, None)"
+            )
+        return uted.uted(
+            nodes1, adj1, nodes2, adj2, delta=delta_tmp
+        ) / norm_dict[norm]([norm1, norm2])
 
-        return uted.uted(nodes1, adj1, nodes2, adj2, delta=delta_tmp) / max(
-            tree1.get_norm(), tree2.get_norm()
+    @staticmethod
+    def __plot_nodes(
+        hier, selected_nodes, color, size, ax, default_color="black", **kwargs
+    ):
+        """
+        Private method that plots the nodes of the tree.
+        """
+        hier_unselected = np.array(
+            [v for k, v in hier.items() if k not in selected_nodes]
         )
+        if hier_unselected.any():
+            ax.scatter(
+                *hier_unselected.T,
+                s=size,
+                zorder=10,
+                color=default_color,
+                **kwargs,
+            )
+        if selected_nodes.intersection(hier.keys()):
+            hier_selected = np.array(
+                [v for k, v in hier.items() if k in selected_nodes]
+            )
+            ax.scatter(
+                *hier_selected.T, s=size, zorder=10, color=color, **kwargs
+            )
+
+    @staticmethod
+    def __plot_edges(
+        hier,
+        lnks_tms,
+        selected_edges,
+        color,
+        ax,
+        default_color="black",
+        **kwargs,
+    ):
+        """
+        Private method that plots the edges of the tree.
+        """
+        x, y = [], []
+        for pred, succs in lnks_tms["links"].items():
+            for succ in succs:
+                if pred not in selected_edges or succ not in selected_edges:
+                    x.extend((hier[succ][0], hier[pred][0], None))
+                    y.extend((hier[succ][1], hier[pred][1], None))
+        ax.plot(x, y, linewidth=0.3, zorder=0.1, c=default_color, **kwargs)
+        x, y = [], []
+        for pred, succs in lnks_tms["links"].items():
+            for succ in succs:
+                if pred in selected_edges and succ in selected_edges:
+                    x.extend((hier[succ][0], hier[pred][0], None))
+                    y.extend((hier[succ][1], hier[pred][1], None))
+        ax.plot(x, y, linewidth=0.3, zorder=0.2, c=color, **kwargs)
 
     def draw_tree_graph(
         self,
         hier,
         lnks_tms,
-        selected_cells=None,
-        color="magenta",
+        selected_nodes=None,
+        selected_edges=None,
+        color_of_nodes="magenta",
+        color_of_edges=None,
+        size=10,
         ax=None,
-        figure=None,
+        default_color="black",
         **kwargs,
     ):
-        if selected_cells is None:
-            selected_cells = []
+        """Function to plot the tree graph.
+
+        Args:
+            hier (dict): Dictinary that contains the positions of all nodes.
+            lnks_tms (dict): 2 dictionaries: 1 contains all links from start of life cycle to end of life cycle and
+                                             the succesors of each cell.
+                                             1 contains the length of each life cycle.
+            selected_nodes (list|set, optional): Which cells are to be selected (Painted with a different color). Defaults to None.
+            selected_edges (list|set, optional): Which edges are to be selected (Painted with a different color). Defaults to None.
+            color_of_nodes (str, optional): Color of selected nodes. Defaults to "magenta".
+            color_of_edges (_type_, optional): Color of selected edges. Defaults to None.
+            size (int, optional): Size of the nodes. Defaults to 10.
+            ax (_type_, optional): Plot the graph on existing ax. Defaults to None.
+            figure (_type_, optional): _description_. Defaults to None.
+            default_color (str, optional): Default color of nodes. Defaults to "black".
+
+        Returns:
+            figure, ax: The matplotlib figure and ax object.
+        """
+        if selected_nodes is None:
+            selected_nodes = []
+        if selected_edges is None:
+            selected_edges = []
         if ax is None:
             figure, ax = plt.subplots()
         else:
             ax.clear()
-        if not isinstance(selected_cells, set):
-            selected_cells = set(selected_cells)
-        hier_unselected = {
-            k: v for k, v in hier.items() if k not in selected_cells
-        }
-        hier_selected = {k: v for k, v in hier.items() if k in selected_cells}
-        unselected = np.array(tuple(hier_unselected.values()))
-        x = []
-        y = []
-        if hier_unselected:
-            for pred, succs in lnks_tms["links"].items():
-                if pred not in selected_cells:
-                    for succ in succs:
-                        x.extend((hier[succ][0], hier[pred][0], None))
-                        y.extend((hier[succ][1], hier[pred][1], None))
-            ax.plot(x, y, c="black", linewidth=0.3, zorder=0.5, **kwargs)
-            ax.scatter(
-                *unselected.T,
-                s=0.1,
-                c="black",
-                zorder=1,
-                **kwargs,
-            )
-        if selected_cells:
-            selected = np.array(tuple(hier_selected.values()))
-            x = []
-            y = []
-            for pred, succs in lnks_tms["links"].items():
-                if pred in selected_cells:
-                    for succ in succs:
-                        x.extend((hier[succ][0], hier[pred][0], None))
-                        y.extend((hier[succ][1], hier[pred][1], None))
-            ax.plot(x, y, c=color, linewidth=0.3, zorder=0.4, **kwargs)
-            ax.scatter(
-                selected.T[0],
-                selected.T[1],
-                s=0.1,
-                c=color,
-                zorder=0.9,
-                **kwargs,
-            )
+        if not isinstance(selected_nodes, set):
+            selected_nodes = set(selected_nodes)
+        if not isinstance(selected_edges, set):
+            selected_edges = set(selected_edges)
+        self.__plot_nodes(
+            hier,
+            selected_nodes,
+            color_of_nodes,
+            size=size,
+            ax=ax,
+            default_color=default_color,
+            **kwargs,
+        )
+        if not color_of_edges:
+            color_of_edges = color_of_nodes
+        self.__plot_edges(
+            hier,
+            lnks_tms,
+            selected_edges,
+            color_of_edges,
+            ax,
+            default_color=default_color,
+            **kwargs,
+        )
         ax.get_yaxis().set_visible(False)
         ax.get_xaxis().set_visible(False)
-        return figure, ax
+        return ax.get_figure(), ax
 
     def to_simple_graph(self, node=None, start_time: int = None):
         """Generates a dictionary of graphs where the keys are the index of the graph and
@@ -1862,8 +1837,8 @@ class lineageTree(lineageTreeLoaders):
         figsize=(10, 15),
         dpi=100,
         fontsize=15,
-        figure=None,
         axes=None,
+        vert_gap=1,
         **kwargs,
     ):
         """Plots all lineages.
@@ -1893,14 +1868,24 @@ class lineageTree(lineageTreeLoaders):
             )
         pos = {
             i: hierarchical_pos(
-                g, g["root"], ycenter=-int(self.time[g["root"]])
+                g,
+                g["root"],
+                ycenter=-int(self.time[g["root"]]),
+                vert_gap=vert_gap,
             )
             for i, g in graphs.items()
         }
-        ncols = int(len(graphs) // nrows) + (+np.sign(len(graphs) % nrows))
-        figure, axes = plt.subplots(
-            figsize=figsize, nrows=nrows, ncols=ncols, dpi=dpi, sharey=True
-        )
+        if axes is None:
+            ncols = int(len(graphs) // nrows) + (+np.sign(len(graphs) % nrows))
+            figure, axes = plt.subplots(
+                figsize=figsize, nrows=nrows, ncols=ncols, dpi=dpi, sharey=True
+            )
+        else:
+            figure, axes = axes.flatten()[0].get_figure(), axes
+            if len(axes.flatten()) < len(graphs):
+                raise Exception(
+                    f"Not enough axes, they should be at least {len(graphs)}."
+                )
         flat_axes = axes.flatten()
         ax2root = {}
         min_width, min_height = float("inf"), float("inf")
@@ -1938,9 +1923,17 @@ class lineageTree(lineageTreeLoaders):
                 },
             )
         [figure.delaxes(ax) for ax in axes.flatten() if not ax.has_data()]
-        return figure, axes, ax2root
+        return axes.flatten()[0].get_figure(), axes, ax2root
 
-    def plot_node(self, node, figsize=(4, 7), dpi=150, vert_gap=2, **kwargs):
+    def plot_node(
+        self,
+        node,
+        figsize=(4, 7),
+        dpi=150,
+        vert_gap=2,
+        ax=None,
+        **kwargs,
+    ):
         """Plots the subtree spawn by a node.
 
         Args:
@@ -1951,7 +1944,10 @@ class lineageTree(lineageTreeLoaders):
         if len(graph) > 1:
             raise Warning("Please enter only one node")
         graph = graph[0]
-        figure, ax = plt.subplots(nrows=1, ncols=1, figsize=figsize, dpi=dpi)
+        if not ax:
+            figure, ax = plt.subplots(
+                nrows=1, ncols=1, figsize=figsize, dpi=dpi
+            )
         self.draw_tree_graph(
             hier=hierarchical_pos(
                 graph,
@@ -1962,7 +1958,7 @@ class lineageTree(lineageTreeLoaders):
             lnks_tms=graph,
             ax=ax,
         )
-        return figure, ax
+        return ax.get_figure(), ax
 
     # def DTW(self, t1, t2, max_w=None, start_delay=None, end_delay=None,
     #         metric='euclidian', **kwargs):
@@ -2615,6 +2611,7 @@ class lineageTree(lineageTreeLoaders):
             warnings.warn(
                 "Error: not possible to show alignment in PCA projection !",
                 UserWarning,
+                stacklevel=2,
             )
 
         return distance, fig
@@ -2636,6 +2633,7 @@ class lineageTree(lineageTreeLoaders):
         reorder: bool = False,
         xml_attributes: tuple = None,
         name: str = None,
+        time_resolution: Union[int, None] = None,
     ):
         """
         TODO: complete the doc
@@ -2646,13 +2644,22 @@ class lineageTree(lineageTreeLoaders):
             file_format (str): either - path format to TGMM xmls
                                       - path to the MaMuT xml
                                       - path to the binary file
-            tb (int): first time point (necessary for TGMM xmls only)
-            te (int): last time point (necessary for TGMM xmls only)
-            z_mult (float): z aspect ratio if necessary (usually only for TGMM xmls)
-            file_type (str): type of input file. Accepts:
+            tb (int, optional):first time point (necessary for TGMM xmls only)
+            te (int, optional): last time point (necessary for TGMM xmls only)
+            z_mult (float, optional):z aspect ratio if necessary (usually only for TGMM xmls)
+            file_type (str, optional):type of input file. Accepts:
                 'TGMM, 'ASTEC', MaMuT', 'TrackMate', 'csv', 'celegans', 'binary'
                 default is 'binary'
+            delim (str, optional): _description_. Defaults to ",".
+            eigen (bool, optional): _description_. Defaults to False.
+            shape (tuple, optional): _description_. Defaults to None.
+            raw_size (tuple, optional): _description_. Defaults to None.
+            reorder (bool, optional): _description_. Defaults to False.
+            xml_attributes (tuple, optional): _description_. Defaults to None.
+            name (str, optional): The name of the dataset. Defaults to None.
+            time_resolution (Union[int, None], optional): Time resolution in mins (If time resolution is smaller than one minute input the time in ms). Defaults to None.
         """
+
         self.name = name
         self.time_nodes = {}
         self.time_edges = {}
@@ -2664,6 +2671,8 @@ class lineageTree(lineageTreeLoaders):
         self.pos = {}
         self.time_id = {}
         self.time = {}
+        if time_resolution is not None:
+            self._time_resolution = time_resolution
         self.kdtrees = {}
         self.spatial_density = {}
         if file_type and file_format:
@@ -2698,6 +2707,8 @@ class lineageTree(lineageTreeLoaders):
                 self.read_from_ASTEC(file_format, eigen)
             elif file_type == "csv":
                 self.read_from_csv(file_format, z_mult, link=1, delim=delim)
+            elif file_type == "bao":
+                self.read_C_elegans_bao(file_format)
             elif file_format and file_format.endswith(".lT"):
                 with open(file_format, "br") as f:
                     tmp = pkl.load(f)
@@ -2708,15 +2719,5 @@ class lineageTree(lineageTreeLoaders):
             if self.name is None:
                 try:
                     self.name = Path(file_format).stem
-                except:
+                except TypeError:
                     self.name = Path(file_format[0]).stem
-        if [] in self.successor.values():
-            successors = list(self.successor.keys())
-            for succ in successors:
-                if self[succ] == []:
-                    self.successor.pop(succ)
-        if [] in self.predecessor.values():
-            predecessors = list(self.predecessor.keys())
-            for succ in predecessors:
-                if self[succ] == []:
-                    self.predecessor.pop(succ)