LineageTree 1.7.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

LineageTree/lineageTree.py

@@ -2,625 +2,538 @@
 # 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)
+
+from __future__ import annotations
+
+import importlib.metadata
 import os
 import pickle as pkl
 import struct
 import warnings
-from collections.abc import Iterable
-from functools import partial
+from collections.abc import Callable, Iterable, Sequence
+from functools import partial, wraps
 from itertools import combinations
 from numbers import Number
-from pathlib import Path
-from typing import TextIO, Union
-
-from .loaders import lineageTreeLoaders
-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.",
-        stacklevel=2,
-    )
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Literal
+
+import matplotlib.colors as mcolors
 import matplotlib.pyplot as plt
 import numpy as np
+import svgwrite
+from edist import uted
+from matplotlib import colormaps
+from matplotlib.collections import LineCollection
+from packaging.version import Version
 from scipy.interpolate import InterpolatedUnivariateSpline
-from scipy.spatial import Delaunay, distance
-from scipy.spatial import cKDTree as KDTree
+from scipy.sparse import dok_array
+from scipy.spatial import Delaunay, KDTree, distance
 
+from .tree_approximation import TreeApproximationTemplate, tree_style
 from .utils import (
-    create_links_and_cycles,
+    convert_style_to_number,
+    create_links_and_chains,
     hierarchical_pos,
 )
 
+if TYPE_CHECKING:
+    from edist.alignment import Alignment
+
+
+class dynamic_property(property):
+    def __init__(
+        self, fget=None, fset=None, fdel=None, doc=None, protected_name=None
+    ):
+        super().__init__(fget, fset, fdel, doc)
+        self.protected_name = protected_name
+
+    def __set_name__(self, owner, name):
+        self.name = name
+        if self.protected_name is None:
+            self.protected_name = f"_{name}"
+        if not hasattr(owner, "_protected_dynamic_properties"):
+            owner._protected_dynamic_properties = []
+        owner._protected_dynamic_properties.append(self.protected_name)
+        if not hasattr(owner, "_dynamic_properties"):
+            owner._dynamic_properties = []
+        owner._dynamic_properties += [name, self.protected_name]
+        setattr(owner, self.protected_name, None)
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+        instance._has_been_reset = False
+        if getattr(instance, self.protected_name) is None:
+            value = super().__get__(instance, owner)
+            setattr(instance, self.protected_name, value)
+            return value
+        else:
+            return getattr(instance, self.protected_name)
+
+
+class lineageTree:
+    norm_dict = {"max": max, "sum": sum, None: lambda x: 1}
+
+    def modifier(wrapped_func):
+        @wraps(wrapped_func)
+        def raising_flag(self, *args, **kwargs):
+            should_reset = (
+                not hasattr(self, "_has_been_reset")
+                or not self._has_been_reset
+            )
+            out_func = wrapped_func(self, *args, **kwargs)
+            if should_reset:
+                for prop in self._protected_dynamic_properties:
+                    self.__dict__[prop] = None
+                self._has_been_reset = True
+            return out_func
+
+        return raising_flag
+
+    def __check_cc_cycles(self, n: int) -> tuple[bool, set[int]]:
+        """Check if the connected component of a given node `n` has a cycle.
+
+        Returns
+        -------
+        bool
+            True if the tree has cycles, False otherwise.
+        set of int
+            The set of nodes that have been checked.
+        """
+        to_do = [n]
+        no_cycle = True
+        already_done = set()
+        while to_do and no_cycle:
+            current = to_do.pop(-1)
+            if current not in already_done:
+                already_done.add(current)
+            else:
+                no_cycle = False
+            to_do.extend(self._successor[current])
+        to_do = list(self._predecessor[n])
+        while to_do and no_cycle:
+            current = to_do.pop(-1)
+            if current not in already_done:
+                already_done.add(current)
+            else:
+                no_cycle = False
+            to_do.extend(self._predecessor[current])
+        return not no_cycle, already_done
+
+    def __check_for_cycles(self) -> bool:
+        """Check if the tree has cycles.
+
+        Returns
+        -------
+        bool
+            True if the tree has cycles, False otherwise.
+        """
+        to_do = set(self.nodes)
+        found_cycle = False
+        while to_do and not found_cycle:
+            current = to_do.pop()
+            found_cycle, done = self.__check_cc_cycles(current)
+            to_do.difference_update(done)
+        return found_cycle
 
-class lineageTree(lineageTreeLoaders):
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         if isinstance(other, lineageTree):
-            return other.successor == self.successor
-        return False
+            return (
+                other._successor == self._successor
+                and other._predecessor == self._predecessor
+                and other._time == self._time
+            )
+        else:
+            return False
 
-    def get_next_id(self):
-        """Computes the next authorized id.
+    def get_next_id(self) -> int:
+        """Computes the next authorized id and assign it.
 
-        Returns:
-            int: 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 == []:
+        if not hasattr(self, "max_id") or (self.max_id == -1 and self.nodes):
+            self.max_id = max(self.nodes) if len(self.nodes) else 0
+        if not hasattr(self, "next_id") or 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(
+    @modifier
+    def add_chain(
         self,
-        pred: int,
+        node: 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.
+        downstream: bool,
+        pos: Callable | None = None,
+    ) -> int:
+        """Adds a chain 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 True will create a branch that goes forwards in time otherwise backwards.
-        Returns:
-            (int): Id of the first node of the sublineage.
+        Parameters
+        ----------
+        node : int
+            Id of the successor (predecessor if `downstream==False`)
+        length : int
+            The length of the new chain.
+        downstream : bool, default=True
+            If `True` will create a chain that goes forwards in time otherwise backwards.
+        pos : np.ndarray, optional
+            The new position of the chain. Defaults to None.
+
+        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
+            return node
+        if length < 1:
+            raise ValueError("Length cannot be <1")
+        if downstream:
+            for _ in range(int(length)):
+                old_node = node
+                node = self._add_node(pred=[old_node])
+                self._time[node] = self._time[old_node] + 1
         else:
-            for _ in range(length):
-                _next = self.add_node(
-                    self.time[pred] + 1,
-                    succ=pred,
-                    pos=self.pos[original],
-                    reverse=False,
+            if self._predecessor[node]:
+                raise Warning("The node already has a predecessor.")
+            if self._time[node] - length < self.t_b:
+                raise Warning(
+                    "A node cannot created outside the lower bound of the dataset. (It is possible to change it by lT.t_b = int(...))"
                 )
-                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
+            for _ in range(int(length)):
+                old_node = node
+                node = self._add_node(succ=[old_node])
+                self._time[node] = self._time[old_node] - 1
+        return node
+
+    @modifier
+    def add_root(self, t: int, pos: list | None = None) -> int:
+        """Adds a root to a specific timepoint.
+
+        Parameters
+        ----------
+        t :int
+            The timepoint the node is going to be added.
+        pos : list
+            The position of the new node.
+        Returns
+        -------
+        int
+            The id of the new root.
+        """
+        C_next = self.get_next_id()
+        self._successor[C_next] = ()
+        self._predecessor[C_next] = ()
+        self._time[C_next] = t
+        self.pos[C_next] = pos if isinstance(pos, list) else []
+        self._changed_roots = True
+        return C_next
 
-        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(
+    def _add_node(
         self,
-        t: int = None,
-        succ: int = None,
-        pos: np.ndarray = None,
-        nid: int = None,
-        reverse: bool = False,
+        succ: list | None = None,
+        pred: list | None = None,
+        pos: np.ndarray | None = None,
+        nid: int | None = None,
     ) -> int:
-        """Adds a node to the lineageTree and update it accordingly.
-
-        Args:
-            t (int): int, time to which to add the node
-            succ (int): id of the node the new node is a successor to
-            pos ([float, ]): list of three floats representing the 3D
-                spatial position of the node
-            nid (int): id value of the new node, to be used carefully,
-                if None is provided the new id is automatically computed.
-            reverse (bool): True if in this lineageTree the predecessors
-                are the successors and reciprocally.
-                This is there for bacward compatibility, should be left at False.
-        Returns:
-            int: id of the new node.
-        """
+        """Adds a node to the LineageTree object that is either a successor or a predecessor of another node.
+        Does not handle time! You cannot enter both a successor and a predecessor.
+
+        Parameters
+        ----------
+        succ : list
+            list of ids of the nodes the new node is a successor to
+        pred : list
+            list of ids of the nodes the new node is a predecessor to
+        pos : np.ndarray, optional
+            position of the new node
+        nid : int, optional
+            id value of the new node, to be used carefully,
+            if None is provided the new id is automatically computed.
+
+        Returns
+        -------
+        int
+            id of the new node.
+        """
+        if not succ and not pred:
+            raise Warning(
+                "Please enter a successor or a predecessor, otherwise use the add_roots() function."
+            )
         C_next = self.get_next_id() if nid is None else nid
-        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)
-        elif succ is not None:
-            self.predecessor.setdefault(succ, []).append(C_next)
-            self.successor.setdefault(C_next, []).append(succ)
-        self.nodes.add(C_next)
-        self.pos[C_next] = pos
-        self.time[C_next] = t
+        if succ:
+            self._successor[C_next] = succ
+            for suc in succ:
+                self._predecessor[suc] = (C_next,)
+        else:
+            self._successor[C_next] = ()
+        if pred:
+            self._predecessor[C_next] = pred
+            self._successor[pred[0]] = self._successor.setdefault(
+                pred[0], ()
+            ) + (C_next,)
+        else:
+            self._predecessor[C_next] = ()
+        if isinstance(pos, list):
+            self.pos[C_next] = pos
         return C_next
 
-    def remove_nodes(self, group: Union[int, set, list]):
+    @modifier
+    def remove_nodes(self, group: int | set | list) -> None:
         """Removes a group of nodes from the LineageTree
 
-        Args:
-            group (set|list|int): One or more nodes that are to be removed.
+        Parameters
+        ----------
+        group : set of int or list of int or int
+            One or more nodes that are to be removed.
         """
-        if isinstance(group, int):
+        if isinstance(group, int | float):
             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)
+        group = self.nodes.intersection(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:
-            node (int): Any node of the branch to be modified/
-            new_length (int): The new length of the tree.
-        """
-        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 - 1,
-                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 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 1", stacklevel=2)
-            self._time_resolution = 10
+            for attr in self.__dict__:
+                attr_value = self.__getattribute__(attr)
+                if isinstance(attr_value, dict) and attr not in [
+                    "successor",
+                    "predecessor",
+                    "_successor",
+                    "_predecessor",
+                    "_time",
+                ]:
+                    attr_value.pop(node, ())
+            if self._predecessor.get(node):
+                self._successor[self._predecessor[node][0]] = tuple(
+                    set(
+                        self._successor[self._predecessor[node][0]]
+                    ).difference(group)
+                )
+            for p_node in self._successor.get(node, []):
+                self._predecessor[p_node] = ()
+            self._predecessor.pop(node, ())
+            self._successor.pop(node, ())
 
     @property
-    def roots(self):
-        if not hasattr(self, "_roots"):
-            self._roots = set(self.nodes).difference(self.predecessor)
-        return self._roots
+    def successor(self) -> MappingProxyType[int, tuple[int]]:
+        """The successor of the tree."""
+        if not hasattr(self, "_protected_successor"):
+            self._protected_successor = MappingProxyType(self._successor)
+        return self._protected_successor
 
     @property
-    def edges(self):
-        return {(k, vi) for k, v in self.successor.items() for vi in v}
+    def predecessor(self) -> MappingProxyType[int, tuple[int]]:
+        """The predecessor of the tree."""
+        if not hasattr(self, "_protected_predecessor"):
+            self._protected_predecessor = MappingProxyType(self._predecessor)
+        return self._protected_predecessor
 
     @property
-    def leaves(self):
-        return set(self.predecessor).difference(self.successor)
+    def time(self) -> MappingProxyType[int, int]:
+        """The time of the tree."""
+        if not hasattr(self, "_protected_time"):
+            self._protected_time = MappingProxyType(self._time)
+        return self._protected_time
+
+    @dynamic_property
+    def t_b(self) -> int:
+        """The first timepoint of the tree."""
+        return min(self._time.values())
+
+    @dynamic_property
+    def t_e(self) -> int:
+        """The last timepoint of the tree."""
+        return max(self._time.values())
+
+    @dynamic_property
+    def nodes(self) -> frozenset[int]:
+        """Nodes of the tree"""
+        return frozenset(self._successor.keys())
+
+    @dynamic_property
+    def depth(self) -> dict[int, int]:
+        """The depth of each node in the tree."""
+        _depth = {}
+        for leaf in self.leaves:
+            _depth[leaf] = 1
+            while leaf in self._predecessor and self._predecessor[leaf]:
+                parent = self._predecessor[leaf][0]
+                current_depth = _depth.get(parent, 0)
+                _depth[parent] = max(_depth[leaf] + 1, current_depth)
+                leaf = parent
+        for root in self.roots - set(_depth):
+            _depth[root] = 1
+        return _depth
+
+    @dynamic_property
+    def roots(self) -> frozenset[int]:
+        """Set of roots of the tree"""
+        return frozenset({s for s, p in self._predecessor.items() if p == ()})
+
+    @dynamic_property
+    def leaves(self) -> frozenset[int]:
+        """Set of leaves"""
+        return frozenset({p for p, s in self._successor.items() if s == ()})
+
+    @dynamic_property
+    def edges(self) -> tuple[tuple[int, int]]:
+        """Set of edges"""
+        return tuple((p, si) for p, s in self._successor.items() for si in s)
 
     @property
-    def labels(self):
+    def labels(self) -> dict[int, str]:
+        """The labels of the nodes."""
         if not hasattr(self, "_labels"):
-            if hasattr(self, "cell_name"):
+            if hasattr(self, "node_name"):
                 self._labels = {
-                    i: self.cell_name.get(i, "Unlabeled") for i in self.roots
+                    i: self.node_name.get(i, "Unlabeled") for i in self.roots
                 }
             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
 
-    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")
-        f.write("define VERTEX %d\n" % (nb_points * 2))
-        f.write("define EDGE %d\n" % nb_points)
-        f.write("define POINT %d\n" % ((length) * nb_points))
-        f.write("Parameters {\n")
-        f.write('\tContentType "HxSpatialGraph"\n')
-        f.write("}\n")
-
-        f.write("VERTEX { float[3] VertexCoordinates } @1\n")
-        f.write("EDGE { int[2] EdgeConnectivity } @2\n")
-        f.write("EDGE { int NumEdgePoints } @3\n")
-        f.write("POINT { float[3] EdgePointCoordinates } @4\n")
-        f.write("VERTEX { float Vcolor } @5\n")
-        f.write("VERTEX { int Vbool } @6\n")
-        f.write("EDGE { float Ecolor } @7\n")
-        f.write("VERTEX { int Vbool2 } @8\n")
-
-    def write_to_am(
-        self,
-        path_format: str,
-        t_b: int = None,
-        t_e: int = None,
-        length: int = 5,
-        manual_labels: dict = None,
-        default_label: int = 5,
-        new_pos: np.ndarray = None,
-    ):
-        """Writes a lineageTree into an Amira readable data (.am format).
-
-        Args:
-            path_format (str): path to the output. It should contain 1 %03d where the time step will be entered
-            t_b (int): first time point to write (if None, min(LT.to_take_time) is taken)
-            t_e (int): last time point to write (if None, max(LT.to_take_time) is taken)
-                note, if there is no 'to_take_time' attribute, self.time_nodes
-                is considered instead (historical)
-            length (int): length of the track to print (how many time before).
-            manual_labels ({id: label, }): dictionary that maps cell ids to
-            default_label (int): default value for the manual label
-            new_pos ({id: [x, y, z]}): dictionary that maps a 3D position to a cell ID.
-                if new_pos == None (default) then self.pos is considered.
-        """
-        if not hasattr(self, "to_take_time"):
-            self.to_take_time = self.time_nodes
-        if t_b is None:
-            t_b = min(self.to_take_time.keys())
-        if t_e is None:
-            t_e = max(self.to_take_time.keys())
-        if new_pos is None:
-            new_pos = self.pos
-
-        if manual_labels is None:
-            manual_labels = {}
-        for t in range(t_b, t_e + 1):
-            with open(path_format % t, "w") as f:
-                nb_points = len(self.to_take_time[t])
-                self._write_header_am(f, nb_points, length)
-                points_v = {}
-                for C in self.to_take_time[t]:
-                    C_tmp = C
-                    positions = []
-                    for _ in range(length):
-                        C_tmp = self.predecessor.get(C_tmp, [C_tmp])[0]
-                        positions.append(new_pos[C_tmp])
-                    points_v[C] = positions
-
-                f.write("@1\n")
-                for C in self.to_take_time[t]:
-                    f.write("{:f} {:f} {:f}\n".format(*tuple(points_v[C][0])))
-                    f.write("{:f} {:f} {:f}\n".format(*tuple(points_v[C][-1])))
-
-                f.write("@2\n")
-                for i, _ in enumerate(self.to_take_time[t]):
-                    f.write("%d %d\n" % (2 * i, 2 * i + 1))
-
-                f.write("@3\n")
-                for _ in self.to_take_time[t]:
-                    f.write("%d\n" % (length))
-
-                f.write("@4\n")
-                for C in self.to_take_time[t]:
-                    for p in points_v[C]:
-                        f.write("{:f} {:f} {:f}\n".format(*tuple(p)))
-
-                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"{0:f}\n")
-
-                f.write("@6\n")
-                for C in self.to_take_time[t]:
-                    f.write(
-                        "%d\n"
-                        % (
-                            int(
-                                manual_labels.get(C, default_label)
-                                != default_label
-                            )
-                        )
-                    )
-                    f.write("%d\n" % (0))
-
-                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.write("@8\n")
-                for _ in self.to_take_time[t]:
-                    f.write("%d\n" % (1))
-                    f.write("%d\n" % (0))
-                f.close()
+    @property
+    def time_resolution(self) -> float:
+        if not hasattr(self, "_time_resolution"):
+            self._time_resolution = 0
+        return self._time_resolution / 10
 
