dataeval 0.76.1__py3-none-any.whl → 0.82.0__py3-none-any.whl

dataeval/utils/_array.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+__all__ = []
+
+import logging
+import warnings
+from importlib import import_module
+from types import ModuleType
+from typing import Any, Iterable, Iterator, Literal, TypeVar, overload
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from dataeval._log import LogMessage
+from dataeval.typing import ArrayLike
+
+_logger = logging.getLogger(__name__)
+
+_MODULE_CACHE = {}
+
+T = TypeVar("T", ArrayLike, np.ndarray, torch.Tensor)
+_np_dtype = TypeVar("_np_dtype", bound=np.generic)
+
+
+def _try_import(module_name) -> ModuleType | None:
+    if module_name in _MODULE_CACHE:
+        return _MODULE_CACHE[module_name]
+
+    try:
+        module = import_module(module_name)
+    except ImportError:  # pragma: no cover
+        _logger.log(logging.INFO, f"Unable to import {module_name}.")
+        module = None
+
+    _MODULE_CACHE[module_name] = module
+    return module
+
+
+def as_numpy(array: ArrayLike | None) -> NDArray[Any]:
+    """Converts an ArrayLike to Numpy array without copying (if possible)"""
+    return to_numpy(array, copy=False)
+
+
+def to_numpy(array: ArrayLike | None, copy: bool = True) -> NDArray[Any]:
+    """Converts an ArrayLike to new Numpy array"""
+    if array is None:
+        return np.ndarray([])
+
+    if isinstance(array, np.ndarray):
+        return array.copy() if copy else array
+
+    if array.__class__.__module__.startswith("tensorflow"):  # pragma: no cover - removed tf from deps
+        tf = _try_import("tensorflow")
+        if tf and tf.is_tensor(array):
+            _logger.log(logging.INFO, "Converting Tensorflow array to NumPy array.")
+            return array.numpy().copy() if copy else array.numpy()  # type: ignore
+
+    if array.__class__.__module__.startswith("torch"):
+        torch = _try_import("torch")
+        if torch and isinstance(array, torch.Tensor):
+            _logger.log(logging.INFO, "Converting PyTorch array to NumPy array.")
+            numpy = array.detach().cpu().numpy().copy() if copy else array.detach().cpu().numpy()  # type: ignore
+            _logger.log(logging.DEBUG, LogMessage(lambda: f"{str(array)} -> {str(numpy)}"))
+            return numpy
+
+    return np.array(array) if copy else np.asarray(array)
+
+
+def to_numpy_iter(iterable: Iterable[ArrayLike]) -> Iterator[NDArray[Any]]:
+    """Yields an iterator of numpy arrays from an ArrayLike"""
+    for array in iterable:
+        yield to_numpy(array)
+
+
+@overload
+def ensure_embeddings(
+    embeddings: T,
+    dtype: torch.dtype,
+    unit_interval: Literal[True, False, "force"] = False,
+) -> torch.Tensor: ...
+
+
+@overload
+def ensure_embeddings(
+    embeddings: T,
+    dtype: type[_np_dtype],
+    unit_interval: Literal[True, False, "force"] = False,
+) -> NDArray[_np_dtype]: ...
+
+
+@overload
+def ensure_embeddings(
+    embeddings: T,
+    dtype: None,
+    unit_interval: Literal[True, False, "force"] = False,
+) -> T: ...
+
+
+def ensure_embeddings(
+    embeddings: T,
+    dtype: type[_np_dtype] | torch.dtype | None = None,
+    unit_interval: Literal[True, False, "force"] = False,
+) -> torch.Tensor | NDArray[_np_dtype] | T:
+    """
+    Validates the embeddings array and converts it to the specified type
+
+    Parameters
+    ----------
+    embeddings : ArrayLike
+        Embeddings array
+    dtype : numpy dtype or torch dtype or None, default None
+        The desired dtype of the output array, None to skip conversion
+    unit_interval : bool or "force", default False
+        Whether to validate or force the embeddings to unit interval
+
+    Returns
+    -------
+        Converted embeddings array
+
+    Raises
+    ------
+    ValueError
+        If the embeddings array is not 2D
+    ValueError
+        If the embeddings array is not unit interval [0, 1]
+    """
+    if isinstance(dtype, torch.dtype):
+        arr = torch.as_tensor(embeddings, dtype=dtype)
+    else:
+        arr = (
+            embeddings.detach().cpu().numpy().astype(dtype)
+            if isinstance(embeddings, torch.Tensor)
+            else np.asarray(embeddings, dtype=dtype)
+        )
+
+    if arr.ndim != 2:
+        raise ValueError(f"Expected a 2D array, but got a {arr.ndim}D array.")
+
+    if unit_interval:
+        arr_min, arr_max = arr.min(), arr.max()
+        if arr_min < 0 or arr_max > 1:
+            if unit_interval == "force":
+                warnings.warn("Embeddings are not unit interval [0, 1]. Forcing to unit interval.")
+                arr = (arr - arr_min) / (arr_max - arr_min)
+            else:
+                raise ValueError("Embeddings must be unit interval [0, 1].")
+
+    if dtype is None:
+        return embeddings
+    else:
+        return arr
+
+
+def flatten(array: ArrayLike) -> NDArray[Any]:
+    """
+    Flattens input array from (N, ... ) to (N, -1) where all samples N have all data in their last dimension
+
+    Parameters
+    ----------
+    X : NDArray, shape - (N, ... )
+        Input array
+
+    Returns
+    -------
+    NDArray, shape - (N, -1)
+    """
+    nparr = as_numpy(array)
+    return nparr.reshape((nparr.shape[0], -1))

