Clustering

concrete_fhe_toolkit.ml.clustering

Clustering task namespace: models, trainers, and metrics.

sklearn-style entry point for unsupervised grouping of encrypted samples.

FHEKMeans

Bases: FHEModel

Encrypted K-Means Inference Model (cluster assignment).

Assigns encrypted samples to the nearest public centroid. Train with FHEKMeansTrainer.fit_encrypted or provide centroids directly.

Parameters:

Name	Type	Description	Default
centroids	list	The public list of cluster centroids.	required
max_distance	int	Upper bound on the squared distance from any sample to any centroid (sizes the argmin reduction).	required

Example

model = FHEKMeans(centroids=[[0,0], [10,10]], max_distance=200)
model.compile(dummy_inputset, batch_size=1)

Source code in src/concrete_fhe_toolkit/ml/classes.py

class FHEKMeans(FHEModel):
    """Encrypted K-Means Inference Model (cluster assignment).

    Assigns encrypted samples to the nearest public centroid. Train with
    ``FHEKMeansTrainer.fit_encrypted`` or provide centroids directly.

    Args:
        centroids (list): The public list of cluster centroids.
        max_distance (int): Upper bound on the squared distance from any
            sample to any centroid (sizes the argmin reduction).

    Example:
        ```python
        model = FHEKMeans(centroids=[[0,0], [10,10]], max_distance=200)
        model.compile(dummy_inputset, batch_size=1)
        ```
    """

    def __init__(self, centroids: Any, *, max_distance: Any):
        super().__init__()
        """Initialize the object."""
        from .models import nearest_centroid_inference

        self._nearest_centroid_inference = nearest_centroid_inference
        self.centroids = [list(centroid) for centroid in centroids]
        self.max_distance = max_distance

    def _circuit_logic(self, features: Any) -> Any:
        return self._nearest_centroid_inference(
            features,
            self.centroids,
            max_distance=self.max_distance,
        )

FHEKMeansTrainer

Bases: FHETrainer

Hybrid encrypted k-means training (fixed iterations).

Each iteration compiles one circuit with the current public centroids: it assigns every encrypted sample to its nearest centroid and returns per-cluster sums and counts. Only those aggregates are decrypted; the centroid update (division) happens clear-side, and the new centroids become the next iteration's public constants.

Parameters:

Name	Type	Description	Default
initial_centroids	List[List[int]]	Public, data-independent starting centroids.	required
min_value	int	The lower bound for the encrypted inputs, used to dimension the FHE circuit.	required
max_value	int	The upper bound for the encrypted inputs, used to dimension the FHE circuit.	required
n_iterations	int	Fixed number of Lloyd iterations.	5
simulate	bool	Run circuits in simulation.	False
configuration	Optional[Configuration]	Optional FHE compiler configuration object.	None

Example

trainer = FHEKMeansTrainer(initial_centroids=[[0, 0]], min_value=0, max_value=1)

Source code in src/concrete_fhe_toolkit/ml/trainers.py