-    def _get_height(self, c: int, done: dict):
-        """Recursively computes the height of a cell within a tree * a space factor.
+    @time_resolution.setter
+    def time_resolution(self, time_resolution) -> None:
+        if time_resolution is not None and time_resolution > 0:
+            self._time_resolution = int(time_resolution * 10)
+        else:
+            warnings.warn("Time resolution set to default 0", stacklevel=2)
+            self._time_resolution = 0
+
+    def __setstate__(self, state):
+        if "_successor" not in state:
+            state["_successor"] = state["successor"]
+        if "_predecessor" not in state:
+            state["_predecessor"] = state["predecessor"]
+        if "_time" not in state:
+            state["_time"] = state["time"]
+        self.__dict__.update(state)
+
+    def _get_height(self, c: int, done: dict) -> float:
+        """Recursively computes the height of a node within a tree times a space factor.
         This function is specific to the function write_to_svg.
 
-        Args:
-            c (int): id of a cell in a lineage tree from which the height will be computed from
-            done ({int: [int, int]}): a dictionary that maps a cell id to its vertical and horizontal position
-        Returns:
-            float:
+        Parameters
+        ----------
+        c : int
+            id of a node in a lineage tree from which the height will be computed from
+        done : dict mapping int to list of two int
+            a dictionary that maps a node id to its vertical and horizontal position
+
+        Returns
+        -------
+        float
+            the height of the node `c`
         """
         if c in done:
             return done[c][0]
         else:
             P = np.mean(
-                [self._get_height(di, done) for di in self.successor[c]]
+                [self._get_height(di, done) for di in self._successor[c]]
             )
-            done[c] = [P, self.vert_space_factor * self.time[c]]
+            done[c] = [P, self.vert_space_factor * self._time[c]]
             return P
 
     def write_to_svg(
         self,
         file_name: str,
-        roots: list = None,
+        roots: list | None = None,
         draw_nodes: bool = True,
         draw_edges: bool = True,
-        order_key: callable = None,
+        order_key: Callable | None = None,
         vert_space_factor: float = 0.5,
         horizontal_space: float = 1,
-        node_size: callable = None,
-        stroke_width: callable = None,
+        node_size: Callable | str | None = None,
+        stroke_width: Callable | None = None,
         factor: float = 1.0,
-        node_color: callable = None,
-        stroke_color: callable = None,
-        positions: dict = None,
-        node_color_map: callable = None,
-        normalize: bool = True,
-    ):
-        ##### remove background? default True background value? default 1
-
+        node_color: Callable | str | None = None,
+        stroke_color: Callable | None = None,
+        positions: dict | None = None,
+        node_color_map: Callable | str | None = None,
+    ) -> None:
         """Writes the lineage tree to an SVG file.
         Node and edges coloring and size can be provided.
 
-        Args:
-            file_name (str): filesystem filename valid for `open()`
-            roots ([int, ...]): list of node ids to be drawn. If `None` all the nodes will be drawn. Default `None`
-            draw_nodes (bool): wether to print the nodes or not, default `True`
-            draw_edges (bool): wether to print the edges or not, default `True`
-            order_key (callable): function that would work for the attribute `key=` for the `sort`/`sorted` function
-            vert_space_factor (float): the vertical position of a node is its time. `vert_space_factor` is a
-                               multiplier to space more or less nodes in time
-            horizontal_space (float): space between two consecutive nodes
-            node_size (callable | str): a function that maps a node id to a `float` value that will determine the
-                       radius of the node. The default function return the constant value `vertical_space_factor/2.1`
-                       If a string is given instead and it is a property of the tree,
-                       the the size will be mapped according to the property
-            stroke_width (callable): a function that maps a node id to a `float` value that will determine the
-                          width of the daughter edge.  The default function return the constant value `vertical_space_factor/2.1`
-            factor (float): scaling factor for nodes positions, default 1
-            node_color (callable | str): a function that maps a node id to a triplet between 0 and 255.
-                        The triplet will determine the color of the node. If a string is given instead and it is a property
-                        of the tree, the the color will be mapped according to the property
-            node_color_map (callable | str): the name of the colormap to use to color the nodes, or a colormap function
-            stroke_color (callable): a function that maps a node id to a triplet between 0 and 255.
-                          The triplet will determine the color of the stroke of the inward edge.
-            positions ({int: [float, float], ...}): dictionary that maps a node id to a 2D position.
-                       Default `None`. If provided it will be used to position the nodes.
+        Parameters
+        ----------
+        file_name : str
+            filesystem filename valid for `open()`
+        roots : list of int, defaults to `self.roots`
+            list of node ids to be drawn. If `None` or not provided all the nodes will be drawn. Default `None`
+        draw_nodes : bool, default True
+            wether to print the nodes or not
+        draw_edges : bool, default True
+            wether to print the edges or not
+        order_key : Callable, optional
+            function that would work for the attribute `key=` for the `sort`/`sorted` function
+        vert_space_factor : float, default=0.5
+            the vertical position of a node is its time. `vert_space_factor` is a
+            multiplier to space more or less nodes in time
+        horizontal_space : float, default=1
+            space between two consecutive nodes
+        node_size : Callable or str, optional
+            a function that maps a node id to a `float` value that will determine the
+            radius of the node. The default function return the constant value `vertical_space_factor/2.1`
+            If a string is given instead and it is a property of the tree,
+            the the size will be mapped according to the property
+        stroke_width : Callable, optional
+            a function that maps a node id to a `float` value that will determine the
+            width of the daughter edge.  The default function return the constant value `vertical_space_factor/2.1`
+        factor : float, default=1.0
+            scaling factor for nodes positions, default 1
+        node_color : Callable or str, optional
+            a function that maps a node id to a triplet between 0 and 255.
+            The triplet will determine the color of the node. If a string is given instead and it is a property
+            of the tree, the the color will be mapped according to the property
+        node_color_map : Callable or str, optional
+            the name of the colormap to use to color the nodes, or a colormap function
+        stroke_color : Callable, optional
+            a function that maps a node id to a triplet between 0 and 255.
+            The triplet will determine the color of the stroke of the inward edge.
+        positions : dict mapping int to list of two float, optional
+            dictionary that maps a node id to a 2D position.
+            Default `None`. If provided it will be used to position the nodes.
         """
-        import svgwrite
 
         def normalize_values(v, nodes, _range, shift, mult):
             min_ = np.percentile(v, 1)
             max_ = np.percentile(v, 99)
             values = _range * ((v - min_) / (max_ - min_)) + shift
-            values_dict_nodes = dict(zip(nodes, values))
+            values_dict_nodes = dict(zip(nodes, values, strict=True))
             return lambda x: values_dict_nodes[x] * mult
 
         if roots is None:
             roots = self.roots
             if hasattr(self, "image_label"):
-                roots = [cell for cell in roots if self.image_label[cell] != 1]
+                roots = [node for node in roots if self.image_label[node] != 1]
 
         if node_size is None:
 
             def node_size(x):
                 return vert_space_factor / 2.1
 
-        elif isinstance(node_size, str) and node_size in self.__dict__:
-            values = np.array([self[node_size][c] for c in self.nodes])
+        else:
+            values = np.array(
+                [self._successor[node_size][c] for c in self.nodes]
+            )
             node_size = normalize_values(
                 values, self.nodes, 0.5, 0.5, vert_space_factor / 2.1
             )
@@ -635,18 +548,19 @@ class lineageTree(lineageTreeLoaders):
                 return 0, 0, 0
 
         elif isinstance(node_color, str) and node_color in self.__dict__:
-            if isinstance(node_color_map, str):
-                from matplotlib import colormaps
+            from matplotlib import colormaps
 
-                if node_color_map in colormaps:
-                    node_color_map = colormaps[node_color_map]
-                else:
-                    node_color_map = colormaps["viridis"]
-            values = np.array([self[node_color][c] for c in self.nodes])
+            if node_color_map in colormaps:
+                cmap = colormaps[node_color_map]
+            else:
+                cmap = colormaps["viridis"]
+            values = np.array(
+                [self._successor[node_color][c] for c in self.nodes]
+            )
             normed_vals = normalize_values(values, self.nodes, 1, 0, 1)
 
             def node_color(x):
-                return [k * 255 for k in node_color_map(normed_vals(x))[:-1]]
+                return [k * 255 for k in cmap(normed_vals(x))[:-1]]
 
         coloring_edges = stroke_color is not None
         if not coloring_edges:
@@ -655,24 +569,25 @@ class lineageTree(lineageTreeLoaders):
                 return 0, 0, 0
 
         elif isinstance(stroke_color, str) and stroke_color in self.__dict__:
-            if isinstance(node_color_map, str):
-                from matplotlib import colormaps
+            from matplotlib import colormaps
 
-                if node_color_map in colormaps:
-                    node_color_map = colormaps[node_color_map]
-                else:
-                    node_color_map = colormaps["viridis"]
-            values = np.array([self[stroke_color][c] for c in self.nodes])
+            if node_color_map in colormaps:
+                cmap = colormaps[node_color_map]
+            else:
+                cmap = colormaps["viridis"]
+            values = np.array(
+                [self._successor[stroke_color][c] for c in self.nodes]
+            )
             normed_vals = normalize_values(values, self.nodes, 1, 0, 1)
 
             def stroke_color(x):
-                return [k * 255 for k in node_color_map(normed_vals(x))[:-1]]
+                return [k * 255 for k in cmap(normed_vals(x))[:-1]]
 
         prev_x = 0
         self.vert_space_factor = vert_space_factor
         if order_key is not None:
             roots.sort(key=order_key)
-        treated_cells = []
+        treated_nodes = []
 
         pos_given = positions is not None
         if not pos_given:
@@ -683,25 +598,26 @@ class lineageTree(lineageTreeLoaders):
                         [0.0, 0.0],
                     ]
                     * len(self.nodes),
-                )
+                    strict=True,
+                ),
             )
         for _i, r in enumerate(roots):
             r_leaves = []
             to_do = [r]
             while len(to_do) != 0:
                 curr = to_do.pop(0)
-                treated_cells += [curr]
-                if curr in self.successor:
+                treated_nodes += [curr]
+                if not self._successor[curr]:
                     if order_key is not None:
-                        to_do += sorted(self.successor[curr], key=order_key)
+                        to_do += sorted(self._successor[curr], key=order_key)
                     else:
-                        to_do += self.successor[curr]
+                        to_do += self._successor[curr]
                 else:
                     r_leaves += [curr]
             r_pos = {
                 leave: [
                     prev_x + horizontal_space * (1 + j),
-                    self.vert_space_factor * self.time[leave],
+                    self.vert_space_factor * self._time[leave],
                 ]
                 for j, leave in enumerate(r_leaves)
             }
@@ -716,12 +632,12 @@ class lineageTree(lineageTreeLoaders):
             size=factor * np.max(list(positions.values()), axis=0),
         )
         if draw_edges and not draw_nodes and not coloring_edges:
-            to_do = set(treated_cells)
+            to_do = set(treated_nodes)
             while len(to_do) > 0:
                 curr = to_do.pop()
-                c_cycle = self.get_cycle(curr)
-                x1, y1 = positions[c_cycle[0]]
-                x2, y2 = positions[c_cycle[-1]]
+                c_chain = self.get_chain_of_node(curr)
+                x1, y1 = positions[c_chain[0]]
+                x2, y2 = positions[c_chain[-1]]
                 dwg.add(
                     dwg.line(
                         (factor * x1, factor * y1),
@@ -729,7 +645,7 @@ class lineageTree(lineageTreeLoaders):
                         stroke=svgwrite.rgb(0, 0, 0),
                     )
                 )
-                for si in self[c_cycle[-1]]:
+                for si in self._successor[c_chain[-1]]:
                     x3, y3 = positions[si]
                     dwg.add(
                         dwg.line(
@@ -738,11 +654,11 @@ class lineageTree(lineageTreeLoaders):
                             stroke=svgwrite.rgb(0, 0, 0),
                         )
                     )
-                to_do.difference_update(c_cycle)
+                to_do.difference_update(c_chain)
         else:
-            for c in treated_cells:
+            for c in treated_nodes:
                 x1, y1 = positions[c]
-                for si in self[c]:
+                for si in self._successor[c]:
                     x2, y2 = positions[si]
                     if draw_edges:
                         dwg.add(
@@ -753,7 +669,7 @@ class lineageTree(lineageTreeLoaders):
                                 stroke_width=svgwrite.pt(stroke_width(si)),
                             )
                         )
-            for c in treated_cells:
+            for c in treated_nodes:
                 x1, y1 = positions[c]
                 if draw_nodes:
                     dwg.add(
@@ -770,49 +686,58 @@ class lineageTree(lineageTreeLoaders):
         fname: str,
         t_min: int = -1,
         t_max: int = np.inf,
-        nodes_to_use: list = None,
+        nodes_to_use: list[int] | None = None,
         temporal: bool = True,
-        spatial: str = None,
+        spatial: str | None = None,
         write_layout: bool = True,
-        node_properties: dict = None,
+        node_properties: dict | None = None,
         Names: bool = False,
-    ):
+    ) -> None:
         """Write a lineage tree into an understable tulip file.
 
