LineageTree 1.4.4__py3-none-any.whl → 1.5.1__py3-none-any.whl

LineageTree/lineageTree.py

@@ -2,35 +2,244 @@
 # This file is subject to the terms and conditions defined in
 # file 'LICENCE', which is part of this source code package.
 # Author: Leo Guignard (leo.guignard...@AT@...gmail.com)
-
 import csv
 import os
 import pickle as pkl
 import struct
+import warnings
 import xml.etree.ElementTree as ET
+from collections.abc import Iterable
 from functools import partial
 from itertools import combinations
 from numbers import Number
-from typing import TextIO
-
+from pathlib import Path
+from typing import TextIO, Union
+
+from .tree_styles import tree_style
+
+try:
+    from edist import uted
+except ImportError:
+    warnings.warn(
+        "No edist installed therefore you will not be able to compute the tree edit distance."
+    )
+import matplotlib.pyplot as plt
+import networkx as nx
 import numpy as np
-from scipy.spatial import Delaunay
+from scipy.interpolate import InterpolatedUnivariateSpline
+from scipy.spatial import Delaunay, distance
 from scipy.spatial import cKDTree as KDTree
 
+from .utils import hierarchy_pos, postions_of_nx
+
 
 class lineageTree:
+    def __eq__(self, other):
+        if isinstance(other, lineageTree):
+            return other.successor == self.successor
+        return False
+
     def get_next_id(self):
         """Computes the next authorized id.
 
         Returns:
             int: next authorized id
         """
+        if self.max_id == -1 and self.nodes:
+            self.max_id = max(self.nodes)
         if self.next_id == []:
             self.max_id += 1
             return self.max_id
         else:
             return self.next_id.pop()
 
+    def complete_lineage(self, nodes: Union[int, set] = None):
+        """Makes all leaf branches longer so that they reach the last timepoint( self.t_e), useful
+        for tree edit distance algorithms.
+
+        Args:
+            nodes (int,set), optional): Which trees should be "completed", if None it will complete the whole dataset. Defaults to None.
+        """
+        if nodes is None:
+            nodes = set(self.roots)
+        elif isinstance(nodes, int):
+            nodes = {nodes}
+        for node in nodes:
+            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)
+
+    ###TODO pos can be callable and stay motionless (copy the position of the succ node, use something like optical flow)
+    def add_branch(
+        self,
+        pred: int,
+        length: int,
+        move_timepoints: bool = True,
+        pos: Union[callable, None] = None,
+        reverse: bool = False,
+    ):
+        """Adds a branch of specific length to a node either as a successor or as a predecessor.
+        If it is placed on top of a tree all the nodes will move timepoints #length down.
+
+        Args:
+            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
+        Returns:
+            (int): Id of the first node of the sublineage.
+        """
+        if length == 0:
+            return pred
+        if self.predecessor.get(pred) and not reverse:
+            raise Warning("Cannot add 2 predecessors to a node")
+        time = self.time[pred]
+        original = pred
+        if not reverse:
+            if move_timepoints:
+                nodes_to_move = set(self.get_sub_tree(pred))
+                new_times = {
+                    node: self.time[node] + length for node in nodes_to_move
+                }
+                for node in nodes_to_move:
+                    old_time = self.time[node]
+                    self.time_nodes[old_time].remove(node)
+                    self.time_nodes.setdefault(old_time + length, set()).add(
+                        node
+                    )
+                self.time.update(new_times)
+                for t in range(length - 1, -1, -1):
+                    _next = self.add_node(
+                        time + t,
+                        succ=pred,
+                        pos=self.pos[original],
+                        reverse=True,
+                    )
+                    pred = _next
+            else:
+                for t in range(length):
+                    _next = self.add_node(
+                        time - t,
+                        succ=pred,
+                        pos=self.pos[original],
+                        reverse=True,
+                    )
+                    pred = _next
+        else:
+            for t in range(length):
+                _next = self.add_node(
+                    time + t, succ=pred, pos=self.pos[original], reverse=False
+                )
+                pred = _next
+            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
+
+    def cut_tree(self, root):
+        """It transforms a lineage that has at least 2 divisions into 2 independent lineages,
+        that spawn from the time point of the first node. (splits a tree into 2)
+
+        Args:
+            root (int): The id of the node, which will be cut.
+
+        Returns:
+            int: The id of the new tree
+        """
+        cycle = self.get_successors(root)
+        last_cell = cycle[-1]
+        if last_cell in self.successor:
+            new_lT = self.successor[last_cell].pop()
+            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
+            )
+            self.roots.add(new_tr)
+            self.labels[new_tr] = f"R-Split {label_of_root}"
+            return new_tr
+        else:
+            raise Warning("No division of the branch")
+
+    def fuse_lineage_tree(
+        self,
+        l1_root: int,
+        l2_root: int,
+        length_l1: int = 0,
+        length_l2: int = 0,
+        length: int = 1,
+    ):
+        """Fuses 2 lineages from the lineagetree object. The 2 lineages that are to be fused can have a longer
+        first node and the node of the resulting lineage can also be longer.
+
+        Args:
+            l1_root (int): Id of the first root
+            l2_root (int): Id of the second root
+            length_l1 (int, optional): The length of the branch that will be added on top of the first lineage. Defaults to 0, which means only one node will be added.
+            length_l2 (int, optional): The length of the branch that will be added on top of the second lineage. Defaults to 0, which means only one node will be added.
+            length (int, optional): The length of the branch that will be added on top of the resulting lineage. Defaults to 1.
+
+        Returns:
+            int: The id of the root of the new lineage.
+        """
+        if self.predecessor.get(l1_root) or self.predecessor.get(l2_root):
+            raise ValueError("Please select 2 roots.")
+        if self.time[l1_root] != self.time[l2_root]:
+            warnings.warn(
+                "Using lineagetrees that do not exist in the same timepoint. The operation will continue"
+            )
+        new_root1 = self.add_branch(l1_root, length_l1)
+        new_root2 = self.add_branch(l2_root, length_l2)
+        next_root1 = self[new_root1][0]
+        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)
+        self.labels[new_branch] = f"Fusion of {new_root1} and {new_root2}"
+        return new_branch
+
+    def copy_lineage(self, root):
+        """
+        Copies the structure of a tree and makes a new with new nodes.
+        Warning does not take into account the predecessor of the root node.
+
+        Args:
+            root (int): The root of the tree to be copied
+
+        Returns:
+            int: The root of the new tree.
+        """
+        new_nodes = {
+            old_node: self.get_next_id()
+            for old_node in self.get_sub_tree(root)
+        }
+        self.nodes.update(new_nodes.values())
+        for old_node, new_node in new_nodes.items():
+            self.time[new_node] = self.time[old_node]
+            succ = self.successor.get(old_node)
+            if succ:
+                self.successor[new_node] = [new_nodes[n] for n in succ]
+            pred = self.predecessor.get(old_node)
+            if pred:
+                self.predecessor[new_node] = [new_nodes[n] for n in pred]
+            self.pos[new_node] = self.pos[old_node] + 0.5
+            self.time_nodes[self.time[old_node]].add(new_nodes[old_node])
+        new_root = new_nodes[root]
+        self.labels[new_root] = f"Copy of {root}"
+        if self.time[new_root] == 0:
+            self.roots.add(new_root)
+        return new_root
+
     def add_node(
         self,
         t: int = None,
@@ -55,109 +264,126 @@ class lineageTree:
             int: id of the new node.
         """
         C_next = self.get_next_id() if nid is None else nid
-        self.time_nodes.setdefault(t, []).append(C_next)
+        self.time_nodes.setdefault(t, set()).add(C_next)
         if succ is not None and not reverse:
             self.successor.setdefault(succ, []).append(C_next)
             self.predecessor.setdefault(C_next, []).append(succ)
-            self.edges.add((succ, C_next))
         elif succ is not None:
             self.predecessor.setdefault(succ, []).append(C_next)
             self.successor.setdefault(C_next, []).append(succ)
-            self.edges.add((C_next, succ))
         self.nodes.add(C_next)
         self.pos[C_next] = pos
-        self.progeny[C_next] = 0
         self.time[C_next] = t
         return C_next
 
-    def remove_track(self, track: list):
-        self.nodes.difference_update(track)
-        times = {self.time[n] for n in track}
-        for t in times:
-            self.time_nodes[t] = list(
-                set(self.time_nodes[t]).difference(track)
-            )
-        for i, c in enumerate(track):
-            self.pos.pop(c)
-            if i != 0:
-                self.predecessor.pop(c)
-            if i < len(track) - 1:
-                self.successor.pop(c)
-            self.time.pop(c)
-
-    def remove_node(self, c: int) -> tuple:
-        """Removes a node and update the lineageTree accordingly
+    def remove_nodes(self, group: Union[int, set, list]):
+        """Removes a group of nodes from the LineageTree
 
         Args:
-            c (int): id of the node to remove
-        """
-        self.nodes.remove(c)
-        self.time_nodes[self.time[c]].remove(c)
-        # self.time_nodes.pop(c, 0)
-        pos = self.pos.pop(c, 0)
-        e_to_remove = [e for e in self.edges if c in e]
-        for e in e_to_remove:
-            self.edges.remove(e)
-        if c in self.roots:
-            self.roots.remove(c)
-        succ = self.successor.pop(c, [])
-        s_to_remove = [s for s, ci in self.successor.items() if c in ci]
-        for s in s_to_remove:
-            self.successor[s].remove(c)
-
-        pred = self.predecessor.pop(c, [])
-        p_to_remove = [s for s, ci in self.predecessor.items() if ci == c]
-        for s in p_to_remove:
-            self.predecessor[s].remove(c)
-
-        self.time.pop(c, 0)
-        self.spatial_density.pop(c, 0)
-
-        self.next_id.append(c)
-        return e_to_remove, succ, s_to_remove, pred, p_to_remove, pos
-
-    def fuse_nodes(self, c1: int, c2: int):
-        """Fuses together two nodes that belong to the same time point
-        and update the lineageTree accordingly.
+            group (set|list|int): One or more nodes that are to be removed.
+        """
+        if isinstance(group, int):
+            group = {group}
+        if isinstance(group, list):
+            group = set(group)
+        group = group.intersection(self.nodes)
+        self.nodes.difference_update(group)
+        times = {self.time.pop(n) for n in group}
+        for t in times:
+            self.time_nodes[t] = set(self.time_nodes[t]).difference(group)
+        for node in group:
+            self.pos.pop(node)
+            if self.predecessor.get(node):
+                pred = self.predecessor[node][0]
+                siblings = self.successor.pop(pred, [])
+                if len(siblings) == 2:
+                    siblings.remove(node)
+                    self.successor[pred] = siblings
+                self.predecessor.pop(node, [])
+            for succ in self.successor.get(node, []):
+                self.predecessor.pop(succ, [])
+            self.successor.pop(node, [])
+            self.labels.pop(node, 0)
+            if node in self.roots:
+                self.roots.remove(node)
+
+    def modify_branch(self, node, new_length):
+        """Changes the length of a branch, so it adds or removes nodes
+        to make the correct length of the cycle.
 
         Args:
-            c1 (int): id of the first node to fuse
-            c2 (int): id of the second node to fuse
+            node (int): Any node of the branch to be modified/
+            new_length (int): The new length of the tree.
         """
-        (
-            e_to_remove,
-            succ,
-            s_to_remove,
-            pred,
-            p_to_remove,
-            c2_pos,
-        ) = self.remove_node(c2)
-        for e in e_to_remove:
-            new_e = [c1] + [other_c for other_c in e if e != c2]
-            self.edges.add(new_e)
-
-        self.successor.setdefault(c1, []).extend(succ)
-        self.predecessor.setdefault(c1, []).extend(pred)
-
-        for s in s_to_remove:
-            self.successor[s].append(c1)
-
-        for p in p_to_remove:
-            self.predecessor[p].append(c1)
-
-        self.pos[c1] = np.mean([self.pos[c1], c2_pos], axis=0)
-        self.progeny[c1] += 1
+        if new_length <= 1:
+            warnings.warn("New length should be more than 1")
+            return None
+        cycle = self.get_cycle(node)
+        length = len(cycle)
+        successors = self.successor.get(cycle[-1])
+        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
+            )
+            if pred:
+                self.successor[pred[0]].remove(node)
+                self.successor[pred[0]].append(new_node)
+        elif self.leaves.intersection(cycle) and new_length < length:
+            self.remove_nodes(cycle[new_length:])
+        elif new_length < length:
+            to_remove = length - new_length
+            last_cell = cycle[new_length - 1]
+            subtree = self.get_sub_tree(cycle[-1])[1:]
+            self.remove_nodes(cycle[new_length:])
+            self.successor[last_cell] = successors
+            if successors:
+                for succ in successors:
+                    self.predecessor[succ] = [last_cell]
+            for node in subtree:
+                if node not in cycle[new_length - 1 :]:
+                    old_time = self.time[node]
+                    self.time[node] = old_time - to_remove
+                    self.time_nodes[old_time].remove(node)
+                    self.time_nodes.setdefault(
+                        old_time - to_remove, set()
+                    ).add(node)
+        elif length < new_length:
+            to_add = new_length - length
+            last_cell = cycle[-1]
+            self.successor.pop(cycle[-2])
+            self.predecessor.pop(last_cell)
+            succ = self.add_branch(
+                last_cell, length=to_add, move_timepoints=True, reverse=False
+            )
+            self.predecessor[succ] = [cycle[-2]]
+            self.successor[cycle[-2]] = [succ]
+            self.time[last_cell] = (
+                self.time[self.predecessor[last_cell][0]] + 1
+            )
+        else:
+            return None
 
     @property
     def roots(self):
         if not hasattr(self, "_roots"):