dataeval/utils/_bin.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+__all__ = []
+
+import warnings
+from typing import Any, Iterable
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.stats import wasserstein_distance as wd
+
+DISCRETE_MIN_WD = 0.054
+CONTINUOUS_MIN_SAMPLE_SIZE = 20
+
+
+def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
+    """
+    Returns columnwise unique counts for discrete data.
+
+    Parameters
+    ----------
+    data : NDArray
+        Array containing integer values for metadata factors
+    min_num_bins : int | None, default None
+        Minimum number of bins for bincount, helps force consistency across runs
+
+    Returns
+    -------
+    NDArray[np.int]
+        Bin counts per column of data.
+    """
+    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
+    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
+    for idx in range(data.shape[1]):
+        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
+
+    return cnt_array
+
+
+def digitize_data(data: list[Any] | NDArray[Any], bins: int | Iterable[float]) -> NDArray[np.intp]:
+    """
+    Digitizes a list of values into a given number of bins.
+
+    Parameters
+    ----------
+    data : list | NDArray
+        The values to be digitized.
+    bins : int | Iterable[float]
+        The number of bins or list of bin edges for the discrete values that data will be digitized into.
+
+    Returns
+    -------
+    NDArray[np.intp]
+        The digitized values
+    """
+
+    if not np.all([np.issubdtype(type(n), np.number) for n in data]):
+        raise TypeError(
+            "Encountered a data value with non-numeric type when digitizing a factor. "
+            "Ensure all occurrences of continuous factors are numeric types."
+        )
+    if isinstance(bins, int):
+        _, bin_edges = np.histogram(data, bins=bins)
+        bin_edges[-1] = np.inf
+        bin_edges[0] = -np.inf
+    else:
+        bin_edges = list(bins)
+    return np.digitize(data, bin_edges)
+
+
+def bin_data(data: NDArray[Any], bin_method: str) -> NDArray[np.int_]:
+    """
+    Bins continuous data through either equal width bins, equal amounts in each bin, or by clusters.
+    """
+    if bin_method == "clusters":
+        bin_edges = bin_by_clusters(data)
+
+    else:
+        counts, bin_edges = np.histogram(data, bins="auto")
+        n_bins = counts.size
+        if counts[counts > 0].min() < 10:
+            counter = 20
+            while counts[counts > 0].min() < 10 and n_bins >= 2 and counter > 0:
+                counter -= 1
+                n_bins -= 1
+                counts, bin_edges = np.histogram(data, bins=n_bins)
+
+        if bin_method == "uniform_count":
+            quantiles = np.linspace(0, 100, n_bins + 1)
+            bin_edges = np.asarray(np.percentile(data, quantiles))
+
+    bin_edges[0] = -np.inf
+    bin_edges[-1] = np.inf
+    return np.digitize(data, bin_edges)
+
+
+def is_continuous(data: NDArray[np.number[Any]], image_indices: NDArray[np.number[Any]]) -> bool:
+    """
+    Determines whether the data is continuous or discrete using the Wasserstein distance.
+
+    Given a 1D sample, we consider the intervals between adjacent points. For a continuous distribution,
+    a point is equally likely to lie anywhere in the interval bounded by its two neighbors. Furthermore,
+    we can put all "between neighbor" locations on the same scale of 0 to 1 by subtracting the smaller
+    neighbor and dividing out the length of the interval. (Duplicates are either assigned to zero or
+    ignored, depending on context). These normalized locations will be much more uniformly distributed
+    for continuous data than for discrete, and this gives us a way to distinguish them. Call this the
+    Normalized Near Neighbor distribution (NNN), defined on the interval [0,1].
+
+    The Wasserstein distance is available in scipy.stats.wasserstein_distance. We can use it to measure
+    how close the NNN is to a uniform distribution over [0,1]. We found that as long as a sample has at
+    least 20 points, and furthermore at least half as many points as there are discrete values, we can
+    reliably distinguish discrete from continuous samples by testing that the Wasserstein distance
+    measured from a uniform distribution is greater or less than 0.054, respectively.
+    """
+    # Check if the metadata is image specific
+    _, data_indices_unsorted = np.unique(data, return_index=True)
+    if data_indices_unsorted.size == image_indices.size:
+        data_indices = np.sort(data_indices_unsorted)
+        if (data_indices == image_indices).all():
+            data = data[data_indices]
+
+    n_examples = len(data)
+
+    if n_examples < CONTINUOUS_MIN_SAMPLE_SIZE:
+        warnings.warn(
+            f"All samples look discrete with so few data points (< {CONTINUOUS_MIN_SAMPLE_SIZE})", UserWarning
+        )
+        return False
+
+    # Require at least 3 unique values before bothering with NNN
+    xu = np.unique(data, axis=None)
+    if xu.size < 3:
+        return False
+
+    Xs = np.sort(data)
+
+    X0, X1 = Xs[0:-2], Xs[2:]  # left and right neighbors
+
+    dx = np.zeros(n_examples - 2)  # no dx at end points
+    gtz = (X1 - X0) > 0  # check for dups; dx will be zero for them
+    dx[np.logical_not(gtz)] = 0.0
+
+    dx[gtz] = (Xs[1:-1] - X0)[gtz] / (X1 - X0)[gtz]  # the core idea: dx is NNN samples.
+
+    shift = wd(dx, np.linspace(0, 1, dx.size))  # how far is dx from uniform, for this feature?
+
+    return shift < DISCRETE_MIN_WD  # if NNN is close enough to uniform, consider the sample continuous.
+
+
+def bin_by_clusters(data: NDArray[np.number[Any]]) -> NDArray[np.float64]:
+    """
+    Bins continuous data by using the Clusterer to identify clusters
+    and incorporates outliers by adding them to the nearest bin.
+    """
+    # Delay load numba compiled functions
+    from dataeval.utils._clusterer import cluster
+
+    # Create initial clusters
+    c = cluster(data)
+
+    # Create bins from clusters
+    bin_edges = np.zeros(c.clusters.max() + 2)
+    for group in range(c.clusters.max() + 1):
+        points = np.nonzero(c.clusters == group)[0]
+        bin_edges[group] = data[points].min()
+
+    # Get the outliers
+    outliers = np.nonzero(c.clusters == -1)[0]
+
+    # Identify non-outlier neighbors
+    nbrs = c.k_neighbors[outliers]
+    nbrs = np.where(np.isin(nbrs, outliers), -1, nbrs)
+
+    # Find the nearest non-outlier neighbor for each outlier
+    nn = np.full(outliers.size, -1, dtype=np.int32)
+    for row in range(outliers.size):
+        non_outliers = nbrs[row, nbrs[row] != -1]
+        if non_outliers.size > 0:
+            nn[row] = non_outliers[0]
+
+    # Group outliers by their neighbors
+    unique_nnbrs, same_nbr, counts = np.unique(nn, return_inverse=True, return_counts=True)
+
+    # Adjust bin_edges based on each unique neighbor group
+    extend_bins = []
+    for i, nnbr in enumerate(unique_nnbrs):
+        outlier_indices = np.nonzero(same_nbr == i)[0]
+        min2add = data[outliers[outlier_indices]].min()
+        if counts[i] >= 4:
+            extend_bins.append(min2add)
+        else:
+            if min2add < data[nnbr]:
+                clusters = c.clusters[nnbr]
+                bin_edges[clusters] = min2add
+    if extend_bins:
+        bin_edges = np.concatenate([bin_edges, extend_bins])
+
+    bin_edges = np.sort(bin_edges)
+    return bin_edges