-        Args:
-            fname (str): path to the tulip file to create
-            t_min (int): minimum time to consider, default -1
-            t_max (int): maximum time to consider, default np.inf
-            nodes_to_use ([int, ]): list of nodes to show in the graph,
-                          default *None*, then self.nodes is used
-                          (taking into account *t_min* and *t_max*)
-            temporal (bool): True if the temporal links should be printed, default True
-            spatial (str): Build spatial edges from a spatial neighbourhood graph.
-                The graph has to be computed before running this function
-                'ball': neighbours at a given distance,
-                'kn': k-nearest neighbours,
-                'GG': gabriel graph,
-                None: no spatial edges are writen.
-                Default None
-            write_layout (bool): True, write the spatial position as layout,
-                                   False, do not write spatial positionm
-                                   default True
-            node_properties ({`p_name`, [{id, p_value}, default]}): a dictionary of properties to write
-                                                To a key representing the name of the property is
-                                                paired a dictionary that maps a cell id to a property
-                                                and a default value for this property
-            Names (bool): Only works with ASTEC outputs, True to sort the cells by their names
+        Parameters
+        ----------
+        fname : str
+            path to the tulip file to create
+        t_min : int, default=-1
+            minimum time to consider
+        t_max : int, default=np.inf
+            maximum time to consider
+        nodes_to_use : list of int, optional
+            list of nodes to show in the graph,
+            if `None` then self.nodes is used
+            (taking into account `t_min` and `t_max`)
+        temporal : bool, default=True
+            True if the temporal links should be printed
+        spatial : str, optional
+            Build spatial edges from a spatial neighbourhood graph.
+            The graph has to be computed before running this function
+            'ball': neighbours at a given distance,
+            'kn': k-nearest neighbours,
+            'GG': gabriel graph,
+            None: no spatial edges are writen.
+            Default None
+        write_layout : bool, default=True
+            write the spatial position as layout if True
+            do not write spatial position otherwise
+        node_properties : dict mapping str to list of dict of properties and its default value, optional
+            a dictionary of properties to write
+            To a key representing the name of the property is
+            paired a dictionary that maps a node id to a property
+            and a default value for this property
+        Names : bool, default=True
+            Only works with ASTEC outputs, True to sort the nodes by their names
         """
 
         def format_names(names_which_matter):
-            """Return an ensured formated cell names"""
+            """Return an ensured formated node names"""
             tmp = {}
             for k, v in names_which_matter.items():
                 tmp[k] = (
                     v.split(".")[0][0]
-                    + "%02d" % int(v.split(".")[0][1:])
+                    + "{:02d}".format(int(v.split(".")[0][1:]))
                     + "."
-                    + "%04d" % int(v.split(".")[1][:-1])
+                    + "{:04d}".format(int(v.split(".")[1][:-1]))
                     + v.split(".")[1][-1]
                 )
             return tmp
@@ -839,20 +764,20 @@ class lineageTree(lineageTreeLoaders):
             if not nodes_to_use:
                 if t_max != np.inf or t_min > -1:
                     nodes_to_use = [
-                        n for n in self.nodes if t_min < self.time[n] <= t_max
+                        n for n in self.nodes if t_min < self._time[n] <= t_max
                     ]
                     edges_to_use = []
                     if temporal:
                         edges_to_use += [
                             e
                             for e in self.edges
-                            if t_min < self.time[e[0]] < t_max
+                            if t_min < self._time[e[0]] < t_max
                         ]
                     if spatial:
                         edges_to_use += [
                             e
                             for e in s_edges
-                            if t_min < self.time[e[0]] < t_max
+                            if t_min < self._time[e[0]] < t_max
                         ]
                 else:
                     nodes_to_use = list(self.nodes)
@@ -866,12 +791,12 @@ class lineageTree(lineageTreeLoaders):
                 nodes_to_use = set(nodes_to_use)
                 if temporal:
                     for n in nodes_to_use:
-                        for d in self.successor.get(n, []):
+                        for d in self._successor[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
+                        e for e in s_edges if t_min < self._time[e[0]] < t_max
                     ]
             nodes_to_use = set(nodes_to_use)
             if Names:
@@ -889,12 +814,12 @@ class lineageTree(lineageTreeLoaders):
                 for k, v in node_properties[Names][0].items():
                     if (
                         len(
-                            self.successor.get(
-                                self.predecessor.get(k, [-1])[0], []
+                            self._successor.get(
+                                self._predecessor.get(k, [-1])[0], ()
                             )
                         )
                         != 1
-                        or self.time[k] == t_min + 1
+                        or self._time[k] == t_min + 1
                     ):
                         tmp_names[k] = v
                 node_properties[Names][0] = tmp_names
@@ -924,7 +849,7 @@ class lineageTree(lineageTreeLoaders):
             f.write('\t(default "0" "0")\n')
             for n in nodes_to_use:
                 f.write(
-                    "\t(node " + str(n) + ' "' + str(self.time[n]) + '")\n'
+                    "\t(node " + str(n) + ' "' + str(self._time[n]) + '")\n'
                 )
             f.write(")\n")
 
@@ -953,10 +878,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(
@@ -971,7 +896,9 @@ class lineageTree(lineageTreeLoaders):
             f.write(")")
             f.close()
 
-    def to_binary(self, fname: str, starting_points: list = None):
+    def to_binary(
+        self, fname: str, starting_points: list[int] | None = None
+    ) -> None:
         """Writes the lineage tree (a forest) as a binary structure
         (assuming it is a binary tree, it would not work for *n* ary tree with 2 < *n*).
         The binary file is composed of 3 sequences of numbers and
@@ -986,36 +913,37 @@ class lineageTree(lineageTreeLoaders):
         The *time_sequence* is stored as a list of unsigned short (0 -> 2^(8*2)-1)
         The *pos_sequence* is stored as a list of double.
 
-        Args:
-            fname (str): name of the binary file
-            starting_points ([int, ]): list of the roots to be written.
-                If None, all roots are written, default value, None
+        Parameters
+        ----------
+        fname : str
+            name of the binary file
+        starting_points : list of int, optional
+            list of the roots to be written.
+            If `None`, all roots are written, default value, None
         """
         if starting_points is None:
-            starting_points = [
-                c for c in self.successor if self.predecessor.get(c, []) == []
-            ]
+            starting_points = list(self.roots)
         number_sequence = [-1]
         pos_sequence = []
         time_sequence = []
         for c in starting_points:
-            time_sequence.append(self.time.get(c, 0))
+            time_sequence.append(self._time.get(c, 0))
             to_treat = [c]
-            while to_treat != []:
+            while to_treat:
                 curr_c = to_treat.pop()
                 number_sequence.append(curr_c)
                 pos_sequence += list(self.pos[curr_c])
-                if self[curr_c] == []:
+                if self._successor[curr_c] == ():
                     number_sequence.append(-1)
-                elif len(self.successor[curr_c]) == 1:
-                    to_treat += self.successor[curr_c]
+                elif len(self._successor[curr_c]) == 1:
+                    to_treat += self._successor[curr_c]
                 else:
                     number_sequence.append(-2)
-                    to_treat += self.successor[curr_c]
+                    to_treat += self._successor[curr_c]
         remaining_nodes = set(self.nodes) - set(number_sequence)
 
         for c in remaining_nodes:
-            time_sequence.append(self.time.get(c, 0))
+            time_sequence.append(self._time.get(c, 0))
             number_sequence.append(c)
             pos_sequence += list(self.pos[c])
             number_sequence.append(-1)
@@ -1030,204 +958,270 @@ class lineageTree(lineageTreeLoaders):
 
             f.close()
 
-    def read_from_binary(self, fname: str):
-        """
-        Reads a binary lineageTree file name.
-        Format description: see self.to_binary
+    def write(self, fname: str) -> None:
+        """Write a lineage tree on disk as an .lT file.
 
-        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.
-
-        Args:
-            fname (str): path to and name of the file to save
+        Parameters
+        ----------
+        fname : str
+            path to and name of the file to save
         """
         if os.path.splitext(fname)[-1] != ".lT":
             fname = os.path.extsep.join((fname, "lT"))
+        if hasattr(self, "_protected_predecessor"):
+            del self._protected_predecessor
+        if hasattr(self, "_protected_successor"):
+            del self._protected_successor
+        if hasattr(self, "_protected_time"):
+            del self._protected_time
         with open(fname, "bw") as f:
             pkl.dump(self, f)
             f.close()
 
     @classmethod
-    def load(clf, fname: str, rm_empty_lists=True):
-        """
-        Loading a lineage tree from a ".lT" file.
+    def load(clf, fname: str):
+        """Loading a lineage tree from a '.lT' file.
 
-        Args:
-            fname (str): path to and name of the file to read
+        Parameters
+        ----------
+        fname : str
+            path to and name of the file to read
 
-        Returns:
-            (lineageTree): loaded file
+        Returns
+        -------
+        LineageTree
+            loaded file
         """
         with open(fname, "br") as f:
             lT = pkl.load(f)
             f.close()
+        if not hasattr(lT, "__version__") or Version(lT.__version__) < Version(
+            "2.0.0"
+        ):
+            properties = {
+                prop_name: prop
+                for prop_name, prop in lT.__dict__.items()
+                if isinstance(prop, dict)
+                and prop_name
+                not in [
+                    "successor",
+                    "predecessor",
+                    "time",
+                    "_successor",
+                    "_predecessor",
+                    "_time",
+                    "pos",
+                    "labels",
+                ]
+                + lineageTree._dynamic_properties
+                + lineageTree._protected_dynamic_properties
+            }
+            print("_comparisons" in properties)
+            lT = lineageTree(
+                successor=lT._successor,
+                time=lT._time,
+                pos=lT.pos,
+                name=lT.name if hasattr(lT, "name") else None,
+                **properties,
+            )
         if not hasattr(lT, "time_resolution"):
-            lT.time_resolution = None
-        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")
+            lT.time_resolution = 1
+
         return lT
 
-    def get_idx3d(self, t: int) -> tuple:
-        """Get a 3d kdtree for the dataset at time *t* .
-        The  kdtree is stored in *self.kdtrees[t]*
-
-        Args:
-            t (int): time
-        Returns:
-            (kdtree, [int, ]): the built kdtree and
-                the correspondancy list,
-                If the query in the kdtree gives you the value i,
-                then it corresponds to the id in the tree to_check_self[i]
+    def get_predecessors(
+        self,
+        x: int,
+        depth: int | None = None,
+        start_time: int | None = None,
+        end_time: int | None = None,
+    ) -> list[int]:
+        """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.
+
+        Parameters
+        ----------
+        x : int
+            id of the node to compute
+        depth : int
+            maximum number of predecessors to return
+
+        Returns
+        -------
+        list of int
+            list of ids, the last id is `x`
+        """
+        if not start_time:
+            start_time = self.t_b
+        if not end_time:
+            end_time = self.t_e
+        unconstrained_chain = [x]
+        chain = [x] if start_time <= self._time[x] <= end_time else []
+        acc = 0
+        while (
+            acc != depth
+            and start_time < self._time[unconstrained_chain[0]]
+            and (
+                self._predecessor[unconstrained_chain[0]] != ()
+                and (  # Please dont change very important even if it looks weird.
+                    len(
+                        self._successor[
+                            self._predecessor[unconstrained_chain[0]][0]
+                        ]
+                    )
+                    == 1
+                )
+            )
+        ):
+            unconstrained_chain.insert(
+                0, self._predecessor[unconstrained_chain[0]][0]
+            )
+            acc += 1
+            if start_time <= self._time[unconstrained_chain[0]] <= end_time:
+                chain.insert(0, unconstrained_chain[0])
+
+        return chain
+
+    def get_successors(
+        self, x: int, depth: int | None = None, end_time: int | None = None
+    ) -> list[int]:
+        """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.
+
+        Parameters
+        ----------
+        x : int
+            id of the node to compute
+        depth : int, optional
+            maximum number of predecessors to return
+        end_time : int, optional
+            maximum time to consider
+
+        Returns
+        -------
+        list of int
+            list of ids, the first id is `x`
+        """
+        if end_time is None:
+            end_time = self.t_e
+        chain = [x]
+        acc = 0
+        while (
+            len(self._successor[chain[-1]]) == 1
+            and acc != depth
+            and self._time[chain[-1]] < end_time
+        ):
+            chain += self._successor[chain[-1]]
+            acc += 1
+
+        return chain
+
+    def get_chain_of_node(
+        self,
+        x: int,
+        depth: int | None = None,
+        depth_pred: int | None = None,
+        depth_succ: int | None = None,
+        end_time: int | None = None,
+    ) -> list[int]:
+        """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 chain is returned.
+
+        Parameters
+        ----------
+        x : int
+            id of the node to compute
+        depth : int, optional
+            maximum number of predecessors and successor to return
+        depth_pred : int, optional
+            maximum number of predecessors to return
+        depth_succ : int, optional
+            maximum number of successors to return
+
+        Returns
+        -------
+        list of int
+            list of node ids
         """
-        to_check_self = list(self.time_nodes[t])
+        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, end_time=end_time)[
+            :-1
+        ] + self.get_successors(x, depth_succ, end_time=end_time)
+
+    @dynamic_property
+    def all_chains(self) -> list[list[int]]:
+        """List of all chains in the tree, ordered in depth-first search."""
+        return self._compute_all_chains()
+
+    @dynamic_property
+    def time_nodes(self):
+        _time_nodes = {}
+        for c, t in self._time.items():
+            _time_nodes.setdefault(t, set()).add(c)
+        return _time_nodes
+
+    def m(self, i, j):
+        if (i, j) not in self._tmp_parenting:
+            if i == j:  # the distance to the node itself is 0
+                self._tmp_parenting[(i, j)] = 0
+                self._parenting[i, j] = self._tmp_parenting[(i, j)]
+            elif not self._predecessor[
+                j
+            ]:  # j and i are note connected so the distance if inf
+                self._tmp_parenting[(i, j)] = np.inf
+            else:  # the distance between i and j is the distance between i and pred(j) + 1
+                self._tmp_parenting[(i, j)] = (
+                    self.m(i, self._predecessor[j][0]) + 1
+                )
+                self._parenting[i, j] = self._tmp_parenting[(i, j)]
+                self._parenting[j, i] = -self._tmp_parenting[(i, j)]
+        return self._tmp_parenting[(i, j)]
+
+    @property
+    def parenting(self):
+        if not hasattr(self, "_parenting"):
+            self._parenting = dok_array((max(self.nodes) + 1,) * 2)
+            self._tmp_parenting = {}
+            for i, j in combinations(self.nodes, 2):
+                if self._time[j] < self.time[i]:
+                    i, j = j, i
+                self._tmp_parenting[(i, j)] = self.m(i, j)
+            del self._tmp_parenting
+        return self._parenting
+
+    def get_idx3d(self, t: int) -> tuple[KDTree, np.ndarray]:
+        """Get a 3d kdtree for the dataset at time `t`.
+        The  kdtree is stored in `self.kdtrees[t]` and returned.
+        The correspondancy list is also returned.
+
+        Parameters
+        ----------
+        t : int
+            time
+
+        Returns
+        -------
+        KDTree
+            The KDTree corresponding to the lineage tree at time `t`
+        np.ndarray
+            The correspondancy list in the KDTree.
+            If the query in the kdtree gives you the value `i`,
+            then it corresponds to the id in the tree `to_check_self[i]`
+        """
+        to_check_self = list(self.nodes_at_t(t=t))
+
+        if not hasattr(self, "kdtrees"):
+            self.kdtrees = {}
+
         if t not in self.kdtrees:
             data_corres = {}
             data = []
