Add leaf labelisation from num regions

doc/source/python/tree_algorithm.rst

@@ -10,6 +10,7 @@ Algorithm for trees
     binary_labelisation_from_markers
     filter_binary_partition_tree
     labelisation_hierarchy_supervertices
+    labelisation_horizontal_cut_from_num_regions
     labelisation_horizontal_cut_from_threshold
     reconstruct_leaf_data
     sort_hierarchy_with_altitudes
@@ -21,6 +22,8 @@ Algorithm for trees
 
 .. autofunction:: higra.labelisation_hierarchy_supervertices
 
+.. autofunction:: higra.labelisation_horizontal_cut_from_num_regions
+
 .. autofunction:: higra.labelisation_horizontal_cut_from_threshold
 
 .. autofunction:: higra.reconstruct_leaf_data

higra/algo/tree.py

@@ -34,7 +34,7 @@ def reconstruct_leaf_data(tree, altitudes, deleted_nodes, leaf_graph=None):
 @hg.argument_helper(hg.CptHierarchy)
 def labelisation_horizontal_cut_from_threshold(tree, altitudes, threshold, leaf_graph=None):
     """
-    Labelize tree leaves according to an horizontal cut in the tree.
+    Labelize tree leaves according to an horizontal cut of the tree given by its altitude.
 
     Two leaves are in the same region (ie. have the same label) if
     the altitude of their lowest common ancestor is strictly greater
@@ -55,6 +55,44 @@ def labelisation_horizontal_cut_from_threshold(tree, altitudes, threshold, leaf_
     return leaf_labels
 
 
+@hg.argument_helper(hg.CptHierarchy)
+def labelisation_horizontal_cut_from_num_regions(tree, altitudes, num_regions, mode="at_least", leaf_graph=None):
+    """
+    Labelize tree leaves according to an horizontal cut of the tree given by its number of regions.
+
+    If :attr:`mode` is ``"at_least"`` (default), the the smallest horizontal cut having at least the given number of
+    regions is considered.
+    If :attr:`mode` is ``"at_most"``, the the largest horizontal cut having at most the given number of
+    regions is considered.
+
+    :param tree: input tree (deduced from :class:`~higra.CptHierarchy`)
+    :param altitudes: node altitudes of the input tree
+    :param num_regions: a number of regions
+    :param mode: ``"at_least"`` or ``"at_most"``
+    :param leaf_graph: graph of the tree leaves (optional, deduced from :class:`~higra.CptHierarchy`)
+    :return: Leaf labels
+    """
+
+    num_regions = int(num_regions)
+
+    if mode == "at_least":
+        modeb = True
+    elif mode == "at_most":
+        modeb = False
+    else:
+        raise ValueError("Incorrect mode")
+
+    hc = hg.HorizontalCutExplorer(tree, altitudes)
+    cut = hc.horizontal_cut_from_num_regions(num_regions, modeb)
+
+    leaf_labels = cut.labelisation_leaves(tree)
+
+    if leaf_graph is not None:
+        leaf_labels = hg.delinearize_vertex_weights(leaf_labels, leaf_graph)
+
+    return leaf_labels
+
+
 @hg.argument_helper(hg.CptHierarchy)
 def labelisation_hierarchy_supervertices(tree, altitudes, leaf_graph=None, handle_rag=True):
     """

test/python/test_algo/test_tree.py

@@ -55,6 +55,36 @@ def test_labelisation_horizontal_cut(self):
         self.assertTrue(hg.is_in_bijection(ref_t1, output_t1))
         self.assertTrue(hg.is_in_bijection(ref_t2, output_t2))
 
+    def test_labelisation_horizontal_cut_num_regions(self):
+        g = hg.get_4_adjacency_graph((1, 11))
+        tree = hg.Tree((11, 11, 11, 12, 12, 16, 13, 13, 13, 14, 14, 17, 16, 15, 15, 18, 17, 18, 18))
+        hg.CptHierarchy.link(tree, g)
+        altitudes = np.asarray((0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 3, 1, 2, 3))
+
+        ref_labels = (
+            np.asarray((1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1)),
+            np.asarray((1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3)),
+            np.asarray((0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3)),
+            np.asarray((0, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8))
+        )
+
+        k_cuts = (1, 3, 4, 9)
+        for i in range(4):
+            labels = hg.labelisation_horizontal_cut_from_num_regions(tree, altitudes, k_cuts[i])
+            self.assertTrue(hg.is_in_bijection(labels, ref_labels[i]))
+
+        # cuts with at least the given number of regions
+        k_cuts = (1, 2, 4, 5)
+        for i in range(4):
+            labels = hg.labelisation_horizontal_cut_from_num_regions(tree, altitudes, k_cuts[i])
+            self.assertTrue(hg.is_in_bijection(labels, ref_labels[i]))
+
+        # cuts with at most the given number of regions
+        k_cuts = (2, 3, 8, 20)
+        for i in range(4):
+            labels = hg.labelisation_horizontal_cut_from_num_regions(tree, altitudes, k_cuts[i], "at_most")
+            self.assertTrue(hg.is_in_bijection(labels, ref_labels[i]))
+
     def test_labelisation_hierarchy_supervertices(self):
         tree = hg.Tree(np.asarray((5, 5, 6, 6, 6, 7, 7, 7)))