class FHEKMeansTrainer(FHETrainer):
    """Hybrid encrypted k-means training (fixed iterations).

    Each iteration compiles one circuit with the current public centroids:
    it assigns every encrypted sample to its nearest centroid and returns
    per-cluster sums and counts. Only those aggregates are decrypted; the
    centroid update (division) happens clear-side, and the new centroids
    become the next iteration's public constants.

    Args:
        initial_centroids (List[List[int]]): Public, data-independent starting centroids.
        min_value (int): The lower bound for the encrypted inputs, used to dimension the FHE circuit.
        max_value (int): The upper bound for the encrypted inputs, used to dimension the FHE circuit.
        n_iterations (int): Fixed number of Lloyd iterations.
        simulate (bool): Run circuits in simulation.
        configuration (Optional[fhe.Configuration]): Optional FHE compiler configuration object.

    Example:
        ```python
        trainer = FHEKMeansTrainer(initial_centroids=[[0, 0]], min_value=0, max_value=1)
        ```
    """

    def __init__(
        self,
        initial_centroids: List[List[int]],
        *,
        min_value: int,
        max_value: int,
        n_iterations: int = 5,
        simulate: bool = False,
        configuration: Optional[fhe.Configuration] = None,
    ) -> None:
        super().__init__(simulate=simulate, configuration=configuration)
        """Initialize the object."""
        if not initial_centroids:
            raise ValueError("initial_centroids must contain at least one centroid")
        self.initial_centroids = [list(centroid) for centroid in initial_centroids]
        self.minimum, self.maximum = validate_bounds(min_value, max_value)
        self.n_iterations = validate_integer("n_iterations", n_iterations, minimum=1)

    def _step_circuit(self, centroids: Any, n_samples: int, n_features: int, max_distance: int) -> Any:
        n_clusters = len(centroids)
        arg_min = make_argmin(n_clusters, 0, max_distance)

        def assign_and_aggregate(X_train: Any) -> Any:
            """Assign and aggregate values."""
            flags = []
            for row in range(n_samples):
                sample = [X_train[row][column] for column in range(n_features)]
                distances = [
                    euclidean_distance_squared(centroid, sample)
                    for centroid in centroids
                ]
                nearest = arg_min(distances)
                flags.append(
                    [equal(cluster, nearest) for cluster in range(n_clusters)]
                )

            outputs = []
            for cluster in range(n_clusters):
                count: Any = 0
                for row in range(n_samples):
                    count = count + flags[row][cluster]
                outputs.append(count)
            for cluster in range(n_clusters):
                for column in range(n_features):
                    total: Any = 0
                    for row in range(n_samples):
                        total = total + flags[row][cluster] * X_train[row][column]
                    outputs.append(total)
            return fhe.array(outputs)

        return assign_and_aggregate

    def fit_encrypted(self, X_train: List[List[int]]) -> Any:
        """Fits the model securely on encrypted training data.

        Args:
            X_train (List[List[int]]): Encrypted training features.

        Returns:
            Any: The trained FHEKMeans model.
        """
        n_samples = len(X_train)
        if n_samples == 0:
            raise ValueError("X_train must contain at least one sample")
        n_features = len(X_train[0])
        span = self.maximum - self.minimum
        max_distance = max(1, n_features * span * span)
        X_array = np.array(X_train, dtype=np.int64)

        centroids = [list(centroid) for centroid in self.initial_centroids]
        n_clusters = len(centroids)

        for _ in range(self.n_iterations):
            circuit_fn = self._step_circuit(
                centroids, n_samples, n_features, max_distance
            )
            flat = list(
                np.array(
                    self._run_circuit(
                        circuit_fn,
                        {"X_train": "encrypted"},
                        [X_array],
                        (X_array,),
                    )
                ).astype(int)
            )
            counts = flat[:n_clusters]
            sums = flat[n_clusters:]
            new_centroids = []
            for cluster in range(n_clusters):
                if counts[cluster] == 0:
                    new_centroids.append(list(centroids[cluster]))
                    continue
                new_centroids.append(
                    [
                        int(
                            _pymath.floor(
                                sums[cluster * n_features + column]
                                / counts[cluster]
                                + 0.5
                            )
                        )
                        for column in range(n_features)
                    ]
                )
            centroids = new_centroids

        return FHEKMeans(centroids, max_distance=max_distance)

fit_encrypted

fit_encrypted(X_train: List[List[int]]) -> Any

Fits the model securely on encrypted training data.

Parameters:

Name	Type	Description	Default
X_train	List[List[int]]	Encrypted training features.	required

Returns:

Name	Type	Description
Any	Any	The trained FHEKMeans model.

Source code in src/concrete_fhe_toolkit/ml/trainers.py

def fit_encrypted(self, X_train: List[List[int]]) -> Any:
    """Fits the model securely on encrypted training data.

    Args:
        X_train (List[List[int]]): Encrypted training features.

    Returns:
        Any: The trained FHEKMeans model.
    """
    n_samples = len(X_train)
    if n_samples == 0:
        raise ValueError("X_train must contain at least one sample")
    n_features = len(X_train[0])
    span = self.maximum - self.minimum
    max_distance = max(1, n_features * span * span)
    X_array = np.array(X_train, dtype=np.int64)

    centroids = [list(centroid) for centroid in self.initial_centroids]
    n_clusters = len(centroids)

    for _ in range(self.n_iterations):
        circuit_fn = self._step_circuit(
            centroids, n_samples, n_features, max_distance
        )
        flat = list(
            np.array(
                self._run_circuit(
                    circuit_fn,
                    {"X_train": "encrypted"},
                    [X_array],
                    (X_array,),
                )
            ).astype(int)
        )
        counts = flat[:n_clusters]
        sums = flat[n_clusters:]
        new_centroids = []
        for cluster in range(n_clusters):
            if counts[cluster] == 0:
                new_centroids.append(list(centroids[cluster]))
                continue
            new_centroids.append(
                [
                    int(
                        _pymath.floor(
                            sums[cluster * n_features + column]
                            / counts[cluster]
                            + 0.5
                        )
                    )
                    for column in range(n_features)
                ]
            )
        centroids = new_centroids

    return FHEKMeans(centroids, max_distance=max_distance)

nearest_centroid_inference

nearest_centroid_inference(sample: Union[ndarray, List[Any]], centroids: Union[ndarray, List[List[Any]]], labels: Optional[Union[ndarray, List[Any]]] = None, *, max_distance: int = 15) -> Any