@@ -1240,16 +1234,21 @@ class lineageTree(lineageTreeLoaders):
             idx3d = self.kdtrees[t]
         return idx3d, np.array(to_check_self)
 
-    def get_gabriel_graph(self, t: int) -> dict:
-        """Build the Gabriel graph of the given graph for time point `t`
-        The Garbiel graph is then stored in self.Gabriel_graph and returned
-        *WARNING: the graph is not recomputed if already computed. even if nodes were added*.
+    def get_gabriel_graph(self, t: int) -> dict[int, set[int]]:
+        """Build the Gabriel graph of the given graph for time point `t`.
+        The Garbiel graph is then stored in `self.Gabriel_graph` and returned.
+
+        .. warning:: the graph is not recomputed if already computed, even if the point cloud has changed
 
-        Args:
-            t (int): time
-        Returns:
-            {int, set([int, ])}: a dictionary that maps a node to
-                the set of its neighbors
+        Parameters
+        ----------
+        t : int
+            time
+
+        Returns
+        -------
+        dict of int to set of int
+            A dictionary that maps a node to the set of its neighbors
         """
         if not hasattr(self, "Gabriel_graph"):
             self.Gabriel_graph = {}
@@ -1294,178 +1293,110 @@ class lineageTree(lineageTreeLoaders):
 
         return self.Gabriel_graph[t]
 
-    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.
-
-        Args:
-            x (int): id of the node to compute
-            depth (int): maximum number of predecessors to return
-        Returns:
-            [int, ]: list of ids, the last id is `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[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
-            )
-        ):
-            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, 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.
-
-        Args:
-            x (int): id of the node to compute
-            depth (int): maximum number of predecessors to return
-        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[cycle[-1]]) == 1
-            and acc != depth
-            and self.time[cycle[-1]] < end_time
-        ):
-            cycle += self.successor[cycle[-1]]
-            acc += 1
-
-        return cycle
+    def get_all_chains_of_subtree(
+        self, node: int, end_time: int | None = None
+    ) -> list[list[int]]:
+        """Computes all the chains of the subtree spawn by a given node.
+        Similar to get_all_chains().
 
-    def get_cycle(
-        self,
-        x: int,
-        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
-            depth (int): maximum number of predecessors and successor to return
-            depth_pred (int): maximum number of predecessors to return
-            depth_succ (int): maximum number of successors to return
-        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, end_time=end_time)[
-            :-1
-        ] + self.get_successors(x, depth_succ, end_time=end_time)
-
-    @property
-    def all_tracks(self):
-        if not hasattr(self, "_all_tracks"):
-            self._all_tracks = self.get_all_tracks()
-        return self._all_tracks
+        Parameters
+        ----------
+        node : int
+            The node from which we want to get its chains.
+        end_time : int, optional
+            The time at which we want to stop the chains.
 
-    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
+        Returns
+        -------
+        list of list of int
+            list of chains
         """
         if not end_time:
             end_time = self.t_e
-        branches = [self.get_successors(node)]
-        to_do = list(self[branches[0][-1]])
+        chains = [self.get_successors(node)]
+        to_do = list(self._successor[chains[0][-1]])
         while to_do:
             current = to_do.pop()
-            track = self.get_successors(current, end_time=end_time)
-            # 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:
-        """Computes all the tracks of a given lineage tree,
-        stores it in `self.all_tracks` and returns it.
-
-        Returns:
-            ([[int, ...], ...]): list of lists containing track cell ids
-        """
-        if not hasattr(self, "_all_tracks") or force_recompute:
-            self._all_tracks = []
-            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.extend(self[track[-1]])
-        return self._all_tracks
-
-    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
+            chain = self.get_successors(current, end_time=end_time)
+            if self._time[chain[-1]] <= end_time:
+                chains += [chain]
+                to_do += self._successor[chain[-1]]
+        return chains
+
+    def _compute_all_chains(self) -> list[list[int]]:
+        """Computes all the chains of a given lineage tree,
+        stores it in `self.all_chains` and returns it.
+
+        Returns
+        -------
+        list of list of int
+            list of chains
+        """
+        all_chains = []
+        to_do = sorted(self.roots, key=self.time.get, reverse=True)
+        while len(to_do) != 0:
+            current = to_do.pop()
+            chain = self.get_chain_of_node(current)
+            all_chains += [chain]
+            to_do.extend(self._successor[chain[-1]])
+        return all_chains
+
+    def __get_chains(  # TODO: Probably should be removed, might be used by DTW. Might also be a @dynamic_property
+        self, nodes: Iterable | int | None = None
+    ) -> dict[int, list[list[int]]]:
+        """Returns all the chains in the subtrees spawned by each of the given nodes.
+
+        Parameters
+        ----------
+        nodes : Iterable or int, optional
+            id or Iterable of ids of the nodes to be computed, if `None` all roots are used
+
+        Returns
+        -------
+        dict mapping int to list of Chain
+            dictionary mapping the node ids to a list of chains
+        """
+        all_chains = self.all_chains
+        if nodes is None:
+            nodes = self.roots
+        if not isinstance(nodes, Iterable):
+            nodes = [nodes]
+        output_chains = {}
+        for n in nodes:
+            starting_node = self.get_predecessors(n)[0]
+            found = False
+            done = False
+            starting_time = self.time[n]
+            i = 0
+            current_chain = []
+            while not done and i < len(all_chains):
+                curr_found = all_chains[i][0] == starting_node
+                found = found or curr_found
+                if found:
+                    done = (
+                        self.time[all_chains[i][0]] <= starting_time
+                    ) and not curr_found
+                    if not done:
+                        if curr_found:
+                            current_chain.append(self.get_successors(n))
+                        else:
+                            current_chain.append(all_chains[i])
+                i += 1
+            output_chains[n] = current_chain
+        return output_chains
 
-    def find_leaves(self, roots: Union[int, set, list, tuple]):
+    def find_leaves(self, roots: int | Iterable) -> set[int]:
         """Finds the leaves of a tree spawned by one or more nodes.
 
-        Args:
-            roots (Union[int,set,list,tuple]): The roots of the trees.
+        Parameters
+        ----------
+        roots : int or Iterable
+            The roots of the trees spawning the leaves
 
-        Returns:
-            set: The leaves of one or more trees.
+        Returns
+        -------
+        set
+            The leaves of one or more trees.
         """
         if not isinstance(roots, Iterable):
             to_do = [roots]
@@ -1474,29 +1405,34 @@ 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[curr]
+            if not succ:
                 leaves.add(curr)
-            else:
-                to_do += succ
+            to_do += succ
         return leaves
 
-    def get_sub_tree(
+    def get_subtree_nodes(
         self,
-        x: Union[int, Iterable],
-        end_time: Union[int, None] = None,
+        x: int | Iterable,
+        end_time: 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.
+    ) -> list[int]:
+        """Computes the list of nodes from the subtree spawned by *x*
+        The default output order is Breadth First Traversal.
         Unless preorder is `True` in that case the order is
-        Depth first traversal preordered.
+        Depth First Traversal (DFT) preordered.
+
+        Parameters
+        ----------
+        x : int
+            id of root node
+        preorder : bool, default=False
+            if True the output preorder is DFT
 
-        Args:
-            x (int): id of root node
-            preorder (bool): if True the output is preorder DFT
-        Returns:
-            ([int, ...]): the ordered list of node ids
+        Returns
+        -------
+        list of int
+            the ordered list of node ids
         """
         if not end_time:
             end_time = self.t_e
@@ -1504,233 +1440,258 @@ class lineageTree(lineageTreeLoaders):
             to_do = [x]
         elif isinstance(x, Iterable):
             to_do = list(x)
-        sub_tree = []
+        subtree = []
         while to_do:
             curr = to_do.pop()
-            succ = self.successor.get(curr, [])
-            if succ and end_time < self.time.get(curr, end_time):
+            succ = self._successor[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]
-        return sub_tree
+                subtree += [curr]
+        return subtree
 
     def compute_spatial_density(
-        self, t_b: int = None, t_e: int = None, th: float = 50
-    ) -> dict:
-        """Computes the spatial density of cells between `t_b` and `t_e`.
-        The spatial density is computed as follow:
-        #cell/(4/3*pi*th^3)
-        The results is stored in self.spatial_density is returned.
-
-        Args:
-            t_b (int): starting time to look at, default first time point
-            t_e (int): ending time to look at, default last time point
-            th (float): size of the neighbourhood
-        Returns:
-            {int, float}: dictionary that maps a cell id to its spatial density
-        """
+        self, t_b: int | None = None, t_e: int | None = None, th: float = 50
+    ) -> dict[int, float]:
+        """Computes the spatial density of nodes between `t_b` and `t_e`.
+        The results is stored in `self.spatial_density` and returned.
+
+        Parameters
+        ----------
+        t_b : int, optional
+            starting time to look at, default first time point
+        t_e : int, optional
+            ending time to look at, default last time point
+        th : float, default=50
+            size of the neighbourhood
+
+        Returns
+        -------
+        dict mapping int to float
+            dictionary that maps a node id to its spatial density
+        """
+        if not hasattr(self, "spatial_density"):
+            self.spatial_density = {}
         s_vol = 4 / 3.0 * np.pi * th**3
-        time_range = set(range(t_b, t_e + 1)).intersection(self.time_nodes)
+        if t_b is None:
+            t_b = self.t_b
+        if t_e is None:
+            t_e = self.t_e
+        time_range = set(range(t_b, t_e)).intersection(self._time.values())
         for t in time_range:
             idx3d, nodes = self.get_idx3d(t)
             nb_ni = [
                 (len(ni) - 1) / s_vol
                 for ni in idx3d.query_ball_tree(idx3d, th)
             ]
-            self.spatial_density.update(dict(zip(nodes, nb_ni)))
+            self.spatial_density.update(dict(zip(nodes, nb_ni, strict=True)))
         return self.spatial_density
 
-    def compute_k_nearest_neighbours(self, k: int = 10) -> dict:
+    def compute_k_nearest_neighbours(self, k: int = 10) -> dict[int, set[int]]:
         """Computes the k-nearest neighbors
         Writes the output in the attribute `kn_graph`
         and returns it.
 
-        Args:
-            k (float): number of nearest neighours
-        Returns:
-            {int, set([int, ...])}: dictionary that maps
-                a cell id to its `k` nearest neighbors
+        Parameters
+        ----------
+        k : float
+            number of nearest neighours
+
+        Returns
+        -------
+        dict mapping int to set of int
+            dictionary that maps
+            a node id to its `k` nearest neighbors
         """
         self.kn_graph = {}
-        for t, nodes in self.time_nodes.items():
-            use_k = k if k < len(nodes) else len(nodes)
-            idx3d, nodes = self.get_idx3d(t)
-            pos = [self.pos[c] for c in nodes]
-            _, neighbs = idx3d.query(pos, use_k)
-            out = dict(zip(nodes, [set(nodes[ni[1:]]) for ni in neighbs]))
-            self.kn_graph.update(out)
+        for t in set(self._time.values()):
+            nodes = self.nodes_at_t(t)
+            if 1 < len(nodes):
+                use_k = k if k < len(nodes) else len(nodes)
+                idx3d, nodes = self.get_idx3d(t)
+                pos = [self.pos[c] for c in nodes]
+                _, neighbs = idx3d.query(pos, use_k)
+                out = dict(
+                    zip(
+                        nodes,
+                        map(set, nodes[neighbs]),
+                        strict=True,
+                    )
+                )
+                self.kn_graph.update(out)
+            else:
+                n = nodes.pop
+                self.kn_graph.update({n: {n}})
         return self.kn_graph
 
-    def compute_spatial_edges(self, th: int = 50) -> dict:
+    def compute_spatial_edges(self, th: int = 50) -> dict[int, set[int]]:
         """Computes the neighbors at a distance `th`
         Writes the output in the attribute `th_edge`
         and returns it.
 
-        Args:
-            th (float): distance to consider neighbors
-        Returns:
-            {int, set([int, ...])}: dictionary that maps
-                a cell id to its neighbors at a distance `th`
+        Parameters
+        ----------
+        th : float, default=50
+            distance to consider neighbors
+
+        Returns
+        -------
+        dict mapping int to set of int
+            dictionary that maps a node id to its neighbors at a distance `th`
         """
         self.th_edges = {}
-        for t, _ in self.time_nodes.items():
+        for t in set(self._time.values()):
+            nodes = self.nodes_at_t(t)
             idx3d, nodes = self.get_idx3d(t)
             neighbs = idx3d.query_ball_tree(idx3d, th)
-            out = dict(zip(nodes, [set(nodes[ni]) for ni in neighbs]))
+            out = dict(
+                zip(nodes, [set(nodes[ni]) for ni in neighbs], strict=True)
+            )
             self.th_edges.update(
                 {k: v.difference([k]) for k, v in out.items()}
             )
         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 = max(self.time_nodes, key=lambda x: len(self.time_nodes[x]))
-        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`
+    def get_ancestor_at_t(self, n: int, time: int | None = None) -> int:
+        """Find the id of the ancestor of a give node `n`
         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
+        If there is no ancestor, returns `None`
+        If time is None return the root of the subtree that spawns
         the node n.
 
-        Args:
-            n (int): node for which to look the ancestor
-            time (int): time at which the ancestor has to be found.
-                If `None` the ancestor at the first time point
-                will be found (default `None`)
-
-        Returns:
-            (int): the id of the ancestor at time `time`,
-                `-1` if it does not exist
+        Parameters
+        ----------
+        n : int
+            node for which to look the ancestor
+        time : int, optional
+            time at which the ancestor has to be found.
+            If `None` the ancestor at the first time point
+            will be found.
+
+        Returns
+        -------
+        int
+            the id of the ancestor at time `time`,
+            `-1` if there is no ancestor.
         """
         if n not in self.nodes:
-            return
+            return -1
         if time is None:
             time = self.t_b
         ancestor = n
         while (
-            time < self.time.get(ancestor, -1) and ancestor in self.predecessor
+            time < self._time.get(ancestor, self.t_b - 1)
+            and self._predecessor[ancestor]
         ):
-            ancestor = self.predecessor.get(ancestor, [-1])[0]
-        return ancestor
+            ancestor = self._predecessor[ancestor][0]
+        if self._time.get(ancestor, self.t_b - 1) == time:
+            return ancestor
+        else:
+            return -1
 
-    def get_labelled_ancestor(self, node: int):
-        """Finds the first labelled ancestor and returns its ID otherwise returns None
+    def get_labelled_ancestor(self, node: int) -> int:
+        """Finds the first labelled ancestor and returns its ID otherwise returns -1
 
-        Args:
-            node (int): The id of the node
+        Parameters
+        ----------
+        node : int
+            The id of the node
 
-        Returns:
-            [None,int]: Returns the first ancestor found that has a label otherwise
-            None.
+        Returns
+        -------
+        int
+            Returns the first ancestor found that has a label otherwise `-1`.
         """
         if node not in self.nodes:
-            return None
+            return -1
         ancestor = node
         while (
-            self.t_b <= self.time.get(ancestor, self.t_b - 1)
+            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
+            ancestor = self._predecessor.get(ancestor, [-1])[0]
+        return -1
+
+    def get_ancestor_with_attribute(self, node: int, attribute: str) -> int:
+        """General purpose function to help with searching the first ancestor that has an attribute.
+        Similar to get_labeled_ancestor and may make it redundant.
+
+        Parameters
+        ----------
+        node : int
+            The id of the node
+
+        Returns
+        -------
+        int
+            Returns the first ancestor found that has an attribute otherwise `-1`.
+        """
+        attr_dict = self.__getattribute__(attribute)
+        if not isinstance(attr_dict, dict):
+            raise ValueError("Please select a dict attribute")
+        if node not in self.nodes:
+            return -1
+        if node in attr_dict:
+            return node
+        if node in self.roots:
+            return -1
+        ancestor = (node,)
+        while ancestor and ancestor != [-1]:
+            ancestor = ancestor[0]
+            if ancestor in attr_dict:
+                return ancestor
+            ancestor = self._predecessor.get(ancestor, [-1])
+        return -1
 
     def unordered_tree_edit_distances_at_time_t(
         self,
         t: int,
-        end_time: int = None,
-        style="simple",
+        end_time: int | None = None,
+        style: (
+            Literal["simple", "full", "downsampled", "normalized_simple"]
+            | type[TreeApproximationTemplate]
+        ) = "simple",
         downsample: int = 2,
-        normalize: bool = True,
+        norm: Literal["max", "sum", None] = "max",
         recompute: bool = False,
-    ) -> dict:
-        """
-        Compute all the pairwise unordered tree edit distances from Zhang 996 between the trees spawned at time `t`
-
-        Args:
-            t (int): time to look at
-            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.
-            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
+    ) -> dict[tuple[int, int], float]:
+        """Compute all the pairwise unordered tree edit distances from Zhang 996 between the trees spawned at time `t`
+
+        Parameters
+        ----------
+        t : int
+            time to look at
+        end_time : int
+            The final time point the comparison algorithm will take into account.
+            If None all nodes will be taken into account.
+        style : {"simple", "full", "downsampled", "normalized_simple"} or TreeApproximationTemplate subclass, default="simple"
+            Which tree approximation is going to be used for the comparisons.
+        downsample : int, default=2
+            The downsample factor for the downsampled tree approximation.
+            Used only when `style="downsampled"`.
+        norm : {"max", "sum"}, default="max"
+            The normalization method to use.
+        recompute : bool, default=False
+            If True, forces to recompute the distances
+
+        Returns
+        -------
+        dict mapping a tuple of tuple that contains 2 ints to float
+            a dictionary that maps a pair of node ids at time `t` to their unordered tree edit distance
         """
         if not hasattr(self, "uted"):
             self.uted = {}
         elif t in self.uted and not recompute:
             return self.uted[t]
         self.uted[t] = {}
-        roots = self.time_nodes[t]
+        roots = self.nodes_at_t(t=t)
         for n1, n2 in combinations(roots, 2):
             key = tuple(sorted((n1, n2)))
             self.uted[t][key] = self.unordered_tree_edit_distance(
@@ -1739,37 +1700,132 @@ class lineageTree(lineageTreeLoaders):
                 end_time=end_time,
                 style=style,
                 downsample=downsample,
-                normalize=normalize,
+                norm=norm,
             )
         return self.uted[t]
 
-    def unordered_tree_edit_distance(
+    def __calculate_distance_of_sub_tree(
+        self,
+        node1: int,
+        node2: int,
+        alignment: Alignment,
+        corres1: dict[int, int],
+        corres2: dict[int, int],
+        delta_tmp: Callable,
+        norm: Callable,
+        norm1: int | float,
+        norm2: int | float,
+    ) -> float:
+        """Calculates the distance of the subtree of each node matched in a comparison.
+        DOES NOT CALCULATE THE DISTANCE FROM SCRATCH BUT USING THE ALIGNMENT.
+        TODO ITS BOUND TO CHANGE
+        Parameters
+        ----------
+        node1 : int
+            The root of the first subtree
+        node2 : int
+            The root of the second subtree
+        alignment : Alignment
+            The alignment of the subtree
+        corres1 : dict
+            The correspndance dictionary of the first lineage
+        corres2 : dict
+            The correspondance dictionary of the second lineage
+        delta_tmp : Callable
+            The delta function for the comparisons
+        norm : Callable
+            How should the lineages be normalized
+        norm1 : int or float
+            The result of the normalization of the first tree
+        norm2 : int or float
+            The result of the normalization of the second tree
+
+        Returns
+        -------
+        float
+            The result of the comparison of the subtree
+        """
+        sub_tree_1 = set(self.get_subtree_nodes(node1))
+        sub_tree_2 = set(self.get_subtree_nodes(node2))
+        res = 0
+        for m in alignment:
+            if (
+                corres1.get(m._left, -1) in sub_tree_1
+                or corres2.get(m._right, -1) in sub_tree_2
+            ):
+                res += delta_tmp(
+                    m._left if m._left != -1 else None,
+                    m._right if m._right != -1 else None,
+                )
+        return res / norm([norm1, norm2])
+
+    def clear_comparisons(self):
+        self._comparisons.clear()
+
+    def __unordereded_backtrace(
         self,
         n1: int,
         n2: int,
-        end_time: int = None,
-        norm: Union["max", "sum", None] = "max",
-        style="simple",
+        end_time: int | None = None,
+        norm: Literal["max", "sum", None] = "max",
+        style: (
+            Literal["simple", "normalized_simple", "full", "downsampled"]
+            | type[TreeApproximationTemplate]
+        ) = "simple",
         downsample: int = 2,
-    ) -> float:
+    ) -> dict[
+        str,
+        Alignment
+        | tuple[TreeApproximationTemplate, TreeApproximationTemplate],
+    ]:
         """
-        Compute the unordered tree edit distance from Zhang 1996 between the trees spawned
+        Compute the unordered tree edit backtrace 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).
-        The distance is normed by the function norm that takes the two list of nodes
-        spawned by the trees `n1` and `n2`.
 
-        Args:
-            n1 (int): id of the first node to compare
-            n2 (int): id of the second node to compare
-            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
-        """
-
-        tree = tree_style[style].value
+        Parameters
+        ----------
+        n1 : int
+            id of the first node to compare
+        n2 : int
+            id of the second node to compare
+        end_time : int
+            The final time point the comparison algorithm will take into account.
+            If None all nodes will be taken into account.
+        norm : {"max", "sum"}, default="max"
+            The normalization method to use.
+        style : {"simple", "full", "downsampled", "normalized_simple"} or TreeApproximationTemplate subclass, default="simple"
+            Which tree approximation is going to be used for the comparisons.
+        downsample : int, default=2
+            The downsample factor for the downsampled tree approximation.
+            Used only when `style="downsampled"`.
+
+        Returns
+        -------
+        dict mapping str to Alignment or tuple of [TreeApproximationTemplate, TreeApproximationTemplate]
+            - 'alignment'
+                The alignment between the nodes by the subtrees spawned by the nodes n1,n2 and the normalization function.
+            - 'trees'
+                A list of the two trees that have been mapped to each other.
+        """
+
+        parameters = (
+            end_time,
+            convert_style_to_number(style=style, downsample=downsample),
+        )
+        n1, n2 = sorted([n1, n2])
+        self._comparisons.setdefault(parameters, {})
+        if len(self._comparisons) > 100:
+            warnings.warn(
+                "More than 100 comparisons are saved, use clear_comparisons() to delete them.",
+                stacklevel=2,
+            )
+        if isinstance(style, str):
+            tree = tree_style[style].value
+        elif issubclass(style, TreeApproximationTemplate):
+            tree = style
+        else:
+            raise ValueError("Please use a valid approximation.")
         tree1 = tree(
             lT=self,
             downsample=downsample,
@@ -1798,7 +1854,11 @@ class lineageTree(lineageTreeLoaders):
             corres2,
         ) = tree2.edist
         if len(nodes1) == len(nodes2) == 0:
-            return 0
+            self._comparisons[parameters][(n1, n2)] = {
+                "alignment": (),
+                "trees": (),
+            }
+            return self._comparisons[parameters][(n1, n2)]
         delta_tmp = partial(
             delta,
             corres1=corres1,
@@ -1806,127 +1866,538 @@ 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:
+        btrc = uted.uted_backtrace(nodes1, adj1, nodes2, adj2, delta=delta_tmp)
+
+        self._comparisons[parameters][(n1, n2)] = {
+            "alignment": btrc,
+            "trees": (tree1, tree2),
+        }
+        return self._comparisons[parameters][(n1, n2)]
+
+    def plot_tree_distance_graphs(
+        self,
+        n1: int,
+        n2: int,
+        end_time: int | None = None,
+        norm: Literal["max", "sum", None] = "max",
+        style: (
+            Literal["simple", "normalized_simple", "full", "downsampled"]
+            | type[TreeApproximationTemplate]
+        ) = "simple",
+        downsample: int = 2,
+        colormap: str = "cool",
+        default_color: str = "black",
+        size: float = 10,
+        lw: float = 0.3,
+        ax: list[plt.Axes] | None = None,
+    ) -> tuple[plt.figure, plt.Axes]:
+        """
+        Plots the subtrees compared and colors them according to the quality of the matching of their subtree.
+
+        Parameters
+        ----------
+        n1 : int
+            id of the first node to compare
+        n2 : int
+            id of the second node to compare
+        end_time : int
+            The final time point the comparison algorithm will take into account.
+            If None all nodes will be taken into account.
+        norm : {"max", "sum"}, default="max"
+            The normalization method to use.
+        style : {"simple", "full", "downsampled", "normalized_simple} or TreeApproximationTemplate subclass, default="simple"
+            Which tree approximation is going to be used for the comparisons.
+        downsample : int, default=2
+            The downsample factor for the downsampled tree approximation.
+            Used only when `style="downsampled"`.
+        colormap : str, default="cool"
+            The colormap used for matched nodes, defaults to "cool"
+        default_color : str
+            The color of the unmatched nodes, defaults to "black"
+        size : float
+            The size of the nodes, defaults to 10
+        lw : float
+            The width of the edges, defaults to 0.3
+        ax : np.ndarray, optional
+            The axes used, if not provided another set of axes is produced, defaults to None
+
+        Returns
+        -------
+        plt.Figure
+             The figure of the plot
+        plt.Axes
+             The axes of the plot
+        """
+        parameters = (
+            end_time,
+            convert_style_to_number(style=style, downsample=downsample),
+        )
+        n1, n2 = sorted([n1, n2])
+        self._comparisons.setdefault(parameters, {})
+        if self._comparisons[parameters].get((n1, n2)):
+            tmp = self._comparisons[parameters][(n1, n2)]
+        else:
+            tmp = self.__unordereded_backtrace(
+                n1, n2, end_time, norm, style, downsample
+            )
+        btrc: Alignment = tmp["alignment"]
+        tree1, tree2 = tmp["trees"]
+        _, times1 = tree1.tree
+        _, times2 = tree2.tree
+        (
+            *_,
+            corres1,
+        ) = tree1.edist
+        (
+            *_,
+            corres2,
+        ) = tree2.edist
+        delta_tmp = partial(
+            tree1.delta,
+            corres1=corres1,
+            corres2=corres2,
+            times1=times1,
+            times2=times2,
+        )
+
+        if norm not in self.norm_dict:
+            raise Warning(
+                "Select a viable normalization method (max, sum, None)"
+            )
+        matched_right = []
+        matched_left = []
+        colors = {}
+        if style not in ("full", "downsampled"):
+            for m in btrc:
+                if m._left != -1 and m._right != -1:
+                    cyc1 = self.get_chain_of_node(corres1[m._left])
+                    if len(cyc1) > 1:
+                        node_1, *_, l_node_1 = cyc1
+                        matched_left.append(node_1)
+                        matched_left.append(l_node_1)
+                    elif len(cyc1) == 1:
+                        node_1 = l_node_1 = cyc1.pop()
+                        matched_left.append(node_1)
+
+                    cyc2 = self.get_chain_of_node(corres2[m._right])
+                    if len(cyc2) > 1:
+                        node_2, *_, l_node_2 = cyc2
+                        matched_right.append(node_2)
+                        matched_right.append(l_node_2)
+
+                    elif len(cyc2) == 1:
+                        node_2 = l_node_2 = cyc2.pop()
+                        matched_right.append(node_2)
+
+                    colors[node_1] = self.__calculate_distance_of_sub_tree(
+                        node_1,
+                        node_2,
+                        btrc,
+                        corres1,
+                        corres2,
+                        delta_tmp,
+                        self.norm_dict[norm],
+                        tree1.get_norm(node_1),
+                        tree2.get_norm(node_2),
+                    )
+                    colors[node_2] = colors[node_1]
+                    colors[l_node_1] = colors[node_1]
+                    colors[l_node_2] = colors[node_2]
+        else:
+            for m in btrc:
+                if m._left != -1 and m._right != -1:
+                    node_1 = corres1[m._left]
+                    node_2 = corres2[m._right]
+
+                    if (
+                        self.get_chain_of_node(node_1)[0] == node_1
+                        or self.get_chain_of_node(node_2)[0] == node_2
+                        and (node_1 not in colors or node_2 not in colors)
+                    ):
+                        matched_left.append(node_1)
+                        l_node_1 = self.get_chain_of_node(node_1)[-1]
+                        matched_left.append(l_node_1)
+                        matched_right.append(node_2)
+                        l_node_2 = self.get_chain_of_node(node_2)[-1]
+                        matched_right.append(l_node_2)
+                        colors[node_1] = self.__calculate_distance_of_sub_tree(
+                            node_1,
+                            node_2,
+                            btrc,
+                            corres1,
+                            corres2,
+                            delta_tmp,
+                            self.norm_dict[norm],
+                            tree1.get_norm(node_1),
+                            tree2.get_norm(node_2),
+                        )
+                        colors[l_node_1] = colors[node_1]
+                        colors[node_2] = colors[node_1]
+                        colors[l_node_2] = colors[node_1]
+        if ax is None:
+            fig, ax = plt.subplots(nrows=1, ncols=2, sharey=True)
+        cmap = colormaps[colormap]
+        c_norm = mcolors.Normalize(0, 1)
+        colors = {c: cmap(c_norm(v)) for c, v in colors.items()}
+        self.plot_subtree(
+            self.get_ancestor_at_t(n1),
+            end_time=end_time,
+            size=size,
+            selected_nodes=matched_left,
+            color_of_nodes=colors,
+            selected_edges=matched_left,
+            color_of_edges=colors,
+            default_color=default_color,
+            lw=lw,
+            ax=ax[0],
+        )
+        self.plot_subtree(
+            self.get_ancestor_at_t(n2),
+            end_time=end_time,
+            size=size,
+            selected_nodes=matched_right,
+            color_of_nodes=colors,
+            selected_edges=matched_right,
+            color_of_edges=colors,
+            default_color=default_color,
+            lw=lw,
+            ax=ax[1],
+        )
+        return ax[0].get_figure(), ax
+
+    def labelled_mappings(
+        self,
+        n1: int,
+        n2: int,
+        end_time: int | None = None,
+        norm: Literal["max", "sum", None] = "max",
+        style: (
+            Literal["simple", "normalized_simple", "full", "downsampled"]
+            | type[TreeApproximationTemplate]
+        ) = "simple",
+        downsample: int = 2,
+    ) -> dict[str, list[str]]:
+        """
+        Returns the labels or IDs of all the nodes in the subtrees compared.
+
+
+        Parameters
+        ----------
+        n1 : int
+            id of the first node to compare
+        n2 : int
+            id of the second node to compare
+        end_time : int, optional
+            The final time point the comparison algorithm will take into account.
+            If None or not provided all nodes will be taken into account.
+        norm : {"max", "sum"}, default="max"
+            The normalization method to use, defaults to 'max'.
+        style : {"simple", "full", "downsampled", "normalized_simple} or TreeApproximationTemplate subclass, default="simple"
+            Which tree approximation is going to be used for the comparisons, defaults to 'simple'.
+        downsample : int, default=2
+            The downsample factor for the downsampled tree approximation.
+            Used only when `style="downsampled"`.
+
+        Returns
+        -------
+        dict mapping str to list of str
+            - 'matched' The labels of the matched nodes of the alignment.
+            - 'unmatched' The labels of the unmatched nodes of the alginment.
+        """
+        parameters = (
+            end_time,
+            convert_style_to_number(style=style, downsample=downsample),
+        )
+        n1, n2 = sorted([n1, n2])
+        self._comparisons.setdefault(parameters, {})
+        if self._comparisons[parameters].get((n1, n2)):
+            tmp = self._comparisons[parameters][(n1, n2)]
+        else:
+            tmp = self.__unordereded_backtrace(
+                n1, n2, end_time, norm, style, downsample
+            )
+        btrc = tmp["alignment"]
+        tree1, tree2 = tmp["trees"]
+
+        (
+            *_,
+            corres1,
+        ) = tree1.edist
+        (
+            *_,
+            corres2,
+        ) = tree2.edist
+
+        if norm not in self.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])
+        matched = []
+        unmatched = []
+        if style not in ("full", "downsampled"):
+            for m in btrc:
+                if m._left != -1 and m._right != -1:
+                    cyc1 = self.get_chain_of_node(corres1[m._left])
+                    if len(cyc1) > 1:
+                        node_1, *_ = cyc1
+                    elif len(cyc1) == 1:
+                        node_1 = cyc1.pop()
+                    cyc2 = self.get_chain_of_node(corres2[m._right])
+                    if len(cyc2) > 1:
+                        node_2, *_ = cyc2
+                    elif len(cyc2) == 1:
+                        node_2 = cyc2.pop()
+                    matched.append(
+                        (
+                            self.labels.get(node_1, node_1),
+                            self.labels.get(node_2, node_2),
+                        )
+                    )
+
+                else:
+                    if m._left != -1:
+                        node_1 = self.get_chain_of_node(
+                            corres1.get(m._left, "-")
+                        )[0]
+                    else:
+                        node_1 = self.get_chain_of_node(
+                            corres2.get(m._right, "-")
+                        )[0]
+                    unmatched.append(self.labels.get(node_1, node_1))
+        else:
+            for m in btrc:
+                if m._left != -1 and m._right != -1:
+                    node_1 = corres1[m._left]
+                    node_2 = corres2[m._right]
+                    matched.append(
+                        (
+                            self.labels.get(node_1, node_1),
+                            self.labels.get(node_2, node_2),
+                        )
+                    )
+                else:
+                    if m._left != -1:
+                        node_1 = corres1[m._left]
+                    else:
+                        node_1 = corres2[m._right]
+                    unmatched.append(self.labels.get(node_1, node_1))
+        return {"matched": matched, "unmatched": unmatched}
+
+    def unordered_tree_edit_distance(
+        self,
+        n1: int,
+        n2: int,
+        end_time: int | None = None,
+        norm: Literal["max", "sum", None] = "max",
+        style: (
+            Literal["simple", "normalized_simple", "full", "downsampled"]
+            | type[TreeApproximationTemplate]
+        ) = "simple",
+        downsample: int = 2,
+        return_norms: bool = False,
+    ) -> float | tuple[float, tuple[float, float]]:
+        """
+        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).
+        The distance is normed by the function norm that takes the two list of nodes
+        spawned by the trees `n1` and `n2`.
+
+        Parameters
+        ----------
+        n1 : int
+            id of the first node to compare
+        n2 : int
+            id of the second node to compare
+        end_time : int, optional
+            The final time point the comparison algorithm will take into account.
+            If None or not provided all nodes will be taken into account.
+        norm : {"max", "sum"}, default="max"
+            The normalization method to use, defaults to 'max'.
+        style : {"simple", "normalized_simple", "full", "downsampled"} or TreeApproximationTemplate subclass, default="simple"
+            Which tree approximation is going to be used for the comparisons.
+        downsample : int, default=2
+            The downsample factor for the downsampled tree approximation.
+            Used only when `style="downsampled"`.
+
+        Returns
+        -------
+        float
+            The normalized unordered tree edit distance between `n1` and `n2`
+        """
+        parameters = (
+            end_time,
+            convert_style_to_number(style=style, downsample=downsample),
+        )
+        n1, n2 = sorted([n1, n2])
+        self._comparisons.setdefault(parameters, {})
+        if self._comparisons[parameters].get((n1, n2)):
+            tmp = self._comparisons[parameters][(n1, n2)]
+        else:
+            tmp = self.__unordereded_backtrace(
+                n1, n2, end_time, norm, style, downsample
+            )
+        btrc = tmp["alignment"]
+        tree1, tree2 = tmp["trees"]
+        _, times1 = tree1.tree
+        _, times2 = tree2.tree
+        (
+            nodes1,
+            adj1,
+            corres1,
+        ) = tree1.edist
+        (
+            nodes2,
+            adj2,
+            corres2,
+        ) = tree2.edist
+        delta_tmp = partial(
+            tree1.delta,
+            corres1=corres1,
+            corres2=corres2,
+            times1=times1,
+            times2=times2,
+        )
+
+        if norm not in self.norm_dict:
+            raise ValueError(
+                "Select a viable normalization method (max, sum, None)"
+            )
+        cost = btrc.cost(nodes1, nodes2, delta_tmp)
+        norm_values = (tree1.get_norm(n1), tree2.get_norm(n2))
+        if return_norms:
+            return cost, norm_values
+        return cost / self.norm_dict[norm](norm_values)
 
     @staticmethod
     def __plot_nodes(
-        hier, selected_nodes, color, size, ax, default_color="black", **kwargs
-    ):
+        hier: dict,
+        selected_nodes: set,
+        color: str | dict | list,
+        size: int | float,
+        ax: plt.Axes,
+        default_color: str = "black",
+        **kwargs,
+    ) -> None:
         """
         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
-            )
+
+        if isinstance(color, dict):
+            color = [color.get(k, default_color) for k in hier]
+        elif isinstance(color, str | list):
+            color = [
+                color if node in selected_nodes else default_color
+                for node in hier
+            ]
+        hier_pos = np.array(list(hier.values()))
+        ax.scatter(*hier_pos.T, s=size, zorder=10, color=color, **kwargs)
 
     @staticmethod
     def __plot_edges(
-        hier,
-        lnks_tms,
-        selected_edges,
-        color,
-        ax,
-        default_color="black",
+        hier: dict,
+        lnks_tms: dict,
+        selected_edges: Iterable,
+        color: str | dict | list,
+        lw: float,
+        ax: plt.Axes,
+        default_color: str = "black",
         **kwargs,
-    ):
+    ) -> None:
         """
         Private method that plots the edges of the tree.
         """
-        x, y = [], []
+        if isinstance(color, dict):
+            selected_edges = color.keys()
+        lines = []
+        c = []
         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)
+            for suc in succs:
+                lines.append(
+                    [
+                        [hier[suc][0], hier[suc][1]],
+                        [hier[pred][0], hier[pred][1]],
+                    ]
+                )
+                if pred in selected_edges:
+                    if isinstance(color, str | list):
+                        c.append(color)
+                    elif isinstance(color, dict):
+                        c.append(color[pred])
+                else:
+                    c.append(default_color)
+        lc = LineCollection(lines, colors=c, linewidth=lw, **kwargs)
+        ax.add_collection(lc)
 
     def draw_tree_graph(
         self,
-        hier,
-        lnks_tms,
-        selected_nodes=None,
-        selected_edges=None,
-        color_of_nodes="magenta",
-        color_of_edges=None,
-        size=10,
-        ax=None,
-        figure=None,
-        default_color="black",
+        hier: dict[int, tuple[int, int]],
+        lnks_tms: dict[str, dict[int, list | int]],
+        selected_nodes: list | set | None = None,
+        selected_edges: list | set | None = None,
+        color_of_nodes: str | dict = "magenta",
+        color_of_edges: str | dict = "magenta",
+        size: int | float = 10,
+        lw: float = 0.3,
+        ax: plt.Axes | None = None,
+        default_color: str = "black",
         **kwargs,
-    ):
+    ) -> tuple[plt.Figure, plt.Axes]:
         """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.
