flash1dkmeans 0.0.1__py3-none-any.whl

flash1dkmeans/__init__.py

@@ -0,0 +1 @@
+from .main import kmeans_1d

flash1dkmeans/config.py

@@ -0,0 +1,9 @@
+"""config.py
+Contains some configuration variables for the flash1dkmeans package.
+"""
+
+import numpy as np
+
+PREFIX_SUM_DTYPE = np.float64
+ARRAY_INDEX_DTYPE = np.int32
+LABEL_DTYPE = np.int32

flash1dkmeans/k_cluster.py

@@ -0,0 +1,341 @@
+"""n_cluster.py
+Optimized kmeans for 1D data with n clusters.
+
+This module contains an optimized kmeans for 1D data with n clusters.
+
+The following is the main function:
+- flash_1d_kmeans_n_cluster
+
+Inputs must be sorted in ascending order, no default values are provided, and no error checking is done.
+This is because this module is intended to be used internally.
+"""
+
+import numba
+from .utils import query_prefix_sum, query_weights_prefix_sum
+import numpy as np
+from .config import ARRAY_INDEX_DTYPE
+
+
+#@numba.njit(cache=True)
+def flash_1d_kmeans_k_cluster(
+        X,
+        n_clusters,
+        max_iter,
+        is_weighted,
+        weights_prefix_sum, weighted_X_prefix_sum,
+        weighted_X_squared_prefix_sum,
+        n_local_trials,  # This value is 2 + int(np.log(n_clusters)) in sklearn
+        start_idx,
+        stop_idx,
+):
+    """An optimized kmeans for 1D data with n clusters.
+    Exploits the fact that the data is 1D to optimize the calculations.
+    Time complexity: O(n_clusters * log(n_clusters) * log(len(X)) * n_clusters + max_iter * log(len(X)) * n_clusters)
+                      = O(k ^ 2 * log(k) * log(n) + max_iter * log(n) * k)
+
+    Args:
+        X: np.ndarray
+            The input data. Should be sorted in ascending order.
+        n_clusters: int
+            The number of clusters to generate
+        max_iter: int
+            The maximum number of iterations to run
+        is_weighted: bool
+            Whether the data is weighted. If True, the weighted versions of the arrays should be provided.
+        weights_prefix_sum: np.ndarray
+            The prefix sum of the weights. Should be None if the data is unweighted.
+        weighted_X_prefix_sum: np.ndarray
+            The prefix sum of the weighted X
+        weighted_X_squared_prefix_sum: np.ndarray
+            The prefix sum of the weighted X squared
+        n_local_trials: int
+            The number of local trials to run when choosing the next centroid
+        start_idx: int
+            The start index of the range to consider
+        stop_idx: int
+            The stop index of the range to consider
+
+    Returns:
+        centroids: np.ndarray
+            The centroids of the clusters
+        cluster_borders: np.ndarray
+            The borders of the clusters
+    """
+    if is_weighted:
+        assert weights_prefix_sum is not None, "weights_prefix_sum must be provided for weighted data"
+    else:
+        assert weights_prefix_sum is None, "weights_prefix_sum must not be provided for unweighted data"
+
+    cluster_borders = np.empty(n_clusters + 1, dtype=ARRAY_INDEX_DTYPE)
+    cluster_borders[0] = start_idx
+    cluster_borders[-1] = stop_idx
+
+    centroids = _kmeans_plusplus(
+        X, n_clusters, is_weighted,
+        weights_prefix_sum, weighted_X_prefix_sum,
+        weighted_X_squared_prefix_sum,
+        n_local_trials, start_idx, stop_idx,
+    )
+    sorted_centroids = np.sort(centroids)
+
+    for _ in range(max_iter):
+        new_cluster_borders = _centroids_to_cluster_borders(X, sorted_centroids, start_idx, stop_idx)
+
+        if np.array_equal(cluster_borders, new_cluster_borders):
+            break
+
+        cluster_borders[:] = new_cluster_borders
+        for i in range(n_clusters):
+            cluster_start = cluster_borders[i]
+            cluster_end = cluster_borders[i + 1]
+
+            if cluster_end < cluster_start:
+                raise ValueError("Cluster end is less than cluster start")
+
+            if cluster_start == cluster_end:
+                continue
+
+            cluster_weighted_X_sum = query_prefix_sum(weighted_X_prefix_sum, cluster_start, cluster_end)
+            cluster_weight_sum = query_weights_prefix_sum(is_weighted, weights_prefix_sum, cluster_start, cluster_end)
+
+            if cluster_weight_sum == 0:
+                # if the sum of the weights is zero, we set the centroid to the mean of the cluster
+                sorted_centroids[i] = X[cluster_start:cluster_end].mean()
+            else:
+                sorted_centroids[i] = cluster_weighted_X_sum / cluster_weight_sum
+
+    return sorted_centroids, cluster_borders
+
+
+@numba.njit(cache=True)
+def _rand_choice_prefix_sum(arr, prob_prefix_sum, start_idx, stop_idx):
+    """Randomly choose an element from arr according to the probability distribution given by prob_prefix_sum
+    Time complexity: O(log(n))
+
+    Args:
+        arr: np.ndarray
+            The array to choose from
+        prob_prefix_sum: np.ndarray
+            The prefix sum of the probability distribution
+        start_idx: int
+            The start index of the range to consider
+        stop_idx: int
+            The stop index of the range to consider
+
+    Returns:
+        The chosen element
+    """
+    total_prob = query_prefix_sum(prob_prefix_sum, start_idx, stop_idx)
+    selector = np.random.random_sample() * total_prob
+
+    # Because we are using start_idx as the base, but the prefix sum is calculated from 0,
+    # we need to adjust the selector if start_idx is not 0.
+    adjusted_selector = selector + prob_prefix_sum[start_idx - 1] if start_idx > 0 else selector
+
+    # Search for the index of the selector in the prefix sum, and add start_idx to get the index in the original array
+    idx = np.searchsorted(prob_prefix_sum[start_idx:stop_idx], adjusted_selector) + start_idx
+
+    return arr[idx]
+
+
+#@numba.njit(cache=True)
+def _centroids_to_cluster_borders(X, sorted_centroids, start_idx, stop_idx):
+    """Converts the centroids to cluster borders.
+    The cluster borders are where the clusters are divided.
+    The centroids must be sorted.
+
+    Time complexity: O(log(len(X)) * len(centroids))
+
+    Args:
+        X: np.ndarray
+            The input data. Should be sorted in ascending order.
+        sorted_centroids: np.ndarray
+            The sorted centroids
+        start_idx: int
+            The start index of the range to consider
+        stop_idx: int
+            The stop index of the range to consider
+
+    Returns:
+        np.ndarray: The cluster borders
+    """
+    midpoints = (sorted_centroids[:-1] + sorted_centroids[1:]) / 2
+    cluster_borders = np.empty(len(sorted_centroids) + 1, dtype=ARRAY_INDEX_DTYPE)
+    cluster_borders[0] = start_idx
+    cluster_borders[-1] = stop_idx
+    cluster_borders[1:-1] = np.searchsorted(X[start_idx:stop_idx], midpoints) + start_idx
+    return cluster_borders
+
+
+#@numba.njit(cache=True)
+def _calculate_inertia(sorted_centroids, centroid_ranges,
+                       is_weighted, weights_prefix_sum, weighted_X_prefix_sum, weighted_X_squared_prefix_sum,
+                       stop_idx):
+    """Calculates the inertia of the clusters given the centroids.
+    The inertia is the sum of the squared distances of each sample to the closest centroid.
+    The calculations are done efficiently using prefix sums.
+
+    Time complexity: O(len(centroids))
+
+    Args:
+        sorted_centroids: np.ndarray
+            The centroids of the clusters
+        centroid_ranges: np.ndarray
+            The borders of the clusters
+        is_weighted: bool
+            Whether the data is weighted. If True, the weighted versions of the arrays should be provided.
+        weights_prefix_sum: np.ndarray
+            The prefix sum of the weights. Should be None if the data is unweighted.
+        weighted_X_prefix_sum: np.ndarray
+            The prefix sum of the weighted X
+        weighted_X_squared_prefix_sum: np.ndarray
+            The prefix sum of the weighted X squared
+        stop_idx: int
+            The stop index of the range to consider
+    """
+    # inertia = sigma_i(w_i * abs(x_i - c)^2) = sigma_i(w_i * (x_i^2 - 2 * x_i * c + c^2))
+    #         = sigma_i(w_i * x_i^2) - 2 * c * sigma_i(w_i * x_i) + c^2 * sigma_i(w_i)
+    #         = sigma_i(weighted_X_squared) - 2 * c * sigma_i(weighted_X) + c^2 * sigma_i(weight)
+    #  Note that the centroid c is the CLOSEST centroid to x_i, so the above calculation must be done for each cluster
+
+    inertia = 0
+    for i in range(len(sorted_centroids)):
+        start = centroid_ranges[i]
+        end = centroid_ranges[i + 1]
+
+        if start >= stop_idx:
+            break
+        if end >= stop_idx:
+            end = stop_idx
+
+        if start == end:
+            continue
+
+        cluster_weighted_X_squared_sum = query_prefix_sum(weighted_X_squared_prefix_sum, start, end)
+        cluster_weighted_X_sum = query_prefix_sum(weighted_X_prefix_sum, start, end)
+        cluster_weight_sum = query_weights_prefix_sum(is_weighted, weights_prefix_sum, start, end)
+
+        inertia += (cluster_weighted_X_squared_sum - 2 * sorted_centroids[i] * cluster_weighted_X_sum +
+                    sorted_centroids[i] ** 2 * cluster_weight_sum)
+
+    return inertia
+
+
+#@numba.njit(cache=True)
+def _rand_choice_centroids(X, centroids,
+                           is_weighted, weights_prefix_sum, weighted_X_prefix_sum, weighted_X_squared_prefix_sum,
+                           sample_size, start_idx, stop_idx):
+    """Randomly choose sample_size elements from X, weighted by the distance to the closest centroid.
+    The weighted logic is implemented efficiently by utilizing the _calculate_inertia function.
+
+    Time complexity: O(sample_size * log(len(X)) * len(centroids))
+
+    Args:
+        X: np.ndarray
+            The input data. Should be sorted in ascending order.
+        centroids: np.ndarray
+            The centroids of the clusters
+        is_weighted: bool
+            Whether the data is weighted. If True, the weighted versions of the arrays should be provided.
+        weights_prefix_sum: np.ndarray
+            The prefix sum of the weights. Should be None if the data is unweighted.
+        weighted_X_prefix_sum: np.ndarray
+            The prefix sum of the weighted X
+        weighted_X_squared_prefix_sum: np.ndarray
+            The prefix sum of the weighted X squared
+        sample_size: int
+            The number of samples to choose
+        start_idx: int
+            The start index of the range to consider
+        stop_idx: int
+            The stop index of the range to consider
+
+    Returns:
+        np.ndarray: The chosen samples
+    """
+    sorted_centroids = np.sort(centroids)
+    cluster_borders = _centroids_to_cluster_borders(X, sorted_centroids, start_idx, stop_idx)
+    total_inertia = _calculate_inertia(sorted_centroids, cluster_borders,
+                                       is_weighted, weights_prefix_sum, weighted_X_prefix_sum,
+                                       weighted_X_squared_prefix_sum, stop_idx)
+    selectors = np.random.random_sample(sample_size) * total_inertia
+    results = np.empty(sample_size, dtype=centroids.dtype)
+
+    for i in range(sample_size):
+        selector = selectors[i]
+        floor = start_idx + 1
+        ceiling = stop_idx
+        while floor < ceiling:
+            stop_idx_cand = (floor + ceiling) // 2
+            inertia = _calculate_inertia(sorted_centroids, cluster_borders,
+                                         is_weighted, weights_prefix_sum, weighted_X_prefix_sum,
+                                         weighted_X_squared_prefix_sum, stop_idx_cand)
+            if inertia < selector:
+                floor = stop_idx_cand + 1
+            else:
+                ceiling = stop_idx_cand
+        results[i] = X[floor - 1]
+
+    return results
+
+
+#@numba.njit(cache=True)
+def _kmeans_plusplus(X, n_clusters,
+                     is_weighted, weights_prefix_sum, weighted_X_prefix_sum, weighted_X_squared_prefix_sum,
+                     n_local_trials, start_idx, stop_idx):
+    """An optimized version of the kmeans++ initialization algorithm for 1D data.
+    The algorithm is optimized for 1D data and utilizes prefix sums for efficient calculations.
+
+    Time complexity: O(n_clusters * log(n_clusters) * log(len(X)) * n_clusters)
+                     = O(k ^ 2 * log(k) * log(n))
+
+    Args:
+        X: np.ndarray
+            The input data
+        n_clusters: int
+            The number of clusters to choose
+        is_weighted: bool
+            Whether the data is weighted. If True, the weighted versions of the arrays should be provided.
+        weights_prefix_sum: np.ndarray
+            The prefix sum of the weights. Should be None if the data is unweighted.
+        weighted_X_prefix_sum: np.ndarray
+            The prefix sum of the weighted X
+        weighted_X_squared_prefix_sum: np.ndarray
+            The prefix sum of the weighted X squared
+
+    Returns:
+        np.ndarray: The chosen centroids
+    """
+    centroids = np.empty(n_clusters, dtype=X.dtype)
+
+    # First centroid is chosen randomly according to sample_weight
+    if is_weighted:
+        centroids[0] = _rand_choice_prefix_sum(X, weights_prefix_sum, start_idx, stop_idx)
+    else:
+        centroids[0] = X[np.random.randint(start_idx, stop_idx)]
+
+    for c_id in range(1, n_clusters):
+        # Choose the next centroid randomly according to the weighted distances
+        # Sample n_local_trials candidates and choose the best one
+        centroid_candidates = _rand_choice_centroids(
+            X, centroids[:c_id],
+            is_weighted, weights_prefix_sum, weighted_X_prefix_sum,
+            weighted_X_squared_prefix_sum, n_local_trials,
+            start_idx, stop_idx
+        )
+
+        best_inertia = np.inf
+        best_centroid = None
+        for i in range(len(centroid_candidates)):
+            centroids[c_id] = centroid_candidates[i]
+            sorted_centroids = np.sort(centroids[:c_id + 1])
+            centroid_ranges = _centroids_to_cluster_borders(X, sorted_centroids, start_idx, stop_idx)
+            inertia = _calculate_inertia(sorted_centroids, centroid_ranges,
+                                         is_weighted, weights_prefix_sum, weighted_X_prefix_sum,
+                                         weighted_X_squared_prefix_sum, stop_idx)
+            if inertia < best_inertia:
+                best_inertia = inertia
+                best_centroid = centroid_candidates[i]
+        centroids[c_id] = best_centroid
+
+    return centroids

