New labelling system

src/lineagetree/_core/_navigation.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 import warnings
+from .._labelling import Labels
 
 if TYPE_CHECKING:
     from ..lineage_tree import LineageTree
@@ -297,7 +298,9 @@ def get_ancestor_at_t(lT: LineageTree, n: int, time: int | None = None) -> int:
         return -1
 
 
-def get_labelled_ancestor(lT: LineageTree, node: int) -> int:
+def get_labelled_ancestor(
+    lT: LineageTree, node: int, labels: str | None = None
+) -> int:
     """Finds the first labelled ancestor and returns its ID otherwise returns -1
 
     Parameters
@@ -315,8 +318,13 @@ def get_labelled_ancestor(lT: LineageTree, node: int) -> int:
     if node not in lT.nodes:
         return -1
     ancestor = node
+    if labels is not None and labels not in lT.labelling.list_of_labels:
+        raise ValueError("Label set not defined.")
+    labs = (
+        getattr(lT.labelling, labels) if labels else lT.labelling.default_dict
+    )
     while lT.t_b <= lT._time.get(ancestor, lT.t_b - 1) and ancestor != -1:
-        if ancestor in lT.labels:
+        if ancestor in labs:
             return ancestor
         ancestor = lT._predecessor.get(ancestor, [-1])[0]
     return -1
@@ -340,8 +348,13 @@ def get_ancestor_with_attribute(
     int
         Returns the first ancestor found that has an attribute otherwise `-1`.
     """
-    attr_dict = lT.__getattribute__(attribute)
-    if not isinstance(attr_dict, dict):
+    if attribute in lT.labelling.list_of_labels:
+        attr_dict = getattr(lT.labelling, attribute, None)
+    else:
+        attr_dict = getattr(lT, attribute, None)
+    if attr_dict is None:
+        raise ValueError(f"No label or property with name {attribute}")
+    if not isinstance(attr_dict, dict | Labels):
         raise ValueError("Please select a dict attribute")
     if node not in lT.nodes:
         return -1
