add branch detection

JelmerBot · JelmerBot · commit d40ab87e9074 · 2024-12-28T10:06:14.000+01:00
diff --git a/fast_hdbscan/__init__.py b/fast_hdbscan/__init__.py
@@ -1,4 +1,5 @@
 from .hdbscan import HDBSCAN, fast_hdbscan
+from .branches import BranchDetector, find_branch_sub_clusters
 
 # Force JIT compilation on import
 import numpy as np
@@ -7,4 +8,4 @@
 HDBSCAN(allow_single_cluster=True).fit(random_data)
 HDBSCAN(cluster_selection_method="leaf").fit(random_data)
 
-__all__ = ["HDBSCAN", "fast_hdbscan"]
+__all__ = ["HDBSCAN", "fast_hdbscan", "BranchDetector", "find_branch_sub_clusters"]
diff --git a/fast_hdbscan/branches.py b/fast_hdbscan/branches.py
@@ -0,0 +1,136 @@
+import numpy as np
+from .sub_clusters import SubClusterDetector, find_sub_clusters
+
+
+def compute_centrality(data, probabilities, *args):
+    points = args[-1]
+    cluster_data = data[points, :]
+    centroid = np.average(cluster_data, weights=probabilities[points], axis=0)
+    return 1 / np.linalg.norm(cluster_data - centroid[None, :], axis=1)
+
+
+def apply_branch_threshold(
+    labels,
+    branch_labels,
+    probabilities,
+    cluster_probabilities,
+    cluster_points,
+    linkage_trees,
+    label_sides_as_branches=False,
+):
+    running_id = 0
+    min_branch_count = 1 if label_sides_as_branches else 2
+    for pts, tree in zip(cluster_points, linkage_trees):
+        unique_branch_labels = np.unique(branch_labels[pts])
+        has_noise = int(unique_branch_labels[0] == -1)
+        num_branches = len(unique_branch_labels) - has_noise
+        if num_branches <= min_branch_count and tree is not None:
+            labels[pts] = running_id
+            probabilities[pts] = cluster_probabilities[pts]
+            running_id += 1
+            continue
+        else:
+            branch_labels[pts] = np.where(
+                branch_labels[pts] < 0, num_branches, branch_labels[pts]
+            )
+            labels[pts] = branch_labels[pts] + running_id
+            running_id += num_branches + has_noise
+
+
+def find_branch_sub_clusters(
+    clusterer,
+    cluster_labels=None,
+    cluster_probabilities=None,
+    *,
+    min_branch_size=None,
+    max_branch_size=None,
+    allow_single_branch=None,
+    branch_selection_method=None,
+    branch_selection_epsilon=0.0,
+    branch_selection_persistence=0.0,
+    label_sides_as_branches=False,
+):
+    result = find_sub_clusters(
+        clusterer,
+        cluster_labels,
+        cluster_probabilities,
+        lens_callback=compute_centrality,
+        min_cluster_size=min_branch_size,
+        max_cluster_size=max_branch_size,
+        allow_single_cluster=allow_single_branch,
+        cluster_selection_method=branch_selection_method,
+        cluster_selection_epsilon=branch_selection_epsilon,
+        cluster_selection_persistence=branch_selection_persistence,
+    )
+    apply_branch_threshold(
+        result[0],
+        result[4],
+        result[1],
+        result[3],
+        result[-1],
+        label_sides_as_branches=label_sides_as_branches,
+    )
+    return result
+
+
+class BranchDetector(SubClusterDetector):
+    """
+    Performs a flare-detection post-processing step to detect branches within
+    clusters [1]_.
+
+    For each cluster, a graph is constructed connecting the data points based on
+    their mutual reachability distances. Each edge is given a centrality value
+    based on how far it lies from the cluster's center. Then, the edges are
+    clustered as if that centrality was a distance, progressively removing the
+    'center' of each cluster and seeing how many branches remain.
+
+    References
+    ----------
+    .. [1] Bot, D. M., Peeters, J., Liesenborgs J., & Aerts, J. (2023, November).
+       FLASC: A Flare-Sensitive Clustering Algorithm: Extending HDBSCAN* for
+       Detecting Branches in Clusters. arXiv:2311.15887.
+    """
+
+    def __init__(
+        self,
+        *,
+        min_branch_size=None,
+        max_branch_size=None,
+        allow_single_branch=None,
+        branch_selection_method=None,
+        branch_selection_epsilon=0.0,
+        branch_selection_persistence=0.0,
+        label_sides_as_branches=False,
+    ):
+        super().__init__(
+            min_cluster_size=min_branch_size,
+            max_cluster_size=max_branch_size,
+            allow_single_cluster=allow_single_branch,
+            cluster_selection_method=branch_selection_method,
+            cluster_selection_epsilon=branch_selection_epsilon,
+            cluster_selection_persistence=branch_selection_persistence,
+        )
+        self.label_sides_as_branches = label_sides_as_branches
+
+    def fit(self, clusterer, labels=None, probabilities=None):
+        super().fit(clusterer, labels, probabilities, compute_centrality)
+        apply_branch_threshold(
+            self.labels_,
+            self.sub_cluster_labels_,
+            self.probabilities_,
+            self.cluster_probabilities_,
+            self.cluster_points_,
+            self.linkage_trees_,
+            label_sides_as_branches=self.label_sides_as_branches,
+        )
+        self.branch_labels_ = self.sub_cluster_labels_
+        self.branch_probabilities_ = self.sub_cluster_probabilities_
+        self.centralities_ = self.lens_values_
+        return self
+
+    @property
+    def approximation_graph_(self):
+        """See :class:`~hdbscan.plots.ApproximationGraph` for documentation."""
+        return super()._make_approximation_graph(
+            lens_name="centrality", sub_cluster_name="branch"
+        )
diff --git a/fast_hdbscan/tests/test_branches.py b/fast_hdbscan/tests/test_branches.py
@@ -0,0 +1,173 @@
+"""
+Test for the branches module.
+"""
+
+import pytest
+import numpy as np
+from sklearn.exceptions import NotFittedError
+from fast_hdbscan import HDBSCAN, BranchDetector, find_branch_sub_clusters
+
+
+def make_branches(points_per_branch=30):
+    # Control points for line segments that merge three clusters
+    p0 = (0.13, -0.26)
+    p1 = (0.24, -0.12)
+    p2 = (0.32, 0.1)
+    p3 = (0.13, 0.1)
+
+    # Noisy points along lines between three clusters
+    return np.concatenate(
+        [
+            np.column_stack(
+                (
+                    np.linspace(p_start[0], p_end[0], points_per_branch),
+                    np.linspace(p_start[1], p_end[1], points_per_branch),
+                )
+            )
+            + np.random.normal(size=(points_per_branch, 2), scale=0.01)
+            for p_start, p_end in [(p0, p1), (p1, p2), (p1, p3)]
+        ]
+    )
+
+
+np.random.seed(0)
+X = np.concatenate(
+    (
+        make_branches(),
+        make_branches()[:60] + np.array([0.3, 0]),
+    )
+)
+c = HDBSCAN(min_samples=5, min_cluster_size=10).fit(X)
+
+
+def check_detected_groups(c, n_clusters=3, n_branches=6, overridden=False):
+    """Checks branch_detector output for main invariants."""
+    noise_mask = c.labels_ == -1
+    assert np.all(np.unique(c.labels_[~noise_mask]) == np.arange(n_branches))
+    assert np.all(np.unique(c.cluster_labels_[~noise_mask]) == np.arange(n_clusters))
+    assert (c.branch_labels_[noise_mask] == 0).all()
+    assert (c.branch_probabilities_[noise_mask] == 1.0).all()
+    assert (c.probabilities_[noise_mask] == 0.0).all()
+    assert (c.cluster_probabilities_[noise_mask] == 0.0).all()
+    if not overridden:
+        assert len(c.cluster_points_) == n_clusters
+        for condensed_tree, linkage_tree in zip(c._condensed_trees, c._linkage_trees):
+            assert linkage_tree is not None
+            assert condensed_tree is not None
+
+
+# --- Detecting Branches
+
+
+def test_attributes():
+    def check_attributes():
+        b = BranchDetector().fit(c)
+        check_detected_groups(b, n_clusters=2, n_branches=5)
+        assert len(b.linkage_trees_) == 2
+        assert len(b.condensed_trees_) == 2
+        assert isinstance(b.condensed_trees_[0], CondensedTree)
+        assert isinstance(b.linkage_trees_[0], SingleLinkageTree)
+        assert isinstance(b.approximation_graph_, ApproximationGraph)
+
+    try:
+        from hdbscan.plots import ApproximationGraph, CondensedTree, SingleLinkageTree
+
+        check_attributes()
+    except ImportError:
+        pass
+
+
+def test_selection_method():
+    b = BranchDetector(branch_selection_method="eom").fit(c)
+    check_detected_groups(b, n_clusters=2, n_branches=5)
+
+    b = BranchDetector(branch_selection_method="leaf").fit(c)
+    check_detected_groups(b, n_clusters=2, n_branches=5)
+
+
+def test_min_branch_size():
+    b = BranchDetector(min_branch_size=7).fit(c)
+    labels, counts = np.unique(b.labels_[b.branch_probabilities_ > 0], return_counts=True)
+    assert (counts[labels >= 0] >= 7).all()
+    check_detected_groups(b, n_clusters=2, n_branches=5)
+
+
+def test_label_sides_as_branches():
+    b = BranchDetector(label_sides_as_branches=True).fit(c)
+    check_detected_groups(b, n_clusters=2, n_branches=6)
+
+
+def test_max_branch_size():
+    b = BranchDetector(label_sides_as_branches=True, max_branch_size=25).fit(c)
+    check_detected_groups(b, n_clusters=2, n_branches=4)
+
+
+def test_override_cluster_labels():
+    X_missing = X.copy()
+    X_missing[60:80] = np.nan
+    c = HDBSCAN(min_cluster_size=5).fit(X_missing)
+    split_y = c.labels_.copy()
+    split_y[split_y == 1] = 0
+    split_y[split_y == 2] = 1
+    b = BranchDetector(label_sides_as_branches=True).fit(c, split_y)
+    check_detected_groups(b, n_clusters=2, n_branches=5, overridden=True)
+    assert b._condensed_trees[0] is None
+    assert b._linkage_trees[0] is None
+
+
+def test_allow_single_branch_with_filters():
+    # Without persistence, find 6 branches
+    b = BranchDetector(
+        min_branch_size=5,
+        branch_selection_method="leaf",
+    ).fit(c)
+    unique_labels = np.unique(b.labels_)
+    assert len(unique_labels) == 5
+
+    # Adding persistence removes the branches
+    b = BranchDetector(
+        min_branch_size=5,
+        branch_selection_method="leaf",
+        branch_selection_persistence=0.15,
+    ).fit(c)
+    unique_labels = np.unique(b.labels_)
+    assert len(unique_labels) == 2
+
+    # Adding epsilon removes some branches
+    b = BranchDetector(
+        min_branch_size=5,
+        branch_selection_method="leaf",
+        branch_selection_epsilon=1 / 0.002,
+    ).fit(c)
+    unique_labels = np.unique(b.labels_)
+    assert len(unique_labels) == 2
+
+
+def test_badargs():
+    c_nofit = HDBSCAN(min_cluster_size=5)
+
+    with pytest.raises(TypeError):
+        find_branch_sub_clusters("fail")
+    with pytest.raises(TypeError):
+        find_branch_sub_clusters(None)
+    with pytest.raises(NotFittedError):
+        find_branch_sub_clusters(c_nofit)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, min_branch_size=-1)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, min_branch_size=0)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, min_branch_size=1)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, min_branch_size=2.0)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, min_branch_size="fail")
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, branch_selection_persistence=-0.1)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(c, branch_selection_epsilon=-0.1)
+    with pytest.raises(ValueError):
+        find_branch_sub_clusters(
+            c,
+            branch_selection_method="something_else",
+        )