dataeval/utils/_clusterer.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+__all__ = []
+
+import warnings
+from dataclasses import dataclass
+
+import numba
+import numpy as np
+from numpy.typing import NDArray
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", category=FutureWarning)
+    from fast_hdbscan.cluster_trees import (
+        CondensedTree,
+        cluster_tree_from_condensed_tree,
+        condense_tree,
+        ds_find,
+        ds_rank_create,
+        ds_union_by_rank,
+        extract_eom_clusters,
+        get_cluster_label_vector,
+        get_point_membership_strength_vector,
+        mst_to_linkage_tree,
+    )
+
+from dataeval.typing import ArrayLike
+from dataeval.utils._array import flatten, to_numpy
+from dataeval.utils._fast_mst import calculate_neighbor_distances, minimum_spanning_tree
+
+
+@numba.njit(parallel=True, locals={"i": numba.types.int32})
+def compare_links_to_cluster_std(mst, clusters):
+    cluster_ids = np.unique(clusters)
+    cluster_grouping = np.full(mst.shape[0], -1, dtype=np.int16)
+
+    for i in numba.prange(mst.shape[0]):
+        cluster_id = clusters[np.int32(mst[i, 0])]
+        if cluster_id == clusters[np.int32(mst[i, 1])]:
+            cluster_grouping[i] = np.int16(cluster_id)
+
+    overall_mean = mst.T[2].mean()
+    order_mag = np.floor(np.log10(overall_mean)) if overall_mean > 0 else 0
+    compare_mag = -3 if order_mag >= 0 else order_mag - 3
+
+    exact_dup = np.full((mst.shape[0], 2), -1, dtype=np.int32)
+    exact_dups_index = np.nonzero(mst[:, 2] < 10**compare_mag)[0]
+    exact_dup[exact_dups_index] = mst[exact_dups_index, :2]
+
+    near_dup = np.full((mst.shape[0], 2), -1, dtype=np.int32)
+    for i in range(cluster_ids.size):
+        cluster_links = np.nonzero(cluster_grouping == cluster_ids[i])[0]
+        cluster_std = mst[cluster_links, 2].std()
+
+        near_dups = np.nonzero(mst[cluster_links, 2] < cluster_std)[0]
+        near_dups_index = cluster_links[near_dups]
+        near_dup[near_dups_index] = mst[near_dups_index, :2]
+
+    exact_idx = np.nonzero(exact_dup.T[0] != -1)[0]
+    near_dup[exact_idx] = np.full((exact_idx.size, 2), -1, dtype=np.int32)
+    near_idx = np.nonzero(near_dup.T[0] != -1)[0]
+
+    return exact_dup[exact_idx], near_dup[near_idx]
+
+
+@dataclass
+class ClusterData:
+    clusters: NDArray[np.intp]
+    mst: NDArray[np.double]
+    linkage_tree: NDArray[np.double]
+    condensed_tree: CondensedTree
+    membership_strengths: NDArray[np.double]
+    k_neighbors: NDArray[np.int32]
+    k_distances: NDArray[np.double]
+
+
+def cluster(data: ArrayLike) -> ClusterData:
+    single_cluster = False
+    cluster_selection_epsilon = 0.0
+    # cluster_selection_method = "eom"
+
+    x = flatten(to_numpy(data))
+    samples, features = x.shape  # Due to flatten(), we know shape has a length of 2
+    if samples < 2:
+        raise ValueError(f"Data should have at least 2 samples; got {samples}")
+    if features < 1:
+        raise ValueError(f"Samples should have at least 1 feature; got {features}")
+
+    num_samples = len(x)
+    min_num = int(num_samples * 0.05)
+    min_cluster_size: int = min(max(5, min_num), 100)
+
+    max_neighbors = min(25, num_samples - 1)
+    kneighbors, kdistances = calculate_neighbor_distances(x, max_neighbors)
+    unsorted_mst: NDArray[np.double] = minimum_spanning_tree(x, kneighbors, kdistances)
+    mst: NDArray[np.double] = unsorted_mst[np.argsort(unsorted_mst.T[2])]
+    linkage_tree: NDArray[np.double] = mst_to_linkage_tree(mst)
+    condensed_tree: CondensedTree = condense_tree(linkage_tree, min_cluster_size, None)
+
+    cluster_tree = cluster_tree_from_condensed_tree(condensed_tree)
+
+    selected_clusters = extract_eom_clusters(condensed_tree, cluster_tree, allow_single_cluster=single_cluster)
+
+    # Uncomment if cluster_selection_method is made a parameter
+    # if cluster_selection_method != "eom":
+    #     selected_clusters = extract_leaves(condensed_tree, allow_single_cluster=single_cluster)
+
+    # Uncomment if cluster_selection_epsilon is made a parameter
+    # if len(selected_clusters) > 1 and cluster_selection_epsilon > 0.0:
+    #     selected_clusters = cluster_epsilon_search(
+    #         selected_clusters,
+    #         cluster_tree,
+    #         min_persistence=cluster_selection_epsilon,
+    #     )
+
+    clusters = get_cluster_label_vector(
+        condensed_tree,
+        selected_clusters,
+        cluster_selection_epsilon,
+        n_samples=x.shape[0],
+    )
+
+    membership_strengths = get_point_membership_strength_vector(condensed_tree, selected_clusters, clusters)
+
+    return ClusterData(clusters, mst, linkage_tree, condensed_tree, membership_strengths, kneighbors, kdistances)
+
+
+def sorted_union_find(index_groups):
+    """Merges and sorts groups of indices that share any common index"""
+    groups = [[np.int32(x) for x in range(0)] for y in range(0)]
+    uniques, inverse = np.unique(index_groups, return_inverse=True)
+    inverse = inverse.flatten()
+    disjoint_set = ds_rank_create(uniques.size)
+    cluster_points = np.empty(uniques.size, dtype=np.uint32)
+    for i in range(index_groups.shape[0]):
+        point, nbr = np.int32(inverse[i * 2]), np.int32(inverse[i * 2 + 1])
+        ds_union_by_rank(disjoint_set, point, nbr)
+    for i in range(uniques.size):
+        cluster_points[i] = ds_find(disjoint_set, i)
+    for i in range(uniques.size):
+        dups = np.nonzero(cluster_points == i)[0]
+        if dups.size > 0:
+            groups.append(uniques[dups].tolist())
+    return sorted(groups)