flash1dkmeans/main.py

@@ -0,0 +1,300 @@
+from .k_cluster import flash_1d_kmeans_k_cluster
+from .two_cluster import flash_1d_kmeans_two_cluster
+from .config import LABEL_DTYPE, PREFIX_SUM_DTYPE
+import numpy as np
+import logging
+
+
+def kmeans_1d(
+        X,
+        n_clusters,
+        max_iter=300,
+        is_sorted=False,
+        sample_weights=None,
+        n_local_trials=None,
+        return_cluster_borders=False,
+        weights_prefix_sum=None,
+        weighted_X_prefix_sum=None,
+        weighted_X_squared_prefix_sum=None,
+        start_idx=None,
+        stop_idx=None,
+):
+    """An optimized kmeans for 1D data.
+    Utilizes a binary search to find the optimal division points for 2 clusters,
+    and an optimized kmeans++ initialization for more than 2 clusters.
+
+    Exploits the fact that 1D data can be sorted.
+
+    Time complexity:
+        2 clusters:
+            O(log(n)) (+ (O(n) for prefix sum calculation if not provided))
+        n clusters:
+            O(n_clusters * 2 + log(n_clusters) * log(n) + max_iter * log(n) * n_clusters)
+            (+ (O(n) for prefix sum calculation if not provided))
+
+    Args:
+        X: np.ndarray or list
+            The input data. Should be sorted in ascending order if is_sorted is False.
+        n_clusters: int
+            The number of clusters.
+        max_iter: int
+            The maximum number of iterations. Only relevant for n_clusters > 2.
+        is_sorted: bool
+            Whether the data is already sorted.
+        sample_weights: np.ndarray or list or None
+            The sample weights. If None, all samples are weighted equally.
+        n_local_trials: int
+            The number of local trials for kmeans++ initialization. Only relevant for n_clusters > 2.
+        return_cluster_borders: bool
+            Whether to return the cluster border indices instead of the labels.
+
+        In the case of repeatedly calling kmeans on different ranges of the same data, the following arguments
+        can be provided to avoid redundant calculations. Make sure that the data is sorted if these are provided.
+
+        weights_prefix_sum: np.ndarray or list or None
+            The prefix sum of the sample weights. Should be None if the data is unweighted.
+        weighted_X_prefix_sum: np.ndarray or list or None
+            The prefix sum of (the weighted) X.
+        weighted_X_squared_prefix_sum: np.ndarray or list or None
+            The prefix sum of (weighted) X squared.
+        start_idx: int
+            The start index of the range to consider.
+        stop_idx: int
+            The stop index of the range to consider.
+
+    Returns:
+        centroids: np.ndarray
+            The centroids of the clusters.
+        labels: np.ndarray
+            The labels of the samples. Only returned if return_cluster_borders is False.
+        cluster_borders: np.ndarray
+            The borders of the clusters. Only returned if return_cluster_borders is True.
+    """
+    # -------------- Input validation --------------
+    if weighted_X_prefix_sum is not None:
+        assert weighted_X_squared_prefix_sum is not None, \
+            "If weighted_X_prefix_sum is provided, weighted_X_squared_prefix_sum should also be provided."
+        # Note that weights_prefix_sum may be None if the data is unweighted
+
+    # Check if user provided both sample_weights and weights_prefix_sum
+    if sample_weights is not None and weights_prefix_sum is not None:
+        raise ValueError("Both sample_weights and weights_prefix_sum cannot be provided. Please provide only one.")
+
+    # Check if data is weighted
+    if sample_weights is not None or weights_prefix_sum is not None:
+        is_weighted = True
+    else:
+        is_weighted = False
+
+    # Check that data is sorted if weighted_X_prefix_sum is provided
+    if weighted_X_prefix_sum is not None and not is_sorted:
+        raise ValueError("The data must be sorted if weighted_X_prefix_sum is provided.")
+
+    # Check if all lengths are the same
+    if sample_weights is not None:
+        assert len(X) == len(sample_weights), "X and sample_weights must have the same length"
+    if weights_prefix_sum is not None:
+        assert len(X) == len(weights_prefix_sum), "X and weights_prefix_sum must have the same length"
+    if weighted_X_prefix_sum is not None:
+        assert len(X) == len(weighted_X_prefix_sum), "X and weighted_X_prefix_sum must have the same length"
+    if weighted_X_squared_prefix_sum is not None:
+        assert len(X) == len(weighted_X_squared_prefix_sum), \
+            "X and weighted_X_squared_prefix_sum must have the same length"
+
+    # Check that return_cluster_borders is called with is_sorted=True, warn the user otherwise
+    if return_cluster_borders:  # User requested cluster borders
+        if not is_sorted:  # However, the user did not provide sorted data
+            logging.warning(  # Warn the user that the data will be sorted
+                "The returned cluster borders will be indexed in the sorted order of X, not the original order.")
+
+    # Check if start_idx and stop_idx are given with sorted data
+    if not is_sorted:
+        assert start_idx is None and stop_idx is None, "start_idx and stop_idx are only valid if the data is sorted."
+
+    # Check some values
+    assert n_clusters >= 2, "n_clusters must be at least 2"
+    assert max_iter >= 0, "max_iter must be non-negative"
+    assert n_local_trials is None or n_local_trials > 0, "n_local_trials must be positive"
+
+    # -------------- Convert all arrays to numpy arrays --------------
+    X = np.asarray(X)
+
+    if sample_weights is not None:
+        sample_weights = np.asarray(sample_weights)
+
+    if weights_prefix_sum is not None:
+        weights_prefix_sum = np.asarray(weights_prefix_sum)
+
+    if weighted_X_prefix_sum is not None:
+        weighted_X_prefix_sum = np.asarray(weighted_X_prefix_sum)
+
+    if weighted_X_squared_prefix_sum is not None:
+        weighted_X_squared_prefix_sum = np.asarray(weighted_X_squared_prefix_sum)
+
+    # -------------- Setting up defaults --------------
+
+    if start_idx is None:
+        start_idx = 0
+    if stop_idx is None:
+        stop_idx = len(X)
+
+    assert 0 <= start_idx < stop_idx <= len(X), "Invalid start_idx and stop_idx"
+
+    if n_local_trials is None:
+        n_local_trials = 2 + int(np.log(n_clusters))
+
+    # -------------- Sorting --------------
+
+    # Sort the data if it is not already sorted
+    # Note that we checked that start_idx and stop_idx are not given if the data is not sorted
+    if not is_sorted:
+        sorted_indices = np.argsort(X)
+        sorted_X = X[sorted_indices]
+        sorted_sample_weights = None if sample_weights is None else sample_weights[sorted_indices]
+    else:
+        sorted_indices = None
+        sorted_X = X
+        sorted_sample_weights = sample_weights
+
+    # -------------- Main algorithm --------------
+
+    centroids, cluster_borders = _sorted_kmeans_1d(
+        sorted_X,
+        n_clusters,
+        max_iter,
+        is_weighted,
+        sorted_sample_weights,
+        n_local_trials,
+        weights_prefix_sum,
+        weighted_X_prefix_sum,
+        weighted_X_squared_prefix_sum,
+        start_idx,
+        stop_idx,
+    )
+
+    # -------------- Post-processing --------------
+
+    if return_cluster_borders:
+        # We checked that the data is sorted if cluster borders are requested, so no post-processing is needed
+        return centroids, cluster_borders
+    else:
+        # Convert cluster borders to labels
+        labels = np.zeros(len(X), dtype=LABEL_DTYPE)
+        for i in range(n_clusters):
+            labels[cluster_borders[i]:cluster_borders[i + 1]] = i
+
+        # Unsort the labels if the data was not sorted
+        if not is_sorted:
+            assert sorted_indices is not None, "sorted_indices should be available here"
+            labels = labels[np.argsort(sorted_indices)]
+
+        return centroids, labels
+
+
+def _sorted_kmeans_1d(
+        sorted_X,  # caller should ensure that X is sorted
+        n_clusters,  # caller should ensure n_clusters >= 2
+        max_iter,  # caller should ensure max_iter >= 0
+        is_weighted,
+        sample_weights,
+        n_local_trials,
+        weights_prefix_sum,
+        weighted_X_prefix_sum,
+        weighted_X_squared_prefix_sum,
+        start_idx,  # caller should ensure 0 <= start_idx < stop_idx <= len(X)
+        stop_idx,  # caller should ensure 0 <= start_idx < stop_idx <= len(X)
+):
+    """The main algorithm for flash_1d_kmeans.
+    This function assumes that the input data is sorted, and all input is assumed to be valid.
+
+    Time complexity:
+        2 clusters:
+            O(log(n)) (+ (O(n) for prefix sum calculation if not provided))
+        n clusters:
+            O(n_clusters * 2 + log(n_clusters) * log(n) + max_iter * log(n) * n_clusters)
+            (+ (O(n) for prefix sum calculation if not provided))
+
+    Args:
+        sorted_X: np.ndarray
+            The input data. Should be sorted in ascending order.
+        n_clusters: int
+            The number of clusters.
+        max_iter: int
+            The maximum number of iterations.
+        is_weighted: bool
+            Whether the data is weighted.
+        sample_weights: np.ndarray or list or None
+            The sample weights.
+        n_local_trials: int
+            The number of local trials for kmeans++ initialization.
+        weights_prefix_sum: np.ndarray or list or None
+            The prefix sum of the sample weights.
+        weighted_X_prefix_sum: np.ndarray or list or None
+            The prefix sum of (the weighted) X.
+        weighted_X_squared_prefix_sum: np.ndarray or list or None
+            The prefix sum of (the squared weighted) X.
+        start_idx: int
+            The start index of the range to consider.
+        stop_idx: int
+            The stop index of the range to consider.
+
+    Returns:
+        centroids: np.ndarray
+            The centroids of the clusters.
+        cluster_borders: np.ndarray
+            The borders of the clusters.
+    """
+    if n_clusters == 2:
+        # If weighted_X_prefix_sum is not provided, neither is weights_prefix_sum (checked by the caller)
+        # They will be generated from sample_weights if needed
+        if weighted_X_prefix_sum is None:
+            sorted_X_casted = sorted_X.astype(PREFIX_SUM_DTYPE)
+            if is_weighted:
+                weights_casted = sample_weights.astype(PREFIX_SUM_DTYPE)
+                weights_prefix_sum = np.cumsum(weights_casted)
+
+                weighted_X_prefix_sum = np.cumsum(sorted_X_casted * weights_casted)
+            else:
+                weights_prefix_sum = None  # weights_prefix_sum is not needed
+                weighted_X_prefix_sum = np.cumsum(sorted_X_casted)  # no weighting needed
+        else:
+            pass  # weighted_X_prefix_sum and weights_prefix_sum are already provided
+        centroids, cluster_borders = flash_1d_kmeans_two_cluster(
+            sorted_X,
+            is_weighted,
+            weighted_X_prefix_sum,
+            weights_prefix_sum,
+            start_idx,
+            stop_idx,
+        )
+    else:
+        # If weighted_X_prefix_sum is not provided, neither is weights_prefix_sum and weighted_X_squared_prefix_sum
+        # They will be generated from sample_weights if needed
+        if weighted_X_prefix_sum is None:
+            sorted_X_casted = sorted_X.astype(PREFIX_SUM_DTYPE)
+            if is_weighted:
+                weights_casted = sample_weights.astype(PREFIX_SUM_DTYPE)
+                weights_prefix_sum = np.cumsum(weights_casted)
+
+                weighted_X_casted = sorted_X_casted * weights_casted
+                weighted_X_prefix_sum = np.cumsum(weighted_X_casted)
+                weighted_X_squared_prefix_sum = np.cumsum(weighted_X_casted * sorted_X_casted)
+            else:
+                weights_prefix_sum = None  # weights_prefix_sum is not needed
+                weighted_X_prefix_sum = np.cumsum(sorted_X_casted)
+                weighted_X_squared_prefix_sum = np.cumsum(sorted_X_casted ** 2)
+
+        centroids, cluster_borders = flash_1d_kmeans_k_cluster(
+            sorted_X,
+            n_clusters,
+            max_iter,
+            is_weighted,
+            weights_prefix_sum,
+            weighted_X_prefix_sum,
+            weighted_X_squared_prefix_sum,
+            n_local_trials,
+            start_idx,
+            stop_idx,
+        )
+
+    return centroids, cluster_borders

