Example Using Custom and External KNN Graph

Hello,

I tried to find an example of using external KNN graphs with UMAP. I wasn't able to find any so I created one and hope people find this snippet of code helpful. It's largely based on the internal APIs and functions within UMAP itself, but it works well for our purposes. An example of how to use the code is at the end and here is the example data from the Iris dataset: approximate-knng.txt

Any feedback would be welcomed.

Take care!

from umap.umap_ import compute_membership_strengths, smooth_knn_dist, make_epochs_per_sample, simplicial_set_embedding, find_ab_params
import scipy

   
# Load graphs created by NGT
def load_knn_graph(ngt_graph):
    knn_indices = []
    knn_dist = []

    nodes_with_no_edges = set()
    with open(ngt_graph, "r") as f:
        while True:
            line = f.readline().strip("\n").strip()
            if not line:
                break
            token = line.split("\t")
            start_node = int(token[0])
            end_nodes = token[1:]

            if end_nodes == []:
                nodes_with_no_edges.add(start_node)
                continue

            knn_entry = []
            knn_entry_dist = []
            for ith in range(0, len(end_nodes), 2):
                try:
                    end_node = int(end_nodes[ith])
                except Exception as e:
                    nodes_with_no_edges.add(start_node)
                    continue
                knn_entry.append(end_node)
                weight = float(end_nodes[ith+1])
                knn_entry_dist.append(weight)
            knn_indices.append(knn_entry)
            knn_dist.append(knn_entry_dist)


    return np.array(knn_indices), np.array(knn_dist)
    
    
def knn_fuzzy_simplicial_set(knn_indices, knn_dists, local_connectivity, set_op_mix_ratio, apply_set_operations):

    n_samples = knn_indices.shape[0]
    n_neighbors = knn_indices.shape[1]

    knn_dists = knn_dist.astype(np.float32)

    sigmas, rhos = smooth_knn_dist(
        knn_dists, float(n_neighbors), local_connectivity=float(local_connectivity),
    )

    rows, cols, vals = compute_membership_strengths(
        knn_indices, knn_dists, sigmas, rhos
    )

    result = scipy.sparse.coo_matrix(
        (vals, (rows, cols)), shape=(n_samples, n_samples)
    )
    result.eliminate_zeros()

    if apply_set_operations:
        transpose = result.transpose()

        prod_matrix = result.multiply(transpose)

        result = (
            set_op_mix_ratio * (result + transpose - prod_matrix)
            + (1.0 - set_op_mix_ratio) * prod_matrix
        )

    result.eliminate_zeros()

    return result, sigmas, rhos


def transform(graph, metric="l2", n_components = 2, n_epochs = 500, spread=1.0, min_dist = 0.0, initial_alpha=1.0, negative_sample_rate=5, repulsion_strength=7):
    graph = graph.tocoo()
    graph.sum_duplicates()
    n_vertices = graph.shape[1]

    if n_epochs <= 0:
        # For smaller datasets we can use more epochs
        if graph.shape[0] <= 10000:
            n_epochs = 500
        else:
            n_epochs = 200

    graph.data[graph.data < (graph.data.max() / float(n_epochs))] = 0.0
    graph.eliminate_zeros()

    epochs_per_sample = make_epochs_per_sample(graph.data, n_epochs)

    head = graph.row
    tail = graph.col
    weight = graph.data

    a, b = find_ab_params(spread, min_dist)
    
    emebedding = simplicial_set_embedding(None, graph, n_components=2, initial_alpha=1.0, a=a, b=b, gamma=repulsion_strength, negative_sample_rate=negative_sample_rate, random_state=np.random, metric=metric, metric_kwds=None, verbose=True, parallel=True, n_epochs=n_epochs, init="random")
    
    return emebedding

An example usage is a follows:

import numpy as np

knn_indices, knn_dist = load_knn_graph("approximate-knng.txt")

knn_indices = knn_indices - 1

local_connectivity = 1 
apply_set_operations = True
set_op_mix_ratio=1.0       

graph, sigmas, rhos = knn_fuzzy_simplicial_set(knn_indices, knn_dist, local_connectivity, set_op_mix_ratio, apply_set_operations)

emebedding = transform(graph, metric="l2", n_components = 2, n_epochs = 500, spread=1.0, min_dist = 0.0, initial_alpha=1.0, negative_sample_rate=5, repulsion_strength=7)

Start by reviewing the proposed Python example and the UMAP entry points it uses, including smooth_knn_dist, compute_membership_strengths, simplicial_set_embedding, and find_ab_params. Check whether the external KNN graph workflow still matches the current APIs and determine where an example would belong; done means the example is validated and its intended documentation location and scope are clear.