+        Parameters
+        ----------
+        hier : dict mapping int to tuple of int
+            Dictionary that contains the positions of all nodes.
+        lnks_tms : dict mapping string to dictionaries mapping int to list or int
+            - 'links' : conatains the hierarchy of the nodes (only start and end of each chain)
+            - 'times' : contains the distance between the  start and the end of each chain.
+        selected_nodes : list or set, optional
+            Which nodes are to be selected (Painted with a different color, according to 'color_'of_nodes')
+        selected_edges : list or set, optional
+            Which edges are to be selected (Painted with a different color, according to 'color_'of_edges')
+        color_of_nodes : str, default="magenta"
+            Color of selected nodes
+        color_of_edges : str, default="magenta"
+            Color of selected edges
+        size : int, default=10
+            Size of the nodes, defaults to 10
+        lw : float, default=0.3
+            The width of the edges of the tree graph, defaults to 0.3
+        ax : plt.Axes, optional
+            Plot the graph on existing ax. If not provided or None a new ax is going to be created.
+        default_color : str, default="black"
+            Default color of nodes
+
+        Returns
+        -------
+        plt.Figure
+            The matplotlib figure
+        plt.Axes
+            The matplotlib ax
         """
         if selected_nodes is None:
             selected_nodes = []
         if selected_edges is None:
             selected_edges = []
         if ax is None:
-            figure, ax = plt.subplots()
+            _, ax = plt.subplots()
         else:
             ax.clear()
         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 0 < size:
+            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(
@@ -1934,62 +2405,107 @@ class lineageTree(lineageTreeLoaders):
             lnks_tms,
             selected_edges,
             color_of_edges,
+            lw,
             ax,
             default_color=default_color,
             **kwargs,
         )
+        ax.autoscale()
+        plt.draw()
         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):
+    def _create_dict_of_plots(
+        self,
+        node: int | Iterable[int] | None = None,
+        start_time: int | None = None,
+        end_time: int | None = None,
+    ) -> dict[int, dict]:
         """Generates a dictionary of graphs where the keys are the index of the graph and
-        the values are the graphs themselves which are produced by create_links_and _cycles
-
-        Args:
-            node (_type_, optional): The id of the node/nodes to produce the simple graphs. Defaults to None.
-            start_time (int, optional): Important only if there are no nodes it will produce the graph of every
-            root that starts before or at start time. Defaults to None.
-
-        Returns:
-            (dict): The keys are just index values  0-n and the values are the graphs produced.
+        the values are the graphs themselves which are produced by `create_links_and_chains`
+
+        Parameters
+        ----------
+        node : int or Iterable of int, optional
+            The id of the node/nodes to produce the simple graphs, if not provided or None will
+            calculate the dicts for every root that starts before 'start_time'
+        start_time : int, optional
+            Important only if there are no nodes it will produce the graph of every
+            root that starts before or at start time. If not provided or None the 'start_time' defaults to the start of the dataset.
+        end_time : int, optional
+            The last timepoint to be considered, if not provided or None the last timepoint of the
+            dataset (t_e) is considered.
+
+        Returns
+        -------
+        dict mapping int to dict
+            The keys are just index values 0-n and the values are the graphs produced.
         """
         if start_time is None:
             start_time = self.t_b
+        if end_time is None:
+            end_time = self.t_e
         if node is None:
             mothers = [
-                root for root in self.roots if self.time[root] <= start_time
+                root for root in self.roots if self._time[root] <= start_time
             ]
+        elif isinstance(node, Iterable):
+            mothers = node
         else:
-            mothers = node if isinstance(node, (list, set)) else [node]
+            mothers = [node]
         return {
-            i: create_links_and_cycles(self, mother)
+            i: create_links_and_chains(self, mother, end_time=end_time)
             for i, mother in enumerate(mothers)
         }
 
     def plot_all_lineages(
         self,
-        nodes: list = None,
-        last_time_point_to_consider: int = None,
-        nrows=2,
-        figsize=(10, 15),
-        dpi=100,
-        fontsize=15,
-        figure=None,
-        axes=None,
+        nodes: list | None = None,
+        last_time_point_to_consider: int | None = None,
+        nrows: int = 2,
+        figsize: tuple[int, int] = (10, 15),
+        dpi: int = 100,
+        fontsize: int = 15,
+        axes: plt.Axes | None = None,
+        vert_gap: int = 1,
         **kwargs,
-    ):
+    ) -> tuple[plt.Figure, plt.Axes, dict[plt.Axes, int]]:
         """Plots all lineages.
 
-        Args:
-            last_time_point_to_consider (int, optional): Which timepoints and upwards are the graphs to be plotted.
-                                                        For example if start_time is 10, then all trees that begin
-                                                        on tp 10 or before are calculated. Defaults to None, where
-                                                        it will plot all the roots that exist on self.t_b.
-            nrows (int):  How many rows of plots should be printed.
-            kwargs: args accepted by matplotlib
+        Parameters
+        ----------
+        nodes : list, optional
+            The nodes spawning the graphs to be plotted.
+        last_time_point_to_consider : int, optional
+            Which timepoints and upwards are the graphs to be plotted.
+            For example if start_time is 10, then all trees that begin
+            on tp 10 or before are calculated. Defaults to None, where
+            it will plot all the roots that exist on `self.t_b`.
+        nrows : int, default=2
+            How many rows of plots should be printed.
+        figsize : tuple, default=(10, 15)
+            The size of the figure.
+        dpi : int, default=100
+            The dpi of the figure.
+        fontsize : int, default=15
+            The fontsize of the labels.
+        axes : plt.Axes, optional
+            The axes to plot the graphs on. If None or not provided new axes are going to be created.
+        vert_gap : int, default=1
+            space between the nodes, defaults to 1
+        **kwargs:
+            kwargs accepted by matplotlib.pyplot.plot, matplotlib.pyplot.scatter
+
+        Returns
+        -------
+        plt.Figure
+            The figure
+        plt.Axes
+            The axes
+        dict of plt.Axes to int
+            A dictionary that maps the axes to the root of the tree.
         """
-
         nrows = int(nrows)
         if last_time_point_to_consider is None:
             last_time_point_to_consider = self.t_b
@@ -1998,22 +2514,33 @@ class lineageTree(lineageTreeLoaders):
             raise Warning("Number of rows has to be at least 1")
         if nodes:
             graphs = {
-                i: self.to_simple_graph(node) for i, node in enumerate(nodes)
+                i: self._create_dict_of_plots(node)
+                for i, node in enumerate(nodes)
             }
         else:
-            graphs = self.to_simple_graph(
+            graphs = self._create_dict_of_plots(
                 start_time=last_time_point_to_consider
             )
         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")
@@ -2051,127 +2578,149 @@ 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_subtree(
+        self,
+        node: int,
+        end_time: int | None = None,
+        figsize: tuple[int, int] = (4, 7),
+        dpi: int = 150,
+        vert_gap: int = 2,
+        selected_nodes: list | None = None,
+        selected_edges: list | None = None,
+        color_of_nodes: str | dict = "magenta",
+        color_of_edges: str | dict = "magenta",
+        size: int | float = 10,
+        lw: float = 0.1,
+        default_color: str = "black",
+        ax: plt.Axes | None = None,
+    ) -> tuple[plt.Figure, plt.Axes]:
         """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 matplotlib
-        """
-        graph = self.to_simple_graph(node)
+        Parameters
+        ----------
+        node : int
+            The id of the node that is going to be plotted.
+        end_time : int, optional
+            The last timepoint to be considered, if None or not provided the last timepoint of the dataset (t_e) is considered.
+        figsize : tuple of 2 ints, default=(4,7)
+            The size of the figure, deafults to (4,7)
+        vert_gap : int, default=2
+            The verical gap of a node when it divides, defaults to 2.
+        dpi : int, default=150
+            The dpi of the figure, defaults to 150
+        selected_nodes : list, optional
+            The nodes that are selected by the user to be colored in a different color, defaults to None
+        selected_edges : list, optional
+            The edges that are selected by the user to be colored in a different color, defaults to None
+        color_of_nodes : str, default="magenta"
+            The color of the nodes to be colored, except the default colored ones, defaults to "magenta"
+        color_of_edges : str, default="magenta"
+            The color of the edges to be colored, except the default colored ones, defaults to "magenta"
+        size : int, default=10
+            The size of the nodes, defaults to 10
+        lw : float, default=0.1
+            The widthe of the edges of the tree graph, defaults to 0.1
+        default_color : str, default="black"
+            The default color of nodes and edges, defaults to "black"
+        ax : plt.Axes, optional
+            The ax where the plot is going to be applied, if not provided or None new axes will be created.
+
+        Returns
+        -------
+        plt.Figure
+            The matplotlib figure
+        plt.Axes
+            The matplotlib axes
+
+        Raises
+        ------
+        Warning
+            If more than one nodes are received
+        """
+        graph = self._create_dict_of_plots(node, end_time=end_time)
         if len(graph) > 1:
-            raise Warning("Please enter only one node")
+            raise Warning(
+                "Please use lT.plot_all_lineages(nodes) for plotting multiple nodes."
+            )
         graph = graph[0]
-        figure, ax = plt.subplots(nrows=1, ncols=1, figsize=figsize, dpi=dpi)
+        if not ax:
+            _, ax = plt.subplots(nrows=1, ncols=1, figsize=figsize, dpi=dpi)
         self.draw_tree_graph(
             hier=hierarchical_pos(
                 graph,
                 graph["root"],
                 vert_gap=vert_gap,
-                ycenter=-int(self.time[node]),
+                ycenter=-int(self._time[node]),
             ),
+            selected_edges=selected_edges,
+            selected_nodes=selected_nodes,
+            color_of_edges=color_of_edges,
+            color_of_nodes=color_of_nodes,
+            default_color=default_color,
+            size=size,
+            lw=lw,
             lnks_tms=graph,
             ax=ax,
-            **kwargs,
         )
-        return figure, ax
-
-    # def DTW(self, t1, t2, max_w=None, start_delay=None, end_delay=None,
-    #         metric='euclidian', **kwargs):
-    #     """ Computes the dynamic time warping distance between the tracks t1 and t2
-
-    #         Args:
-    #             t1 ([int, ]): list of node ids for the first track
-    #             t2 ([int, ]): list of node ids for the second track
-    #             w (int): maximum wapring allowed (default infinite),
-    #                      if w=1 then the DTW is the distance between t1 and t2
-    #             start_delay (int): maximum number of time points that can be
-    #                                skipped at the beginning of the track
-    #             end_delay (int): minimum number of time points that can be
-    #                              skipped at the beginning of the track
-    #             metric (str): str or callable, optional The distance metric to use.
-    #                           Default='euclidean'. Refer to the documentation for
-    #                           scipy.spatial.distance.cdist. Some examples:
-    #                           'braycurtis', 'canberra', 'chebyshev', 'cityblock', 'correlation',
-    #                           'cosine', 'dice', 'euclidean', 'hamming', 'jaccard', 'kulsinski',
-    #                           'mahalanobis', 'matching', 'minkowski', 'rogerstanimoto', 'russellrao',
-    #                           'seuclidean', 'sokalmichener', 'sokalsneath', 'sqeuclidean',
-    #                           'wminkowski', 'yule'
-    #             **kwargs (dict): Extra arguments to `metric`: refer to each metric
-    #                              documentation in scipy.spatial.distance (optional)
-
-    #         Returns:
-    #             float: the dynamic time warping distance between the two tracks
-    #     """
-    #     from scipy.sparse import
-    #     pos_t1 = [self.pos[ti] for ti in t1]
-    #     pos_t2 = [self.pos[ti] for ti in t2]
-    #     distance_matrix = np.zeros((len(t1), len(t2))) + np.inf
-
-    #     c = distance.cdist(exp_data, num_data, metric=metric, **kwargs)
-
-    #     d = np.zeros(c.shape)
-    #     d[0, 0] = c[0, 0]
-    #     n, m = c.shape
-    #     for i in range(1, n):
-    #         d[i, 0] = d[i-1, 0] + c[i, 0]
-    #     for j in range(1, m):
-    #         d[0, j] = d[0, j-1] + c[0, j]
-    #     for i in range(1, n):
-    #         for j in range(1, m):
-    #             d[i, j] = c[i, j] + min((d[i-1, j], d[i, j-1], d[i-1, j-1]))
-    #     return d[-1, -1], d
-
-    def __getitem__(self, item):
-        if isinstance(item, str):
-            return self.__dict__[item]
-        elif np.issubdtype(type(item), np.integer):
-            return self.successor.get(item, [])
-        else:
-            raise KeyError(
-                "Only integer or string are valid key for lineageTree"
-            )
+        return ax.get_figure(), ax
 
-    def get_cells_at_t_from_root(self, r: [int, list], t: int = None) -> list:
+    def nodes_at_t(
+        self,
+        t: int,
+        r: int | Iterable[int] | None = None,
+    ) -> list:
         """
-        Returns the list of cells at time `t` that are spawn by the node(s) `r`.
+        Returns the list of nodes at time `t` that are spawn by the node(s) `r`.
 
-            Args:
-                r (int | list): id or list of ids of the spawning node
-                t (int): target time, if None goes as far as possible
-                        (default None)
+        Parameters
+        ----------
+        t : int
+            target time, if `None` goes as far as possible
+        r : int or Iterable of int, optional
+            id or list of ids of the spawning node
 
-            Returns:
-                (list) list of nodes at time `t` spawned by `r`
+        Returns
+        -------
+        list
+            list of nodes at time `t` spawned by `r`
         """
-        if not isinstance(r, list):
+        if not r and r != 0:
+            r = {root for root in self.roots if self.time[root] <= t}
+        if isinstance(r, int):
             r = [r]
+        if t is None:
+            t = self.t_e
         to_do = list(r)
         final_nodes = []
         while len(to_do) > 0:
             curr = to_do.pop()
-            for _next in self[curr]:
-                if self.time[_next] < t:
+            for _next in self._successor[curr]:
+                if self._time[_next] < t:
                     to_do.append(_next)
-                elif self.time[_next] == t:
+                elif self._time[_next] == t:
                     final_nodes.append(_next)
         if not final_nodes:
             return list(r)
         return final_nodes
 
     @staticmethod
-    def __calculate_diag_line(dist_mat: np.ndarray) -> (float, float):
+    def __calculate_diag_line(dist_mat: np.ndarray) -> tuple[float, float]:
         """
         Calculate the line that centers the band w.
 
-            Args:
-                dist_mat (matrix): distance matrix obtained by the function calculate_dtw
+        Parameters
+        ----------
+        dist_mat : np.ndarray
+            distance matrix obtained by the function calculate_dtw
 
-            Returns:
-                (float) Slope
-                (float) intercept of the line
+        Returns
+        -------
+        float
+            The slope of the curve
+        float
+            The intercept of the curve
         """
         i, j = dist_mat.shape
         x1 = max(0, i - j) / 2
@@ -2191,22 +2740,33 @@ class lineageTree(lineageTreeLoaders):
         fast: bool = False,
         w: int = 0,
         centered_band: bool = True,
-    ) -> (((int, int), ...), np.ndarray):
+    ) -> tuple[list[int], np.ndarray, float]:
         """
         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
+        Parameters
+        ----------
+        dist_mat : np.ndarray
+            distance matrix obtained by the function calculate_dtw
+        start_d : int, default=0
+            start delay
+        back_d : int, default=0
+            end delay
+        fast : bool, default=False
+            if `True`, the algorithm will use a faster version but might not find the optimal alignment
+        w : int, default=0
+            window constrain
+        centered_band : bool, default=True
+            if `True`, the band will be centered around the diagonal
+
+        Returns
+        -------
+        tuple of tuples of int
+            Aligment path
+        np.ndarray
+            cost matrix
+        float
+            optimal cost
         """
         N, M = dist_mat.shape
         w_limit = max(w, abs(N - M))  # Calculate the Sakoe-Chiba band width
@@ -2327,7 +2887,6 @@ class lineageTree(lineageTreeLoaders):
 
         # 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
 