dataeval/utils/_fast_mst.py

@@ -0,0 +1,189 @@
+# Adapted from fast_hdbscan python module
+# Original Authors: Leland McInnes <https://github.com/TutteInstitute/fast_hdbscan>
+# Adapted for DataEval by Ryan Wood
+# License: BSD 2-Clause
+
+__all__ = []
+
+import warnings
+
+import numba
+import numpy as np
+from sklearn.neighbors import NearestNeighbors
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", category=FutureWarning)
+    from fast_hdbscan.disjoint_set import ds_find, ds_rank_create
+
+
+@numba.njit()
+def _ds_union_by_rank(disjoint_set, point, nbr):
+    y = ds_find(disjoint_set, point)
+    x = ds_find(disjoint_set, nbr)
+
+    if x == y:
+        return 0
+
+    if disjoint_set.rank[x] < disjoint_set.rank[y]:
+        x, y = y, x
+
+    disjoint_set.parent[y] = x
+    if disjoint_set.rank[x] == disjoint_set.rank[y]:
+        disjoint_set.rank[x] += 1
+    return 1
+
+
+@numba.njit(locals={"i": numba.types.uint32, "nbr": numba.types.uint32, "dist": numba.types.float32})
+def _init_tree(n_neighbors, n_distance):
+    # Initial graph to hold tree connections
+    tree = np.zeros((n_neighbors.size - 1, 3), dtype=np.float32)
+    disjoint_set = ds_rank_create(n_neighbors.size)
+    cluster_points = np.empty(n_neighbors.size, dtype=np.uint32)
+
+    int_tree = 0
+    for i in range(n_neighbors.size):
+        nbr = n_neighbors[i]
+        connect = _ds_union_by_rank(disjoint_set, i, nbr)
+        if connect == 1:
+            dist = n_distance[i]
+            tree[int_tree] = (np.float32(i), np.float32(nbr), dist)
+            int_tree += 1
+
+    for i in range(cluster_points.size):
+        cluster_points[i] = ds_find(disjoint_set, i)
+
+    return tree, int_tree, disjoint_set, cluster_points
+
+
+@numba.njit(locals={"i": numba.types.uint32, "nbr": numba.types.uint32})
+def _update_tree_by_distance(tree, int_tree, disjoint_set, n_neighbors, n_distance):
+    cluster_points = np.empty(n_neighbors.size, dtype=np.uint32)
+    sort_dist = np.argsort(n_distance)
+    dist_sorted = n_distance[sort_dist]
+    nbrs_sorted = n_neighbors[sort_dist]
+    points = np.arange(n_neighbors.size)
+    point_sorted = points[sort_dist]
+
+    for i in range(n_neighbors.size):
+        point = point_sorted[i]
+        nbr = nbrs_sorted[i]
+        connect = _ds_union_by_rank(disjoint_set, point, nbr)
+        if connect == 1:
+            dist = dist_sorted[i]
+            tree[int_tree] = (np.float32(point), np.float32(nbr), dist)
+            int_tree += 1
+
+    for i in range(cluster_points.size):
+        cluster_points[i] = ds_find(disjoint_set, i)
+
+    return tree, int_tree, disjoint_set, cluster_points
+
+
+@numba.njit(locals={"i": numba.types.uint32})
+def _cluster_edges(tracker, last_idx, cluster_distances):
+    cluster_ids = np.unique(tracker)
+    edge_points = []
+    for idx in range(cluster_ids.size):
+        cluster_points = np.nonzero(tracker == cluster_ids[idx])[0]
+        cluster_size = cluster_points.size
+        cluster_mean = cluster_distances[: last_idx + 1, cluster_points].mean()
+        cluster_std = cluster_distances[: last_idx + 1, cluster_points].std()
+        threshold = cluster_mean + cluster_std
+        points_mean = np.empty_like(cluster_points, dtype=np.float32)
+        for i in range(cluster_size):
+            points_mean[i] = cluster_distances[: last_idx + 1, cluster_points[i]].mean()
+        pts_to_add = cluster_points[np.nonzero(points_mean > threshold)[0]]
+        threshold = int(cluster_size * 0.01) if np.floor(np.log10(cluster_size)) > 2 else int(cluster_size * 0.1)
+        threshold = max(10, threshold)
+        if pts_to_add.size > threshold:
+            edge_points.append(pts_to_add)
+        else:
+            edge_points.append(cluster_points)
+    return edge_points
+
+
+def _compute_nn(dataA, dataB, k):
+    distances, neighbors = NearestNeighbors(n_neighbors=k + 1, algorithm="brute").fit(dataA).kneighbors(dataB)
+    neighbors = np.array(neighbors[:, 1 : k + 1], dtype=np.int32)
+    distances = np.array(distances[:, 1 : k + 1], dtype=np.float32)
+    return neighbors, distances
+
+
+def _calculate_cluster_neighbors(data, groups, point_array):
+    """Rerun nearest neighbor based on clusters"""
+    cluster_neighbors = np.zeros(point_array.size, dtype=np.uint32)
+    cluster_nbr_distances = np.full(point_array.size, np.inf, dtype=np.float32)
+
+    for i in range(len(groups)):
+        selectionA = groups[i]
+        groupA = data[selectionA]
+        selectionB = np.concatenate([arr for j, arr in enumerate(groups) if j != i])
+        groupB = data[selectionB]
+        new_neighbors, new_distances = _compute_nn(groupB, groupA, 2)
+        cluster_neighbors[selectionA] = selectionB[new_neighbors[:, 1]]
+        cluster_nbr_distances[selectionA] = new_distances[:, 1]
+
+    return cluster_neighbors, cluster_nbr_distances
+
+
+def minimum_spanning_tree(data, neighbors, distances):
+    # Transpose arrays to get number of samples along a row
+    k_neighbors = neighbors.T.astype(np.uint32).copy()
+    k_distances = distances.T.astype(np.float32).copy()
+
+    # Create cluster merging tracker
+    merge_tracker = np.full((k_neighbors.shape[0] + 1, k_neighbors.shape[1]), -1, dtype=np.int32)
+
+    # Initialize tree
+    tree, int_tree, tree_disjoint_set, merge_tracker[0] = _init_tree(k_neighbors[0], k_distances[0])
+
+    # Loop through all of the neighbors, updating the tree
+    last_idx = 0
+    for i in range(1, k_neighbors.shape[0]):
+        tree, int_tree, tree_disjoint_set, merge_tracker[i] = _update_tree_by_distance(
+            tree, int_tree, tree_disjoint_set, k_neighbors[i], k_distances[i]
+        )
+        last_idx = i
+        if (merge_tracker[i] == merge_tracker[i - 1]).all():
+            last_idx -= 1
+            break
+
+    # Identify final clusters
+    cluster_ids = np.unique(merge_tracker[last_idx])
+    if cluster_ids.size > 1:
+        # Determining the edge points
+        edge_points = _cluster_edges(merge_tracker[last_idx], last_idx, k_distances)
+
+        # Run nearest neighbor again between clusters to reach single cluster
+        additional_neighbors, additional_distances = _calculate_cluster_neighbors(
+            data, edge_points, merge_tracker[last_idx]
+        )
+
+        # Update clusters
+        last_idx += 1
+        tree, int_tree, tree_disjoint_set, merge_tracker[last_idx] = _update_tree_by_distance(
+            tree, int_tree, tree_disjoint_set, additional_neighbors, additional_distances
+        )
+
+    return tree
+
+
+def calculate_neighbor_distances(data: np.ndarray, k: int = 10):
+    # Have the potential to add in other distance calculations - supported calculations:
+    # https://github.com/lmcinnes/pynndescent/blob/master/pynndescent/pynndescent_.py#L524
+    try:
+        from pynndescent import NNDescent
+
+        max_descent = 30 if k <= 20 else k + 16
+        index = NNDescent(
+            data,
+            metric="euclidean",
+            n_neighbors=max_descent,
+        )
+        neighbors, distances = index.neighbor_graph
+    except ImportError:
+        distances, neighbors = NearestNeighbors(n_neighbors=k + 1, algorithm="brute").fit(data).kneighbors(data)
+
+    neighbors = np.array(neighbors[:, 1 : k + 1], dtype=np.int32)
+    distances = np.array(distances[:, 1 : k + 1], dtype=np.float32)
+    return neighbors, distances