-            self._roots = set(self.successor).difference(self.predecessor)
+            self._roots = set(self.nodes).difference(self.predecessor)
         return self._roots
 
+    @property
+    def edges(self):
+        return {(k, vi) for k, v in self.successor.items() for vi in v}
+
     @property
     def leaves(self):
         return set(self.predecessor).difference(self.successor)
 
+    @property
+    def labels(self):
+        if not hasattr(self, "_labels"):
+            self._labels = {i: "Unlabeled" for i in self.roots}
+        return self._labels
+
     def _write_header_am(self, f: TextIO, nb_points: int, length: int):
         """Header for Amira .am files"""
         f.write("# AmiraMesh 3D ASCII 2.0\n")
@@ -470,7 +696,7 @@ class lineageTree:
                         stroke=svgwrite.rgb(0, 0, 0),
                     )
                 )
-                for si in self.successor.get(c_cycle[-1], []):
+                for si in self[c_cycle[-1]]:
                     x3, y3 = positions[si]
                     dwg.add(
                         dwg.line(
@@ -483,7 +709,7 @@ class lineageTree:
         else:
             for c in treated_cells:
                 x1, y1 = positions[c]
-                for si in self.successor.get(c, []):
+                for si in self[c]:
                     x2, y2 = positions[si]
                     if draw_edges:
                         dwg.add(
@@ -535,7 +761,7 @@ class lineageTree:
         start_time = times_to_consider[0]
         for t in times_to_consider:
             for id_mother in self.time_nodes[t]:
-                ids_daughters = self.successor.get(id_mother, [])
+                ids_daughters = self[id_mother]
                 new_ids_daughters = ids_daughters.copy()
                 for _ in range(sampling - 1):
                     tmp = []
@@ -659,12 +885,12 @@ class lineageTree:
                         edges_to_use += list(s_edges)
             else:
                 edges_to_use = []
+                nodes_to_use = set(nodes_to_use)
                 if temporal:
-                    edges_to_use += [
-                        e
-                        for e in self.edges
-                        if e[0] in nodes_to_use and e[1] in nodes_to_use
-                    ]
+                    for n in nodes_to_use:
+                        for d in self.successor.get(n, []):
+                            if d in nodes_to_use:
+                                edges_to_use.append((n, d))
                 if spatial:
                     edges_to_use += [
                         e for e in s_edges if t_min < self.time[e[0]] < t_max
@@ -787,7 +1013,6 @@ class lineageTree:
         self.time_edges = {}
         unique_id = 0
         self.nodes = set()
-        self.edges = set()
         self.successor = {}
         self.predecessor = {}
         self.pos = {}
@@ -826,7 +1051,6 @@ class lineageTree:
                 M = corres[pred]
                 self.predecessor[C] = [M]
                 self.successor.setdefault(M, []).append(C)
-                self.edges.add((M, C))
                 self.time_edges.setdefault(t, set()).add((M, C))
                 self.lin.setdefault(lin_id, []).append(C)
                 self.C_lin[C] = lin_id
@@ -962,8 +1186,9 @@ class lineageTree:
             if "cell_fate" in tmp_data:
                 self.fates[unique_id] = tmp_data["cell_fate"].get(n, "")
             if "cell_barycenter" in tmp_data:
-                self.pos[unique_id] = tmp_data["cell_barycenter"].get(n, np.zeros(3))
-
+                self.pos[unique_id] = tmp_data["cell_barycenter"].get(
+                    n, np.zeros(3)
+                )
 
             unique_id += 1
         if do_surf:
@@ -982,9 +1207,7 @@ class lineageTree:
                 self.successor[new_id] = [
                     self.pkl2lT[ni] for ni in lt[n] if ni in self.pkl2lT
                 ]
-                self.edges.update(
-                    [(new_id, ni) for ni in self.successor[new_id]]
-                )
+
                 for ni in self.successor[new_id]:
                     self.time_edges.setdefault(t - 1, set()).add((new_id, ni))
 
@@ -993,31 +1216,43 @@ class lineageTree:
         self.max_id = unique_id
 
         # do this in the end of the process, skip lineage tree and whatever is stored already
-        pre_treated_prop = [
+        discard = {
             "cell_volume",
             "cell_fate",
             "cell_barycenter",
             "cell_contact_surface",
             "cell_lineage",
-        ]
+            "all_cells",
+            "cell_history",
+            "problematic_cells",
+            "cell_labels_in_time",
+        }
+        self.specific_properties = []
         for prop_name, prop_values in tmp_data.items():
-            if not (prop_name in pre_treated_prop or hasattr(self, prop_name)):
+            if not (prop_name in discard or hasattr(self, prop_name)):
                 if isinstance(prop_values, dict):
                     dictionary = {
-                        self.pkl2lT.get(k, -1): v for k, v in prop_values.items()
+                        self.pkl2lT.get(k, -1): v
+                        for k, v in prop_values.items()
                     }
                     # is it a regular dictionary or a dictionary with dictionaries inside?
                     for key, value in dictionary.items():
                         if isinstance(value, dict):
                             # rename all ids from old to new
                             dictionary[key] = {
-                                self.pkl2lT.get(k, -1): v for k, v in value.items()
+                                self.pkl2lT.get(k, -1): v
+                                for k, v in value.items()
                             }
                     self.__dict__[prop_name] = dictionary
+                    self.specific_properties.append(prop_name)
                 # is any of this necessary? Or does it mean it anyways does not contain
                 # information about the id and a simple else: is enough?
-                elif isinstance(prop_values, (list, set, np.ndarray)):
+                elif (
+                    isinstance(prop_values, (list, set, np.ndarray))
+                    and prop_name not in []
+                ):
                     self.__dict__[prop_name] = prop_values
+                    self.specific_properties.append(prop_name)
 
             # what else could it be?
 
@@ -1129,7 +1364,6 @@ class lineageTree:
                         p = None
                     self.predecessor.setdefault(c, []).append(p)
                     self.successor.setdefault(p, []).append(c)
-                    self.edges.add((p, c))
                     self.time_edges.setdefault(t - 1, set()).add((p, c))
             self.max_id = unique_id
 
@@ -1217,7 +1451,6 @@ class lineageTree:
                         p = None
                     self.predecessor.setdefault(c, []).append(p)
                     self.successor.setdefault(p, []).append(c)
-                    self.edges.add((p, c))
                     self.time_edges.setdefault(t - 1, set()).add((p, c))
             self.max_id = unique_id
 
@@ -1241,7 +1474,6 @@ class lineageTree:
         self.time_edges = {}
         unique_id = 0
         self.nodes = set()
-        self.edges = set()
         self.successor = {}
         self.predecessor = {}
         self.pos = {}
@@ -1301,7 +1533,6 @@ class lineageTree:
                             M = self.time_id[(t - 1, M_id)]
                             self.successor.setdefault(M, []).append(C)
                             self.predecessor.setdefault(C, []).append(M)
-                            self.edges.add((M, C))
                             self.time_edges[t].add((M, C))
                         else:
                             if M_id != -1:
@@ -1338,7 +1569,6 @@ class lineageTree:
 
         mr = MastodonReader(path)
         spots, links = mr.read_tables()
-        mr.read_tags(spots, links)
 
         self.node_name = {}
 
@@ -1358,7 +1588,6 @@ class lineageTree:
             target = e.target_idx
             self.predecessor.setdefault(target, []).append(source)
             self.successor.setdefault(source, []).append(target)
-            self.edges.add((source, target))
             self.time_edges.setdefault(self.time[source], set()).add(
                 (source, target)
             )
@@ -1393,14 +1622,13 @@ class lineageTree:
             self.nodes.add(unique_id)
             self.time[unique_id] = t
             self.node_name[unique_id] = spot[1]
-            self.pos[unique_id] = np.array([x, y, z])
+            self.pos[unique_id] = np.array([x, y, z], dtype=float)
 
         for link in links:
             source = int(float(link[4]))
             target = int(float(link[5]))
             self.predecessor.setdefault(target, []).append(source)
             self.successor.setdefault(source, []).append(target)
-            self.edges.add((source, target))
             self.time_edges.setdefault(self.time[source], set()).add(
                 (source, target)
             )
@@ -1455,23 +1683,24 @@ class lineageTree:
                     if attr in self.xml_attributes:
                         self.__dict__[attr][cell_id] = eval(cell.attrib[attr])
 
-        self.edges = set()
         tracks = {}
         self.successor = {}
         self.predecessor = {}
         self.track_name = {}
         for track in AllTracks:
             if "TRACK_DURATION" in track.attrib:
-                t_id, _ = int(track.attrib["TRACK_ID"]), float(
-                    track.attrib["TRACK_DURATION"]
+                t_id, _ = (
+                    int(track.attrib["TRACK_ID"]),
+                    float(track.attrib["TRACK_DURATION"]),
                 )
             else:
                 t_id = int(track.attrib["TRACK_ID"])
             t_name = track.attrib["name"]
             tracks[t_id] = []
             for edge in track:
-                s, t = int(edge.attrib["SPOT_SOURCE_ID"]), int(
-                    edge.attrib["SPOT_TARGET_ID"]
+                s, t = (
+                    int(edge.attrib["SPOT_SOURCE_ID"]),
+                    int(edge.attrib["SPOT_TARGET_ID"]),
                 )
                 if s in self.nodes and t in self.nodes:
                     if self.time[s] > self.time[t]:
@@ -1481,7 +1710,6 @@ class lineageTree:
                     self.track_name[s] = t_name
                     self.track_name[t] = t_name
                     tracks[t_id].append((s, t))
-                    self.edges.add((s, t))
         self.t_b = min(self.time_nodes.keys())
         self.t_e = max(self.time_nodes.keys())
 
@@ -1519,7 +1747,7 @@ class lineageTree:
                 curr_c = to_treat.pop()
                 number_sequence.append(curr_c)
                 pos_sequence += list(self.pos[curr_c])
-                if self.successor.get(curr_c, []) == []:
+                if self[curr_c] == []:
                     number_sequence.append(-1)
                 elif len(self.successor[curr_c]) == 1:
                     to_treat += self.successor[curr_c]
@@ -1681,7 +1909,6 @@ class lineageTree:
         self.time_edges = time_edges
         self.pos = pos
         self.nodes = set(nodes)
-        self.edges = set(edges)
         self.t_b = min(time_nodes.keys())
         self.t_e = max(time_nodes.keys())
         self.is_root = is_root
@@ -1701,7 +1928,7 @@ class lineageTree:
             f.close()
 
     @classmethod
-    def load(clf, fname: str):
+    def load(clf, fname: str, rm_empty_lists=True):
         """
         Loading a lineage tree from a ".lT" file.
 
@@ -1714,6 +1941,18 @@ class lineageTree:
         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")
         return lT
 
     def get_idx3d(self, t: int) -> tuple:
@@ -1795,7 +2034,9 @@ class lineageTree:
 
         return self.Gabriel_graph[t]
 
-    def get_predecessors(self, x: int, depth: int = None) -> list:
+    def get_predecessors(
+        self, x: int, depth: int = None, start_time: int = None, end_time=None
+    ) -> list:
         """Computes the predecessors of the node `x` up to
         `depth` predecessors or the begining of the life of `x`.
         The ordered list of ids is returned.
@@ -1806,20 +2047,34 @@ class lineageTree:
         Returns:
             [int, ]: list of ids, the last id is `x`
         """
-        cycle = [x]
+        if not start_time:
+            start_time = self.t_b
+        if not end_time:
+            end_time = self.t_e
+        unconstrained_cycle = [x]
+        cycle = [x] if start_time <= self.time[x] <= end_time else []
         acc = 0
         while (
-            len(
-                self.successor.get(self.predecessor.get(cycle[0], [-1])[0], [])
-            )
+            len(self[self.predecessor.get(unconstrained_cycle[0], [-1])[0]])
             == 1
             and acc != depth
+            and start_time
+            <= self.time.get(
+                self.predecessor.get(unconstrained_cycle[0], [-1])[0], -1
+            )
         ):
-            cycle.insert(0, self.predecessor[cycle[0]][0])
+            unconstrained_cycle.insert(
+                0, self.predecessor[unconstrained_cycle[0]][0]
+            )
             acc += 1
+            if start_time <= self.time[unconstrained_cycle[0]] <= end_time:
+                cycle.insert(0, unconstrained_cycle[0])
+
         return cycle
 
-    def get_successors(self, x: int, depth: int = None) -> list:
+    def get_successors(
+        self, x: int, depth: int = None, end_time: int = None
+    ) -> list:
         """Computes the successors of the node `x` up to
         `depth` successors or the end of the life of `x`.
         The ordered list of ids is returned.
@@ -1830,11 +2085,18 @@ class lineageTree:
         Returns:
             [int, ]: list of ids, the first id is `x`
         """
+        if end_time is None:
+            end_time = self.t_e
         cycle = [x]
         acc = 0
-        while len(self.successor.get(cycle[-1], [])) == 1 and acc != depth:
+        while (
+            len(self[cycle[-1]]) == 1
+            and acc != depth
+            and self.time[cycle[-1]] < end_time
+        ):
             cycle += self.successor[cycle[-1]]
             acc += 1
+
         return cycle
 
     def get_cycle(
@@ -1843,12 +2105,14 @@ class lineageTree:
         depth: int = None,
         depth_pred: int = None,
         depth_succ: int = None,
+        end_time: int = None,
     ) -> list:
         """Computes the predecessors and successors of the node `x` up to
         `depth_pred` predecessors plus `depth_succ` successors.
         If the value `depth` is provided and not None,
         `depth_pred` and `depth_succ` are overwriten by `depth`.
         The ordered list of ids is returned.
+        If all `depth` are None, the full cycle is returned.
 
         Args:
             x (int): id of the node to compute
@@ -1858,11 +2122,13 @@ class lineageTree:
         Returns:
             [int, ]: list of ids
         """
+        if end_time is None:
+            end_time = self.t_e
         if depth is not None:
             depth_pred = depth_succ = depth
-        return self.get_predecessors(x, depth_pred)[:-1] + self.get_successors(
-            x, depth_succ
-        )
+        return self.get_predecessors(x, depth_pred, end_time=end_time)[
+            :-1
+        ] + self.get_successors(x, depth_succ, end_time=end_time)
 
     @property
     def all_tracks(self):
@@ -1870,6 +2136,29 @@ class lineageTree:
             self._all_tracks = self.get_all_tracks()
         return self._all_tracks
 
+    def get_all_branches_of_node(
+        self, node: int, end_time: int = None
+    ) -> list:
+        """Computes all the tracks of the subtree spawn by a given node.
+        Similar to get_all_tracks().
+
+        Args:
+            node (int, optional): The node that we want to get its branches.
+
+        Returns:
+            ([[int, ...], ...]): list of lists containing track cell ids
+        """
+        if not end_time:
+            end_time = self.t_e
+        branches = [self.get_successors(node)]
+        to_do = self[branches[0][-1]].copy()
+        while to_do:
+            current = to_do.pop()
+            track = self.get_cycle(current, end_time=end_time)
+            branches += [track]
+            to_do.extend(self[track[-1]])
+        return branches
+
     def get_all_tracks(self, force_recompute: bool = False) -> list:
         """Computes all the tracks of a given lineage tree,
         stores it in `self.all_tracks` and returns it.
@@ -1877,17 +2166,42 @@ class lineageTree:
         Returns:
             ([[int, ...], ...]): list of lists containing track cell ids
         """
-        if not hasattr(self, "_all_tracks"):
+        if not hasattr(self, "_all_tracks") or force_recompute:
             self._all_tracks = []
-            to_do = set(self.nodes)
+            to_do = list(self.roots)
             while len(to_do) != 0:
                 current = to_do.pop()
                 track = self.get_cycle(current)
                 self._all_tracks += [track]
-                to_do -= set(track)
+                to_do.extend(self[track[-1]])
         return self._all_tracks
 
-    def get_sub_tree(self, x: int, preorder: bool = False) -> list:
+    def get_tracks(self, roots: list = None) -> list:
+        """Computes the tracks given by the list of nodes `roots` and returns it.
+
+        Args:
+            roots (list): list of ids of the roots to be computed
+        Returns:
+            ([[int, ...], ...]): list of lists containing track cell ids
+        """
+        if roots is None:
+            return self.get_all_tracks(force_recompute=True)
+        else:
+            tracks = []
+            to_do = list(roots)
+            while len(to_do) != 0:
+                current = to_do.pop()
+                track = self.get_cycle(current)
+                tracks.append(track)
+                to_do.extend(self[track[-1]])
+            return tracks
+
+    def get_sub_tree(
+        self,
+        x: Union[int, Iterable],
+        end_time: Union[int, None] = None,
+        preorder: bool = False,
+    ) -> list:
         """Computes the list of cells from the subtree spawned by *x*
         The default output order is breadth first traversal.
         Unless preorder is `True` in that case the order is
@@ -1899,16 +2213,24 @@ class lineageTree:
         Returns:
             ([int, ...]): the ordered list of node ids
         """
-        to_do = [x]
+        if not end_time:
+            end_time = self.t_e
+        if not isinstance(x, Iterable):
+            to_do = [x]
+        elif isinstance(x, Iterable):
+            to_do = list(x)
         sub_tree = []
-        while len(to_do) > 0:
-            curr = to_do.pop(0)
+        while to_do:
+            curr = to_do.pop()
             succ = self.successor.get(curr, [])
+            if succ and end_time < self.time.get(curr, end_time):
+                succ = []
+                continue
             if preorder:
                 to_do = succ + to_do
             else:
                 to_do += succ
-            sub_tree += [curr]
+                sub_tree += [curr]
         return sub_tree
 
     def compute_spatial_density(
@@ -1979,6 +2301,70 @@ class lineageTree:
             )
         return self.th_edges
 
+    def main_axes(self, time: int = None):
+        """Finds the main axes for a timepoint.
+        If none will select the timepoint with the highest amound of cells.
+
+        Args:
+            time (int, optional): The timepoint to find the main axes.
+                                  If None will find the timepoint
+                                  with the largest number of cells.
+
+        Returns:
+            list: A list that contains the array of eigenvalues and eigenvectors.
+        """
+        if time is None:
+            time = np.argmax(
+                [len(self.time_nodes[t]) for t in range(int(self.t_e))]
+            )
+        pos = np.array([self.pos[node] for node in self.time_nodes[time]])
+        pos = pos - np.mean(pos, axis=0)
+        cov = np.cov(np.array(pos).T)
+        eig_val, eig_vec = np.linalg.eig(cov)
+        srt = np.argsort(eig_val)[::-1]
+        self.eig_val, self.eig_vec = eig_val[srt], eig_vec[:, srt]
+        return eig_val[srt], eig_vec[:, srt]
+
+    def scale_embryo(self, scale=1000):
+        """Scale the embryo using their eigenvalues.
+
+        Args:
+            scale (int, optional): The resulting scale you want to achieve. Defaults to 1000.
+
+        Returns:
+            float: The scale factor.
+        """
+        eig = self.main_axes()[0]
+        return scale / (np.sqrt(eig[0]))
+
+    @staticmethod
+    def __rodrigues_rotation_matrix(vector1, vector2=(0, 1, 0)):
+        """Calculates the rodrigues matrix of a dataset. It should use vectors from the find_main_axes(eigenvectors) function of LineagTree.
+        Uses the Rodrigues rotation formula.
+
+        Args:
+            vector1 (list|np.array): The vector that should be rotated to be aligned to the second vector
+            vector2 (list|np.array, optional): The second vector. Defaults to [1,0,0].
+
+        Returns:
+            np.array: The rotation matrix.
+        """
+        vector1 = vector1 / np.linalg.norm(vector1)
+        vector2 = vector2 / np.linalg.norm(vector2)
+        if vector1 @ vector2 == 1:
+            return np.eye(3)
+        angle = np.arccos(vector1 @ vector2)
+        axis = np.cross(vector1, vector2)
+        axis = axis / np.linalg.norm(axis)
+        K = np.array(
+            [
+                [0, -axis[2], axis[1]],
+                [axis[2], 0, -axis[0]],
+                [-axis[1], axis[0], 0],
+            ]
+        )
+        return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
+
     def get_ancestor_at_t(self, n: int, time: int = None):
         """
         Find the id of the ancestor of a give node `n`
@@ -2005,62 +2391,27 @@ class lineageTree:
             ancestor = self.predecessor.get(ancestor, [-1])[0]
         return ancestor
 
-    def get_simple_tree(self, r: int, time_resolution: int = 1) -> tuple:
-        """
-        Get a "simple" version of the tree spawned by the node `r`
-        This simple version is just one node per cell (as opposed to
-        one node per cell per time-point). The life time duration of
-        a cell `c` is stored in `self.cycle_time` and return by this
-        function
+    def get_labelled_ancestor(self, node: int):
+        """Finds the first labelled ancestor and returns its ID otherwise returns None
 
         Args:
-            r (int): root of the tree to spawn
-            time_resolution (float): the time between two consecutive time points
+            node (int): The id of the node
 
         Returns:
-            (dict) {m (int): [d1 (int), d2 (int)]}: a adjacency dictionnary
-                where the ids are the ids of the cells in the original tree
-                at their first time point (except for the cell `r` if it was
-                not the first time point).
-            (dict) {m (int): duration (float)}: life time duration of the cell `m`
-        """
-        if not hasattr(self, "cycle_time"):
-            self.cycle_time = {}
-        out_dict = {}
-        to_do = [r]
-        while to_do:
-            current = to_do.pop()
-            cycle = self.get_successors(current)
-            _next = self.successor.get(cycle[-1], [])
-            if _next:
-                out_dict[current] = _next
-            to_do.extend(_next)
-            self.cycle_time[current] = len(cycle) * time_resolution
-        return out_dict, self.cycle_time
-
-    @staticmethod
-    def __edist_format(adj_dict: dict):
-        inv_adj = {vi: k for k, v in adj_dict.items() for vi in v}
-        roots = set(adj_dict).difference(inv_adj)
-        nid2list = {}
-        list2nid = {}
-        nodes = []
-        adj_list = []
-        curr_id = 0
-        for r in roots:
-            to_do = [r]
-            while to_do:
-                curr = to_do.pop(0)
-                nid2list[curr] = curr_id
-                list2nid[curr_id] = curr
-                nodes.append(curr_id)
-                to_do = adj_dict.get(curr, []) + to_do
-                curr_id += 1
-            adj_list = [
-                [nid2list[d] for d in adj_dict.get(list2nid[_id], [])]
-                for _id in nodes
-            ]
-        return nodes, adj_list, list2nid
+            [None,int]: Returns the first ancestor found that has a label otherwise
+            None.
+        """
+        if node not in self.nodes:
+            return None
+        ancestor = node
+        while (
+            self.t_b <= self.time.get(ancestor, self.t_b - 1)
+            and ancestor != -1
+        ):
+            if ancestor in self.labels:
+                return ancestor
+            ancestor = self.predecessor.get(ancestor, [-1])[0]
+        return
 
     def unordered_tree_edit_distances_at_time_t(
         self,
@@ -2068,6 +2419,7 @@ class lineageTree:
         delta: callable = None,
         norm: callable = None,
         recompute: bool = False,
+        end_time: int = None,
     ) -> dict:
         """
         Compute all the pairwise unordered tree edit distances from Zhang 996 between the trees spawned at time `t`
@@ -2079,6 +2431,8 @@ class lineageTree:
                 of the tree spawned by `n1` and the number of nodes
                 of the tree spawned by `n2` as arguments.
             recompute (bool): if True, forces to recompute the distances (default: False)
+            end_time (int): The final time point the comparison algorithm will take into account. If None all nodes
+                            will be taken into account.
 
         Returns:
             (dict) a dictionary that maps a pair of cell ids at time `t` to their unordered tree edit distance
@@ -2092,14 +2446,20 @@ class lineageTree:
         for n1, n2 in combinations(roots, 2):
             key = tuple(sorted((n1, n2)))
             self.uted[t][key] = self.unordered_tree_edit_distance(
-                n1, n2, delta=delta, norm=norm
+                n1, n2, end_time=end_time
             )
         return self.uted[t]
 
     def unordered_tree_edit_distance(
-        self, n1: int, n2: int, delta: callable = None, norm: callable = None
+        self,
+        n1: int,
+        n2: int,
+        end_time: int = None,
+        style="fragmented",
+        node_lengths: tuple = (1, 5, 7),
     ) -> 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).
@@ -2109,48 +2469,178 @@ class lineageTree:
         Args:
             n1 (int): id of the first node to compare
             n2 (int): id of the second node to compare
-            delta (callable): comparison function (see edist doc for more information)
-            norm (callable): norming function that takes the number of nodes
-                of the tree spawned by `n1` and the number of nodes
-                of the tree spawned by `n2` as arguments.
+            tree_style ("mini","simple","fragmented","full"): Which tree approximation is going to be used for the comparisons.
+                                                              Defaults to "fragmented".
 
         Returns:
             (float) The normed unordered tree edit distance
         """
 
-        from edist.uted import uted
-
-        if delta is None or not callable(delta):
+        tree = tree_style[style].value
+        tree1 = tree(
+            lT=self, node_length=node_lengths, end_time=end_time, root=n1
+        )
+        tree2 = tree(
+            lT=self, node_length=node_lengths, end_time=end_time, root=n2
+        )
+        delta = tree1.delta
+        _, times1 = tree1.tree
+        _, times2 = tree2.tree
+        (
+            nodes1,
+            adj1,
+            corres1,
+        ) = tree1.edist
+        (
+            nodes2,
+            adj2,
+            corres2,
+        ) = tree2.edist
+        if len(nodes1) == len(nodes2) == 0:
+            return 0
+        delta_tmp = partial(
+            delta,
+            corres1=corres1,
+            corres2=corres2,
+            times1=times1,
+            times2=times2,
+        )
 
-            def delta(x, y, corres1, corres2, times):
-                if x is None or y is None:
-                    return 1
-                len_x = times[corres1[x]]
-                len_y = times[corres2[y]]
-                return np.abs(len_x - len_y) / (len_x + len_y)
+        return uted.uted(nodes1, adj1, nodes2, adj2, delta=delta_tmp) / max(
+            tree1.get_norm(), tree2.get_norm()
+        )
 
-        if norm is None or not callable(norm):
+    def to_simple_networkx(
+        self, node: Union[int, list, set, tuple] = None, start_time: int = 0
+    ):
+        """
+        Creates a simple networkx tree graph (every branch is a cell lifetime). This function is to be used for producing nx.graph objects(
+        they can be used for visualization or other tasks),
+        so only the start and the end of a branch are calculated, all cells in between are not taken into account.
+        Args:
+            start_time (int): From which timepoints are the graphs to be calculated.
+                                For example if start_time is 10, then all trees that begin
+                                on tp 10 or before are calculated.
+        returns:
+            G : list(nx.Digraph(),...)
+            pos : list(dict(id:position))
+        """
 
-            def norm(x, y):
-                return max(len(x), len(y))
+        if node is None:
+            mothers = [
+                root for root in self.roots if self.time[root] <= start_time
+            ]
+        else:
+            mothers = node if isinstance(node, (list, set)) else [node]
+        graph = {}
+        all_nodes = {}
+        all_edges = {}
+        for mom in mothers:
+            edges = set()
+            nodes = set()
+            for branch in self.get_all_branches_of_node(mom):
+                nodes.update((branch[0], branch[-1]))
+                if len(branch) > 1:
+                    edges.add((branch[0], branch[-1]))
+                for suc in self[branch[-1]]:
+                    edges.add((branch[-1], suc))
+            all_edges[mom] = edges
+            all_nodes[mom] = nodes
+        for i, mother in enumerate(mothers):
+            graph[i] = nx.DiGraph()
+            graph[i].add_nodes_from(all_nodes[mother])
+            graph[i].add_edges_from(all_edges[mother])
+
+        return graph
+
+    def plot_all_lineages(
+        self,
+        starting_point: int = 0,
+        nrows=2,
+        figsize=(10, 15),
+        dpi=70,
+        fontsize=22,
+        figure=None,
+        axes=None,
+        **kwargs,
+    ):
+        """Plots all lineages.
 
-        if norm is False:
+        Args:
+            starting_point (int, optional): Which timepoints and upwards are the graphs to be calculated.
+                                For example if start_time is 10, then all trees that begin
+                                on tp 10 or before are calculated. Defaults to None.
+            nrows (int):  How many rows of plots should be printed.
+            kwargs: args accepted by networkx
+        """
 
-            def norm(*args):
-                return 1
+        nrows = int(nrows)
+        if nrows < 1 or not nrows:
+            nrows = 1
+            raise Warning("Number of rows has to be at least 1")
 
-        simple_tree_1, _ = self.get_simple_tree(n1)
-        simple_tree_2, _ = self.get_simple_tree(n2)
-        nodes1, adj1, corres1 = self.__edist_format(simple_tree_1)
-        nodes2, adj2, corres2 = self.__edist_format(simple_tree_2)
-        if len(nodes1) == len(nodes2) == 0:
-            return 0
-        delta_tmp = partial(
-            delta, corres1=corres1, corres2=corres2, times=self.cycle_time
+        graphs = self.to_simple_networkx(start_time=starting_point)
+        ncols = int(len(graphs) // nrows) + (+np.sign(len(graphs) % nrows))
+        pos = postions_of_nx(self, graphs)
+        figure, axes = plt.subplots(
+            figsize=figsize, nrows=nrows, ncols=ncols, dpi=dpi, sharey=True
         )
-        return uted(nodes1, adj1, nodes2, adj2, delta=delta_tmp) / norm(
-            nodes1, nodes2
+        flat_axes = axes.flatten()
+        ax2root = {}
+        for i, graph in enumerate(graphs.values()):
+            nx.draw_networkx(
+                graph,
+                pos[i],
+                with_labels=False,
+                arrows=False,
+                **kwargs,
+                ax=flat_axes[i],
+            )
+            root = [n for n, d in graph.in_degree() if d == 0][0]
+            label = self.labels.get(root, "Unlabeled")
+            xlim = flat_axes[i].get_xlim()
+            ylim = flat_axes[i].get_ylim()
+            x_pos = (xlim[1]) / 10
+            y_pos = ylim[0] + 15
+            ax2root[flat_axes[i]] = root
+            flat_axes[i].text(
+                x_pos,
+                y_pos,
+                label,
+                fontsize=fontsize,
+                color="black",
+                ha="center",
+                va="center",
+                bbox={
+                    "facecolor": "white",
+                    "edgecolor": "green",
+                    "boxstyle": "round",
+                },
+            )
+        [figure.delaxes(ax) for ax in axes.flatten() if not ax.has_data()]
+        return figure, axes, ax2root
+
+    def plot_node(self, node, figsize=(4, 7), dpi=150, **kwargs):
+        """Plots the subtree spawn by a node.
+
+        Args:
+            node (int): The id of the node that is going to be plotted.
+            kwargs: args accepted by networkx
+        """
+        graph = self.to_simple_networkx(node)
+        if len(graph) > 1:
+            raise Warning("Please enter only one node")
+        graph = graph[list(graph)[0]]
+        figure, ax = plt.subplots(nrows=1, ncols=1)
+        nx.draw_networkx(
+            graph,
+            hierarchy_pos(graph, self, node),
+            with_labels=False,
+            arrows=False,
+            ax=ax,
+            **kwargs,
         )
+        return figure, ax
 
     # def DTW(self, t1, t2, max_w=None, start_delay=None, end_delay=None,
     #         metric='euclidian', **kwargs):
@@ -2231,11 +2721,584 @@ class lineageTree:
                     to_do.append(_next)
                 elif self.time[_next] == t:
                     final_nodes.append(_next)
-        if not final_nodes: return list(r)
+        if not final_nodes:
+            return list(r)
         return final_nodes
 
+    @staticmethod
+    def __calculate_diag_line(dist_mat: np.ndarray) -> (float, float):
+        """
+        Calculate the line that centers the band w.
+
+            Args:
+                dist_mat (matrix): distance matrix obtained by the function calculate_dtw
+
+            Returns:
+                (float) Slope
+                (float) intercept of the line
+        """
+        i, j = dist_mat.shape
+        x1 = max(0, i - j) / 2
+        x2 = (i + min(i, j)) / 2
+        y1 = max(0, j - i) / 2
+        y2 = (j + min(i, j)) / 2
+        slope = (y1 - y2) / (x1 - x2)
+        intercept = y1 - slope * x1
+        return slope, intercept
+
+    # Reference: https://github.com/kamperh/lecture_dtw_notebook/blob/main/dtw.ipynb
+    def __dp(
+        self,
+        dist_mat: np.ndarray,
+        start_d: int = 0,
+        back_d: int = 0,
+        fast: bool = False,
+        w: int = 0,
+        centered_band: bool = True,
+    ) -> (((int, int), ...), np.ndarray):
+        """
+        Find DTW minimum cost between two series using dynamic programming.
+
+            Args:
+                dist_mat (matrix): distance matrix obtained by the function calculate_dtw
+                start_d (int): start delay
+                back_d (int): end delay
+                w (int): window constrain
+                slope (float): to calculate window - givem by the function __calculate_diag_line
+                intercept (flost): to calculate window - givem by the function __calculate_diag_line
+                use_absolute (boolean): if the window constraing is calculate by the absolute difference between points (uncentered)
+
+            Returns:
+                (tuple of tuples) Aligment path
+                (matrix) Cost matrix
+        """
+        N, M = dist_mat.shape
+        w_limit = max(w, abs(N - M))  # Calculate the Sakoe-Chiba band width
+
+        if centered_band:
+            slope, intercept = self.__calculate_diag_line(dist_mat)
+            square_root = np.sqrt((slope**2) + 1)
+
+        # Initialize the cost matrix
+        cost_mat = np.full((N + 1, M + 1), np.inf)
+        cost_mat[0, 0] = 0
+
+        # Fill the cost matrix while keeping traceback information
+        traceback_mat = np.zeros((N, M))
+
+        cost_mat[: start_d + 1, 0] = 0
+        cost_mat[0, : start_d + 1] = 0
+
+        cost_mat[N - back_d :, M] = 0
+        cost_mat[N, M - back_d :] = 0
+
+        for i in range(N):
+            for j in range(M):
+                if fast and not centered_band:
+                    condition = abs(i - j) <= w_limit
+                elif fast:
+                    condition = (
+                        abs(slope * i - j + intercept) / square_root <= w_limit
+                    )
+                else:
+                    condition = True
+
+                if condition:
+                    penalty = [
+                        cost_mat[i, j],  # match (0)
+                        cost_mat[i, j + 1],  # insertion (1)
+                        cost_mat[i + 1, j],  # deletion (2)
+                    ]
+                    i_penalty = np.argmin(penalty)
+                    cost_mat[i + 1, j + 1] = (
+                        dist_mat[i, j] + penalty[i_penalty]
+                    )
+                    traceback_mat[i, j] = i_penalty
+
+        min_index1 = np.argmin(cost_mat[N - back_d :, M])
+        min_index2 = np.argmin(cost_mat[N, M - back_d :])
+
+        if (
+            cost_mat[N, M - back_d + min_index2]
+            < cost_mat[N - back_d + min_index1, M]
+        ):
+            i = N - 1
+            j = M - back_d + min_index2 - 1
+            final_cost = cost_mat[i + 1, j + 1]
+        else:
+            i = N - back_d + min_index1 - 1
+            j = M - 1
+            final_cost = cost_mat[i + 1, j + 1]
+
+        path = [(i, j)]
+
+        while (
+            start_d != 0
+            and ((start_d < i and j > 0) or (i > 0 and start_d < j))
+        ) or (start_d == 0 and (i > 0 or j > 0)):
+            tb_type = traceback_mat[i, j]
+            if tb_type == 0:
+                # Match
+                i -= 1
+                j -= 1
+            elif tb_type == 1:
+                # Insertion
+                i -= 1
+            elif tb_type == 2:
+                # Deletion
+                j -= 1
+
+            path.append((i, j))
+
+        # Strip infinity edges from cost_mat before returning
+        cost_mat = cost_mat[1:, 1:]
+        return path[::-1], cost_mat, final_cost
+
+    # Reference: https://github.com/nghiaho12/rigid_transform_3D
+    @staticmethod
+    def __rigid_transform_3D(A, B):
+        assert A.shape == B.shape
+
+        num_rows, num_cols = A.shape
+        if num_rows != 3:
+            raise Exception(
+                f"matrix A is not 3xN, it is {num_rows}x{num_cols}"
+            )
+
+        num_rows, num_cols = B.shape
+        if num_rows != 3:
+            raise Exception(
+                f"matrix B is not 3xN, it is {num_rows}x{num_cols}"
+            )
+
+        # find mean column wise
+        centroid_A = np.mean(A, axis=1)
+        centroid_B = np.mean(B, axis=1)
+
+        # ensure centroids are 3x1
+        centroid_A = centroid_A.reshape(-1, 1)
+        centroid_B = centroid_B.reshape(-1, 1)
+
+        # subtract mean
+        Am = A - centroid_A
+        Bm = B - centroid_B
+
+        H = Am @ np.transpose(Bm)
+
+        # find rotation
+        U, S, Vt = np.linalg.svd(H)
+        R = Vt.T @ U.T
+
+        # special reflection case
+        if np.linalg.det(R) < 0:
+            # print("det(R) < R, reflection detected!, correcting for it ...")
+            Vt[2, :] *= -1
+            R = Vt.T @ U.T
+
+        t = -R @ centroid_A + centroid_B
+
+        return R, t
+
+    def __interpolate(
+        self, track1: list, track2: list, threshold: int
+    ) -> (np.ndarray, np.ndarray):
+        """
+        Interpolate two series that have different lengths
+
+            Args:
+                track1 (list): list of nodes of the first cell cycle to compare
+                track2 (list): list of nodes of the second cell cycle to compare
+                threshold (int): set a maximum number of points a track can have
+
+            Returns:
+                (list of list) x, y, z postions for track1
+                (list of list) x, y, z postions for track2
+        """
+        inter1_pos = []
+        inter2_pos = []
+
+        track1_pos = np.array([self.pos[c_id] for c_id in track1])
+        track2_pos = np.array([self.pos[c_id] for c_id in track2])
+
+        # Both tracks have the same length and size below the threshold - nothing is done
+        if len(track1) == len(track2) and (
+            len(track1) <= threshold or len(track2) <= threshold
+        ):
+            return track1_pos, track2_pos
+        # Both tracks have the same length but one or more sizes are above the threshold
+        elif len(track1) > threshold or len(track2) > threshold:
+            sampling = threshold
+        # Tracks have different lengths and the sizes are below the threshold
+        else:
+            sampling = max(len(track1), len(track2))
+
+        for pos in range(3):
+            track1_interp = InterpolatedUnivariateSpline(
+                np.linspace(0, 1, len(track1_pos[:, pos])),
+                track1_pos[:, pos],
+                k=1,
+            )
+            inter1_pos.append(track1_interp(np.linspace(0, 1, sampling)))
+
+            track2_interp = InterpolatedUnivariateSpline(
+                np.linspace(0, 1, len(track2_pos[:, pos])),
+                track2_pos[:, pos],
+                k=1,
+            )
+            inter2_pos.append(track2_interp(np.linspace(0, 1, sampling)))
+
+        return np.column_stack(inter1_pos), np.column_stack(inter2_pos)
+
+    def calculate_dtw(
+        self,
+        nodes1: int,
+        nodes2: int,
+        threshold: int = 1000,
+        regist: bool = True,
+        start_d: int = 0,
+        back_d: int = 0,
+        fast: bool = False,
+        w: int = 0,
+        centered_band: bool = True,
+        cost_mat_p: bool = False,
+    ) -> (float, tuple, np.ndarray, np.ndarray, np.ndarray):
+        """
+        Calculate DTW distance between two cell cycles
+
+            Args:
+                nodes1 (int): node to compare distance
+                nodes2 (int): node to compare distance
+                threshold: set a maximum number of points a track can have
+                regist (boolean): Rotate and translate trajectories
+                start_d (int): start delay
+                back_d (int): end delay
+                fast (boolean): True if the user wants to run the fast algorithm with window restrains
+                w (int): window size
+                centered_band (boolean): if running the fast algorithm, True if the windown is centered
+                cost_mat_p (boolean): True if print the not normalized cost matrix
+
+            Returns:
+                (float) DTW distance
+                (tuple of tuples) Aligment path
+                (matrix) Cost matrix
+                (list of lists) pos_cycle1: rotated and translated trajectories positions
+                (list of lists) pos_cycle2: rotated and translated trajectories positions
+        """
+        nodes1_cycle = self.get_cycle(nodes1)
+        nodes2_cycle = self.get_cycle(nodes2)
+
+        interp_cycle1, interp_cycle2 = self.__interpolate(
+            nodes1_cycle, nodes2_cycle, threshold
+        )
+
+        pos_cycle1 = np.array([self.pos[c_id] for c_id in nodes1_cycle])
+        pos_cycle2 = np.array([self.pos[c_id] for c_id in nodes2_cycle])
+
+        if regist:
+            R, t = self.__rigid_transform_3D(
+                np.transpose(interp_cycle1), np.transpose(interp_cycle2)
+            )
+            pos_cycle1 = np.transpose(np.dot(R, pos_cycle1.T) + t)
+
+        dist_mat = distance.cdist(pos_cycle1, pos_cycle2, "euclidean")
+
+        path, cost_mat, final_cost = self.__dp(
+            dist_mat,
+            start_d,
+            back_d,
+            w=w,
+            fast=fast,
+            centered_band=centered_band,
+        )
+        cost = final_cost / len(path)
+
+        if cost_mat_p:
+            return cost, path, cost_mat, pos_cycle1, pos_cycle2
+        else:
+            return cost, path
+
+    def plot_dtw_heatmap(
+        self,
+        nodes1: int,
+        nodes2: int,
+        threshold: int = 1000,
+        regist: bool = True,
+        start_d: int = 0,
+        back_d: int = 0,
+        fast: bool = False,
+        w: int = 0,
+        centered_band: bool = True,
+    ) -> (float, plt.figure):
+        """
+        Plot DTW cost matrix between two cell cycles in heatmap format
+
+            Args:
+                nodes1 (int): node to compare distance
+                nodes2 (int): node to compare distance
+                start_d (int): start delay
+                back_d (int): end delay
+                fast (boolean): True if the user wants to run the fast algorithm with window restrains
+                w (int): window size
+                centered_band (boolean): if running the fast algorithm, True if the windown is centered
+
+            Returns:
+                (float) DTW distance
+                (figure) Heatmap of cost matrix with opitimal path
+        """
+        cost, path, cost_mat, pos_cycle1, pos_cycle2 = self.calculate_dtw(
+            nodes1,
+            nodes2,
+            threshold,
+            regist,
+            start_d,
+            back_d,
+            fast,
+            w,
+            centered_band,
+            cost_mat_p=True,
+        )
+
+        fig = plt.figure(figsize=(8, 6))
+        ax = fig.add_subplot(1, 1, 1)
+        im = ax.imshow(
+            cost_mat, cmap="viridis", origin="lower", interpolation="nearest"
+        )
+        plt.colorbar(im)
+        ax.set_title("Heatmap of DTW Cost Matrix")
+        ax.set_xlabel("Tree 1")
+        ax.set_ylabel("tree 2")
+        x_path, y_path = zip(*path)
+        ax.plot(y_path, x_path, color="black")
+
+        return cost, fig
+
+    @staticmethod
+    def __plot_2d(
+        pos_cycle1,
+        pos_cycle2,
+        nodes1,
+        nodes2,
+        ax,
+        x_idx,
+        y_idx,
+        x_label,
+        y_label,
+    ):
+        ax.plot(
+            pos_cycle1[:, x_idx],
+            pos_cycle1[:, y_idx],
+            "-",
+            label=f"root = {nodes1}",
+        )
+        ax.plot(
+            pos_cycle2[:, x_idx],
+            pos_cycle2[:, y_idx],
+            "-",
+            label=f"root = {nodes2}",
+        )
+        ax.set_xlabel(x_label)
+        ax.set_ylabel(y_label)
+
+    def plot_dtw_trajectory(
+        self,
+        nodes1: int,
+        nodes2: int,
+        threshold: int = 1000,
+        regist: bool = True,
+        start_d: int = 0,
+        back_d: int = 0,
+        fast: bool = False,
+        w: int = 0,
+        centered_band: bool = True,
+        projection: str = None,
+        alig: bool = False,
+    ) -> (float, plt.figure):
+        """
+        Plots DTW trajectories aligment between two cell cycles in 2D or 3D
+
+            Args:
+                nodes1 (int): node to compare distance
+                nodes2 (int): node to compare distance
+                threshold (int): set a maximum number of points a track can have
+                regist (boolean): Rotate and translate trajectories
+                start_d (int): start delay
+                back_d (int): end delay
+                w (int): window size
+                fast (boolean): True if the user wants to run the fast algorithm with window restrains
+                centered_band (boolean): if running the fast algorithm, True if the windown is centered
+                projection (string): specify which 2D to plot ->
+                    '3d' : for the 3d visualization
+                    'xy' or None (default) : 2D projection of axis x and y
+                    'xz' : 2D projection of axis x and z
+                    'yz' : 2D projection of axis y and z
+                    'pca' : PCA projection
+                alig (boolean): True to show alignment on plot
+
+            Returns:
+                (float) DTW distance
+                (figue) Trajectories Plot
+        """
+        (
+            distance,
+            alignment,
+            cost_mat,
+            pos_cycle1,
+            pos_cycle2,
+        ) = self.calculate_dtw(
+            nodes1,
+            nodes2,
+            threshold,
+            regist,
+            start_d,
+            back_d,
+            fast,
+            w,
+            centered_band,
+            cost_mat_p=True,
+        )
+
+        fig = plt.figure(figsize=(10, 6))
+
+        if projection == "3d":
+            ax = fig.add_subplot(1, 1, 1, projection="3d")
+        else:
+            ax = fig.add_subplot(1, 1, 1)
+
+        if projection == "3d":
+            ax.plot(
+                pos_cycle1[:, 0],
+                pos_cycle1[:, 1],
+                pos_cycle1[:, 2],
+                "-",
+                label=f"root = {nodes1}",
+            )
+            ax.plot(
+                pos_cycle2[:, 0],
+                pos_cycle2[:, 1],
+                pos_cycle2[:, 2],
+                "-",
+                label=f"root = {nodes2}",
+            )
+            ax.set_ylabel("y position")
+            ax.set_xlabel("x position")
+            ax.set_zlabel("z position")
+        else:
+            if projection == "xy" or projection == "yx" or projection is None:
+                self.__plot_2d(
+                    pos_cycle1,
+                    pos_cycle2,
+                    nodes1,
+                    nodes2,
+                    ax,
+                    0,
+                    1,
+                    "x position",
+                    "y position",
+                )
+            elif projection == "xz" or projection == "zx":
+                self.__plot_2d(
+                    pos_cycle1,
+                    pos_cycle2,
+                    nodes1,
+                    nodes2,
+                    ax,
+                    0,
+                    2,
+                    "x position",
+                    "z position",
+                )
+            elif projection == "yz" or projection == "zy":
+                self.__plot_2d(
+                    pos_cycle1,
+                    pos_cycle2,
+                    nodes1,
+                    nodes2,
+                    ax,
+                    1,
+                    2,
+                    "y position",
+                    "z position",
+                )
+            elif projection == "pca":
+                try:
+                    from sklearn.decomposition import PCA
+                except ImportError:
+                    Warning(
+                        "scikit-learn is not installed, the PCA orientation cannot be used. You can install scikit-learn with pip install"
+                    )
+
+                # Apply PCA
+                pca = PCA(n_components=2)
+                pca.fit(np.vstack([pos_cycle1, pos_cycle2]))
+                pos_cycle1_2d = pca.transform(pos_cycle1)
+                pos_cycle2_2d = pca.transform(pos_cycle2)
+
+                ax.plot(
+                    pos_cycle1_2d[:, 0],
+                    pos_cycle1_2d[:, 1],
+                    "-",
+                    label=f"root = {nodes1}",
+                )
+                ax.plot(
+                    pos_cycle2_2d[:, 0],
+                    pos_cycle2_2d[:, 1],
+                    "-",
+                    label=f"root = {nodes2}",
+                )
+
+                # Set axis labels
+                axes = ["x", "y", "z"]
+                x_label = axes[np.argmax(np.abs(pca.components_[0]))]
+                y_label = axes[np.argmax(np.abs(pca.components_[1]))]
+                x_percent = 100 * (
+                    np.max(np.abs(pca.components_[0]))
+                    / np.sum(np.abs(pca.components_[0]))
+                )
+                y_percent = 100 * (
+                    np.max(np.abs(pca.components_[1]))
+                    / np.sum(np.abs(pca.components_[1]))
+                )
+                ax.set_xlabel(f"{x_percent:.0f}% of {x_label} position")
+                ax.set_ylabel(f"{y_percent:.0f}% of {y_label} position")
+            else:
+                raise ValueError(
+                    """Error: available projections are:
+                        '3d' : for the 3d visualization
+                        'xy' or None (default) : 2D projection of axis x and y
+                        'xz' : 2D projection of axis x and z
+                        'yz' : 2D projection of axis y and z
+                        'pca' : PCA projection"""
+                )
+
+        connections = [[pos_cycle1[i], pos_cycle2[j]] for i, j in alignment]
+
+        for connection in connections:
+            xyz1 = connection[0]
+            xyz2 = connection[1]
+            x_pos = [xyz1[0], xyz2[0]]
+            y_pos = [xyz1[1], xyz2[1]]
+            z_pos = [xyz1[2], xyz2[2]]
+
+            if alig and projection != "pca":
+                if projection == "3d":
+                    ax.plot(x_pos, y_pos, z_pos, "k--", color="grey")
+                else:
+                    ax.plot(x_pos, y_pos, "k--", color="grey")
+
+        ax.set_aspect("equal")
+        ax.legend()
+        fig.tight_layout()
+
+        if alig and projection == "pca":
+            warnings.warn(
+                "Error: not possible to show alignment in PCA projection !",
+                UserWarning,
+            )
+
+        return distance, fig
+
     def first_labelling(self):
-        self.labels =  {i:"Enter_Label" for i in self.time_nodes[0]}
+        self.labels = {i: "Unlabeled" for i in self.time_nodes[0]}
 
     def __init__(
         self,
@@ -2268,12 +3331,12 @@ class lineageTree:
                 'TGMM, 'ASTEC', MaMuT', 'TrackMate', 'csv', 'celegans', 'binary'
                 default is 'binary'
         """
+        self.name = name
         self.time_nodes = {}
         self.time_edges = {}
         self.max_id = -1
         self.next_id = []
         self.nodes = set()
-        self.edges = set()
         self.successor = {}
         self.predecessor = {}
         self.pos = {}
@@ -2281,40 +3344,57 @@ class lineageTree:
         self.time = {}
         self.kdtrees = {}
         self.spatial_density = {}
-        self.progeny = {}
-        self.labels = {}
-        if xml_attributes is None:
-            self.xml_attributes = []
-        else:
-            self.xml_attributes = xml_attributes
-        file_type = file_type.lower()
-        if file_type == "tgmm":
-            self.read_tgmm_xml(file_format, tb, te, z_mult)
-            self.t_b = tb
-            self.t_e = te
-        elif file_type == "mamut" or file_type == "trackmate":
-            self.read_from_mamut_xml(file_format)
-        elif file_type == "celegans":
-            self.read_from_txt_for_celegans(file_format)
-        elif file_type == "celegans_cao":
-            self.read_from_txt_for_celegans_CAO(
-                file_format, reorder=reorder, shape=shape, raw_size=raw_size
-            )
-        elif file_type == "mastodon":
-            if isinstance(file_format, list) and len(file_format) == 2:
-                self.read_from_mastodon_csv(file_format)
+        if file_type and file_format:
+            if xml_attributes is None:
+                self.xml_attributes = []
             else:
-                if isinstance(file_format, list):
-                    file_format = file_format[0]
-                self.read_from_mastodon(file_format, name)
-        elif file_type == "astec":
-            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_format and file_format.endswith(".lT"):
-            with open(file_format, "br") as f:
-                tmp = pkl.load(f)
-                f.close()
-            self.__dict__.update(tmp.__dict__)
-        elif file_format is not None:
-            self.read_from_binary(file_format)
+                self.xml_attributes = xml_attributes
+            file_type = file_type.lower()
+            if file_type == "tgmm":
+                self.read_tgmm_xml(file_format, tb, te, z_mult)
+                self.t_b = tb
+                self.t_e = te
+            elif file_type == "mamut" or file_type == "trackmate":
+                self.read_from_mamut_xml(file_format)
+            elif file_type == "celegans":
+                self.read_from_txt_for_celegans(file_format)
+            elif file_type == "celegans_cao":
+                self.read_from_txt_for_celegans_CAO(
+                    file_format,
+                    reorder=reorder,
+                    shape=shape,
+                    raw_size=raw_size,
+                )
+            elif file_type == "mastodon":
+                if isinstance(file_format, list) and len(file_format) == 2:
+                    self.read_from_mastodon_csv(file_format)
+                else:
+                    if isinstance(file_format, list):
+                        file_format = file_format[0]
+                    self.read_from_mastodon(file_format, name)
+            elif file_type == "astec":
+                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_format and file_format.endswith(".lT"):
+                with open(file_format, "br") as f:
+                    tmp = pkl.load(f)
+                    f.close()
+                self.__dict__.update(tmp.__dict__)
+            elif file_format is not None:
+                self.read_from_binary(file_format)
+            if self.name is None:
+                try:
+                    self.name = Path(file_format).stem
+                except:
+                    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)