flash1dkmeans/two_cluster.py

@@ -0,0 +1,157 @@
+"""two_cluster.py
+Optimized kmeans for 1D data with 2 clusters.
+
+This module contains an optimized kmeans for 1D data with 2 clusters.
+
+The following is the main function:
+- flash_1d_kmeans_two_cluster
+
+Inputs must be sorted in ascending order, no default values are provided, and no error checking is done.
+This is because this module is intended to be used internally.
+"""
+
+import numpy as np
+import numba
+from .utils import query_prefix_sum, query_weights_prefix_sum
+from .config import PREFIX_SUM_DTYPE, ARRAY_INDEX_DTYPE
+
+
+#@numba.njit(cache=True)
+def flash_1d_kmeans_two_cluster(
+        X,
+        is_weighted,
+        weighted_X_prefix_sum,
+        sample_weight_prefix_sum,
+        start_idx,
+        stop_idx,
+):
+    """An optimized kmeans for 1D data with 2 clusters.
+    Utilizes a binary search to find the optimal division point.
+    Time complexity: O(log(n))
+
+    Args:
+        X: np.ndarray
+            The input data. Should be sorted in ascending order.
+        is_weighted: bool
+            Whether the data is weighted. If True, the weighted versions of the arrays should be provided.
+        weighted_X_prefix_sum: np.ndarray
+            The prefix sum of (the weighted) X.
+        sample_weight_prefix_sum: np.ndarray
+            The prefix sum of the sample weights. Should be None if the data is unweighted.
+        start_idx: int
+            The start index of the range to consider.
+        stop_idx: int
+            The stop index of the range to consider.
+
+    Returns:
+        centroids: np.ndarray
+            The centroids of the two clusters, shape (2,)
+        cluster_borders: np.ndarray
+            The borders of the two clusters, shape (3,)
+
+    WARNING: X should be sorted in ascending order before calling this function.
+    """
+    if is_weighted:
+        assert sample_weight_prefix_sum is not None, "sample_weight_prefix_sum must be provided for weighted data"
+    else:
+        assert sample_weight_prefix_sum is None, "sample_weight_prefix_sum must not be provided for unweighted data"
+
+    size = stop_idx - start_idx
+    centroids = np.empty(2, dtype=X.dtype)
+    cluster_borders = np.empty(3, dtype=ARRAY_INDEX_DTYPE)
+    cluster_borders[0] = start_idx
+    cluster_borders[2] = stop_idx
+    # Remember to set cluster_borders[1] as the division point
+
+    if size == 1:
+        centroids[0], centroids[1] = X[start_idx], X[start_idx]
+        cluster_borders[1] = start_idx + 1
+        return centroids, cluster_borders
+
+    if size == 2:
+        centroids[0], centroids[1] = X[start_idx], X[start_idx + 1]
+        cluster_borders[1] = start_idx + 1
+        return centroids, cluster_borders
+
+    # Now we know that there are at least 3 elements
+
+    if is_weighted:
+        # If the sum of the sample weight in the range is 0, we assume that the data is unweighted
+        if query_weights_prefix_sum(True, sample_weight_prefix_sum, start_idx, stop_idx) == 0:
+            is_weighted = False
+            # We need to recalculate the prefix sum, as previously it would have been all zeros
+            weighted_X_prefix_sum = np.cumsum(X, dtype=PREFIX_SUM_DTYPE)
+            sample_weight_prefix_sum = None
+        else:
+            # Check if there is only one nonzero sample weight
+            total_weight = query_weights_prefix_sum(True, sample_weight_prefix_sum, start_idx, stop_idx)
+            sample_weight_prefix_sum_within_range = sample_weight_prefix_sum[start_idx:stop_idx]
+            final_increase_idx = np.searchsorted(sample_weight_prefix_sum_within_range,
+                                                 sample_weight_prefix_sum_within_range[-1])
+            final_increase_amount = query_weights_prefix_sum(True, sample_weight_prefix_sum,
+                                                             start_idx + final_increase_idx,
+                                                             start_idx + final_increase_idx + 1)
+            if total_weight == final_increase_amount:
+                # If there is only one nonzero sample weight, we need to return the corresponding weight as the centroid
+                # and set all elements to the left cluster
+                nonzero_weight_index = start_idx + final_increase_idx
+                centroids[0], centroids[1] = X[nonzero_weight_index], X[nonzero_weight_index]
+                cluster_borders[1] = stop_idx
+                return centroids, cluster_borders
+
+    # Now we know that there are at least 3 elements and at least 2 nonzero weights
+
+    # KMeans with 2 clusters on 1D data is equivalent to finding a division point.
+    # The division point can be found by doing a binary search on the prefix sum.
+
+    # We will do a search for the division point,
+    # where we search for the optimum number of elements in the first cluster
+    # We don't want empty clusters, so we set the floor and ceiling to start_idx + 1 and stop_idx - 1
+    floor = start_idx + 1
+    ceiling = stop_idx - 1
+    left_centroid = None
+    right_centroid = None
+    division_point = None
+
+    while floor + 1 < ceiling:
+        division_point = (floor + ceiling) // 2
+        # If the left cluster has no weight, we need to move the floor up
+        left_weight_sum = query_weights_prefix_sum(is_weighted, sample_weight_prefix_sum, start_idx, division_point)
+        if left_weight_sum == 0:
+            floor = division_point
+            continue
+        right_weight_sum = query_weights_prefix_sum(is_weighted, sample_weight_prefix_sum, division_point, stop_idx)
+        # If the right cluster has no weight, we need to move the ceiling down
+        if right_weight_sum == 0:
+            ceiling = division_point
+            continue
+
+        left_centroid = query_prefix_sum(weighted_X_prefix_sum, start_idx, division_point) / left_weight_sum
+        right_centroid = query_prefix_sum(weighted_X_prefix_sum, division_point, stop_idx) / right_weight_sum
+
+        new_division_point_value = (left_centroid + right_centroid) / 2
+        if X[division_point - 1] <= new_division_point_value:
+            if new_division_point_value <= X[division_point]:
+                # The new division point matches the previous one, so we can stop
+                break
+            else:
+                floor = division_point
+        else:
+            ceiling = division_point
+
+    # initialize variables in case the loop above does not run through
+    if left_centroid is None:
+        division_point = (floor + ceiling) // 2
+        left_centroid = (query_prefix_sum(weighted_X_prefix_sum, start_idx, division_point) /
+                         query_weights_prefix_sum(is_weighted, sample_weight_prefix_sum, start_idx, division_point))
+    if right_centroid is None:
+        division_point = (floor + ceiling) // 2
+        right_centroid = (query_prefix_sum(weighted_X_prefix_sum, division_point, stop_idx) /
+                          query_weights_prefix_sum(is_weighted, sample_weight_prefix_sum, division_point, stop_idx))
+
+    # avoid using lists to allow numba.njit
+    centroids[0] = left_centroid
+    centroids[1] = right_centroid
+
+    cluster_borders[1] = division_point
+    return centroids, cluster_borders

