diff --git a/src/lineagetree/distance/approximations.py b/src/lineagetree/distance/approximations.py new file mode 100644 index 0000000..e513774 --- /dev/null +++ b/src/lineagetree/distance/approximations.py @@ -0,0 +1,635 @@ +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 .delta import ( + delta_normalized_difference, + delta_nd_norm, + delta_difference, + delta_binary, +) + +if TYPE_CHECKING: + from lineagetree import LineageTree + from .approximations import TreeApproximationTemplate + + +@dataclass(frozen=True) +class TreeSpecs: + """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 + 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: + """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 + tree_specs: TreeSpecs | 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())) + if isinstance(sample, Iterable): + length = len(sample) + default_value = [0] * length + 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} + 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 str(self.tree_specs) + + 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__( + self, + delta: Callable | None = None, + ): + self.delta = delta if delta else self.__class__.default_delta # :( + + @abstractmethod + def approximation( + self, + lT: LineageTree, + 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 + """ + + +class ReducedTreeTimed(TreeApproximationTemplate): + default_delta = delta_normalized_difference + + def approximation( + self, + lT: LineageTree, + 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: + 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, TreeSpecs(hash(lT), 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: LineageTree, + root: int, + 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 + 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, TreeSpecs(hash(lT), 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 + 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 + 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, TreeSpecs(hash(lT), root, end_time) + ) + + +class ResampledTree(TreeApproximationTemplate): ### TODO + """ + 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 + 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 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, TreeSpecs(hash(lT), 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: 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 + 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] + 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, TreeSpecs(hash(lT), root, end_time) + ) diff --git a/src/lineagetree/distance/comparator.py b/src/lineagetree/distance/comparator.py new file mode 100644 index 0000000..ce77bf5 --- /dev/null +++ b/src/lineagetree/distance/comparator.py @@ -0,0 +1,471 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from warnings import warn +from multiprocessing import pool +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): + ( + distance, + app_tree1, + app_tree2, + cache, + ) = args + res = distance.compute_distance_parallel(app_tree1, app_tree2, cache) + + return res + + +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, + tree_distance: TreeDistanceTemplate, + approximator: TreeApproximationTemplate | None = None, + ): + + 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.approximator.delta) + self.labels = {} + self.trees = {} + + @property + def _cached_distances( + self, + ) -> 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 + + @property + 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 + + 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: + f_tree = self._cached_approximations[new_key] + else: + f_tree = self.approximator.approximation(*tree) + self._cached_approximations[new_key] = f_tree + else: + f_tree = tree + return f_tree + + def compare( + self, + tree1: tuple[LineageTree, int, int] | ApproximatedTree, + tree2: tuple[LineageTree, int, int] | ApproximatedTree, + 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( + tree1, tree2, self._cached_distances + ) / self.tree_distance.get_norm(tree1, tree2, norm) + + return distance + + def p_compare( + self, + *trees, + 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)) + combinations = list(itertools.combinations(app_trees, 2)) + processes = [ + (self.tree_distance, comb1, comb2, self._cached_distances) + for comb1, comb2 in combinations + ] + distances = {} + + ch_size = ( + 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] + mid = len(processes_b_s) // 2 + high = processes_b_s[:mid] + low = processes_b_s[mid:] + new_proc = [] + for h, l in zip(high, low): + new_proc.append(h) + new_proc.append(l) + + if len(high) != len(low): + new_proc.append(low[-1]) + assert len(new_proc) == len(processes) + 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 + ), + total=len(combinations), + desc="UTED distances", + ): + 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() + } + + @property + 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: 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() + + 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, norm=norm, **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] + + plot = np.ix_(order, order) + _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) + 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, + 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: 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, + roots, + ) + else: + new_roots = lT.time_nodes[time] + trees = [(lT, r, end_time) for r in new_roots] + return self.plot_clustermap( + *trees, norm=norm, n_processors=n_processors, kwargs=kwargs + ) + + def plot_tree_distance_graph( + self, + tree1, + tree2, + ): ... + + +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 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..9926d49 --- /dev/null +++ b/src/lineagetree/distance/distance_calculator.py @@ -0,0 +1,255 @@ +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 + from .approximations import ApproximatedTree + from lineagetree import LineageTree + + +class TreeDistanceTemplate(ABC): + """The template class to create new distance metrics""" + + def compute_distance_parallel( + self, + approximated_tree1: ApproximatedTree | LineageTree, + 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." + ) + return { + frozenset( + {approximated_tree1, approximated_tree2} + ): self.compute_distance( + approximated_tree1, approximated_tree2, backtrace + ) + } + + @abstractmethod + def compute_distance( + self, + approximated_tree1: ApproximatedTree | LineageTree, + approximated_tree2: ApproximatedTree | LineageTree, + backtrace: Alignment = None, + ) -> float: ... + + @abstractmethod + def get_norm( + self, + tree1: ApproximatedTree, + tree2: ApproximatedTree, + norm_type: {"max", "sum", "tuple"} = "sum", + ) -> float: ... + + +class UnorderedTreeEditDistance(TreeDistanceTemplate): + """Specific implementation for Unordered Tree Edit Distance.""" + + backtrace_based = True + + def __init__(self, delta): + super().__init__() + self.delta = delta + self.backtrace = {} + + def compute_distance_parallel( + self, + 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) + ) + align = self._compute_uted_backtrace( + approximated_tree1, approximated_tree2, backtrace + ) + + return ( + {key: align}, + { + key: self.compute_distance( + approximated_tree1, approximated_tree2, backtrace + ) + }, + ) + + def _compute_uted_backtrace( + self, + approximated_tree1: ApproximatedTree, + approximated_tree2: ApproximatedTree, + cache: dict[frozenset[tuple], Alignment] | None = None, + ) -> 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 + """ + key = frozenset({approximated_tree1, 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: + 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( + self, + approximated_tree1: ApproximatedTree, + approximated_tree2: ApproximatedTree, + btrc: dict[frozenset[tuple], Alignment] | None = None, + ) -> 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. + """ + 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) + + def reconstruct_backtrace(self, tree1, tree2, backtrace): ... + + def get_norm( + self, + tree1: ApproximatedTree, + tree2: ApproximatedTree, + norm_type: {"max", "sum", "tuple"} | Callable = "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 + ) + 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 _: + raise ValueError( + f"Invalid value for `norm_type`. Got {norm_type}," + "expected 'max', 'sum' or 'tuple'" + ) diff --git a/src/lineagetree/lineage_tree.py b/src/lineagetree/lineage_tree.py index 3699928..22162d3 100644 --- a/src/lineagetree/lineage_tree.py +++ b/src/lineagetree/lineage_tree.py @@ -9,7 +9,7 @@ from collections.abc import Iterable, Sequence from packaging.version import Version import numpy as np - +import hashlib from ._core.utils import CompatibleUnpickler from ._mixins.properties_mixin import PropertiesMixin from ._mixins.modifier_mixin import ModifierMixin @@ -42,6 +42,12 @@ def __eq__(self, other) -> bool: else: return False + def __hash__(self) -> int: + 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: state["_successor"] = state["successor"] diff --git a/tests/test_lineageTree.py b/tests/test_lineageTree.py index b349042..5be1e45 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 ( @@ -798,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