add lensed clusters from knn-mst union

Author: JelmerBot

fast_hdbscan/boruvka.py

@@ -4,26 +4,33 @@
 from .disjoint_set import ds_rank_create, ds_find, ds_union_by_rank
 from .numba_kdtree import parallel_tree_query, rdist, point_to_node_lower_bound_rdist
 
-@numba.njit(locals={"i": numba.types.int64})
-def merge_components(disjoint_set, candidate_neighbors, candidate_neighbor_distances, point_components):
-    component_edges = {np.int64(0): (np.int64(0), np.int64(1), np.float32(0.0)) for i in range(0)}
+
+@numba.njit(locals={"parent": numba.types.int32})
+def select_components(candidate_distances, candidate_neighbors, point_components):
+    component_edges = {np.int64(0): (np.int32(0), np.int32(1), np.float32(0.0)) for i in range(0)}
 
     # Find the best edges from each component
-    for i in range(candidate_neighbors.shape[0]):
-        from_component = np.int64(point_components[i])
+    for parent, (distance, neighbor, from_component) in enumerate(
+        zip(candidate_distances, candidate_neighbors, point_components)
+    ):
         if from_component in component_edges:
-            if candidate_neighbor_distances[i] < component_edges[from_component][2]:
-                component_edges[from_component] = (np.int64(i), np.int64(candidate_neighbors[i]), candidate_neighbor_distances[i])
+            if distance < component_edges[from_component][2]:
+                component_edges[from_component] = (parent, neighbor, distance)
         else:
-            component_edges[from_component] = (np.int64(i), np.int64(candidate_neighbors[i]), candidate_neighbor_distances[i])
+            component_edges[from_component] = (parent, neighbor, distance)
+
+    return component_edges
+
 
+@numba.njit()
+def merge_components(disjoint_set, component_edges):
     result = np.empty((len(component_edges), 3), dtype=np.float64)
     result_idx = 0
 
     # Add the best edges to the edge set and merge the relevant components
     for edge in component_edges.values():
-        from_component = ds_find(disjoint_set, np.int32(edge[0]))
-        to_component = ds_find(disjoint_set, np.int32(edge[1]))
+        from_component = ds_find(disjoint_set, edge[0])
+        to_component = ds_find(disjoint_set, edge[1])
         if from_component != to_component:
             result[result_idx] = (np.float64(edge[0]), np.float64(edge[1]), np.float64(edge[2]))
             result_idx += 1