flash1dkmeans/utils.py

@@ -0,0 +1,37 @@
+import numba
+
+
+#@numba.njit(cache=True)
+def query_prefix_sum(arr_prefix_sum, start, stop):
+    """Returns the sum of elements in the range [start, stop) of arr.
+
+    Args:
+        arr_prefix_sum: The prefix sum of the array arr.
+        start: The start index of the range.
+        stop: The stop index of the range.
+
+    Returns:
+        The sum of elements in the range [start, stop) of arr.
+    """
+    return arr_prefix_sum[stop - 1] - arr_prefix_sum[start - 1] if start > 0 else arr_prefix_sum[stop - 1]
+
+
+#@numba.njit(cache=True)
+def query_weights_prefix_sum(is_weighted, sample_weights, start, stop):
+    """Returns the sum of sample_weights in the range [start, stop).
+    If is_weighted is False, simply returns stop - start.
+
+    Args:
+        is_weighted: bool
+            Whether the data is weighted. Unweighted data is equivalent to weights of all 1.
+        sample_weights: np.ndarray
+            The sample weights.
+        start: int
+            The start index of the range.
+        stop: int
+            The stop index of the range.
+    """
+    if is_weighted:
+        return query_prefix_sum(sample_weights, start, stop)
+    else:
+        return stop - start