@@ -2336,52 +2895,59 @@ class lineageTree(lineageTreeLoaders):
         return R, t
 
     def __interpolate(
-        self, track1: list, track2: list, threshold: int
-    ) -> (np.ndarray, np.ndarray):
+        self, chain1: list, chain2: list, threshold: int
+    ) -> tuple[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
+        Parameters
+        ----------
+        chain1 : list of int
+            list of nodes of the first chain to compare
+        chain2 : list of int
+            list of nodes of the second chain to compare
+        threshold : int
+            set a maximum number of points a chain can have
+
+        Returns
+        -------
+        list of np.ndarray
+            `x`, `y`, `z` postions for `chain1`
+        list of np.ndarray
+            `x`, `y`, `z` postions for `chain2`
         """
         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])
+        chain1_pos = np.array([self.pos[c_id] for c_id in chain1])
+        chain2_pos = np.array([self.pos[c_id] for c_id in chain2])
 
-        # 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
+        # Both chains have the same length and size below the threshold - nothing is done
+        if len(chain1) == len(chain2) and (
+            len(chain1) <= threshold or len(chain2) <= 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:
+            return chain1_pos, chain2_pos
+        # Both chains have the same length but one or more sizes are above the threshold
+        elif len(chain1) > threshold or len(chain2) > threshold:
             sampling = threshold
-        # Tracks have different lengths and the sizes are below the threshold
+        # chains have different lengths and the sizes are below the threshold
         else:
-            sampling = max(len(track1), len(track2))
+            sampling = max(len(chain1), len(chain2))
 
         for pos in range(3):
-            track1_interp = InterpolatedUnivariateSpline(
-                np.linspace(0, 1, len(track1_pos[:, pos])),
-                track1_pos[:, pos],
+            chain1_interp = InterpolatedUnivariateSpline(
+                np.linspace(0, 1, len(chain1_pos[:, pos])),
+                chain1_pos[:, pos],
                 k=1,
             )
-            inter1_pos.append(track1_interp(np.linspace(0, 1, sampling)))
+            inter1_pos.append(chain1_interp(np.linspace(0, 1, sampling)))
 
-            track2_interp = InterpolatedUnivariateSpline(
-                np.linspace(0, 1, len(track2_pos[:, pos])),
-                track2_pos[:, pos],
+            chain2_interp = InterpolatedUnivariateSpline(
+                np.linspace(0, 1, len(chain2_pos[:, pos])),
+                chain2_pos[:, pos],
                 k=1,
             )
-            inter2_pos.append(track2_interp(np.linspace(0, 1, sampling)))
+            inter2_pos.append(chain2_interp(np.linspace(0, 1, sampling)))
 
         return np.column_stack(inter1_pos), np.column_stack(inter2_pos)
 
@@ -2397,46 +2963,66 @@ class lineageTree(lineageTreeLoaders):
         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
+    ) -> (
+        tuple[float, tuple, np.ndarray, np.ndarray, np.ndarray]
+        | tuple[float, tuple]
+    ):
         """
-        nodes1_cycle = self.get_cycle(nodes1)
-        nodes2_cycle = self.get_cycle(nodes2)
-
-        interp_cycle1, interp_cycle2 = self.__interpolate(
-            nodes1_cycle, nodes2_cycle, threshold
+        Calculate DTW distance between two chains
+
+        Parameters
+        ----------
+        nodes1 : int
+            node to compare distance
+        nodes2 : int
+            node to compare distance
+        threshold : int, default=1000
+            set a maximum number of points a chain can have
+        regist : bool, default=True
+            Rotate and translate trajectories
+        start_d : int, default=0
+            start delay
+        back_d : int, default=0
+            end delay
+        fast : bool, default=False
+            if `True`, the algorithm will use a faster version but might not find the optimal alignment
+        w : int, default=0
+            window size
+        centered_band : bool, default=True
+            when running the fast algorithm, `True` if the windown is centered
+        cost_mat_p : bool, default=False
+            True if print the not normalized cost matrix
+
+        Returns
+        -------
+        float
+            DTW distance
+        tuple of tuples
+            Aligment path
+        matrix
+            Cost matrix
+        list of lists
+            rotated and translated trajectories positions
+        list of lists
+            rotated and translated trajectories positions
+        """
+        nodes1_chain = self.get_chain_of_node(nodes1)
+        nodes2_chain = self.get_chain_of_node(nodes2)
+
+        interp_chain1, interp_chain2 = self.__interpolate(
+            nodes1_chain, nodes2_chain, 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])
+        pos_chain1 = np.array([self.pos[c_id] for c_id in nodes1_chain])
+        pos_chain2 = np.array([self.pos[c_id] for c_id in nodes2_chain])
 
         if regist:
             R, t = self.__rigid_transform_3D(
-                np.transpose(interp_cycle1), np.transpose(interp_cycle2)
+                np.transpose(interp_chain1), np.transpose(interp_chain2)
             )
-            pos_cycle1 = np.transpose(np.dot(R, pos_cycle1.T) + t)
+            pos_chain1 = np.transpose(np.dot(R, pos_chain1.T) + t)
 
-        dist_mat = distance.cdist(pos_cycle1, pos_cycle2, "euclidean")
+        dist_mat = distance.cdist(pos_chain1, pos_chain2, "euclidean")
 
         path, cost_mat, final_cost = self.__dp(
             dist_mat,
@@ -2449,7 +3035,7 @@ class lineageTree(lineageTreeLoaders):
         cost = final_cost / len(path)
 
         if cost_mat_p:
-            return cost, path, cost_mat, pos_cycle1, pos_cycle2
+            return cost, path, cost_mat, pos_chain1, pos_chain2
         else:
             return cost, path
 
@@ -2464,24 +3050,39 @@ class lineageTree(lineageTreeLoaders):
         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(
+    ) -> tuple[float, plt.Figure]:
+        """
+        Plot DTW cost matrix between two chains in heatmap format
+
+        Parameters
+        ----------
+        nodes1 : int
+            node to compare distance
+        nodes2 : int
+            node to compare distance
+        threshold : int, default=1000
+            set a maximum number of points a chain can have
+        regist : bool, default=True
+            Rotate and translate trajectories
+        start_d : int, default=0
+            start delay
+        back_d : int, default=0
+            end delay
+        fast : bool, default=False
+            if `True`, the algorithm will use a faster version but might not find the optimal alignment
+        w : int, default=0
+            window size
+        centered_band : bool, default=True
+            when running the fast algorithm, `True` if the windown is centered
+
+        Returns
+        -------
+        float
+            DTW distance
+        plt.Figure
+            Heatmap of cost matrix with opitimal path
+        """
+        cost, path, cost_mat, pos_chain1, pos_chain2 = self.calculate_dtw(
             nodes1,
             nodes2,
             threshold,
@@ -2503,32 +3104,32 @@ class lineageTree(lineageTreeLoaders):
         ax.set_title("Heatmap of DTW Cost Matrix")
         ax.set_xlabel("Tree 1")
         ax.set_ylabel("tree 2")
-        x_path, y_path = zip(*path)
+        x_path, y_path = zip(*path, strict=True)
         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,
-    ):
+        pos_chain1: np.ndarray,
+        pos_chain2: np.ndarray,
+        nodes1: list[int],
+        nodes2: list[int],
+        ax: plt.Axes,
+        x_idx: list[int],
+        y_idx: list[int],
+        x_label: str,
+        y_label: str,
+    ) -> None:
         ax.plot(
-            pos_cycle1[:, x_idx],
-            pos_cycle1[:, y_idx],
+            pos_chain1[:, x_idx],
+            pos_chain1[:, y_idx],
             "-",
             label=f"root = {nodes1}",
         )
         ax.plot(
-            pos_cycle2[:, x_idx],
-            pos_cycle2[:, y_idx],
+            pos_chain2[:, x_idx],
+            pos_chain2[:, y_idx],
             "-",
             label=f"root = {nodes2}",
         )
@@ -2546,40 +3147,55 @@ class lineageTree(lineageTreeLoaders):
         fast: bool = False,
         w: int = 0,
         centered_band: bool = True,
-        projection: str = None,
+        projection: Literal["3d", "xy", "xz", "yz", "pca", None] = 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
+    ) -> tuple[float, plt.Figure]:
+        """
+        Plots DTW trajectories aligment between two chains in 2D or 3D
+
+        Parameters
+        ----------
+        nodes1 : int
+            node to compare distance
+        nodes2 : int
+            node to compare distance
+        threshold : int, default=1000
+            set a maximum number of points a chain can have
+        regist : bool, default=True
+            Rotate and translate trajectories
+        start_d : int, default=0
+            start delay
+        back_d : int, default=0
+            end delay
+        w : int, default=0
+            window size
+        fast : bool, default=False
+            True if the user wants to run the fast algorithm with window restrains
+        centered_band : bool, default=True
+            if running the fast algorithm, True if the windown is centered
+        projection : {"3d", "xy", "xz", "yz", "pca"}, optional
+            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 : bool
+            True to show alignment on plot
+
+        Returns
+        -------
+        float
+            DTW distance
+        figure
+            Trajectories Plot
         """
         (
             distance,
             alignment,
             cost_mat,
-            pos_cycle1,
-            pos_cycle2,
+            pos_chain1,
+            pos_chain2,
         ) = self.calculate_dtw(
             nodes1,
             nodes2,
@@ -2602,16 +3218,16 @@ class lineageTree(lineageTreeLoaders):
 
         if projection == "3d":
             ax.plot(
-                pos_cycle1[:, 0],
-                pos_cycle1[:, 1],
-                pos_cycle1[:, 2],
+                pos_chain1[:, 0],
+                pos_chain1[:, 1],
+                pos_chain1[:, 2],
                 "-",
                 label=f"root = {nodes1}",
             )
             ax.plot(
-                pos_cycle2[:, 0],
-                pos_cycle2[:, 1],
-                pos_cycle2[:, 2],
+                pos_chain2[:, 0],
+                pos_chain2[:, 1],
+                pos_chain2[:, 2],
                 "-",
                 label=f"root = {nodes2}",
             )
@@ -2621,8 +3237,8 @@ class lineageTree(lineageTreeLoaders):
         else:
             if projection == "xy" or projection == "yx" or projection is None:
                 self.__plot_2d(
-                    pos_cycle1,
-                    pos_cycle2,
+                    pos_chain1,
+                    pos_chain2,
                     nodes1,
                     nodes2,
                     ax,
@@ -2633,8 +3249,8 @@ class lineageTree(lineageTreeLoaders):
                 )
             elif projection == "xz" or projection == "zx":
                 self.__plot_2d(
-                    pos_cycle1,
-                    pos_cycle2,
+                    pos_chain1,
+                    pos_chain2,
                     nodes1,
                     nodes2,
                     ax,
@@ -2645,8 +3261,8 @@ class lineageTree(lineageTreeLoaders):
                 )
             elif projection == "yz" or projection == "zy":
                 self.__plot_2d(
-                    pos_cycle1,
-                    pos_cycle2,
+                    pos_chain1,
+                    pos_chain2,
                     nodes1,
                     nodes2,
                     ax,
@@ -2660,24 +3276,25 @@ class lineageTree(lineageTreeLoaders):
                     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"
+                        "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)
+                pca.fit(np.vstack([pos_chain1, pos_chain2]))
+                pos_chain1_2d = pca.transform(pos_chain1)
+                pos_chain2_2d = pca.transform(pos_chain2)
 
                 ax.plot(
-                    pos_cycle1_2d[:, 0],
-                    pos_cycle1_2d[:, 1],
+                    pos_chain1_2d[:, 0],
+                    pos_chain1_2d[:, 1],
                     "-",
                     label=f"root = {nodes1}",
                 )
                 ax.plot(
-                    pos_cycle2_2d[:, 0],
-                    pos_cycle2_2d[:, 1],
+                    pos_chain2_2d[:, 0],
+                    pos_chain2_2d[:, 1],
                     "-",
                     label=f"root = {nodes2}",
                 )
@@ -2706,7 +3323,7 @@ class lineageTree(lineageTreeLoaders):
                         'pca' : PCA projection"""
                 )
 
-        connections = [[pos_cycle1[i], pos_cycle2[j]] for i, j in alignment]
+        connections = [[pos_chain1[i], pos_chain2[j]] for i, j in alignment]
 
         for connection in connections:
             xyz1 = connection[0]
@@ -2729,120 +3346,179 @@ class lineageTree(lineageTreeLoaders):
             warnings.warn(
                 "Error: not possible to show alignment in PCA projection !",
                 UserWarning,
+                stacklevel=2,
             )
 
         return distance, fig
 
-    def first_labelling(self):
-        self.labels = {i: "Unlabeled" for i in self.time_nodes[0]}
-
     def __init__(
         self,
-        file_format: str = None,
-        tb: int = None,
-        te: int = None,
-        z_mult: float = 1.0,
-        file_type: str = "",
-        delim: str = ",",
-        eigen: bool = False,
-        shape: tuple = None,
-        raw_size: tuple = None,
-        reorder: bool = False,
-        xml_attributes: tuple = None,
-        name: str = None,
-        time_resolution: Union[int, None] = None,
+        *,
+        successor: dict[int, Sequence] | None = None,
+        predecessor: dict[int, int | Sequence] | None = None,
+        time: dict[int, int] | None = None,
+        starting_time: int | None = None,
+        pos: dict[int, Iterable] | None = None,
+        name: str | None = None,
+        root_leaf_value: Sequence | None = None,
+        **kwargs,
     ):
-        """
-        TODO: complete the doc
-        Main library to build tree graph representation of lineage tree data
-        It can read TGMM, ASTEC, SVF, MaMuT and TrackMate outputs.
-
-        Args:
-            file_format (str): either - path format to TGMM xmls
-                                      - path to the MaMuT xml
-                                      - path to the binary file
-            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.
-        """
+        """Create a lineageTree object from minimal information, without reading from a file.
+        Either `successor` or `predecessor` should be specified.
+
+        Parameters
+        ----------
+        successor : dict mapping int to Iterable
+            Dictionary assigning nodes to their successors.
+        predecessor : dict mapping int to int or Iterable
+            Dictionary assigning nodes to their predecessors.
+        time : dict mapping int to int, optional
+            Dictionary assigning nodes to the time point they were recorded to.
+            Defaults to None, in which case all times are set to `starting_time`.
+        starting_time : int, optional
+            Starting time of the lineage tree. Defaults to 0.
+        pos : dict mapping int to Iterable, optional
+            Dictionary assigning nodes to their positions. Defaults to None.
+        name : str, optional
+            Name of the lineage tree. Defaults to None.
+        root_leaf_value : Iterable, optional
+            Iterable of values of roots' predecessors and leaves' successors in the successor and predecessor dictionaries.
+            Defaults are `[None, (), [], set()]`.
+        **kwargs:
+            Supported keyword arguments are dictionaries assigning nodes to any custom property.
+            The property must be specified for every node, and named differently from lineageTree's own attributes.
+        """
+        self.__version__ = importlib.metadata.version("LineageTree")
+        self.name = str(name) if name is not None else None
+        if successor is not None and predecessor is not None:
+            raise ValueError(
+                "You cannot have both successors and predecessors."
+            )
 
-        self.name = name
-        self.time_nodes = {}
-        self.time_edges = {}
-        self.max_id = -1
-        self.next_id = []
-        self.nodes = set()
-        self.successor = {}
-        self.predecessor = {}
-        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:
-            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,
+        if root_leaf_value is None:
+            root_leaf_value = [None, (), [], set()]
+        elif not isinstance(root_leaf_value, Iterable):
+            raise TypeError(
+                f"root_leaf_value is of type {type(root_leaf_value)}, expected Iterable."
+            )
+        elif len(root_leaf_value) < 1:
+            raise ValueError(
+                "root_leaf_value should have at least one element."
+            )
+        self._successor = {}
+        self._predecessor = {}
+        if successor is not None:
+            for pred, succs in successor.items():
+                if succs in root_leaf_value:
+                    self._successor[pred] = ()
+                else:
+                    if not isinstance(succs, Iterable):
+                        raise TypeError(
+                            f"Successors should be Iterable, got {type(succs)}."
+                        )
+                    if len(succs) == 0:
+                        raise ValueError(
+                            f"{succs} was not declared as a leaf but was found as a successor.\n"
+                            "Please lift the ambiguity."
+                        )
+                    self._successor[pred] = tuple(succs)
+                    for succ in succs:
+                        if succ in self._predecessor:
+                            raise ValueError(
+                                "Node can have at most one predecessor."
+                            )
+                        self._predecessor[succ] = (pred,)
+        elif predecessor is not None:
+            for succ, pred in predecessor.items():
+                if pred in root_leaf_value:
+                    self._predecessor[succ] = ()
+                else:
+                    if isinstance(pred, Sequence):
+                        if len(pred) == 0:
+                            raise ValueError(
+                                f"{pred} was not declared as a leaf but was found as a successor.\n"
+                                "Please lift the ambiguity."
+                            )
+                        if 1 < len(pred):
+                            raise ValueError(
+                                "Node can have at most one predecessor."
+                            )
+                        pred = pred[0]
+                    self._predecessor[succ] = (pred,)
+                    self._successor.setdefault(pred, ())
+                    self._successor[pred] += (succ,)
+        for root in set(self._successor).difference(self._predecessor):
+            self._predecessor[root] = ()
+        for leaf in set(self._predecessor).difference(self._successor):
+            self._successor[leaf] = ()
+
+        if self.__check_for_cycles():
+            raise ValueError(
+                "Cycles were found in the tree, there should not be any."
+            )
+
+        if pos is None:
+            self.pos = {}
+        else:
+            if self.nodes.difference(pos) != set():
+                raise ValueError("Please provide the position of all nodes.")
+            self.pos = {
+                node: np.array(position) for node, position in pos.items()
+            }
+
+        if time is None:
+            if starting_time is None:
+                starting_time = 0
+            if not isinstance(starting_time, int):
+                warnings.warn(
+                    f"Attribute `starting_time` was a `{type(starting_time)}`, has been casted as an `int`.",
+                    stacklevel=2,
                 )
-            elif file_type == "mastodon":
-                if isinstance(file_format, list) and len(file_format) == 2:
-                    self.read_from_mastodon_csv(file_format)
+            self._time = {node: starting_time for node in self.roots}
+            queue = list(self.roots)
+            for node in queue:
+                for succ in self._successor[node]:
+                    self._time[succ] = self._time[node] + 1
+                    queue.append(succ)
+        else:
+            if starting_time is not None:
+                warnings.warn(
+                    "Both `time` and `starting_time` were provided, `starting_time` was ignored.",
+                    stacklevel=2,
+                )
+            self._time = {n: int(time[n]) for n in self.nodes}
+            if self._time != time:
+                if len(self._time) != len(time):
+                    warnings.warn(
+                        "The provided `time` dictionary had keys that were not nodes. "
+                        "They have been removed",
+                        stacklevel=2,
+                    )
                 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)
+                    warnings.warn(
+                        "The provided `time` dictionary had values that were not `int`. "
+                        "These values have been truncated and converted to `int`",
+                        stacklevel=2,
+                    )
+            if self.nodes.symmetric_difference(self._time) != set():
+                raise ValueError(
+                    "Please provide the time of all nodes and only existing nodes."
+                )
+            if not all(
+                self._time[node] < self._time[s]
+                for node, succ in self._successor.items()
+                for s in succ
+            ):
+                raise ValueError(
+                    "Provided times are not strictly increasing. Setting times to default."
+                )
+        # custom properties
+        for name, d in kwargs.items():
+            if name in self.__dict__:
+                warnings.warn(
+                    f"Attribute name {name} is reserved.", stacklevel=2
+                )
+                continue
+            setattr(self, name, d)
+        if not hasattr(self, "_comparisons"):
+            self._comparisons = {}