@@ -34,10 +41,13 @@ def merge_components(disjoint_set, candidate_neighbors, candidate_neighbor_dista
 
 
 @numba.njit(parallel=True)
-def update_component_vectors(tree, disjoint_set, node_components, point_components):
+def update_point_components(disjoint_set, point_components):
     for i in numba.prange(point_components.shape[0]):
         point_components[i] = ds_find(disjoint_set, np.int32(i))
 
+
+@numba.njit()
+def update_node_components(tree, node_components, point_components):
     for i in range(tree.node_data.shape[0] - 1, -1, -1):
         node_info = tree.node_data[i]
 
@@ -272,28 +282,28 @@ def parallel_boruvka(tree, min_samples=10, sample_weights=None):
         expected_neighbors = min_samples / mean_sample_weight
         distances, neighbors = parallel_tree_query(tree, tree.data, k=int(2 * expected_neighbors))
         core_distances = sample_weight_core_distance(distances, neighbors, sample_weights, min_samples)
-        edges = initialize_boruvka_from_knn(neighbors, distances, core_distances, components_disjoint_set)
-        update_component_vectors(tree, components_disjoint_set, node_components, point_components)
     else:
         if min_samples > 1:
             distances, neighbors = parallel_tree_query(tree, tree.data, k=min_samples + 1, output_rdist=True)
             core_distances = distances.T[-1]
-            edges = initialize_boruvka_from_knn(neighbors, distances, core_distances, components_disjoint_set)
-            update_component_vectors(tree, components_disjoint_set, node_components, point_components)
         else:
             core_distances = np.zeros(tree.data.shape[0], dtype=np.float32)
             distances, neighbors = parallel_tree_query(tree, tree.data, k=2)
-            edges = initialize_boruvka_from_knn(neighbors, distances, core_distances, components_disjoint_set)
-            update_component_vectors(tree, components_disjoint_set, node_components, point_components)
 
-    while n_components > 1:
+    edges = [np.empty((0, 3), dtype=np.float64) for _ in range(0)]
+    new_edges = initialize_boruvka_from_knn(neighbors, distances, core_distances, components_disjoint_set)
+    while True:
+        edges.append(new_edges)
+        n_components -= new_edges.shape[0]
+        if n_components == 1:
+            break
+        update_point_components(components_disjoint_set, point_components)
+        update_node_components(tree, node_components, point_components)
         candidate_distances, candidate_indices = boruvka_tree_query(tree, node_components, point_components,
                                                                     core_distances)
-        new_edges = merge_components(components_disjoint_set, candidate_indices, candidate_distances, point_components)
-        update_component_vectors(tree, components_disjoint_set, node_components, point_components)
-
-        edges = np.vstack((edges, new_edges))
-        n_components = np.unique(point_components).shape[0]
+        component_edges = select_components(candidate_distances, candidate_indices, point_components)
+        new_edges = merge_components(components_disjoint_set, component_edges)
 
+    edges = np.vstack(edges)
     edges[:, 2] = np.sqrt(edges.T[2])
     return edges, neighbors[:, 1:], np.sqrt(core_distances)

fast_hdbscan/cluster_trees.py

@@ -160,6 +160,13 @@ def eliminate_branch(branch_node, parent_node, lambda_value, parents, children,
 CondensedTree = namedtuple('CondensedTree', ['parent', 'child', 'lambda_val', 'child_size'])
 
 
+@numba.njit()
+def empty_condensed_tree():
+    parents = np.empty(shape=0, dtype=np.intp)
+    others = np.empty(shape=0, dtype=np.float32)
+    return CondensedTree(parents, parents, others, others)
+
+
 @numba.njit(fastmath=True)
 def condense_tree(hierarchy, min_cluster_size=10, max_cluster_size=np.inf, sample_weights=None):
     root = 2 * hierarchy.shape[0]

fast_hdbscan/core_graph.py

@@ -0,0 +1,219 @@
+import numba
+import numpy as np
+from collections import namedtuple
+
+from .disjoint_set import ds_rank_create
+from .hdbscan import clusters_from_spanning_tree
+from .cluster_trees import empty_condensed_tree
+from .boruvka import merge_components, update_point_components
+
+CoreGraph = namedtuple("CoreGraph", ["weights", "distances", "indices", "indptr"])
+
+
+@numba.njit(parallel=True)
+def knn_mst_union(neighbors, core_distances, min_spanning_tree, lens_values):
+    # List of dictionaries of child: (weight, distance)
+    graph = [
+        {np.int32(0): (np.float64(0.0), np.float64(0.0)) for _ in range(0)}
+        for _ in range(neighbors.shape[0])
+    ]
+
+    # Add knn edges
+    for point in numba.prange(len(core_distances)):
+        children = graph[point]
+        parent_lens = lens_values[point]
+        parent_dist = core_distances[point]
+        for child in neighbors[point]:
+            if child < 0:
+                continue
+            children[child] = (
+                max(parent_lens, lens_values[child]),
+                max(parent_dist, core_distances[child]),
+            )
+
+    # Add non-knn mst edges
+    for parent, child, distance in min_spanning_tree:
+        parent = np.int32(parent)
+        child = np.int32(child)
+        children = graph[parent]
+        if child in children:
+            continue
+        children[child] = (max(lens_values[parent], lens_values[child]), distance)
+
+    return graph
+
+
+@numba.njit(parallel=True)
+def sort_by_lens(graph):
+    for point in numba.prange(len(graph)):
+        graph[point] = {
+            k: v for k, v in sorted(graph[point].items(), key=lambda item: item[1][0])
+        }
+    return graph
+
+
+@numba.njit(parallel=True)
+def apply_lens(core_graph, lens_values):
+    # Apply new lens to the graph
+    for point in numba.prange(len(lens_values)):
+        children = core_graph[point]
+        point_lens = lens_values[point]
+        for child, value in children.items():
+            children[child] = (max(point_lens, lens_values[child]), value[1])
+    return sort_by_lens(core_graph)
+
+
+@numba.njit()
+def flatten_to_csr(graph):
+    # Count children to form indptr
+    num_points = len(graph)
+    indptr = np.empty(num_points + 1, dtype=np.int32)
+    indptr[0] = 0
+    for i, children in enumerate(graph):
+        indptr[i + 1] = indptr[i] + len(children)
+
+    # Flatten children to form indices, weights, and distances
+    weights = np.empty(indptr[-1], dtype=np.float32)
+    distances = np.empty(indptr[-1], dtype=np.float32)
+    indices = np.empty(indptr[-1], dtype=np.int32)
+    for point in numba.prange(num_points):
+        start = indptr[point]
+        children = graph[point]
+        for j, (child, (weight, distance)) in enumerate(children.items()):
+            weights[start + j] = weight
+            distances[start + j] = distance
+            indices[start + j] = child
+
+    # Return as named csr tuple
+    return CoreGraph(weights, distances, indices, indptr)
+
+
+@numba.njit(locals={"parent": numba.types.int32})
+def select_components(graph, point_components):
+    component_edges = {
+        np.int64(0): (np.int32(0), np.int32(1), np.float32(0.0)) for _ in range(0)
+    }
+
+    # Find the best edges from each component
+    for parent, (children, from_component) in enumerate(zip(graph, point_components)):
+        if len(children) == 0:
+            continue
+        neighbor = next(iter(children.keys()))
+        distance = np.float32(children[neighbor][0])
+        if from_component in component_edges:
+            if distance < component_edges[from_component][2]:
+                component_edges[from_component] = (parent, neighbor, distance)
+        else:
+            component_edges[from_component] = (parent, neighbor, distance)
+
+    return component_edges
+
+
+@numba.njit()  # enabling parallel breaks this function
+def update_graph_components(graph, point_components):
+    # deleting from dictionary during iteration breaks in numba.
+    for point in numba.prange(len(graph)):
+        graph[point] = {
+            child: (weight, distance)
+            for child, (weight, distance) in graph[point].items()
+            if point_components[child] != point_components[point]
+        }
+
+
+@numba.njit()
+def minimum_spanning_tree(graph, overwrite=False):
+    """
+    Implements Boruvka on lod-style graph with multiple connected components.
+    """
+    if not overwrite:
+        graph = [children for children in graph]
+
+    disjoint_set = ds_rank_create(len(graph))
+    point_components = np.arange(len(graph))
+    n_components = len(point_components)
+
+    edges_list = [np.empty((0, 3), dtype=np.float64) for _ in range(0)]
+    while n_components > 1:
+        new_edges = merge_components(
+            disjoint_set,
+            select_components(graph, point_components),
+        )
+        if new_edges.shape[0] == 0:
+            break
+
+        edges_list.append(new_edges)
+        update_point_components(disjoint_set, point_components)
+        update_graph_components(graph, point_components)
+        n_components -= new_edges.shape[0]
+
+    counter = 0
+    num_edges = sum([edges.shape[0] for edges in edges_list])
+    result = np.empty((num_edges, 3), dtype=np.float64)
+    for edges in edges_list:
+        result[counter : counter + edges.shape[0]] = edges
+        counter += edges.shape[0]
+    return n_components, point_components, result
+
+
+@numba.njit()
+def core_graph_spanning_tree(neighbors, core_distances, min_spanning_tree, lens):
+    graph = sort_by_lens(
+        knn_mst_union(neighbors, core_distances, min_spanning_tree, lens)
+    )
+    return (*minimum_spanning_tree(graph), flatten_to_csr(graph))
+
+
+def core_graph_clusters(
+    lens,
+    neighbors,
+    core_distances,
+    min_spanning_tree,
+    **kwargs,
+):
+    num_components, component_labels, lensed_mst, graph = core_graph_spanning_tree(
+        neighbors, core_distances, min_spanning_tree, lens
+    )
+    if num_components > 1:
+        for i, label in enumerate(np.unique(component_labels)):
+            component_labels[component_labels == label] = i
+        return (
+            component_labels,
+            np.ones(len(component_labels), dtype=np.float32),
+            np.empty((0, 4)),
+            empty_condensed_tree(),
+            lensed_mst,
+            graph,
+        )
+
+    return (
+        *clusters_from_spanning_tree(lensed_mst, **kwargs),
+        graph,
+    )
+
+
+def core_graph_to_rec_array(graph):
+    result = np.empty(
+        graph.indptr[-1],
+        dtype=[
+            ("parent", np.int32),
+            ("child", np.int32),
+            ("weight", np.float32),
+            ("distance", np.float32),
+        ],
+    )
+    result["parent"] = np.repeat(
+        np.arange(len(graph.indptr) - 1), np.diff(graph.indptr)
+    )
+    result["child"] = graph.indices
+    result["weight"] = graph.weights
+    result["distance"] = graph.distances
+    return result
+
+
+def core_graph_to_edge_list(graph):
+    result = np.empty((graph.indptr[-1], 4), dtype=np.float64)
+    result[:, 0] = np.repeat(np.arange(len(graph.indptr) - 1), np.diff(graph.indptr))
+    result[:, 1] = graph.indices
+    result[:, 2] = graph.weights
+    result[:, 3] = graph.distances
+    return result