flash1dkmeans-0.0.1.dist-info/METADATA

@@ -0,0 +1,128 @@
+Metadata-Version: 2.3
+Name: flash1dkmeans
+Version: 0.0.1
+Summary: An optimized K-means implementation for the one-dimensional case.
+Project-URL: Homepage, https://github.com/SyphonArch/flash1dkmeans
+Project-URL: Issues, https://github.com/SyphonArch/flash1dkmeans/issues
+Author-email: Jake Hyun <jake.hyun@hotmail.com>
+License-File: LICENSE
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Requires-Dist: numba~=0.59.1
+Requires-Dist: numpy~=1.26.4
+Description-Content-Type: text/markdown
+
+# flash1dkmeans
+An optimized K-means implementation for the one-dimensional case
+
+Exploits the fact that one-dimensional data can be sorted.
+
+Numba acceleration is used, so callers can further use parallelization with Numba's `@njit(parallel=True)`
+if multiple instances of the algorithm are run in parallel.
+
+## Features
+
+### Two clusters
+
+Guaranteed to find the optimal solution for two clusters.
+
+### K clusters
+
+Uses the K-means++ initialization algorithm to find the initial centroids.
+Then uses the Lloyd's algorithm to find the final centroids, except with optimizations for the one-dimensional case.
+
+## Time Complexity
+
+- **2 clusters**: O(log(n))
+- (+ (O(n) for prefix sum calculation if not provided))
+- **k clusters**: O(k * 2 + log(k) * log(n) + max_iter * log(n) * k)  
+  (+ (O(n) for prefix sum calculation if not provided))
+
+This is a significant improvement over the standard K-means algorithm, which has a time complexity of O(n * k * max_iter),
+even excluding the time complexity of the K-means++ initialization.
+
+## How fast is it?
+
+TODO: Comparision with sklearn's KMeans
+
+
+## Installation
+```bash
+pip install flash1dkmeans
+```
+
+## Usage
+
+### Basic usage
+```python
+from flash1dkmeans import kmeans_1d
+
+data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
+k = 2
+
+centroids, labels = kmeans_1d(data, k)
+```
+
+### More Options
+```python
+from flash1dkmeans import kmeans_1d
+import numpy as np
+
+data = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0])
+weights = np.random.random_sample(data.shape)
+k = 3
+
+centroids, labels = kmeans_1d(
+    data, k,
+    sample_weights=weights,  # sample weights
+    n_local_trials=(2 + int(np.log(k))),  # number of local tries
+    max_iter=100,  # maximum number of iterations
+)
+```
+
+### Even More Options
+```python
+from flash1dkmeans import kmeans_1d
+import numpy as np
+
+n, k = 1024, 4
+
+# Generate random data
+data = np.random.random_sample(n)
+data = np.sort(data)
+
+# Generate random weights
+weights = np.random.random_sample(data.shape)
+
+# Calculate prefix sums
+weights_prefix_sum = np.cumsum(weights)
+weighted_X_prefix_sum = np.cumsum(data * weights)
+weighted_X_squared_prefix_sum = np.cumsum(data ** 2 * weights)
+
+middle_idx = n // 2
+
+# Providing prefix sums reduces redundant calculations
+# This is useful when the algorithm is run multiple times on different segments of the data
+for start_idx, stop_idx in [(0, middle_idx), (middle_idx, n)]:
+  centroids, cluster_borders = kmeans_1d(
+      data, k,  # Note how the sample weights are not provided when the prefix sums are provided
+      n_local_trials=(2 + int(np.log(k))),  # number of local tries
+      max_iter=100,  # maximum number of iterations
+      return_cluster_borders=True,  # return cluster borders instead of labels
+      weights_prefix_sum=weights_prefix_sum,  # prefix sum of the sample weights, leave empty for unwieghted data
+      weighted_X_prefix_sum=weighted_X_prefix_sum,  # prefix sum of the weighted data
+      weighted_X_squared_prefix_sum=weighted_X_squared_prefix_sum,  # prefix sum of the squared weighted data
+      start_idx=start_idx,  # start index of the data
+      stop_idx=stop_idx,  # stop index of the data
+  )
+```
+
+## Notes
+
+This repository has been created to be used in [Any-Precision-LLM](https://github.com/SNU-ARC/any-precision-llm) project,
+where multiple 1D K-means instances are run in parallel for LLM quantization.
+
+However the algorithm is general and can be used for any 1D K-means problem.
+

flash1dkmeans-0.0.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+flash1dkmeans/__init__.py,sha256=6tJz85chEMKNU3epnPJLxsxyO0huXrnJqCGU3qRdUlE,28
+flash1dkmeans/config.py,sha256=jCGWED26FWxhDtss3TOcn4Jl0AJ1hAZAMIhMHY6H4kk,189
+flash1dkmeans/k_cluster.py,sha256=WY16GQ2iPZkVL3RPS80y1lzrwGmwdYYS9b_uakTavCc,13940
+flash1dkmeans/main.py,sha256=wcq7P-ASS6PmL_cQNcW57mMbEd1ww4AvD8j33WvQirY,12325
+flash1dkmeans/two_cluster.py,sha256=7yMl0DxjzVcmh2_tSJcNigyceq-MqXel76dRCeY7zg4,7119
+flash1dkmeans/utils.py,sha256=awQlNNN7GRQsG4s2H03WcMObH4aCAH6xT36wGSS0G3o,1179
+flash1dkmeans-0.0.1.dist-info/METADATA,sha256=lN-J-4YS7g2XYWqxf10EhPb8elGJAVW00k9Nm2hSeBQ,4076
+flash1dkmeans-0.0.1.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+flash1dkmeans-0.0.1.dist-info/licenses/LICENSE,sha256=4fb3Ics16RwglL3QyYyQvyIkTdTOJ825dCytbUqaDu0,1066
+flash1dkmeans-0.0.1.dist-info/RECORD,,

flash1dkmeans-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.24.2
+Root-Is-Purelib: true
+Tag: py3-none-any

flash1dkmeans-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jake Hyun
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.