@@ -395,106 +408,3 @@ def nodes_at_t(
         elif lT._time[curr] < t:
             to_do.extend(lT.successor[curr])
     return final_nodes
-
-
-def get_available_labels(lT: LineageTree) -> list[str]:
-    """Returns the list all the available label dictionaries
-
-    Parameters
-    ----------
-    lT : LineageTree
-        The LineageTree instance.
-
-    Returns
-    -------
-    list of string
-        list of the names of all the available properties
-        to label the nodes
-    """
-    available_labels = []
-    for prop_name, prop in lT.__dict__.items():
-        if (
-            0 < len(prop_name)
-            and prop_name[0] != "_"
-            and isinstance(prop, dict)
-            and 0 < len(prop)
-            and all(isinstance(k, int) for k in prop.keys())
-            and all(isinstance(v, str) for v in prop.values())
-        ):
-            available_labels.append(prop_name)
-    return available_labels
-
-
-def change_labels(
-    lT: LineageTree,
-    new_labels_name: str | None = None,
-    new_labels_dict: dict[int, str] | None = None,
-    only_first_node_in_chain: bool = False,
-) -> None:
-    """Change the dictionary that serves at labels with
-    the `LineageTree` attribute `new_labels_name`.
-    It has to be a dictionary mapping node id to string.
-
-    If `new_labels_dict` is provided, it will be used to
-    label the cells.
-    If `new_labels_name` is not specified, the labels are reset.
-
-    One can decide to only label the first node of the chain
-    instead of all the nodes of the chain.
-    That can help readability in the napari plugin reLAX.
-
-    Parameters
-    ----------
-    lT : LineageTree
-    new_labels_name : string, optional
-        The name of the dictionary to use
-        (the list of potential dictionaries can be found
-        with `lT.available_labels`)
-        If `new_labels_name` is not provided,
-        the labels are reset to `"Unlabeleld"`
-    new_labels_dict : dictionary mapping integers to strings, optional
-        The new names as a dictionary mapping each named node id to its string label
-        If not provided and lT has a fitting attribute named `new_labels_name`,
-        it will therefore be used
-    only_first_node_in_chain : bool, default=True
-        If `True` only labels the first node of the chains
-    """
-    store_new_labels = True
-    if new_labels_name is not None:
-        lT.labels_name = new_labels_name
-        if new_labels_dict is None:
-            if new_labels_name in lT.__dict__:
-                new_labels_dict = lT.__dict__[new_labels_name]
-                store_new_labels = False
-            else:
-                raise AttributeError(
-                    f"{new_labels_name} is not in the properties of {lT.name}"
-                )
-        if any(not isinstance(v, str) for v in new_labels_dict.values()):
-            raise TypeError(
-                "All values of new_labels dictionary should be `str`"
-            )
-
-        labelled_cells = lT.nodes.intersection(new_labels_dict)
-        if only_first_node_in_chain:
-            labelled_cells = labelled_cells.intersection(
-                {chain[0] for chain in lT.all_chains}
-            )
-
-        if len(labelled_cells) < 1:
-            warnings.warn(
-                "The labeling dictionary does not have any node labels.\n"
-                'Defaulting to the "Unlabeled" labeling'
-            )
-        else:
-            lT._labels = {n: new_labels_dict[n] for n in labelled_cells}
-            if store_new_labels:
-                lT.__dict__[new_labels_name] = lT._labels
-    else:
-        lT.labels_name = ""
-        lT._labels = {
-            root: "Unlabeled"
-            for root in lT.roots
-            for leaf in lT.find_leaves(root)
-            if abs(lT._time[leaf] - lT._time[root]) >= abs(lT.t_e - lT.t_b) / 4
-        }

src/lineagetree/_core/_properties.py

@@ -141,28 +141,6 @@ def edges(lT: LineageTree) -> tuple[tuple[int, int]]:
     return tuple((p, si) for p, s in lT._successor.items() for si in s)
 
 
-@property
-def labels(lT: LineageTree) -> dict[int, str]:
-    """Dictionary that maps a node to its label"""
-    if not hasattr(lT, "_labels"):
-        if hasattr(lT, "node_name"):
-            lT.labels_name = "node_name"
-            lT._labels = {
-                chain[0]: lT.node_name.get(chain[0], "Unlabeled")
-                for chain in lT.all_chains
-            }
-        else:
-            lT.labels_name = ""
-            lT._labels = {
-                root: "Unlabeled"
-                for root in lT.roots
-                for leaf in lT.find_leaves(root)
-                if abs(lT._time[leaf] - lT._time[root])
-                >= abs(lT.t_e - lT.t_b) / 4
-            }
-    return lT._labels
-
-
 @property
 def time_resolution(lT: LineageTree) -> float:
     """Time resolution of the lineage tree"""

src/lineagetree/_io/_loaders.py

@@ -872,7 +872,7 @@ def read_from_mastodon(
     time = dict(zip(nodes, spots[:, 3], strict=True))
     predecessor = {}
     labels, labels_name = {}, []
-
+    label_set = {}
     for succ, pred in zip(links[:, 1], links[:, 0]):
         predecessor[int(succ)] = int(pred)
 
@@ -883,6 +883,13 @@ def read_from_mastodon(
     elif 0 < len(properties):
         labels_name, labels = next(iter(properties.items()))
 
+    for key, prop in tuple(properties.items()):
+        if isinstance(prop, dict) and isinstance(
+            next(iter(prop.values()), None), str
+        ):
+            label_set[key] = prop
+            properties.pop(key)
+
     if not name:
         if isinstance(path, Path):
             tmp_name = path.stem
@@ -891,13 +898,13 @@ def read_from_mastodon(
         if name == "":
             warn(f"Name set to default {tmp_name}", stacklevel=2)
         name = tmp_name
+    properties["label_set"] = label_set
 
     return LineageTree(
         predecessor=predecessor,
         time=time,
         pos=pos,
         labels=labels,
-        labels_name=labels_name,
         name=name,
         **properties,
     )
@@ -994,7 +1001,7 @@ def read_from_mamut_xml(
     nodes = set()
     pos = {}
     time = {}
-    properties["label"] = {}
+    properties["labels"] = {}
 
     for frame in AllSpots:
         t = int(frame.attrib["frame"])
@@ -1009,7 +1016,7 @@ def read_from_mamut_xml(
             nodes.add(cell_id)
             pos[cell_id] = np.array([x, y, z])
             time[cell_id] = t
-            properties["label"][cell_id] = n
+            properties["labels"][cell_id] = str(n)
             if "TISSUE_NAME" in cell.attrib:
                 if "fate" not in properties:
                     properties["fate"] = {}

src/lineagetree/_labelling.py

@@ -0,0 +1,41 @@
+from typing import Any, Iterable, Mapping
+from collections import UserDict
+from .util_types import StaticTypedValueDict
+
+
+class Labels(StaticTypedValueDict):
+    """Subclass of `StaticTypedValueDict` that only allows for string values."""
+
+    def __init__(self, iterable: Iterable) -> None:
+        super().__init__(iterable, str)
+
+
+class Labelling:
+    """Labelling can hold only `Labels` where all their keys are a subset of the nodes of `LineageTree`."""
+
+    default_dict = {}
+
+    def __init__(self, nodes) -> None:
+        if isinstance(nodes, Iterable) and not isinstance(nodes, str):
+            self.__dict__["_nodes"] = set(nodes)
+        else:
+            raise ValueError(
+                f"Labelling expects an Iterable, not {type(nodes)}"
+            )
+
+    def __setattr__(self, name: str, value: dict | Labels) -> None:
+        l = Labels(value)
+        if not set(l).issubset(self._nodes):
+            raise ValueError(
+                "All ids in the labelset should correspond to ids in the LineageTree object."
+            )
+        if (
+            not self.default_dict
+        ):  # if the default dict has not been set set it up with the first label set available.
+            super().__setattr__("default_dict", l)
+        # self.list_of_labels.append(name)
+        super().__setattr__(name, l)
+
+    @property
+    def list_of_labels(self) -> list[str]:
+        return [k for k, v in self.__dict__.items() if isinstance(v, Labels)]

src/lineagetree/_mixins/navigation_mixin.py

@@ -1,10 +1,8 @@
 from .._core._navigation import (
-    change_labels,
     find_leaves,
     get_all_chains_of_subtree,
     get_ancestor_at_t,
     get_ancestor_with_attribute,
-    get_available_labels,
     get_chain_of_node,
     get_labelled_ancestor,
     get_predecessors,
@@ -19,8 +17,6 @@
 class NavigationMixin(metaclass=AutoMethodizeMeta):
     """Mixin for tree navigation operations."""
 
-    get_available_labels = get_available_labels
-    change_labels = change_labels
     find_leaves = find_leaves
     get_all_chains_of_subtree = get_all_chains_of_subtree
     get_ancestor_at_t = get_ancestor_at_t

src/lineagetree/_mixins/properties_mixin.py

@@ -2,7 +2,6 @@
     all_chains,
     depth,
     edges,
-    labels,
     leaves,
     nodes,
     number_of_nodes,
@@ -35,7 +34,6 @@ class PropertiesMixin(metaclass=AutoMethodizeMeta):
     roots = roots
     leaves = leaves
     edges = edges
-    labels = labels
     time_resolution = time_resolution
     all_chains = all_chains
     time_nodes = time_nodes

src/lineagetree/lineage_tree.py

@@ -19,6 +19,7 @@
 from ._mixins.analysis_mixin import AnalysisMixin
 from ._mixins.io_mixin import IOMixin
 from ._core.validation import TreeValidator
+from ._labelling import Labelling, Labels
 
 
 class LineageTree(
@@ -259,9 +260,20 @@ def __init__(
             self.pos = {
                 node: np.array(position) for node, position in pos.items()
             }
-        if "labels" in kwargs:
-            self._labels = kwargs["labels"]
+        if "labelling" in kwargs:  # for loading trees
+            self.labelling = kwargs["labelling"]
+        elif "labels" in kwargs:  # for importing trees
+            _labels = kwargs["labels"]
             kwargs.pop("labels")
+            self.labelling = Labelling(self.nodes)
+            self.labelling.labels = Labels(_labels)
+        else:
+            self.labelling = Labelling(self.nodes)
+        for k, v in tuple(kwargs.items()):
+            if isinstance(v, dict):
+                if all(isinstance(l, str) for l in v.values()):
+                    setattr(self.labelling, k, v)
+                    kwargs.pop(k)
         if time is None:
             if starting_time is None:
                 starting_time = 0

src/lineagetree/plot.py

@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from .lineage_tree import LineageTree
+    from ._labelling import Labels
 
 
 def __plot_nodes(
@@ -223,6 +224,7 @@ def plot_all_lineages(
     fontsize: int = 15,
     axes: plt.Axes | None = None,
     vert_gap: int = 1,
+    labels: str | dict | None = None,
     **kwargs,
 ) -> tuple[plt.Figure, plt.Axes, dict[plt.Axes, int]]:
     """Plots all lineages.
@@ -313,7 +315,14 @@ def plot_all_lineages(
         )
         root = graph["root"]
         ax2root[flat_axes[i]] = root
-        label = lT.labels.get(root, "Unlabeled")
+        if labels is None:
+            label = lT.labelling.default_dict.get(root, "Unlabeled")
+        elif isinstance(labels, dict | Labels):
+            label = labels.get(root, "Unlabelled")
+        elif isinstance(labels, str):
+            if labels not in lT.labelling.list_of_labels:
+                raise ValueError("Label set not defined.")
+            label = getattr(lT.labelling, labels).get(root, "Unlabelled")
         xlim = flat_axes[i].get_xlim()
         ylim = flat_axes[i].get_ylim()
         x_pos = (xlim[0] + xlim[1]) / 2