From 6de00ef4b847baa3a7dfd533965fe4875eee0e34 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Mon, 8 Jun 2026 15:50:00 +0200
Subject: [PATCH 01/14] progress

---
 src/lineagetree/distance/approximations.py    | 554 ++++++++++++++++++
 src/lineagetree/distance/comparator.py        | 112 ++++
 src/lineagetree/distance/delta.py             |  57 ++
 .../distance/distance_calculator.py           | 170 ++++++
 4 files changed, 893 insertions(+)
 create mode 100644 src/lineagetree/distance/approximations.py
 create mode 100644 src/lineagetree/distance/comparator.py
 create mode 100644 src/lineagetree/distance/delta.py
 create mode 100644 src/lineagetree/distance/distance_calculator.py

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
new file mode 100644
index 0000000..ea69da7
--- /dev/null
+++ b/src/lineagetree/distance/approximations.py
@@ -0,0 +1,554 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from collections import deque
+from typing import TYPE_CHECKING, Iterable, Callable
+import numpy as np
+from scipy.interpolate import InterpolatedUnivariateSpline
+from hashlib import sha256
+import pickle
+
+from edist import uted
+
+from .delta import (
+    delta_normalized_difference,
+    delta_nd_norm,
+    delta_difference,
+    delta_binary,
+)
+from warnings import warn
+
+if TYPE_CHECKING:
+    from lineagetree import LineageTree
+    from edist.alignment import Alignment
+    from .approximations import TreeApproximationTemplate
+    from .distance_calculator import TreeDistanceTemplate
+
+
+@dataclass
+class ApproximatedTree:
+    adjacency_dict: dict[int, Iterable[int]]
+    property_dict: dict[int, float | int | list[int | float]] | None
+    tree_specs: tuple | None = None
+
+    def _edist_format(
+        self,
+    ) -> tuple[list, list[list], dict[int, int]]:
+        """
+        A function that transforms an adjacency list into
+        a datastructure understandable by uted
+
+        Returns
+        -------
+        nodes : list of int
+            The list of node ids
+        adj_list: list of list of int
+            The adjacency list where `nodes[adj_list[nodes[0]][0]]`
+            gives you the id of one of the successors of `nodes[0]`.
+        list2nid : dict of int to int
+            A dictionary that maps the new node ids (in `nodes`) to
+            the original ones (from the LineageTree)
+        """
+        if self.property_dict:
+            sample = next(iter(self.property_dict.values()))
+            dict_type = type(sample)
+            if isinstance(sample, Iterable):
+                length = len(sample)
+                default_value = [0] * length
+            elif isinstance(dict_type(), int or float):
+                default_value = 0
+
+        inv_adj = {vi: k for k, v in self.adjacency_dict.items() for vi in v}
+        roots = set(self.adjacency_dict).difference(inv_adj)
+        nid2list = {}
+        list2nid = {}
+        nodes = []
+        adj_list = []
+        curr_id = 0
+        for r in roots:
+            to_do = deque([r])
+            while to_do:
+                curr = to_do.popleft()
+                nid2list[curr] = curr_id
+                list2nid[curr_id] = curr
+                if self.property_dict:
+                    nodes.append(self.property_dict.get(curr, default_value))
+                else:
+                    nodes.append(0)
+                to_do.extendleft(reversed(self.adjacency_dict.get(curr, [])))
+                curr_id += 1
+            adj_list = [
+                [
+                    nid2list[d]
+                    for d in self.adjacency_dict.get(list2nid[_id], [])
+                ]
+                for _id in range(len(nodes))
+            ]
+        return nodes, adj_list, list2nid
+
+    def __post_init__(self):
+        if (
+            self.property_dict is not None
+            and len(set(self.adjacency_dict).difference(self.property_dict))
+            != 0
+        ):
+            raise ValueError(
+                "Mismatch between adjacency_dict keys and property_dict keys.\n"
+                "All nodes have to have a property."
+            )
+        elif self.property_dict is None:
+            self.property_dict = {}
+        self.nodes, self.adjacency_list, self.correspondency = (
+            self._edist_format()
+        )
+
+    def __str__(self):  # For quickly checking how the object was created.
+        return f"Subtree: {self.tree_specs}"
+
+
+class TreeApproximationTemplate(ABC):
+    default_delta = ...
+
+    def __init__(
+        self,
+        delta: Callable | None = None,
+    ):
+        self.delta = delta if delta else self.default_delta
+
+    @abstractmethod
+    def approximation(
+        self,
+        lT: LineageTree,
+        root: int,
+        end_time: int | None = None,
+    ) -> ApproximatedTree: ...
+
+
+class ReducedTreeTimed(TreeApproximationTemplate):
+    default_delta = delta_normalized_difference
+
+    def approximation(
+        self,
+        lT: LineageTree,
+        root: int,
+        end_time: int | None = None,
+    ) -> ApproximatedTree:
+        if end_time is None:
+            end_time = lT.t_e
+        if lT.time_resolution == 0:
+            raise Warning(
+                "Time resolution of `LineageTree` object cannot be `0`."
+            )
+        if lT.time_resolution is None:
+            time_resolution = 1
+        else:
+            time_resolution = lT.time_resolution
+        out_dict = {}
+        final_properties = {}
+        to_do = [root]
+        while to_do:
+            current = to_do.pop()
+            cycle = np.array(lT.get_successors(current))
+            cycle_times = np.array([lT.time[c] for c in cycle])
+            cycle = cycle[cycle_times <= end_time]
+            if cycle.size:
+                _next = list(lT.successor[cycle[-1]])
+                if len(_next) > 1 and lT.time[cycle[-1]] < end_time:
+                    out_dict[current] = _next
+                    to_do.extend(_next)
+                else:
+                    out_dict[current] = []
+            final_properties[current] = len(cycle) * (time_resolution)
+        return ApproximatedTree(
+            out_dict, final_properties, (lT.name, root, end_time)
+        )
+
+
+class ReducedTreeProperties(TreeApproximationTemplate):
+
+    default_delta = delta_nd_norm
+
+    def __init__(self, aggregator: Callable = np.nanmean, delta=None):
+        super().__init__(delta)
+        self.aggregator = aggregator
+
+    def approximation(
+        self,
+        lT,
+        root,
+        end_time=None,
+        properties: dict[int, float | list] | list[str] = None,
+    ):
+        """Build the approximation of a given sub lineage of lineage tree.
+
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        end_time : int
+            The last time point to consider
+        time_resolution : float
+            How much time happens between two consecutive time points
+            This is useful when comparing two different lineage trees
+            that do not have the same time scale
+        properties : dict mapping int to float, optional
+            The dictionary that maps a node id to a given value.
+            *Not* all the nodes need to have a value.
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
+
+        if end_time is None:
+            end_time = lT.t_e
+        if lT.time_resolution == 0:
+            raise Warning(
+                "Time resolution of `LineageTree` object cannot be `0`."
+            )
+        out_dict = {}
+        final_properties = {}
+        if isinstance(properties, list):
+            prop_dicts = [getattr(lT, prop) for prop in properties]
+            properties = {
+                node: [prop.get(node, np.nan) for prop in prop_dicts]
+                for node in lT.nodes
+            }
+
+        default_value = np.zeros_like(next(iter(properties.values())))
+        to_do = [root]
+        while to_do:
+            current = to_do.pop()
+            cycle = np.array(lT.get_successors(current))
+            cycle_times = np.array([lT.time[c] for c in cycle])
+            cycle = cycle[cycle_times <= end_time]
+            if 0 < len(cycle):
+                _next = list(lT.successor[cycle[-1]])
+                if len(_next) > 1 and lT.time[cycle[-1]] < end_time:
+                    out_dict[current] = _next
+                    to_do.extend(_next)
+                else:
+                    out_dict[current] = []
+                final_properties[current] = self.aggregator(
+                    [properties.get(node, np.nan) for node in cycle],
+                    axis=0,
+                )
+                if np.isnan(final_properties[current]).any():
+                    final_properties[current] = default_value
+
+        return ApproximatedTree(
+            out_dict, final_properties, (lT.name, root, end_time)
+        )
+
+
+class DownsampledTree(TreeApproximationTemplate):
+    """
+    An approximator where a lineage tree is downsampled by a given
+    downsampling value `downsample`. One every `downsample` node
+    is kept in the approximated tree.
+    When computing the distance, it only looks at difference in topology,
+    no property dictionary can be provided.
+
+    This approximation is halfway between the precision of the full lineage
+    tree and the rapidity of the SimpleTree.
+
+    .. warning:: This approximation should only be used to compare trees
+    that have the same time resolution. When comparing two trees that have
+    different time resolution, one should use `ResampledTree`
+
+    The default delta for this approximation is simple delta
+    """
+
+    default_delta = delta_binary
+
+    def __init__(self, downsample: int = 2, delta=None):
+        super().__init__(delta)
+        if not isinstance(downsample, int):
+            raise Warning("Please put a valid downsampling value.")
+        self.downsample = downsample
+
+    def approximation(
+        self,
+        lT: LineageTree,
+        root: int,
+        end_time: int | None = None,
+        properties: dict[int, float | list] | list[str] = None,
+    ):
+        """
+        Build the approximation of a given sub lineage of lineage tree.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        end_time : int
+            The last time point to consider
+        downsample : int
+            The downsampling value.
+            One node will be conserved every `downsample` node
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
+
+        if end_time is None:
+            end_time = lT.t_e
+        out_dict = {}
+        if isinstance(properties, list):
+            prop_dicts = [getattr(lT, prop) for prop in properties]
+            properties = {
+                node: [prop.get(node, np.nan) for prop in prop_dicts]
+                for node in lT.nodes
+            }
+        to_do = [root]
+        while to_do:
+            current = to_do.pop()
+            _next = lT.nodes_at_t(
+                r=current,
+                t=lT.time[current] + self.downsample,
+            )
+            if _next == [current]:
+                _next = None
+            if _next and lT.time[_next[0]] <= end_time:
+                out_dict[current] = _next
+                to_do.extend(_next)
+            else:
+                out_dict[current] = []
+        return ApproximatedTree(
+            out_dict, properties, (lT.name, root, end_time)
+        )
+
+
+class ResampledTree(TreeApproximationTemplate):
+    """
+    An approximator that resample lineage trees to a given time resolution.
+    The target time resolution is provided at creation.
+    Each time a new approximated tree is built, its original time resolution
+    has to be provided. Therefore, the resampling can be upsampling or downsampling.
+
+    Moreover, a dictionary mapping node ids to values can be provided.
+    It will be interpolated using a spline interpolation.
+    If no property dictionary is provided then the simple delta will be used.
+
+    This approximation is the go to approximation when comparing
+    different lineage trees.
+
+    The default delta for this approximation is the difference delta
+    """
+
+    default_delta = delta_difference
+
+    def __init__(
+        self,
+        delta=None,
+        target_time_resolution: int = 2,
+        spline_smoothing: int = 3,
+    ):
+        super().__init__(delta)
+        self.spline_smoothing = spline_smoothing
+        self.target_time_resolution = target_time_resolution
+
+    def approximation(
+        self,
+        lT: LineageTree,
+        root: int,
+        end_time: int | None = None,
+        properties: dict[int, float | list] | list[str] = None,
+    ):
+        """
+        Build the approximation of a given sub lineage of lineage tree.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        sampling_property : dict mapping ints to floats, optional
+            A dictionary mapping node ids to their corresponding
+            values for the property considered.
+            If not provided, no property will be attached to the nodes
+            and therefore only the topology of the lineage tree will matter
+        end_time : int
+            The last time point to consider
+        time_resolution : int
+            The time resolution of the lineage tree
+            to approximate. It has to be in the same
+            units than the one provided as `target_time_resolution`.
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
+
+        if lT.time_resolution is None:
+            lT.time_resolution = 1
+        if end_time is None:
+            end_time = lT.t_e
+        if properties is None:
+            properties = {}
+        if isinstance(properties, list):
+            prop_dicts = [getattr(lT, prop) for prop in properties]
+            properties = {
+                node: [prop.get(node, np.nan) for prop in prop_dicts]
+                for node in lT.nodes
+            }
+
+        # All chains will end up existing, at most reduced to 1 if too short
+        chains = lT.get_all_chains_of_subtree(root, end_time=end_time)
+
+        downsampling_rate = self.target_time_resolution / lT.time_resolution
+        out_dict = {}  # adjacency dictionary
+        final_property = {}  # interpolated property dictionary
+        # links remains to add after the first for loop (ie divisions)
+        remaining_links = []
+        temporary_mapping = {}  # Mapping between old and new ids
+
+        next_id = 0  # Initial id value
+        for chain in chains:
+            # First computing the times of the original and new chain
+            # to use for the spline interpolation
+            initial_length = len(chain)
+            target_length = round(initial_length / downsampling_rate)
+            if target_length <= 0:
+                target_length = 1
+            new_chain = np.arange(target_length)
+            initial_time = np.arange(initial_length) * lT.time_resolution
+            new_chain_time = new_chain * self.target_time_resolution
+
+            # Making sure that we don't have missing values in the property values
+            # to avoid weird things with the interpolation
+            initial_property = []
+            smoothing_time = []
+            for i, node in enumerate(chain):
+                if node in properties:
+                    initial_property.append(properties[node])
+                    smoothing_time.append(initial_time[i])
+
+            if self.spline_smoothing < len(smoothing_time):
+                interpolator = InterpolatedUnivariateSpline(
+                    smoothing_time,
+                    initial_property,
+                    k=self.spline_smoothing,
+                )
+                interpolated_property = interpolator(new_chain_time)
+            elif 0 < len(
+                smoothing_time
+            ):  # case when interpolation give too few nodes
+                interpolated_property = [
+                    0,
+                ] * target_length
+            else:
+                interpolated_property = initial_property[:target_length]
+
+            if len(interpolated_property) == 0:
+                interpolated_property = [
+                    0,
+                ] * target_length
+
+            # it can happen that the interpolated property does not
+            # have enough values, so we pad its with the last value
+            if len(interpolated_property) < target_length:
+                interpolated_property.extend(
+                    [
+                        interpolated_property[-1],
+                    ]
+                    * (target_length - len(interpolated_property))
+                )
+
+            new_chain += next_id
+            next_id = new_chain[-1] + 1
+
+            out_dict.update(
+                {
+                    n: [
+                        new_chain[i + 1],
+                    ]
+                    for i, n in enumerate(new_chain[:-1])
+                }
+            )
+
+            # The strict=True ensure that we have
+            # len(new_chain) == len(interpolated_property)
+            final_property.update(
+                zip(new_chain, interpolated_property, strict=True)
+            )
+            if lT.time[chain[-1]] <= end_time:
+                remaining_links.append(chain[-1])
+                temporary_mapping[chain[-1]] = new_chain[-1]
+                temporary_mapping[chain[0]] = new_chain[0]
+
+        for node in remaining_links:
+            out_dict[temporary_mapping[node]] = list(
+                temporary_mapping[nodei] for nodei in lT.successor[node]
+            )
+
+        return ApproximatedTree(
+            out_dict, final_property, (lT.name, root, end_time)
+        )
+
+
+class FullTree(TreeApproximationTemplate):
+    """
+    An approximator do not really approximate.
+    It takes the sub tree to "approximate" as is unless
+    end_time is provided.
+
+    This approximation is the go to approximation when comparing
+    different lineage trees.
+
+    The default delta for this approximation is the unit delta
+    """
+
+    default_delta = delta_binary
+
+    def __init__(self, delta=None):
+        super().__init__(delta)
+
+    def approximation(
+        self,
+        lT,
+        root,
+        end_time: int | None = None,
+        properties: dict[int, float | list] | list[str] = None,
+    ):
+
+        if isinstance(properties, list):
+            prop_dicts = [getattr(lT, prop) for prop in properties]
+            properties = {
+                node: [prop.get(node, np.nan) for prop in prop_dicts]
+                for node in lT.nodes
+            }
+        if end_time is None:
+            end_time = lT.t_e
+        out_dict = {}
+        to_do = [root]
+        while to_do:
+            current = to_do.pop()
+            _next = list(lT.successor[current])
+            if _next:
+                for _n in _next:
+                    if lT.time[_n] <= end_time:
+                        out_dict.setdefault(current, []).append(_n)
+                        to_do.append(_n)
+            else:
+                out_dict[current] = []
+        return ApproximatedTree(
+            out_dict, properties, (lT.name, root, end_time)
+        )
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
new file mode 100644
index 0000000..d0d4e8b
--- /dev/null
+++ b/src/lineagetree/distance/comparator.py
@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from collections import deque
+from typing import TYPE_CHECKING, Iterable, Callable
+import numpy as np
+from scipy.interpolate import InterpolatedUnivariateSpline
+from hashlib import sha256
+import pickle
+
+from edist import uted
+
+from .delta import (
+    delta_normalized_difference,
+    delta_nd_norm,
+    delta_difference,
+    delta_binary,
+)
+from warnings import warn
+
+if TYPE_CHECKING:
+    from lineagetree import LineageTree
+    from edist.alignment import Alignment
+    from .approximations import TreeApproximationTemplate
+    from .distance_calculator import TreeDistanceTemplate
+
+
+class TreeComparator:
+
+    def __init__(
+        self,
+        tree_distance: TreeDistanceTemplate,
+        approximator: TreeApproximationTemplate | None = None,
+    ):
+        self.approximator = approximator() if approximator else lambda x: x
+        if approximator is None:
+            warn(
+                "No approximation was set so the lineageTree object itself will be added to the the Object without any changes."
+            )
+            self.approximator = lambda x: x
+        else:
+            self.approximator = approximator
+
+        self.tree_distance = tree_distance
+        self.labels = {}
+        self.__id_count = 0
+        self.trees = {}
+
+    @property
+    def cached_distances(
+        self,
+    ):
+        if not hasattr(self, "_cache"):
+            self._cache = {}
+        return self._cache
+
+    def _get_next_id(self) -> int:
+        """Provides and id if it is not provided by `add_tree`.
+
+        Returns
+        -------
+        int
+            The generated id.
+        """
+        self.__id_count += 1
+        return self.__id_count - 1
+
+    def add_trees_from_dif_starts(self, lT, starts, emd_time): ...
+
+    def add_multiple_trees(self, *trees): ...
+
+    def add_tree(
+        self,
+        lT: LineageTree,
+        root: int,
+        end_time: int | None = None,
+        name: str = None,
+    ):
+        """Handles adding a tree to the object, by calling the `approximation` method and assigning an id to the tree.
+        This function also saves the labels.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The LineageTree object
+        root : int
+            The root of the tree to be added, can be any nodein the `lT` object.
+        end_time : int | None, optional
+            The final timepoint of the subtree to be taken into account, by default None
+        time_resolution : float, optional
+            The time_resolution, by default None
+        name: str, optional
+            The name of the tree to be added, if None an id will be generated, which will serve as its name, by default None
+        """
+        approximated_tree = self.approximator.approximation(
+            lT, root, end_time, time_resolution
+        )
+        if isinstance(name, str):
+            id = name
+        else:
+            id = self._get_next_id()
+        self.trees.update({id: approximated_tree})
+        self.labels.update({id: lT.labels.get(root, root)})
+
+    def remove_tree(self, tree: int | str):
+        if not isinstance(tree, str):
+            raise Warning(
+                f"Trees are removed form the structure using their names (str)"
+            )
+        if tree in self.trees:
+            self.trees.pop(tree)
diff --git a/src/lineagetree/distance/delta.py b/src/lineagetree/distance/delta.py
new file mode 100644
index 0000000..cec1026
--- /dev/null
+++ b/src/lineagetree/distance/delta.py
@@ -0,0 +1,57 @@
+import numpy as np
+
+
+def delta_normalized_difference(node1, node2):
+    """
+    Delta function for unordered edit distance.
+
+    Returns the absolute difference between `v1` and `v2` normalized by `max(v1, v2)`.
+    Addition and deletion both cost 1.
+    """
+    if node1 is None or node2 is None:
+        return 1
+    else:
+        return np.abs(node1 - node2) / np.max([node1, node2])
+
+
+def delta_nd_norm(node1, node2):
+    """
+    Delta function for unordered edit distance.
+
+    Returns the norm between `np.linalg.norm(v1 - v2)`.
+    Addition and deletion both cost
+    `np.linalg.norm(v1)` or `np.linalg.norm(v2)`.
+    """
+    if node1 is None:
+        return np.linalg.norm(node2)
+    elif node2 is None:
+        return np.linalg.norm(node1)
+    else:
+        return np.linalg.norm(node1 - node2)
+
+
+def delta_difference(node1, node2):
+    """
+    Delta function for unordered edit distance.
+
+    Returns the absolute difference between `v1` and `v2`.
+    Addition and deletion both cost `v1` or `v2` accordingly.
+    """
+    if node1 is None:
+        return node2
+    elif node2 is None:
+        return node1
+    else:
+        return np.abs(node1 - node2)
+
+
+def delta_binary(node1, node2):
+    """
+    Delta function for unordered edit distance.
+
+    Matching costs 0, addition and deletion cost 1
+    """
+    if node1 is None or node2 is None:
+        return 1
+    else:
+        return 0
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
new file mode 100644
index 0000000..ee2c8f3
--- /dev/null
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Iterable, Callable
+import numpy as np
+from scipy.interpolate import InterpolatedUnivariateSpline
+from edist import uted
+
+if TYPE_CHECKING:
+    from lineagetree import LineageTree
+    from edist.alignment import Alignment
+    from .approximations import ApproximatedTree
+
+
+class TreeDistanceTemplate(ABC):
+
+    @abstractmethod
+    def compute_distance(
+        self,
+        approximated_tree1: ApproximatedTree,
+        approximated_tree2: ApproximatedTree,
+        backtrace: Alignment = None,
+    ) -> float: ...
+
+    @abstractmethod
+    def get_norm(
+        self,
+        tree1: ApproximatedTree,
+        tree2: ApproximatedTree,
+        norm_type: {"max", "sum", "tuple"} = "sum",
+    ) -> float: ...
+
+
+class UnorderedTreeEditDistance(TreeDistanceTemplate):
+
+    def __init__(self, delta):
+        super().__init__()
+        self.delta = delta
+        self.backtrace = {}
+
+    def _compute_uted_backtrace(
+        self,
+        approximated_tree1: ApproximatedTree,
+        approximated_tree2: ApproximatedTree,
+        return_serial: bool = False,
+    ) -> Alignment:
+        """
+        Computes the optimal mapping between
+        `approximated_tree1` and `approximated_tree2`
+        according to the delta function predetermined
+
+        Parameters
+        ----------
+        approximated_tree1 : ApproximatedTree
+            The first tree to compare
+        approximated_tree2 : ApproximatedTree
+            The second tree to compare
+
+        Returns
+        -------
+        backtrace : edist.alignment.Alignment
+            The resulting alignment in the edist format
+        """
+        backtrace = uted.uted_backtrace(
+            approximated_tree1.nodes,
+            approximated_tree1.adjacency_list,
+            approximated_tree2.nodes,
+            approximated_tree2.adjacency_list,
+            delta=self.delta,
+        )
+        if return_serial:
+            return (backtrace,)
+        else:
+            return backtrace
+
+    def compute_distance(
+        self,
+        approximated_tree1: ApproximatedTree,
+        approximated_tree2: ApproximatedTree,
+        btrc: dict[frozenset[tuple], Alignment],
+    ) -> float:
+        """
+        Computes the unordered edit distance
+        between two approximated lineage trees.
+
+        It can take as an input the backtrace
+        if it was already computed
+
+        Parameters
+        ----------
+        approximated_tree1 : ApproximatedTree
+            The first tree to compare
+        approximated_tree2 : ApproximatedTree
+            The second tree to compare
+        backtrace : Alignment, optional
+            The precomputed alignement between
+            the two trees
+
+        Returns
+        -------
+        float
+            The unordered tree edit distance
+            between the two trees.
+        """
+        key = frozenset(
+            approximated_tree1.tree_specs, approximated_tree2.tree_specs
+        )
+        if key not in btrc:
+            btrc[key] = self._compute_uted_backtrace(
+                approximated_tree1, approximated_tree2
+            )
+        return btrc[key].cost(
+            approximated_tree1.nodes, approximated_tree2.nodes, self.delta
+        )
+
+    def reconstruct_backtrace(self, tree1, tree2, backtrace): ...
+
+    def get_norm(
+        self,
+        tree1: ApproximatedTree,
+        tree2: ApproximatedTree,
+        norm_type: {"max", "sum", "tuple"} = "sum",
+    ) -> float:
+        """
+        Computes the normalisation value for the
+        unordered tree edit distance between `tree1` and `tree2`
+
+        The normalisation always involve the distance between
+        either of the trees to the empty tree (d(tree, ø)).
+
+        If the norm type is "max" then
+        the max of d(tree1, ø) and d(ø, tree2) is returned
+
+        If the norm type is "sum", then
+        d(tree1, ø) + d(ø, tree2) is returned
+
+        Parameters
+        ----------
+        tree1 : ApproximatedTree
+            The first tree to compute the normalisation
+        tree2 : ApproximatedTree
+            The second tree to compute the normalisation
+        norm_type : {"max", "sum", "tuple"}
+            How to combine the two distances (see above)
+            If "tuple" is provided, return the two raw values
+
+        Returns
+        -------
+        float
+            The normalisation value
+        """
+        distance_to_none1 = uted.uted(
+            tree1.nodes, tree1.adjacency_list, [], [], delta=self.delta
+        )
+        distance_to_none2 = uted.uted(
+            [], [], tree2.nodes, tree2.adjacency_list, delta=self.delta
+        )
+
+        match norm_type.lower():
+            case "max":
+                return max(distance_to_none1, distance_to_none2)
+            case "sum":
+                return distance_to_none1 + distance_to_none2
+            case "tuple":
+                return distance_to_none1, distance_to_none2
+            case _:
+                raise ValueError(
+                    f"Invalid value for `norm_type`. Got {norm_type},"
+                    "expected 'max', 'sum' or 'tuple'"
+                )

From 0f2a740617a4e284f79afb15f620c940f708855c Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Wed, 10 Jun 2026 01:42:18 +0200
Subject: [PATCH 02/14] progress

---
 src/lineagetree/distance/approximations.py    |  14 +-
 src/lineagetree/distance/comparator.py        | 133 +++++++++++++-----
 .../distance/distance_calculator.py           |  36 +++--
 tests/test_lineageTree.py                     |   5 +-
 4 files changed, 117 insertions(+), 71 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index ea69da7..2f75640 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -6,10 +6,7 @@
 from typing import TYPE_CHECKING, Iterable, Callable
 import numpy as np
 from scipy.interpolate import InterpolatedUnivariateSpline
-from hashlib import sha256
-import pickle
 
-from edist import uted
 
 from .delta import (
     delta_normalized_difference,
@@ -17,13 +14,10 @@
     delta_difference,
     delta_binary,
 )
-from warnings import warn
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
-    from edist.alignment import Alignment
     from .approximations import TreeApproximationTemplate
-    from .distance_calculator import TreeDistanceTemplate
 
 
 @dataclass
@@ -175,8 +169,8 @@ def __init__(self, aggregator: Callable = np.nanmean, delta=None):
 
     def approximation(
         self,
-        lT,
-        root,
+        lT: LineageTree,
+        root: int,
         end_time=None,
         properties: dict[int, float | list] | list[str] = None,
     ):
@@ -523,8 +517,8 @@ def __init__(self, delta=None):
 
     def approximation(
         self,
-        lT,
-        root,
+        lT: LineageTree,
+        root: int,
         end_time: int | None = None,
         properties: dict[int, float | list] | list[str] = None,
     ):
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index d0d4e8b..4d318fb 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -1,38 +1,27 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from abc import ABC, abstractmethod
-from collections import deque
-from typing import TYPE_CHECKING, Iterable, Callable
-import numpy as np
-from scipy.interpolate import InterpolatedUnivariateSpline
-from hashlib import sha256
-import pickle
-
-from edist import uted
-
-from .delta import (
-    delta_normalized_difference,
-    delta_nd_norm,
-    delta_difference,
-    delta_binary,
-)
+from typing import TYPE_CHECKING, Iterable
+import inspect
+
 from warnings import warn
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
-    from edist.alignment import Alignment
     from .approximations import TreeApproximationTemplate
     from .distance_calculator import TreeDistanceTemplate
 
 
 class TreeComparator:
 
+    __adding_trees = True  # Variable that allows adding trees to the object or not, the moment a dataset is added no more trees can be added.
+    levels_of_comparison = {}
+
     def __init__(
         self,
         tree_distance: TreeDistanceTemplate,
         approximator: TreeApproximationTemplate | None = None,
     ):
+
         self.approximator = approximator() if approximator else lambda x: x
         if approximator is None:
             warn(
@@ -55,27 +44,99 @@ def cached_distances(
             self._cache = {}
         return self._cache
 
-    def _get_next_id(self) -> int:
+    def _get_next_id(self) -> str:
         """Provides and id if it is not provided by `add_tree`.
 
         Returns
         -------