Assign an encrypted sample to the nearest public centroid (k-means step).

Returns the centroid index, or the matching label when labels is given. max_distance must bound the squared distance to any centroid.

Parameters:

Name	Type	Description	Default
sample	Union[ndarray, List[Any]]	The encrypted sample.	required
centroids	Union[ndarray, List[List[Any]]]	A list of public centroids.	required
labels	Optional[Union[ndarray, List[Any]]]	Optional labels for each centroid.	None
max_distance	int	An upper bound on the squared distance. Defaults to 15.	15

Returns:

Type	Description
Any	The index of the nearest centroid, or its corresponding label.

Example

from concrete_fhe_toolkit.ml.models import nearest_centroid_inference

# Inside an FHE circuit
# cluster_idx = nearest_centroid_inference(
#     enc_sample, public_centroids, max_distance=100
# )

Source code in src/concrete_fhe_toolkit/ml/models.py

def nearest_centroid_inference(
    sample: Union[np.ndarray, List[Any]],
    centroids: Union[np.ndarray, List[List[Any]]],
    labels: Optional[Union[np.ndarray, List[Any]]] = None,
    *,
    max_distance: int = 15,
) -> Any:
    """Assign an encrypted sample to the nearest public centroid (k-means step).

    Returns the centroid index, or the matching label when ``labels`` is
    given. ``max_distance`` must bound the squared distance to any centroid.

    Args:
        sample: The encrypted sample.
        centroids: A list of public centroids.
        labels: Optional labels for each centroid.
        max_distance: An upper bound on the squared distance. Defaults to 15.

    Returns:
        The index of the nearest centroid, or its corresponding label.

    Example:
        ```python
        from concrete_fhe_toolkit.ml.models import nearest_centroid_inference

        # Inside an FHE circuit
        # cluster_idx = nearest_centroid_inference(
        #     enc_sample, public_centroids, max_distance=100
        # )
        ```
    """
    maximum_distance = validate_integer("max_distance", max_distance, minimum=1)
    if not centroids:
        raise ValueError("centroids must contain at least one centroid")
    if labels is not None and len(labels) != len(centroids):
        raise ValueError("labels and centroids must have the same length")

    distances = [
        euclidean_distance_squared(centroid, sample)
        for centroid in centroids
    ]
    arg_min = make_argmin(len(distances), 0, maximum_distance)
    nearest_index = arg_min(distances)

    if labels is None:
        return nearest_index

    prediction: Any = 0
    for index, label in enumerate(labels):
        prediction = prediction + equal(index, nearest_index) * label
    return prediction

inertia

inertia(samples: List[List[Any]], centroids: List[List[int]], *, max_distance: int) -> Any

Sum of squared distances from each sample to its nearest centroid.

The standard k-means quality metric (lower is better). max_distance must bound the squared distance from any sample to any centroid.

Parameters:

Name	Type	Description	Default
samples	List[List[Any]]	The encrypted samples to evaluate.	required
centroids	List[List[int]]	The centroids to compute distance against.	required
max_distance	int	The maximum squared distance bound.	required

Returns:

Name	Type	Description
Any	Any	The sum of squared distances (inertia).

Raises:

Type	Description
ValueError	If centroids is empty.

Example

from concrete_fhe_toolkit.ml.clustering import inertia

# Calculate inertia for two encrypted samples
enc_inertia = inertia([[1, 2], [5, 5]], centroids=[[1, 1], [4, 4]], max_distance=100)

Source code in src/concrete_fhe_toolkit/ml/clustering.py

def inertia(
    samples: List[List[Any]],
    centroids: List[List[int]],
    *,
    max_distance: int,
) -> Any:
    """Sum of squared distances from each sample to its nearest centroid.

    The standard k-means quality metric (lower is better). ``max_distance``
    must bound the squared distance from any sample to any centroid.

    Args:
        samples (List[List[Any]]): The encrypted samples to evaluate.
        centroids (List[List[int]]): The centroids to compute distance against.
        max_distance (int): The maximum squared distance bound.

    Returns:
        Any: The sum of squared distances (inertia).

    Raises:
        ValueError: If centroids is empty.

    Example:
        ```python
        from concrete_fhe_toolkit.ml.clustering import inertia

        # Calculate inertia for two encrypted samples
        enc_inertia = inertia([[1, 2], [5, 5]], centroids=[[1, 1], [4, 4]], max_distance=100)
        ```
    """
    if not centroids:
        raise ValueError("centroids must contain at least one centroid")
    nearest = make_array_minimum(len(centroids), 0, max_distance)
    total: Any = 0
    for sample in samples:
        distances = [
            euclidean_distance_squared(centroid, sample)
            for centroid in centroids
        ]
        total = total + nearest(distances)
    return total