-        int
+        str
             The generated id.
         """
         self.__id_count += 1
-        return self.__id_count - 1
+        return f"LineageTree {self.__id_count - 1}"
 
-    def add_trees_from_dif_starts(self, lT, starts, emd_time): ...
+    def add_trees_one_dataset(
+        self,
+        lT: LineageTree,
+        roots: Iterable,
+        starts: int | list[int],
+        end_time: int | list[int] = None,
+    ):
+        """Creates a dataset using ONLY one `Lineagetree` object for comparison.
+          The user may choose if they want their dataset to have multiple starts
+          or multiple end times but not both.
 
-    def add_multiple_trees(self, *trees): ...
+        Parameters
+        ----------
+        lT : LineageTree
+            The `LineageTree` object.
+        roots : Iterable
+            The roots of the subtrees to be compared.
+        starts : int | list[int]
+            The different starting points for the dataset. If `starts` is a list the `end_time` cannot be a list.
+        end_time : int | list[int], optional
+            The different end times for the dataset. If `end_time` is a list the `starts` cannot be a list. If None the whole subtrees will be taken into account, by default `None`
+        """
+        if not self.__adding_trees:
+            raise Warning("Cannot add any more trees to the Comparator.")
+        if isinstance(starts, Iterable) and isinstance(end_time, Iterable):
+            raise Warning(
+                "Each Comparator can hold either multiple starts or multiple end_times. Not both."
+            )
+        if isinstance(starts, Iterable) and not isinstance(end_time, Iterable):
+            for start in starts:
+                selected_roots = lT.nodes_at_t(start, roots)
+                self._add_trees(lT, selected_roots, end_time, level=start)
+        elif not isinstance(starts, Iterable) and isinstance(
+            end_time, Iterable
+        ):
+            for end in end_time:
+                self._add_trees(lT, roots, end, level=end)
 
-    def add_tree(
+        else:
+            for root in roots:
+                self._add_tree(lT, root, end_time, level=1)
+
+        self.__adding_trees = False
+
+    def _add_trees(
+        self,
+        lT: LineageTree,
+        roots: Iterable,
+        end_time: int | list[int] | None = None,
+        level=1,
+    ):
+        for root in roots:
+            self._add_tree(lT, root, end_time, level=level)
+
+    def add_trees_multiple_datasets(
+        self, *trees: tuple[LineageTree, list[int], int|list[int], int|list[int]]
+    ):
+        """
+        Works the same way as adding trees from one dataset, but its for multiple datasaets.
+
+        Args:
+            *trees: Variable-length collection of tuples, where each tuple contains:
+                - tree (LineageTree): The lineage tree to add.
+                - roots (list[int]): The selected roots of the dataset.
+                - end_time (int): The ending time for the subtrees.
+        """
+        if not self.__adding_trees:
+            raise Warning("Cannot add any more trees to the Comparator.")
+        lTs = [tr[0] for tr in trees]
+        roots = [tr[1] for tr in trees]
+        starting_times = [tr[2] for tr in trees]
+        for s_t in starting_times:
+            
+        end_times = [tr[3] for tr in trees]
+        
+
+    def _add_tree(
         self,
         lT: LineageTree,
         root: int,
         end_time: int | None = None,
-        name: str = None,
+        level: int = None,
     ):
         """Handles adding a tree to the object, by calling the `approximation` method and assigning an id to the tree.
         This function also saves the labels.
@@ -90,23 +151,17 @@ def add_tree(
             The final timepoint of the subtree to be taken into account, by default None
         time_resolution : float, optional
             The time_resolution, by default None
-        name: str, optional
-            The name of the tree to be added, if None an id will be generated, which will serve as its name, by default None
+        level: int, optional
+            Used if the dataset has multiple levels of comparison
         """
-        approximated_tree = self.approximator.approximation(
-            lT, root, end_time, time_resolution
-        )
-        if isinstance(name, str):
-            id = name
+        approximated_tree = self.approximator.approximation(lT, root, end_time)
+        if isinstance(lT.name, str):
+            id = lT.name
         else:
             id = self._get_next_id()
         self.trees.update({id: approximated_tree})
         self.labels.update({id: lT.labels.get(root, root)})
-
-    def remove_tree(self, tree: int | str):
-        if not isinstance(tree, str):
-            raise Warning(
-                f"Trees are removed form the structure using their names (str)"
+        if level is not None:
+            self.levels_of_comparison.setdefault(level, []).append(
+                approximated_tree
             )
-        if tree in self.trees:
-            self.trees.pop(tree)
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index ee2c8f3..57d4137 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -1,13 +1,10 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Iterable, Callable
-import numpy as np
-from scipy.interpolate import InterpolatedUnivariateSpline
+from typing import TYPE_CHECKING
 from edist import uted
 
 if TYPE_CHECKING:
-    from lineagetree import LineageTree
     from edist.alignment import Alignment
     from .approximations import ApproximatedTree
 
@@ -42,7 +39,6 @@ def _compute_uted_backtrace(
         self,
         approximated_tree1: ApproximatedTree,
         approximated_tree2: ApproximatedTree,
-        return_serial: bool = False,
     ) -> Alignment:
         """
         Computes the optimal mapping between
@@ -68,16 +64,13 @@ def _compute_uted_backtrace(
             approximated_tree2.adjacency_list,
             delta=self.delta,
         )
-        if return_serial:
-            return (backtrace,)
-        else:
-            return backtrace
+        return backtrace
 
     def compute_distance(
         self,
         approximated_tree1: ApproximatedTree,
         approximated_tree2: ApproximatedTree,
-        btrc: dict[frozenset[tuple], Alignment],
+        btrc: dict[frozenset[tuple], Alignment] | None = None,
     ) -> float:
         """
         Computes the unordered edit distance
@@ -102,16 +95,21 @@ def compute_distance(
             The unordered tree edit distance
             between the two trees.
         """
-        key = frozenset(
-            approximated_tree1.tree_specs, approximated_tree2.tree_specs
-        )
-        if key not in btrc:
-            btrc[key] = self._compute_uted_backtrace(
-                approximated_tree1, approximated_tree2
+        if btrc:
+            key = frozenset(
+                {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+            )
+            if key not in btrc:
+                btrc[key] = self._compute_uted_backtrace(
+                    approximated_tree1, approximated_tree2
+                )
+            return btrc[key].cost(
+                approximated_tree1.nodes, approximated_tree2.nodes, self.delta
+            )
+        else:
+            uted.uted(
+                approximated_tree1.nodes, approximated_tree2.nodes, self.delta
             )
-        return btrc[key].cost(
-            approximated_tree1.nodes, approximated_tree2.nodes, self.delta
-        )
 
     def reconstruct_backtrace(self, tree1, tree2, backtrace): ...
 
diff --git a/tests/test_lineageTree.py b/tests/test_lineageTree.py
index b349042..35e1dce 100644
--- a/tests/test_lineageTree.py
+++ b/tests/test_lineageTree.py
@@ -636,9 +636,8 @@ def test_spatial_density():
     assert np.isclose(
         lT1.spatial_density(0, th=40)[110832], 7.460387957432594e-06
     )
-    assert lT1.neighbours_in_radius(0, th=40)[110832] == {
-        np.int64(110826)
-    }
+    assert lT1.neighbours_in_radius(0, th=40)[110832] == {np.int64(110826)}
+
 
 def test_k_nearest_neighbours():
     assert (

From a37da03b9ef59d4a35a818625d8764748b5cd87e Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Wed, 10 Jun 2026 16:48:40 +0200
Subject: [PATCH 03/14] .

---
 src/lineagetree/distance/approximations.py    |  11 +-
 src/lineagetree/distance/comparator.py        | 158 ++++--------------
 .../distance/distance_calculator.py           |  57 ++++---
 src/lineagetree/lineage_tree.py               |   3 +
 4 files changed, 79 insertions(+), 150 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 2f75640..6fce9ae 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -128,6 +128,7 @@ def approximation(
         root: int,
         end_time: int | None = None,
     ) -> ApproximatedTree:
+        print(lT, root, end_time)
         if end_time is None:
             end_time = lT.t_e
         if lT.time_resolution == 0:
@@ -155,7 +156,7 @@ def approximation(
                     out_dict[current] = []
             final_properties[current] = len(cycle) * (time_resolution)
         return ApproximatedTree(
-            out_dict, final_properties, (lT.name, root, end_time)
+            out_dict, final_properties, (hash(lT), root, end_time)
         )
 
 
@@ -238,7 +239,7 @@ def approximation(
                     final_properties[current] = default_value
 
         return ApproximatedTree(
-            out_dict, final_properties, (lT.name, root, end_time)
+            out_dict, final_properties, (hash(lT), root, end_time)
         )
 
 
@@ -322,7 +323,7 @@ def approximation(
             else:
                 out_dict[current] = []
         return ApproximatedTree(
-            out_dict, properties, (lT.name, root, end_time)
+            out_dict, properties, (hash(lT), root, end_time)
         )
 
 
@@ -494,7 +495,7 @@ def approximation(
             )
 
         return ApproximatedTree(
-            out_dict, final_property, (lT.name, root, end_time)
+            out_dict, final_property, (hash(lT), root, end_time)
         )
 
 
@@ -544,5 +545,5 @@ def approximation(
             else:
                 out_dict[current] = []
         return ApproximatedTree(
-            out_dict, properties, (lT.name, root, end_time)
+            out_dict, properties, (hash(lT), root, end_time)
         )
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 4d318fb..b8b43c1 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -9,11 +9,12 @@
     from lineagetree import LineageTree
     from .approximations import TreeApproximationTemplate
     from .distance_calculator import TreeDistanceTemplate
+    from .approximations import ApproximatedTree
+    from edist.alignment import Alignment
 
 
 class TreeComparator:
 
-    __adding_trees = True  # Variable that allows adding trees to the object or not, the moment a dataset is added no more trees can be added.
     levels_of_comparison = {}
 
     def __init__(
@@ -33,135 +34,50 @@ def __init__(
 
         self.tree_distance = tree_distance
         self.labels = {}
-        self.__id_count = 0
         self.trees = {}
 
     @property
-    def cached_distances(
+    def _cached_distances(
         self,
-    ):
+    ) -> dict[frozenset[tuple[int, int, int]], Alignment | float | int]:
         if not hasattr(self, "_cache"):
             self._cache = {}
         return self._cache
 
-    def _get_next_id(self) -> str:
-        """Provides and id if it is not provided by `add_tree`.
-
-        Returns
-        -------
-        str
-            The generated id.
-        """
-        self.__id_count += 1
-        return f"LineageTree {self.__id_count - 1}"
-
-    def add_trees_one_dataset(
+    @property
+    def _cached_approximations(
         self,
-        lT: LineageTree,
-        roots: Iterable,
-        starts: int | list[int],
-        end_time: int | list[int] = None,
-    ):
-        """Creates a dataset using ONLY one `Lineagetree` object for comparison.
-          The user may choose if they want their dataset to have multiple starts
-          or multiple end times but not both.
-
-        Parameters
-        ----------
-        lT : LineageTree
-            The `LineageTree` object.
-        roots : Iterable
-            The roots of the subtrees to be compared.
-        starts : int | list[int]
-            The different starting points for the dataset. If `starts` is a list the `end_time` cannot be a list.
-        end_time : int | list[int], optional
-            The different end times for the dataset. If `end_time` is a list the `starts` cannot be a list. If None the whole subtrees will be taken into account, by default `None`
-        """
-        if not self.__adding_trees:
-            raise Warning("Cannot add any more trees to the Comparator.")
-        if isinstance(starts, Iterable) and isinstance(end_time, Iterable):
-            raise Warning(
-                "Each Comparator can hold either multiple starts or multiple end_times. Not both."
-            )
-        if isinstance(starts, Iterable) and not isinstance(end_time, Iterable):
-            for start in starts:
-                selected_roots = lT.nodes_at_t(start, roots)
-                self._add_trees(lT, selected_roots, end_time, level=start)
-        elif not isinstance(starts, Iterable) and isinstance(
-            end_time, Iterable
-        ):
-            for end in end_time:
-                self._add_trees(lT, roots, end, level=end)
-
+    ) -> dict[tuple[int, int, int], ApproximatedTree]:
+        if not hasattr(self, "approximations"):
+            self.__approximations = {}
+        return self.__approximations
+
+    def use_approximation(
+        self, tree: tuple[LineageTree, int, int] | ApproximatedTree
+    ) -> ApproximatedTree:
+        print("peos", tree)
+        if isinstance(tree, tuple):
+            new_key = (hash(tree[0]), *tree[1:])
+            if new_key in self._cached_approximations:
+                f_tree = self._cached_approximations[new_key]
+            else:
+                print("edw", tree)
+                f_tree = self.approximator.approximation(*tree)
+                self._cached_approximations[new_key] = f_tree
         else:
-            for root in roots:
-                self._add_tree(lT, root, end_time, level=1)
+            f_tree = tree
+        return f_tree
 
-        self.__adding_trees = False
-
-    def _add_trees(
+    def compare(
         self,
-        lT: LineageTree,
-        roots: Iterable,
-        end_time: int | list[int] | None = None,
-        level=1,
-    ):
-        for root in roots:
-            self._add_tree(lT, root, end_time, level=level)
-
-    def add_trees_multiple_datasets(
-        self, *trees: tuple[LineageTree, list[int], int|list[int], int|list[int]]
-    ):
-        """
-        Works the same way as adding trees from one dataset, but its for multiple datasaets.
-
-        Args:
-            *trees: Variable-length collection of tuples, where each tuple contains:
-                - tree (LineageTree): The lineage tree to add.
-                - roots (list[int]): The selected roots of the dataset.
-                - end_time (int): The ending time for the subtrees.
-        """
-        if not self.__adding_trees:
-            raise Warning("Cannot add any more trees to the Comparator.")
-        lTs = [tr[0] for tr in trees]
-        roots = [tr[1] for tr in trees]
-        starting_times = [tr[2] for tr in trees]
-        for s_t in starting_times:
-            
-        end_times = [tr[3] for tr in trees]
-        
-
-    def _add_tree(
-        self,
-        lT: LineageTree,
-        root: int,
-        end_time: int | None = None,
-        level: int = None,
-    ):
-        """Handles adding a tree to the object, by calling the `approximation` method and assigning an id to the tree.
-        This function also saves the labels.
-
-        Parameters
-        ----------
-        lT : LineageTree
-            The LineageTree object
-        root : int
-            The root of the tree to be added, can be any nodein the `lT` object.
-        end_time : int | None, optional
-            The final timepoint of the subtree to be taken into account, by default None
-        time_resolution : float, optional
-            The time_resolution, by default None
-        level: int, optional
-            Used if the dataset has multiple levels of comparison
-        """
-        approximated_tree = self.approximator.approximation(lT, root, end_time)
-        if isinstance(lT.name, str):
-            id = lT.name
-        else:
-            id = self._get_next_id()
-        self.trees.update({id: approximated_tree})
-        self.labels.update({id: lT.labels.get(root, root)})
-        if level is not None:
-            self.levels_of_comparison.setdefault(level, []).append(
-                approximated_tree
-            )
+        tree1: tuple[LineageTree, int, int] | ApproximatedTree,
+        tree2: tuple[LineageTree, int, int] | ApproximatedTree,
+        norm="sum",
+    ) -> float | int:
+        tree1 = self.use_approximation(tree1)
+        tree2 = self.use_approximation(tree2)
+        distance = self.tree_distance.compute_distance(
+            tree1, tree2, self._cached_distances
+        ) / self.tree_distance.get_norm(tree1, tree2, norm)
+
+        return distance
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index 57d4137..6563726 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -7,15 +7,19 @@
 if TYPE_CHECKING:
     from edist.alignment import Alignment
     from .approximations import ApproximatedTree
+    from lineagetree import LineageTree
 
 
 class TreeDistanceTemplate(ABC):
+    backtrace_based = (
+        ...
+    )  # If your distance is and edit distance that may provide the alignments between nodes set to `True`, else `False`.
 
     @abstractmethod
     def compute_distance(
         self,
-        approximated_tree1: ApproximatedTree,
-        approximated_tree2: ApproximatedTree,
+        approximated_tree1: ApproximatedTree | LineageTree,
+        approximated_tree2: ApproximatedTree | LineageTree,
         backtrace: Alignment = None,
     ) -> float: ...
 
@@ -30,6 +34,8 @@ def get_norm(
 
 class UnorderedTreeEditDistance(TreeDistanceTemplate):
 
+    backtrace_based = True
+
     def __init__(self, delta):
         super().__init__()
         self.delta = delta
@@ -39,6 +45,7 @@ def _compute_uted_backtrace(
         self,
         approximated_tree1: ApproximatedTree,
         approximated_tree2: ApproximatedTree,
+        cache: dict[frozenset[tuple], Alignment] | None = None,
     ) -> Alignment:
         """
         Computes the optimal mapping between
@@ -57,13 +64,27 @@ def _compute_uted_backtrace(
         backtrace : edist.alignment.Alignment
             The resulting alignment in the edist format
         """
-        backtrace = uted.uted_backtrace(
-            approximated_tree1.nodes,
-            approximated_tree1.adjacency_list,
-            approximated_tree2.nodes,
-            approximated_tree2.adjacency_list,
-            delta=self.delta,
-        )
+        if cache:
+            key = frozenset(
+                {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+            )
+            if key not in cache:
+                cache[key] = uted.uted_backtrace(
+                    approximated_tree1.nodes,
+                    approximated_tree1.adjacency_list,
+                    approximated_tree2.nodes,
+                    approximated_tree2.adjacency_list,
+                    delta=self.delta,
+                )
+                backtrace = cache[key]
+        else:
+            backtrace = uted.uted_backtrace(
+                approximated_tree1.nodes,
+                approximated_tree1.adjacency_list,
+                approximated_tree2.nodes,
+                approximated_tree2.adjacency_list,
+                delta=self.delta,
+            )
         return backtrace
 
     def compute_distance(
@@ -95,21 +116,9 @@ def compute_distance(
             The unordered tree edit distance
             between the two trees.
         """
-        if btrc:
-            key = frozenset(
-                {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
-            )
-            if key not in btrc:
-                btrc[key] = self._compute_uted_backtrace(
-                    approximated_tree1, approximated_tree2
-                )
-            return btrc[key].cost(
-                approximated_tree1.nodes, approximated_tree2.nodes, self.delta
-            )
-        else:
-            uted.uted(
-                approximated_tree1.nodes, approximated_tree2.nodes, self.delta
-            )
+        return self._compute_uted_backtrace(
+            approximated_tree1, approximated_tree2, btrc
+        ).cost(approximated_tree1.nodes, approximated_tree2.nodes, self.delta)
 
     def reconstruct_backtrace(self, tree1, tree2, backtrace): ...
 
diff --git a/src/lineagetree/lineage_tree.py b/src/lineagetree/lineage_tree.py
index 3699928..f50637a 100644
--- a/src/lineagetree/lineage_tree.py
+++ b/src/lineagetree/lineage_tree.py
@@ -42,6 +42,9 @@ def __eq__(self, other) -> bool:
         else:
             return False
 
+    def __hash__(self) -> int:
+        return hash(tuple(self.successor.items()))
+
     def __setstate__(self, state):
         if "_successor" not in state:
             state["_successor"] = state["successor"]

From 944aad72683f37a7ade1b9a596b9df26ac564dac Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Fri, 12 Jun 2026 15:43:25 +0200
Subject: [PATCH 04/14] completed_mp

---
 src/lineagetree/distance/approximations.py    | 41 +++++++--
 src/lineagetree/distance/comparator.py        | 61 +++++++++++---
 .../distance/distance_calculator.py           | 84 +++++++++++++------
 src/lineagetree/lineage_tree.py               | 14 +++-
 4 files changed, 151 insertions(+), 49 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 6fce9ae..4f83aed 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -20,11 +20,30 @@
     from .approximations import TreeApproximationTemplate
 
 
+@dataclass
+class TreeSpecs:
+    lT: int
+    root: int
+    end_time: int
+
+    def __str__(self):
+        return (
+            f"TreeSpecs("
+            f"lt={self.lT}, "
+            f"root={self.root}, "
+            f"end_time={self.end_time}"
+            f")"
+        )
+
+    def __hash__(self):
+        return hash((self.lT, self.root, self.end_time))
+
+
 @dataclass
 class ApproximatedTree:
     adjacency_dict: dict[int, Iterable[int]]
     property_dict: dict[int, float | int | list[int | float]] | None
-    tree_specs: tuple | None = None
+    tree_specs: TreeSpecs | None = None
 
     def _edist_format(
         self,
@@ -50,7 +69,9 @@ def _edist_format(
             if isinstance(sample, Iterable):
                 length = len(sample)
                 default_value = [0] * length
-            elif isinstance(dict_type(), int or float):
+            elif isinstance(sample, int | float):
+                default_value = 0
+            else:
                 default_value = 0
 
         inv_adj = {vi: k for k, v in self.adjacency_dict.items() for vi in v}
@@ -98,7 +119,7 @@ def __post_init__(self):
         )
 
     def __str__(self):  # For quickly checking how the object was created.
-        return f"Subtree: {self.tree_specs}"
+        return self.tree_specs
 
 
 class TreeApproximationTemplate(ABC):
@@ -118,6 +139,9 @@ def approximation(
         end_time: int | None = None,
     ) -> ApproximatedTree: ...
 
+    def delta(self, node1, node2):
+        return self.delta(node1, node2)
+
 
 class ReducedTreeTimed(TreeApproximationTemplate):
     default_delta = delta_normalized_difference
@@ -128,7 +152,6 @@ def approximation(
         root: int,
         end_time: int | None = None,
     ) -> ApproximatedTree:
-        print(lT, root, end_time)
         if end_time is None:
             end_time = lT.t_e
         if lT.time_resolution == 0:
@@ -156,7 +179,7 @@ def approximation(
                     out_dict[current] = []
             final_properties[current] = len(cycle) * (time_resolution)
         return ApproximatedTree(
-            out_dict, final_properties, (hash(lT), root, end_time)
+            out_dict, final_properties, TreeSpecs(hash(lT), root, end_time)
         )
 
 
@@ -239,7 +262,7 @@ def approximation(
                     final_properties[current] = default_value
 
         return ApproximatedTree(
-            out_dict, final_properties, (hash(lT), root, end_time)
+            out_dict, final_properties, TreeSpecs(hash(lT), root, end_time)
         )
 
 
@@ -323,7 +346,7 @@ def approximation(
             else:
                 out_dict[current] = []
         return ApproximatedTree(
-            out_dict, properties, (hash(lT), root, end_time)
+            out_dict, properties, TreeSpecs(hash(lT), root, end_time)
         )
 
 
@@ -495,7 +518,7 @@ def approximation(
             )
 
         return ApproximatedTree(
-            out_dict, final_property, (hash(lT), root, end_time)
+            out_dict, final_property, TreeSpecs(hash(lT), root, end_time)
         )
 
 
@@ -545,5 +568,5 @@ def approximation(
             else:
                 out_dict[current] = []
         return ApproximatedTree(
-            out_dict, properties, (hash(lT), root, end_time)
+            out_dict, properties, TreeSpecs(hash(lT), root, end_time)
         )
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index b8b43c1..7dc4a78 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -2,28 +2,28 @@
 
 from typing import TYPE_CHECKING, Iterable
 import inspect
-
+import pickle
 from warnings import warn
+from multiprocessing import pool
+import types
+from .approximations import ApproximatedTree
+import itertools
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
     from .approximations import TreeApproximationTemplate
     from .distance_calculator import TreeDistanceTemplate
-    from .approximations import ApproximatedTree
     from edist.alignment import Alignment
 
 
 class TreeComparator:
 
-    levels_of_comparison = {}
-
     def __init__(
         self,
         tree_distance: TreeDistanceTemplate,
         approximator: TreeApproximationTemplate | None = None,
     ):
 
-        self.approximator = approximator() if approximator else lambda x: x
         if approximator is None:
             warn(
                 "No approximation was set so the lineageTree object itself will be added to the the Object without any changes."
@@ -32,7 +32,7 @@ def __init__(
         else:
             self.approximator = approximator
 
-        self.tree_distance = tree_distance
+        self.tree_distance = tree_distance(self.approximator.delta)
         self.labels = {}
         self.trees = {}
 
@@ -52,16 +52,14 @@ def _cached_approximations(
             self.__approximations = {}
         return self.__approximations
 
-    def use_approximation(
+    def __use_approximation(
         self, tree: tuple[LineageTree, int, int] | ApproximatedTree
     ) -> ApproximatedTree:
-        print("peos", tree)
         if isinstance(tree, tuple):
             new_key = (hash(tree[0]), *tree[1:])
             if new_key in self._cached_approximations:
                 f_tree = self._cached_approximations[new_key]
             else:
-                print("edw", tree)
                 f_tree = self.approximator.approximation(*tree)
                 self._cached_approximations[new_key] = f_tree
         else:
@@ -72,12 +70,51 @@ def compare(
         self,
         tree1: tuple[LineageTree, int, int] | ApproximatedTree,
         tree2: tuple[LineageTree, int, int] | ApproximatedTree,
-        norm="sum",
+        norm: {"max", "sum", "tuple"} | Callable = "sum",
     ) -> float | int:
-        tree1 = self.use_approximation(tree1)
-        tree2 = self.use_approximation(tree2)
+        if not isinstance(tree1, ApproximatedTree):
+            tree1 = self.__use_approximation(tree1)
+        if not isinstance(tree2, ApproximatedTree):
+            tree2 = self.__use_approximation(tree2)
         distance = self.tree_distance.compute_distance(
             tree1, tree2, self._cached_distances
         ) / self.tree_distance.get_norm(tree1, tree2, norm)
 
         return distance
+
+    def p_compare(self, *trees, norm="sum", n_proccessors: int = 4):
+        app_trees = []
+        for tree in trees:
+            if not isinstance(tree, ApproximatedTree):
+                app_trees.append(self.__use_approximation(tree))
+        print("finished calc approx")
+        combinations = itertools.combinations(app_trees, 2)
+        proccesses = (
+            (self.tree_distance, comb1, comb2, norm, self._cached_distances)
+            for comb1, comb2 in combinations
+        )
+        with pool.Pool(
+            n_proccessors if n_proccessors < len(app_trees) else len(app_trees)
+        ) as p:
+            res = p.map(_worker, proccesses)
+        distances = []
+        for r in res:
+            print(r)
+            self._cached_distances.update(r[0])
+            distances.append(r[1])
+
+        return distances
+
+
+def _worker(args):
+    (
+        distance,
+        app_tree1,
+        app_tree2,
+        norm,
+        cache,
+    ) = args
+    res = distance.compute_distance_parallel(app_tree1, app_tree2, cache)
+    distance = res[1] / distance.get_norm(app_tree1, app_tree2, norm)
+
+    return (res[0], distance)
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index 6563726..39f1b6d 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import Callable
 from typing import TYPE_CHECKING
 from edist import uted
+from warnings import warn
 
 if TYPE_CHECKING:
     from edist.alignment import Alignment
@@ -11,9 +13,19 @@
 
 
 class TreeDistanceTemplate(ABC):
-    backtrace_based = (
-        ...
-    )  # If your distance is and edit distance that may provide the alignments between nodes set to `True`, else `False`.
+
+    def compute_distance_parallel(
+        self,
+        approximated_tree1: ApproximatedTree | LineageTree,
+        approximated_tree2: ApproximatedTree | LineageTree,
+        backtrace: Alignment = None,
+    ) -> dict:
+        warn(
+            "Parallel algorithm not implemented, defaulting to single process."
+        )
+        self.compute_distance(
+            approximated_tree1, approximated_tree2, backtrace
+        )
 
     @abstractmethod
     def compute_distance(
@@ -41,6 +53,26 @@ def __init__(self, delta):
         self.delta = delta
         self.backtrace = {}
 
+    def compute_distance_parallel(
+        self,
+        approximated_tree1: ApproximatedTree,
+        approximated_tree2: ApproximatedTree,
+        backtrace: Alignment = None,
+    ):
+        key = frozenset(
+            {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+        )
+        align = self._compute_uted_backtrace(
+            approximated_tree1, approximated_tree2, backtrace
+        )
+
+        return (
+            {key: align},
+            self.compute_distance(
+                approximated_tree1, approximated_tree2, backtrace
+            ),
+        )
+
     def _compute_uted_backtrace(
         self,
         approximated_tree1: ApproximatedTree,
@@ -64,27 +96,24 @@ def _compute_uted_backtrace(
         backtrace : edist.alignment.Alignment
             The resulting alignment in the edist format
         """
-        if cache:
-            key = frozenset(
-                {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
-            )
-            if key not in cache:
-                cache[key] = uted.uted_backtrace(
-                    approximated_tree1.nodes,
-                    approximated_tree1.adjacency_list,
-                    approximated_tree2.nodes,
-                    approximated_tree2.adjacency_list,
-                    delta=self.delta,
-                )
-                backtrace = cache[key]
-        else:
-            backtrace = uted.uted_backtrace(
-                approximated_tree1.nodes,
-                approximated_tree1.adjacency_list,
-                approximated_tree2.nodes,
-                approximated_tree2.adjacency_list,
-                delta=self.delta,
-            )
+        key = frozenset(
+            {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+        )
+
+        if cache is not None and key in cache:
+            return cache[key]
+
+        backtrace = uted.uted_backtrace(
+            approximated_tree1.nodes,
+            approximated_tree1.adjacency_list,
+            approximated_tree2.nodes,
+            approximated_tree2.adjacency_list,
+            delta=self.delta,
+        )
+
+        if cache is not None:
+            cache[key] = backtrace
+
         return backtrace
 
     def compute_distance(
@@ -126,7 +155,7 @@ def get_norm(
         self,
         tree1: ApproximatedTree,
         tree2: ApproximatedTree,
-        norm_type: {"max", "sum", "tuple"} = "sum",
+        norm_type: {"max", "sum", "tuple"} | Callable = "sum",
     ) -> float:
         """
         Computes the normalisation value for the
@@ -162,12 +191,15 @@ def get_norm(
         distance_to_none2 = uted.uted(
             [], [], tree2.nodes, tree2.adjacency_list, delta=self.delta
         )
-
+        if isinstance(norm_type, Callable):
+            return norm_type(distance_to_none1, distance_to_none2)
         match norm_type.lower():
             case "max":
                 return max(distance_to_none1, distance_to_none2)
             case "sum":
                 return distance_to_none1 + distance_to_none2
+            case "None" | None:
+                return 1
             case "tuple":
                 return distance_to_none1, distance_to_none2
             case _:
diff --git a/src/lineagetree/lineage_tree.py b/src/lineagetree/lineage_tree.py
index f50637a..9dee0db 100644
--- a/src/lineagetree/lineage_tree.py
+++ b/src/lineagetree/lineage_tree.py
@@ -9,7 +9,8 @@
 from collections.abc import Iterable, Sequence
 from packaging.version import Version
 import numpy as np
-
+import hashlib
+import json
 from ._core.utils import CompatibleUnpickler
 from ._mixins.properties_mixin import PropertiesMixin
 from ._mixins.modifier_mixin import ModifierMixin
@@ -21,6 +22,12 @@
 from ._core.validation import TreeValidator
 
 
+def normalize_keys(d):
+    return tuple(
+        sorted((int(k), tuple(int(vv) for vv in v)) for k, v in d.items())
+    )
+
+
 class LineageTree(
     PropertiesMixin,
     ModifierMixin,
@@ -43,7 +50,10 @@ def __eq__(self, other) -> bool:
             return False
 
     def __hash__(self) -> int:
-        return hash(tuple(self.successor.items()))
+        return hash(tuple(self._successor))
+        # succ = normalize_keys(self._successor)
+        # s = json.dumps(succ, sort_keys=True, separators=(",", ":"))
+        # return int(hashlib.sha256(s.encode("utf-8")).hexdigest(), 16)
 
     def __setstate__(self, state):
         if "_successor" not in state:

From b8a5029b036a26d47aecc349ed3109cc758e094f Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Mon, 15 Jun 2026 13:06:29 +0200
Subject: [PATCH 05/14] multiprocessing works super good

---
 src/lineagetree/distance/approximations.py    |  9 +--
 src/lineagetree/distance/comparator.py        | 56 ++++++++++++-------
 .../distance/distance_calculator.py           | 14 +++--
 3 files changed, 50 insertions(+), 29 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 4f83aed..1f8d674 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -119,17 +119,17 @@ def __post_init__(self):
         )
 
     def __str__(self):  # For quickly checking how the object was created.
-        return self.tree_specs
+        return str(self.tree_specs)
 
 
 class TreeApproximationTemplate(ABC):
-    default_delta = ...
+    default_delta: Callable = ...
 
     def __init__(
         self,
         delta: Callable | None = None,
     ):
-        self.delta = delta if delta else self.default_delta
+        self.delta = delta if delta else self.__class__.default_delta  # :(
 
     @abstractmethod
     def approximation(
@@ -139,9 +139,6 @@ def approximation(
         end_time: int | None = None,
     ) -> ApproximatedTree: ...
 
-    def delta(self, node1, node2):
-        return self.delta(node1, node2)
-
 
 class ReducedTreeTimed(TreeApproximationTemplate):
     default_delta = delta_normalized_difference
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 7dc4a78..11f8c65 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -8,6 +8,7 @@
 import types
 from .approximations import ApproximatedTree
 import itertools
+import tqdm
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
@@ -56,7 +57,7 @@ def __use_approximation(
         self, tree: tuple[LineageTree, int, int] | ApproximatedTree
     ) -> ApproximatedTree:
         if isinstance(tree, tuple):
-            new_key = (hash(tree[0]), *tree[1:])
+            new_key = (hash(tree[0]), tree[1], tree[2])
             if new_key in self._cached_approximations:
                 f_tree = self._cached_approximations[new_key]
             else:
@@ -72,36 +73,53 @@ def compare(
         tree2: tuple[LineageTree, int, int] | ApproximatedTree,
         norm: {"max", "sum", "tuple"} | Callable = "sum",
     ) -> float | int:
-        if not isinstance(tree1, ApproximatedTree):
-            tree1 = self.__use_approximation(tree1)
-        if not isinstance(tree2, ApproximatedTree):
-            tree2 = self.__use_approximation(tree2)
+        tree1 = self.__use_approximation(tree1)
+        tree2 = self.__use_approximation(tree2)
         distance = self.tree_distance.compute_distance(
             tree1, tree2, self._cached_distances
         ) / self.tree_distance.get_norm(tree1, tree2, norm)
 
         return distance
 
+    def order_processes(self, proc):
+        nodes1, nodes2 = len(proc[1].nodes), len(proc[2].nodes)
+        return nodes1 + nodes2
+
     def p_compare(self, *trees, norm="sum", n_proccessors: int = 4):
         app_trees = []
-        for tree in trees:
-            if not isinstance(tree, ApproximatedTree):
-                app_trees.append(self.__use_approximation(tree))
-        print("finished calc approx")
-        combinations = itertools.combinations(app_trees, 2)
-        proccesses = (
+        for tree in tqdm.tqdm(trees, desc="Processing Trees: "):
+            app_trees.append(self.__use_approximation(tree))
+        combinations = list(itertools.combinations(app_trees, 2))
+        processes = (
             (self.tree_distance, comb1, comb2, norm, self._cached_distances)
             for comb1, comb2 in combinations
         )
-        with pool.Pool(
-            n_proccessors if n_proccessors < len(app_trees) else len(app_trees)
-        ) as p:
-            res = p.map(_worker, proccesses)
         distances = []
-        for r in res:
-            print(r)
-            self._cached_distances.update(r[0])
-            distances.append(r[1])
+
+        ch_size = (
+            len(combinations) // n_proccessors
+        ) // 4  # Maybe overengineered but it works super good
+        processes_b_s = sorted(processes, key=self.order_processes)
+        processes_s_b = processes_b_s[::-1]
+        mid = len(processes_b_s) // 2
+        high = processes_b_s[:mid]
+        low = processes_s_b[mid:]
+        new_proc = []
+        for h, l in zip(high, low):
+            new_proc.append(h)
+            new_proc.append(l)
+
+        new_proc.extend(high[len(low) :])
+        with pool.Pool(min(n_proccessors, len(app_trees))) as p:
+            for r in tqdm.tqdm(
+                p.imap_unordered(
+                    _worker, new_proc, chunksize=ch_size if ch_size > 0 else 1
+                ),
+                total=len(combinations),
+                desc="UTED distances",
+            ):
+                self._cached_distances.update(r[0])
+                distances.append(r[1])
 
         return distances
 
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index 39f1b6d..92055ec 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -59,8 +59,9 @@ def compute_distance_parallel(
         approximated_tree2: ApproximatedTree,
         backtrace: Alignment = None,
     ):
-        key = frozenset(
-            {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+        key = frozenset({str(approximated_tree1), str(approximated_tree2)})
+        approximated_tree1, approximated_tree2 = sorted(
+            (approximated_tree1, approximated_tree2), key=lambda x: str(x)
         )
         align = self._compute_uted_backtrace(
             approximated_tree1, approximated_tree2, backtrace
@@ -96,8 +97,10 @@ def _compute_uted_backtrace(
         backtrace : edist.alignment.Alignment
             The resulting alignment in the edist format
         """
-        key = frozenset(
-            {approximated_tree1.tree_specs, approximated_tree2.tree_specs}
+        key = frozenset({str(approximated_tree1), str(approximated_tree2)})
+
+        approximated_tree1, approximated_tree2 = sorted(
+            (approximated_tree1, approximated_tree2), key=lambda x: str(x)
         )
 
         if cache is not None and key in cache:
@@ -145,6 +148,9 @@ def compute_distance(
             The unordered tree edit distance
             between the two trees.
         """
+        approximated_tree1, approximated_tree2 = sorted(
+            (approximated_tree1, approximated_tree2), key=lambda x: str(x)
+        )
         return self._compute_uted_backtrace(
             approximated_tree1, approximated_tree2, btrc
         ).cost(approximated_tree1.nodes, approximated_tree2.nodes, self.delta)

From 014108be0d8cb2049c43e4b091f45d86478c110d Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 11:39:52 +0200
Subject: [PATCH 06/14] new hash

---
 src/lineagetree/distance/approximations.py |  1 -
 src/lineagetree/distance/comparator.py     | 63 +++++++++++++++-------
 src/lineagetree/lineage_tree.py            |  8 +--
 3 files changed, 47 insertions(+), 25 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 1f8d674..bb6d4b0 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -65,7 +65,6 @@ def _edist_format(
         """
         if self.property_dict:
             sample = next(iter(self.property_dict.values()))
-            dict_type = type(sample)
             if isinstance(sample, Iterable):
                 length = len(sample)
                 default_value = [0] * length
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 11f8c65..20a8967 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -9,6 +9,10 @@
 from .approximations import ApproximatedTree
 import itertools
 import tqdm
+from pathlib import Path
+import csv
+
+# import seaborn
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
@@ -17,6 +21,20 @@
     from edist.alignment import Alignment
 
 
+def _worker(args):
+    (
+        distance,
+        app_tree1,
+        app_tree2,
+        norm,
+        cache,
+    ) = args
+    res = distance.compute_distance_parallel(app_tree1, app_tree2, cache)
+    distance = res[1] / distance.get_norm(app_tree1, app_tree2, norm)
+
+    return (res[0], distance)
+
+
 class TreeComparator:
 
     def __init__(
@@ -49,9 +67,9 @@ def _cached_distances(
     def _cached_approximations(
         self,
     ) -> dict[tuple[int, int, int], ApproximatedTree]:
-        if not hasattr(self, "approximations"):
-            self.__approximations = {}
-        return self.__approximations
+        if not hasattr(self, "_approximations"):
+            self._approximations = {}
+        return self._approximations
 
     def __use_approximation(
         self, tree: tuple[LineageTree, int, int] | ApproximatedTree
@@ -81,11 +99,23 @@ def compare(
 
         return distance
 
-    def order_processes(self, proc):
+    def _order_processes(self, proc):
         nodes1, nodes2 = len(proc[1].nodes), len(proc[2].nodes)
         return nodes1 + nodes2
 
-    def p_compare(self, *trees, norm="sum", n_proccessors: int = 4):
+    def p_compare(
+        self,
+        *trees,
+        norm="sum",
+        n_proccessors: int = 4,
+        path: Path | None = None,
+    ):
+        if path:
+            for tree in trees:
+                if tree[0].name is None:
+                    raise Warning(
+                        f"All LineageTrees should have a name for saving the comparisons. Missing name: {tree[0]}"
+                    )
         app_trees = []
         for tree in tqdm.tqdm(trees, desc="Processing Trees: "):
             app_trees.append(self.__use_approximation(tree))
@@ -99,7 +129,7 @@ def p_compare(self, *trees, norm="sum", n_proccessors: int = 4):
         ch_size = (
             len(combinations) // n_proccessors
         ) // 4  # Maybe overengineered but it works super good
-        processes_b_s = sorted(processes, key=self.order_processes)
+        processes_b_s = sorted(processes, key=self._order_processes)
         processes_s_b = processes_b_s[::-1]
         mid = len(processes_b_s) // 2
         high = processes_b_s[:mid]
@@ -119,20 +149,13 @@ def p_compare(self, *trees, norm="sum", n_proccessors: int = 4):
                 desc="UTED distances",
             ):
                 self._cached_distances.update(r[0])
-                distances.append(r[1])
-
-        return distances
+                distances.append((r[0].keys(), r[1]))
 
+                # if path: # needs discussion
+                #     with open(path, "w", newline="") as f:
+                #         writer = csv.writer(f)
+                #         writer.writerow((r[0].keys(), r[1]))
 
-def _worker(args):
-    (
-        distance,
-        app_tree1,
-        app_tree2,
-        norm,
-        cache,
-    ) = args
-    res = distance.compute_distance_parallel(app_tree1, app_tree2, cache)
-    distance = res[1] / distance.get_norm(app_tree1, app_tree2, norm)
+        return distances
 
-    return (res[0], distance)
+    def plot_clustermap(self, *trees, **kwargs): ...
diff --git a/src/lineagetree/lineage_tree.py b/src/lineagetree/lineage_tree.py
index 9dee0db..ea7971f 100644
--- a/src/lineagetree/lineage_tree.py
+++ b/src/lineagetree/lineage_tree.py
@@ -50,10 +50,10 @@ def __eq__(self, other) -> bool:
             return False
 
     def __hash__(self) -> int:
-        return hash(tuple(self._successor))
-        # succ = normalize_keys(self._successor)
-        # s = json.dumps(succ, sort_keys=True, separators=(",", ":"))
-        # return int(hashlib.sha256(s.encode("utf-8")).hexdigest(), 16)
+        used_nodes = self.roots.union(self.leaves)
+        return int.from_bytes(
+            hashlib.sha256(str(used_nodes).encode()).digest()[:8], "big"
+        )
 
     def __setstate__(self, state):
         if "_successor" not in state:

From 0f6f8d2e3836b45be999a2c9b8b462429a907e12 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 16:32:30 +0200
Subject: [PATCH 07/14] fix number of comparisons and added clustermap

---
 src/lineagetree/distance/approximations.py    |   5 +-
 src/lineagetree/distance/comparator.py        | 119 +++++++++++++++---
 .../distance/distance_calculator.py           |  12 +-
 tests/test_lineageTree.py                     |   4 +
 4 files changed, 119 insertions(+), 21 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index bb6d4b0..0277881 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -20,7 +20,7 @@
     from .approximations import TreeApproximationTemplate
 
 
-@dataclass
+@dataclass(frozen=True)
 class TreeSpecs:
     lT: int
     root: int
@@ -120,6 +120,9 @@ def __post_init__(self):
     def __str__(self):  # For quickly checking how the object was created.
         return str(self.tree_specs)
 
+    def __hash__(self):
+        return hash(self.tree_specs)
+
 
 class TreeApproximationTemplate(ABC):
     default_delta: Callable = ...
diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 20a8967..ae0e7e9 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -12,8 +12,6 @@
 from pathlib import Path
 import csv
 
-# import seaborn
-
 if TYPE_CHECKING:
     from lineagetree import LineageTree
     from .approximations import TreeApproximationTemplate
@@ -26,13 +24,11 @@ def _worker(args):
         distance,
         app_tree1,
         app_tree2,
-        norm,
         cache,
     ) = args
     res = distance.compute_distance_parallel(app_tree1, app_tree2, cache)
-    distance = res[1] / distance.get_norm(app_tree1, app_tree2, norm)
 
-    return (res[0], distance)
+    return res
 
 
 class TreeComparator:
@@ -120,36 +116,38 @@ def p_compare(
         for tree in tqdm.tqdm(trees, desc="Processing Trees: "):
             app_trees.append(self.__use_approximation(tree))
         combinations = list(itertools.combinations(app_trees, 2))
-        processes = (
-            (self.tree_distance, comb1, comb2, norm, self._cached_distances)
+        processes = [
+            (self.tree_distance, comb1, comb2, self._cached_distances)
             for comb1, comb2 in combinations
-        )
-        distances = []
+        ]
+        distances = {}
 
         ch_size = (
             len(combinations) // n_proccessors
         ) // 4  # Maybe overengineered but it works super good
         processes_b_s = sorted(processes, key=self._order_processes)
-        processes_s_b = processes_b_s[::-1]
+        # processes_s_b = processes_b_s[::-1]
         mid = len(processes_b_s) // 2
         high = processes_b_s[:mid]
-        low = processes_s_b[mid:]
+        low = processes_b_s[mid:]
         new_proc = []
         for h, l in zip(high, low):
             new_proc.append(h)
             new_proc.append(l)
 
-        new_proc.extend(high[len(low) :])
+        if len(high) != len(low):
+            new_proc.append(low[-1])
+        assert len(new_proc) == len(processes)
         with pool.Pool(min(n_proccessors, len(app_trees))) as p:
             for r in tqdm.tqdm(
                 p.imap_unordered(
-                    _worker, new_proc, chunksize=ch_size if ch_size > 0 else 1
+                    _worker, processes, chunksize=ch_size if ch_size > 0 else 1
                 ),
                 total=len(combinations),
                 desc="UTED distances",
             ):
                 self._cached_distances.update(r[0])
-                distances.append((r[0].keys(), r[1]))
+                distances.update(r[1])
 
                 # if path: # needs discussion
                 #     with open(path, "w", newline="") as f:
@@ -158,4 +156,95 @@ def p_compare(
 
         return distances
 
-    def plot_clustermap(self, *trees, **kwargs): ...
+    @property
+    def __get_next_lineage(self):
+        if not hasattr(self, "__id"):
+            self.__id = 0
+        self.__id += 1
+        return f"lT: {self.__id-1}"  # if you start naming your datasets lT 1, 2 .... fuck you.
+
+    def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
+        from scipy.cluster.hierarchy import dendrogram, linkage
+        from scipy.spatial.distance import squareform
+        import numpy as np
+        import matplotlib.pyplot as plt
+
+        if ax is None:
+            fig, ax = plt.subplots()
+
+        lts = {tree[0] for tree in trees}
+        one_lineage = False
+        if len(lts) == 1:
+            one_lineage = True
+        hash_to_name = {hash(lt): lt for lt in lts}
+        res = self.p_compare(*trees, **kwargs)
+        only_trees = set()
+        for tree in res.keys():
+            only_trees.update(tree)
+        matrix = np.zeros((len(trees), len(trees)))
+        for i, t1 in enumerate(only_trees):
+            for j, t2 in enumerate(only_trees):
+                if i != j:
+                    matrix[i, j] = matrix[j, i] = res[frozenset({t1, t2})]
+                else:
+                    matrix[i, j] = 0
+        if one_lineage:
+            labels = [
+                (
+                    hash_to_name[tr.tree_specs.lT].labels.get(
+                        tr.tree_specs.root, tr.tree_specs.root
+                    )
+                )
+                for tr in only_trees
+            ]
+        else:
+            labels = [
+                str(
+                    getattr(
+                        hash_to_name[tr.lT], "name", self.__get_next_lineage
+                    )
+                )
+                + (
+                    hash_to_name[tr.tree_specs.lT].labels.get(
+                        tr.tree_specs.root, tr.tree_specs.root
+                    )
+                )
+                for tr in only_trees
+            ]
+        condensed_dist_matrix = squareform(matrix)
+
+        linkage_data = linkage(condensed_dist_matrix, method="ward")
+        order = dendrogram(linkage_data, no_plot=True)["leaves"]
+        labels = [labels[i] for i in order]
+        print(labels)
+        plot = np.ix_(order, order)
+        plt.imshow(matrix[plot], **kwargs)
+        ax.set_xticks(np.arange(len(labels)), labels=labels)
+        ax.set_yticks(np.arange(len(labels)), labels=labels)
+        ax.tick_params(axis="both", labelsize=10)
+        plt.setp(
+            ax.get_xticklabels(),
+            rotation=45,
+            ha="right",
+            rotation_mode="anchor",
+        )
+        return fig, ax
+
+    def compare_all_trees_that_start_at_t(
+        self, lT: LineageTree, roots, time, end_time
+    ):
+        if roots:
+            new_roots = lT.nodes_at_t(
+                time,
+                roots,
+            )
+        else:
+            new_roots = lT.time_nodes[time]
+        trees = [(lT, r, end_time) for r in new_roots]
+        self.p_compare(*trees)
+
+    def plot_tree_distance_graph(
+        self,
+        tree1,
+        tree2,
+    ): ...
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index 92055ec..84bc086 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -59,7 +59,7 @@ def compute_distance_parallel(
         approximated_tree2: ApproximatedTree,
         backtrace: Alignment = None,
     ):
-        key = frozenset({str(approximated_tree1), str(approximated_tree2)})
+        key = frozenset({approximated_tree1, approximated_tree2})
         approximated_tree1, approximated_tree2 = sorted(
             (approximated_tree1, approximated_tree2), key=lambda x: str(x)
         )
@@ -69,9 +69,11 @@ def compute_distance_parallel(
 
         return (
             {key: align},
-            self.compute_distance(
-                approximated_tree1, approximated_tree2, backtrace
-            ),
+            {
+                key: self.compute_distance(
+                    approximated_tree1, approximated_tree2, backtrace
+                )
+            },
         )
 
     def _compute_uted_backtrace(
@@ -97,7 +99,7 @@ def _compute_uted_backtrace(
         backtrace : edist.alignment.Alignment
             The resulting alignment in the edist format
         """
-        key = frozenset({str(approximated_tree1), str(approximated_tree2)})
+        key = frozenset({approximated_tree1, approximated_tree2})
 
         approximated_tree1, approximated_tree2 = sorted(
             (approximated_tree1, approximated_tree2), key=lambda x: str(x)
diff --git a/tests/test_lineageTree.py b/tests/test_lineageTree.py
index 35e1dce..5be1e45 100644
--- a/tests/test_lineageTree.py
+++ b/tests/test_lineageTree.py
@@ -797,3 +797,7 @@ def test_smoothing():
         new_pos[1552], np.array([462.15385069, 907.17562352, 419.54303692])
     ).all()
     lt.pos = lt.old_pos
+
+
+def test_hash():
+    assert hash(lt) == 7929439451057191037

From 3474db9d839793671145f4e76158b242f33fb7ba Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 17:20:36 +0200
Subject: [PATCH 08/14] .

---
 src/lineagetree/distance/comparator.py        | 38 ++++++++-----------
 .../distance/distance_calculator.py           | 12 ++++--
 2 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index ae0e7e9..a1ae488 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -95,23 +95,12 @@ def compare(
 
         return distance
 
-    def _order_processes(self, proc):
-        nodes1, nodes2 = len(proc[1].nodes), len(proc[2].nodes)
-        return nodes1 + nodes2
-
     def p_compare(
         self,
         *trees,
         norm="sum",
         n_proccessors: int = 4,
-        path: Path | None = None,
     ):
-        if path:
-            for tree in trees:
-                if tree[0].name is None:
-                    raise Warning(
-                        f"All LineageTrees should have a name for saving the comparisons. Missing name: {tree[0]}"
-                    )
         app_trees = []
         for tree in tqdm.tqdm(trees, desc="Processing Trees: "):
             app_trees.append(self.__use_approximation(tree))
@@ -125,7 +114,7 @@ def p_compare(
         ch_size = (
             len(combinations) // n_proccessors
         ) // 4  # Maybe overengineered but it works super good
-        processes_b_s = sorted(processes, key=self._order_processes)
+        processes_b_s = sorted(processes, key=_order_processes)
         # processes_s_b = processes_b_s[::-1]
         mid = len(processes_b_s) // 2
         high = processes_b_s[:mid]
@@ -148,13 +137,10 @@ def p_compare(
             ):
                 self._cached_distances.update(r[0])
                 distances.update(r[1])
-
-                # if path: # needs discussion
-                #     with open(path, "w", newline="") as f:
-                #         writer = csv.writer(f)
-                #         writer.writerow((r[0].keys(), r[1]))
-
-        return distances
+        return {
+            key: dist / self.tree_distance.get_norm(*key, norm)
+            for key, dist in distances.items()
+        }
 
     @property
     def __get_next_lineage(self):
@@ -216,9 +202,10 @@ def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
         linkage_data = linkage(condensed_dist_matrix, method="ward")
         order = dendrogram(linkage_data, no_plot=True)["leaves"]
         labels = [labels[i] for i in order]
-        print(labels)
+
         plot = np.ix_(order, order)
-        plt.imshow(matrix[plot], **kwargs)
+        _pl = plt.imshow(matrix[plot], **kwargs)
+        self.colorbar = fig.colorbar(_pl, ax=ax)
         ax.set_xticks(np.arange(len(labels)), labels=labels)
         ax.set_yticks(np.arange(len(labels)), labels=labels)
         ax.tick_params(axis="both", labelsize=10)
@@ -231,7 +218,7 @@ def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
         return fig, ax
 
     def compare_all_trees_that_start_at_t(
-        self, lT: LineageTree, roots, time, end_time
+        self, lT: LineageTree, time, roots, end_time, norm, n_processors=4
     ):
         if roots:
             new_roots = lT.nodes_at_t(
@@ -241,10 +228,15 @@ def compare_all_trees_that_start_at_t(
         else:
             new_roots = lT.time_nodes[time]
         trees = [(lT, r, end_time) for r in new_roots]
-        self.p_compare(*trees)
+        self.p_compare(*trees, norm=norm, n_proccessors=n_processors)
 
     def plot_tree_distance_graph(
         self,
         tree1,
         tree2,
     ): ...
+
+
+def _order_processes(proc):
+    nodes1, nodes2 = len(proc[1].nodes), len(proc[2].nodes)
+    return nodes1 + nodes2
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index 84bc086..a3dc7ec 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -19,13 +19,17 @@ def compute_distance_parallel(
         approximated_tree1: ApproximatedTree | LineageTree,
         approximated_tree2: ApproximatedTree | LineageTree,
         backtrace: Alignment = None,
-    ) -> dict:
+    ) -> tuple[dict, dict] | dict:
         warn(
             "Parallel algorithm not implemented, defaulting to single process."
         )
-        self.compute_distance(
-            approximated_tree1, approximated_tree2, backtrace
-        )
+        return {
+            frozenset(
+                {approximated_tree1, approximated_tree2}
+            ): self.compute_distance(
+                approximated_tree1, approximated_tree2, backtrace
+            )
+        }
 
     @abstractmethod
     def compute_distance(

From 0451df043f34dd53975d93fcb3938e6c8e062c4d Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 17:33:23 +0200
Subject: [PATCH 09/14] fix compare trees at t

---
 src/lineagetree/distance/comparator.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index a1ae488..7938e42 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -137,6 +137,7 @@ def p_compare(
             ):
                 self._cached_distances.update(r[0])
                 distances.update(r[1])
+
         return {
             key: dist / self.tree_distance.get_norm(*key, norm)
             for key, dist in distances.items()
@@ -218,7 +219,13 @@ def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
         return fig, ax
 
     def compare_all_trees_that_start_at_t(
-        self, lT: LineageTree, time, roots, end_time, norm, n_processors=4
+        self,
+        lT: LineageTree,
+        time,
+        norm="sum",
+        roots=None,
+        end_time=None,
+        n_processors=4,
     ):
         if roots:
             new_roots = lT.nodes_at_t(
@@ -228,7 +235,7 @@ def compare_all_trees_that_start_at_t(
         else:
             new_roots = lT.time_nodes[time]
         trees = [(lT, r, end_time) for r in new_roots]
-        self.p_compare(*trees, norm=norm, n_proccessors=n_processors)
+        return self.p_compare(*trees, norm=norm, n_proccessors=n_processors)
 
     def plot_tree_distance_graph(
         self,

From 0e59b542683ba65bacdd6db14da03042b4dd9edf Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 17:41:04 +0200
Subject: [PATCH 10/14] ruff black

---
 src/lineagetree/distance/comparator.py | 7 +------
 src/lineagetree/lineage_tree.py        | 7 -------
 2 files changed, 1 insertion(+), 13 deletions(-)

diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 7938e42..17036e5 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -1,16 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable
-import inspect
-import pickle
+from typing import TYPE_CHECKING
 from warnings import warn
 from multiprocessing import pool
-import types
 from .approximations import ApproximatedTree
 import itertools
 import tqdm
-from pathlib import Path
-import csv
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
diff --git a/src/lineagetree/lineage_tree.py b/src/lineagetree/lineage_tree.py
index ea7971f..22162d3 100644
--- a/src/lineagetree/lineage_tree.py
+++ b/src/lineagetree/lineage_tree.py
@@ -10,7 +10,6 @@
 from packaging.version import Version
 import numpy as np
 import hashlib
-import json
 from ._core.utils import CompatibleUnpickler
 from ._mixins.properties_mixin import PropertiesMixin
 from ._mixins.modifier_mixin import ModifierMixin
@@ -22,12 +21,6 @@
 from ._core.validation import TreeValidator
 
 
-def normalize_keys(d):
-    return tuple(
-        sorted((int(k), tuple(int(vv) for vv in v)) for k, v in d.items())
-    )
-
-
 class LineageTree(
     PropertiesMixin,
     ModifierMixin,

From 1bb7f6aac436d6350345e260b54c5a2cec7341a6 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 22:01:45 +0200
Subject: [PATCH 11/14] el comparator docs

---
 src/lineagetree/distance/comparator.py | 257 +++++++++++++++++++++++--
 1 file changed, 242 insertions(+), 15 deletions(-)

diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py
index 17036e5..ce77bf5 100644
--- a/src/lineagetree/distance/comparator.py
+++ b/src/lineagetree/distance/comparator.py
@@ -6,12 +6,18 @@
 from .approximations import ApproximatedTree
 import itertools
 import tqdm
+from scipy.cluster.hierarchy import dendrogram, linkage
+from scipy.spatial.distance import squareform
+import numpy as np
+import matplotlib.pyplot as plt
 
 if TYPE_CHECKING:
     from lineagetree import LineageTree
     from .approximations import TreeApproximationTemplate
     from .distance_calculator import TreeDistanceTemplate
     from edist.alignment import Alignment
+    from typing import Callable
+    import matplotlib
 
 
 def _worker(args):
@@ -27,6 +33,11 @@ def _worker(args):
 
 
 class TreeComparator:
+    """Handler class for comparing lineages. The parameters of the initializer
+    are the tree distance algorithm (A class that dictates how the lineages are to be compared)
+    and the approximator (A class that approximates/transforms the lineage to the format appropriate
+    for comparing lineages)
+    """
 
     def __init__(
         self,
@@ -49,7 +60,18 @@ def __init__(
     @property
     def _cached_distances(
         self,
-    ) -> dict[frozenset[tuple[int, int, int]], Alignment | float | int]:
+    ) -> dict[
+        frozenset[tuple[ApproximatedTree, ApproximatedTree]],
+        Alignment | float | int,
+    ]:
+        """This cache is where the distances are saved. For alignment based distances usually the
+        alignment is saved instead.
+
+        Returns
+        -------
+        dict[frozenset[tuple[ApproximatedTree, ApproximatedTree], Alignment | float | int]
+            A dict that contains a frozenset of 2 approximated trees mapped to their distance/alignment.
+        """
         if not hasattr(self, "_cache"):
             self._cache = {}
         return self._cache
@@ -58,6 +80,7 @@ def _cached_distances(
     def _cached_approximations(
         self,
     ) -> dict[tuple[int, int, int], ApproximatedTree]:
+        """The cache of the approximations."""
         if not hasattr(self, "_approximations"):
             self._approximations = {}
         return self._approximations
@@ -65,6 +88,23 @@ def _cached_approximations(
     def __use_approximation(
         self, tree: tuple[LineageTree, int, int] | ApproximatedTree
     ) -> ApproximatedTree:
+        """Helping function that takes a lineagetree a node and an end_time to create
+        tne approximated_tree. If an approximated tree is inserted instead just return that.
+
+        Parameters
+        ----------
+        tree : tuple[LineageTree, int, int] | ApproximatedTree
+               - A tuple ``(lineage_tree, root, end_time)`` containing:
+                    - ``lineage_tree``: a ``LineageTree`` instance.
+                    - ``root``: the root node identifier.
+                    - ``end_time``: the final time step.
+                - An ``ApproximatedTree`` instance.
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree object.
+        """
         if isinstance(tree, tuple):
             new_key = (hash(tree[0]), tree[1], tree[2])
             if new_key in self._cached_approximations:
@@ -80,8 +120,34 @@ def compare(
         self,
         tree1: tuple[LineageTree, int, int] | ApproximatedTree,
         tree2: tuple[LineageTree, int, int] | ApproximatedTree,
-        norm: {"max", "sum", "tuple"} | Callable = "sum",
+        norm: Literal["max", "sum", "tuple"] | Callable = "sum",
     ) -> float | int:
+        """Compares two trees.
+
+        Parameters
+        ----------
+        tree1 : tuple[LineageTree, int, int] | ApproximatedTree
+            tree : tuple[LineageTree, int, int] | ApproximatedTree
+               - A tuple ``(lineage_tree, root, end_time)`` containing:
+                    - ``lineage_tree``: a ``LineageTree`` instance.
+                    - ``root``: the root node identifier.
+                    - ``end_time``: the final time step.
+                - An ``ApproximatedTree`` instance.
+        tree2 : tuple[LineageTree, int, int] | ApproximatedTree
+            tree : tuple[LineageTree, int, int] | ApproximatedTree
+               - A tuple ``(lineage_tree, root, end_time)`` containing:
+                    - ``lineage_tree``: a ``LineageTree`` instance.
+                    - ``root``: the root node identifier.
+                    - ``end_time``: the final time step.
+                - An ``ApproximatedTree`` instance.
+        norm : `max`,`sum` or tuple, by default `sum`
+            How the distances should be normalized, by default "sum"
+
+        Returns
+        -------
+        float | int
+            The distance between 2 trees
+        """
         tree1 = self.__use_approximation(tree1)
         tree2 = self.__use_approximation(tree2)
         distance = self.tree_distance.compute_distance(
@@ -93,9 +159,37 @@ def compare(
     def p_compare(
         self,
         *trees,
-        norm="sum",
-        n_proccessors: int = 4,
+        norm: Literal["max", "sum", "tuple"] | Callable = "sum",
+        n_processors: int = 4,
     ):
+        """Compute pairwise distances between multiple trees in parallel.
+
+        Parameters
+        ----------
+        *trees : tuple[LineageTree, int, int] | ApproximatedTree
+            Trees to compare.
+
+            Each tree can be either:
+
+            - A tuple ``(lineage_tree, root, end_time)``, where:
+                - ``lineage_tree`` is a ``LineageTree`` instance.
+                - ``root`` is the identifier of the root node.
+                - ``end_time`` is the final time step to consider.
+            - An ``ApproximatedTree`` instance.
+
+        norm : {"max", "sum", "tuple"} | Callable, default="sum"
+            Normalization strategy applied to the computed distances.
+            May be one of ``"max"``, ``"sum"``, ``"tuple"``, or a custom
+            callable.
+
+        n_proccessors : int, default=4
+            Number of worker processes used to perform the comparisons.
+
+        Returns
+        -------
+        list[float | int]
+            Distances computed for all requested tree comparisons.
+        """
         app_trees = []
         for tree in tqdm.tqdm(trees, desc="Processing Trees: "):
             app_trees.append(self.__use_approximation(tree))
@@ -107,7 +201,7 @@ def p_compare(
         distances = {}
 
         ch_size = (
-            len(combinations) // n_proccessors
+            len(combinations) // n_processors
         ) // 4  # Maybe overengineered but it works super good
         processes_b_s = sorted(processes, key=_order_processes)
         # processes_s_b = processes_b_s[::-1]
@@ -122,7 +216,7 @@ def p_compare(
         if len(high) != len(low):
             new_proc.append(low[-1])
         assert len(new_proc) == len(processes)
-        with pool.Pool(min(n_proccessors, len(app_trees))) as p:
+        with pool.Pool(min(n_processors, len(app_trees))) as p:
             for r in tqdm.tqdm(
                 p.imap_unordered(
                     _worker, processes, chunksize=ch_size if ch_size > 0 else 1
@@ -139,17 +233,53 @@ def p_compare(
         }
 
     @property
-    def __get_next_lineage(self):
+    def __get_next_lineage(self) -> str:
+        """Returns a name for a nameless dataset."""
         if not hasattr(self, "__id"):
             self.__id = 0
         self.__id += 1
         return f"lT: {self.__id-1}"  # if you start naming your datasets lT 1, 2 .... fuck you.
 
-    def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
-        from scipy.cluster.hierarchy import dendrogram, linkage
-        from scipy.spatial.distance import squareform
-        import numpy as np
-        import matplotlib.pyplot as plt
+    def plot_clustermap(
+        self,
+        *trees,
+        norm: Literal["max", "sum", "tuple"] | Callable = "sum",
+        ax=None,
+        **kwargs,
+    ):
+        """Compute pairwise tree distances using `p_compare` and visualize them
+         as a hierarchical clustermap.
+
+         Parameters
+         ----------
+         *trees : tuple[LineageTree, int, int] | ApproximatedTree
+             Trees to compare.
+
+             Each tree can be either:
+
+             - A tuple ``(lineage_tree, root, end_time)``, where:
+                 - ``lineage_tree`` is a ``LineageTree`` instance.
+                 - ``root`` is the identifier of the root node.
+                 - ``end_time`` is the final time step.
+             - An ``ApproximatedTree`` instance.
+
+         norm : {"max", "sum", "tuple"} | Callable, by default="sum"
+             Normalization method, by default `sum
+
+         ax : matplotlib.axes.Axes, optional
+             Matplotlib axis to draw the clustermap on. If None, a new figure
+             is created, by default `None
+
+         **kwargs :
+             Additional keyword arguments forwarded to the plt.imshow function.
+
+         Returns
+        fig : matplotlib.figure.Figure
+             The matplotlib figure containing the clustermap.
+
+         ax : matplotlib.axes.Axes
+             The axes object associated with the clustermap.
+        """
 
         if ax is None:
             fig, ax = plt.subplots()
@@ -159,7 +289,7 @@ def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
         if len(lts) == 1:
             one_lineage = True
         hash_to_name = {hash(lt): lt for lt in lts}
-        res = self.p_compare(*trees, **kwargs)
+        res = self.p_compare(*trees, norm=norm, **kwargs)
         only_trees = set()
         for tree in res.keys():
             only_trees.update(tree)
@@ -214,14 +344,108 @@ def plot_clustermap(self, *trees, norm, ax=None, **kwargs):
         return fig, ax
 
     def compare_all_trees_that_start_at_t(
+        self,
+        lT: LineageTree,
+        time: int,
+        norm: Literal["max", "sum", "tuple"] | Callable = "sum",
+        roots: list | None = None,
+        end_time: int | None = None,
+        n_processors=4,
+    ) -> dict[frozenset[tuple[ApproximatedTree, ApproximatedTree]], float]:
+        """
+        Compare all trees that originate from a given lineage tree at a specific time.
+
+        This function computes pairwise distances between all subtrees of a given
+        lineage tree that start at a specified time point.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The input lineage tree from which comparisons will be derived.
+
+        time : int
+            The time point at which to extract subtrees for comparison.
+
+        norm : {"max", "sum", "tuple"} | Callable, by default="sum"
+            Normalization method, by default `sum`
+
+        roots : list of int or None, optional
+            Specific root node IDs to include in the comparison.
+            If None, all valid roots at the given time are used.
+
+        end_time : int or None, optional
+            Optional cutoff time for subtree extraction.
+            If None, uses the full available time range.
+
+        n_processors : int, default=4
+            Number of parallel processes used for computing pairwise distances.
+
+        Returns
+        -------
+        dict
+            A dictionary where keys are frozensets of TreeSpecs pairs and
+            values are the corresponding pairwise distances.
+        """
+
+        if roots:
+            new_roots = lT.nodes_at_t(
+                time,
+                roots,
+            )
+        else:
+            new_roots = lT.time_nodes[time]
+        trees = [(lT, r, end_time) for r in new_roots]
+        return self.p_compare(*trees, norm=norm, n_processors=n_processors)
+
+    def clustermap__all_trees_that_start_at_t(
         self,
         lT: LineageTree,
         time,
-        norm="sum",
+        norm: Literal["max", "sum", "tuple"] | Callable = "sum",
         roots=None,
         end_time=None,
         n_processors=4,
+        **kwargs,
     ):
+        """
+        Compute and visualize a clustermap of all tree distances starting at a given time.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            Input lineage tree from which subtrees are extracted.
+
+        time : int
+            Time point at which subtree extraction begins.
+
+        norm : {"max", "sum", "tuple"} or callable, default="sum"
+           Normalization method, by default `sum`
+
+        roots : list, optional
+            Specific root nodes to include in the analysis. If None, all
+            nodes at a specific timepoint are used, by default None
+
+        end_time : int, optional
+            Optional cutoff time for subtree extraction. If None, uses the whole tree, by default None
+
+        n_processors : int, default=4
+            Number of parallel workers used for distance computation.
+
+        **kwargs :
+            Additional keyword arguments passed to matplotlib.pyplot.imshow function
+
+        Returns
+        -------
+        matplotlib.axes.Axes
+            The axes object containing the clustermap visualization.
+
+        Notes
+        -----
+        - Internally computes pairwise distances using a parallelized backend.
+        - The resulting matrix is hierarchically clustered for visualization.
+        - Output ordering is determined by linkage dendrogram reordering.
+        """
+
         if roots:
             new_roots = lT.nodes_at_t(
                 time,
@@ -230,7 +454,9 @@ def compare_all_trees_that_start_at_t(
         else:
             new_roots = lT.time_nodes[time]
         trees = [(lT, r, end_time) for r in new_roots]
-        return self.p_compare(*trees, norm=norm, n_proccessors=n_processors)
+        return self.plot_clustermap(
+            *trees, norm=norm, n_processors=n_processors, kwargs=kwargs
+        )
 
     def plot_tree_distance_graph(
         self,
@@ -240,5 +466,6 @@ def plot_tree_distance_graph(
 
 
 def _order_processes(proc):
+    """Calculates the total number of nodes each process in p_comparisons has to deal with."""
     nodes1, nodes2 = len(proc[1].nodes), len(proc[2].nodes)
     return nodes1 + nodes2

From 845c64676d4156a28b5e9653db48492353bff693 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 22:26:23 +0200
Subject: [PATCH 12/14] docstrings

---
 src/lineagetree/distance/approximations.py    |  2 +-
 .../distance/distance_calculator.py           | 36 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 0277881..7a057aa 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -349,7 +349,7 @@ def approximation(
         )
 
 
-class ResampledTree(TreeApproximationTemplate):
+class ResampledTree(TreeApproximationTemplate):  ### TODO
     """
     An approximator that resample lineage trees to a given time resolution.
     The target time resolution is provided at creation.
diff --git a/src/lineagetree/distance/distance_calculator.py b/src/lineagetree/distance/distance_calculator.py
index a3dc7ec..9926d49 100644
--- a/src/lineagetree/distance/distance_calculator.py
+++ b/src/lineagetree/distance/distance_calculator.py
@@ -13,6 +13,7 @@
 
 
 class TreeDistanceTemplate(ABC):
+    """The template class to create new distance metrics"""
 
     def compute_distance_parallel(
         self,
@@ -20,6 +21,24 @@ def compute_distance_parallel(
         approximated_tree2: ApproximatedTree | LineageTree,
         backtrace: Alignment = None,
     ) -> tuple[dict, dict] | dict:
+        """Computes the distance between 2 trees. Important to implement
+        if parallel processing is gonna be used.
+
+        Parameters
+        ----------
+        approximated_tree1 : ApproximatedTree | LineageTree
+            The first tree to compare
+        approximated_tree2 : ApproximatedTree | LineageTree
+            The secnd tree to compare
+        backtrace : Alignment, optional
+            The cache where the distances/alignments are gonna be saved or extracted from, by default None
+
+        Returns
+        -------
+        tuple[dict, dict] | dict
+            - A dictionary that maps the two trees to a distance
+            - Two dictinaries that map the 2 trees to a distance and their backtrace.
+        """
         warn(
             "Parallel algorithm not implemented, defaulting to single process."
         )
@@ -49,6 +68,7 @@ def get_norm(
 
 
 class UnorderedTreeEditDistance(TreeDistanceTemplate):
+    """Specific implementation for Unordered Tree Edit Distance."""
 
     backtrace_based = True
 
@@ -62,7 +82,21 @@ def compute_distance_parallel(
         approximated_tree1: ApproximatedTree,
         approximated_tree2: ApproximatedTree,
         backtrace: Alignment = None,
-    ):
+    ) -> tuple[dict, dict] | dict:
+        """Performs uted in parallel.
+
+        Parameters
+        ----------
+        approximated_tree1 : ApproximatedTree
+            The first tree to compare
+        approximated_tree2 : ApproximatedTree
+            The second tree to compare
+
+        Returns
+        -------
+        backtrace : edist.alignment.Alignment
+            The resulting alignment in the edist format
+        """
         key = frozenset({approximated_tree1, approximated_tree2})
         approximated_tree1, approximated_tree2 = sorted(
             (approximated_tree1, approximated_tree2), key=lambda x: str(x)

From 135a4a0567aff611dacc7a8779771c9d81ffe84e Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Tue, 16 Jun 2026 23:22:50 +0200
Subject: [PATCH 13/14] docs :)

---
 src/lineagetree/distance/approximations.py | 90 ++++++++++++++++++----
 1 file changed, 77 insertions(+), 13 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 7a057aa..4cf4205 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -22,6 +22,8 @@
 
 @dataclass(frozen=True)
 class TreeSpecs:
+    """Serves as the identifier of an `ApproximatedTree`"""
+
     lT: int
     root: int
     end_time: int
@@ -41,6 +43,8 @@ def __hash__(self):
 
 @dataclass
 class ApproximatedTree:
+    """A tree that has been processed and i ready for comparison should be converted to this class."""
+
     adjacency_dict: dict[int, Iterable[int]]
     property_dict: dict[int, float | int | list[int | float]] | None
     tree_specs: TreeSpecs | None = None
@@ -120,11 +124,16 @@ def __post_init__(self):
     def __str__(self):  # For quickly checking how the object was created.
         return str(self.tree_specs)
 
-    def __hash__(self):
+    def __hash__(self):  # Makes it hashable.
         return hash(self.tree_specs)
 
 
 class TreeApproximationTemplate(ABC):
+    """Template to create approximations
+    One function the `approximation` has to be implemented
+    and there should be a `default_delta` set
+    """
+
     default_delta: Callable = ...
 
     def __init__(
@@ -139,7 +148,26 @@ def approximation(
         lT: LineageTree,
         root: int,
         end_time: int | None = None,
-    ) -> ApproximatedTree: ...
+    ) -> ApproximatedTree:
+        """Build the approximation of a given sub lineage of lineage tree.
+
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        end_time : int
+            The last time point to consider
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
 
 
 class ReducedTreeTimed(TreeApproximationTemplate):
@@ -151,6 +179,25 @@ def approximation(
         root: int,
         end_time: int | None = None,
     ) -> ApproximatedTree:
+        """Build the approximation of a given sub lineage of lineage tree.
+
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        end_time : int
+            The last time point to consider
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
         if end_time is None:
             end_time = lT.t_e
         if lT.time_resolution == 0:
@@ -209,10 +256,6 @@ def approximation(
             to approximate
         end_time : int
             The last time point to consider
-        time_resolution : float
-            How much time happens between two consecutive time points
-            This is useful when comparing two different lineage trees
-            that do not have the same time scale
         properties : dict mapping int to float, optional
             The dictionary that maps a node id to a given value.
             *Not* all the nodes need to have a value.
@@ -310,9 +353,9 @@ def approximation(
             to approximate
         end_time : int
             The last time point to consider
-        downsample : int
-            The downsampling value.
-            One node will be conserved every `downsample` node
+        properties : dict mapping int to float, optional
+            The dictionary that maps a node id to a given value.
+            *Not* all the nodes need to have a value.
 
         Returns
         -------
@@ -402,10 +445,9 @@ def approximation(
             and therefore only the topology of the lineage tree will matter
         end_time : int
             The last time point to consider
-        time_resolution : int
-            The time resolution of the lineage tree
-            to approximate. It has to be in the same
-            units than the one provided as `target_time_resolution`.
+        properties : dict mapping int to float, optional
+            The dictionary that maps a node id to a given value.
+            *Not* all the nodes need to have a value.
 
         Returns
         -------
@@ -545,6 +587,28 @@ def approximation(
         end_time: int | None = None,
         properties: dict[int, float | list] | list[str] = None,
     ):
+        """
+        Build the approximation of a given sub lineage of lineage tree.
+
+        Parameters
+        ----------
+        lT : LineageTree
+            The lineage tree from which to approximate
+        root : int
+            The id of the spawning cell of the sub lineage
+            to approximate
+        end_time : int
+            The last time point to consider
+        properties : dict mapping int to float, optional
+            The dictionary that maps a node id to a given value.
+            *Not* all the nodes need to have a value.
+
+        Returns
+        -------
+        ApproximatedTree
+            The approximated tree that will be used
+            for the computation of the tree edit distance
+        """
 
         if isinstance(properties, list):
             prop_dicts = [getattr(lT, prop) for prop in properties]

From 5bea21e653d0b8b619ed9ea7d9c4b246e7c71c84 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Wed, 17 Jun 2026 10:55:38 +0200
Subject: [PATCH 14/14] .

---
 src/lineagetree/distance/approximations.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py
index 4cf4205..e513774 100644
--- a/src/lineagetree/distance/approximations.py
+++ b/src/lineagetree/distance/approximations.py
@@ -22,7 +22,7 @@
 
 @dataclass(frozen=True)
 class TreeSpecs:
-    """Serves as the identifier of an `ApproximatedTree`"""
+    """Serves as the identifier of an `ApproximatedTree` an should be used when transforming trees to a form that is to be compared."""
 
     lT: int
     root: int
@@ -43,7 +43,7 @@ def __hash__(self):
 
 @dataclass
 class ApproximatedTree:
-    """A tree that has been processed and i ready for comparison should be converted to this class."""
+    """A tree that has been processed and iw ready for comparison should be converted to this class."""
 
     adjacency_dict: dict[int, Iterable[int]]
     property_dict: dict[int, float | int | list[int | float]] | None