Skip to content

Core API

kmeanssa_ng.core

Core abstractions and algorithms for k-means on metric spaces.

Center

Bases: Point

Abstract base class for cluster centers.

A center is a special type of point that can move through the space using two mechanisms: - Brownian motion: Random exploration - Drift: Directed movement toward a target point

This class is used in simulated annealing for k-means clustering.

Source code in kmeanssa_ng/core/abstract.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class Center(Point):
    """Abstract base class for cluster centers.

    A center is a special type of point that can move through the space
    using two mechanisms:
    - Brownian motion: Random exploration
    - Drift: Directed movement toward a target point

    This class is used in simulated annealing for k-means clustering.
    """

    @abstractmethod
    def brownian_motion(self, time_to_travel: float) -> None:
        """Perform random Brownian motion in the space.

        Args:
            time_to_travel: Time parameter controlling the magnitude of motion.
                Typical distance traveled is proportional to sqrt(time_to_travel).
        """
        raise NotImplementedError

    @abstractmethod
    def drift(self, target_point: Point, prop_to_travel: float) -> None:
        """Move toward a target point.

        Args:
            target_point: The point to move toward.
            prop_to_travel: Proportion of the distance to travel (between 0 and 1).
                0 means no movement, 1 means move all the way to target.
        """
        raise NotImplementedError

    def seed_rng(self, rng: np.random.Generator) -> None:
        """Adopt the algorithm's random generator.

        ``SimulatedAnnealing`` seeds every center it drives so that all
        stochastic moves (Brownian steps, tie-breaking in vertex routing)
        draw from one reproducible stream. The default stores the generator
        on ``self._rng``, the attribute the built-in centers read; a center
        with its own noise source must override this method to honor it,
        otherwise its randomness escapes ``random_state`` control.
        """
        self._rng = rng

brownian_motion(time_to_travel) abstractmethod

Perform random Brownian motion in the space.

Parameters:

Name Type Description Default
time_to_travel float

Time parameter controlling the magnitude of motion. Typical distance traveled is proportional to sqrt(time_to_travel).

 required
Source code in kmeanssa_ng/core/abstract.py 
52
53
54
55
56
57
58
59
60
@abstractmethod
def brownian_motion(self, time_to_travel: float) -> None:
    """Perform random Brownian motion in the space.

    Args:
        time_to_travel: Time parameter controlling the magnitude of motion.
            Typical distance traveled is proportional to sqrt(time_to_travel).
    """
    raise NotImplementedError

drift(target_point, prop_to_travel) abstractmethod

Move toward a target point.

Parameters:

Name Type Description Default
target_point Point

The point to move toward.

 required
prop_to_travel float

Proportion of the distance to travel (between 0 and 1). 0 means no movement, 1 means move all the way to target.

 required
Source code in kmeanssa_ng/core/abstract.py 
62
63
64
65
66
67
68
69
70
71
@abstractmethod
def drift(self, target_point: Point, prop_to_travel: float) -> None:
    """Move toward a target point.

    Args:
        target_point: The point to move toward.
        prop_to_travel: Proportion of the distance to travel (between 0 and 1).
            0 means no movement, 1 means move all the way to target.
    """
    raise NotImplementedError

seed_rng(rng)

Adopt the algorithm's random generator.

SimulatedAnnealing seeds every center it drives so that all stochastic moves (Brownian steps, tie-breaking in vertex routing) draw from one reproducible stream. The default stores the generator on self._rng, the attribute the built-in centers read; a center with its own noise source must override this method to honor it, otherwise its randomness escapes random_state control.

Source code in kmeanssa_ng/core/abstract.py 
73
74
75
76
77
78
79
80
81
82
83
def seed_rng(self, rng: np.random.Generator) -> None:
    """Adopt the algorithm's random generator.

    ``SimulatedAnnealing`` seeds every center it drives so that all
    stochastic moves (Brownian steps, tie-breaking in vertex routing)
    draw from one reproducible stream. The default stores the generator
    on ``self._rng``, the attribute the built-in centers read; a center
    with its own noise source must override this method to honor it,
    otherwise its randomness escapes ``random_state`` control.
    """
    self._rng = rng

Lloyd

Implementation of Lloyd's algorithm for k-means clustering.

This class provides a classic iterative implementation of k-means. It is strategy-based, allowing for custom initialization and center update logic.

Attributes:

Name Type Description
points

The list of points to cluster.
k

The number of clusters.
space

The metric space in which the points reside.
update_strategy

The strategy for computing new cluster centers.
Source code in kmeanssa_ng/core/lloyd.py 
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
class Lloyd:
    """Implementation of Lloyd's algorithm for k-means clustering.

    This class provides a classic iterative implementation of k-means.
    It is strategy-based, allowing for custom initialization and center
    update logic.

    Attributes:
        points: The list of points to cluster.
        k: The number of clusters.
        space: The metric space in which the points reside.
        update_strategy: The strategy for computing new cluster centers.
    """

    def __init__(
        self,
        points: list[Point],
        k: int,
        update_strategy: LloydUpdateStrategy,
        random_state: int | np.random.Generator | None = None,
    ):
        """Initialize Lloyd's algorithm.

        Args:
            points: A list of points to be clustered.
            k: The number of clusters.
            update_strategy: The strategy for updating cluster centers.
            random_state: Controls randomness for reproducibility.
        """
        if not points:
            raise ValueError("Input points list cannot be empty.")
        if k <= 0:
            raise ValueError("Number of clusters k must be positive.")

        self.points = points
        self.k = k
        self.space = points[0].space
        self.update_strategy = update_strategy

        if isinstance(random_state, np.random.Generator):
            self._rng = random_state
        else:
            self._rng = np.random.default_rng(random_state)

    @property
    def observations(self) -> list[Point]:
        """Return the points to be clustered.
        This is for compatibility with initialization strategies.
        """
        return self.points

    def run(
        self,
        initialization_strategy: InitializationStrategy | None = None,
        max_iterations: int = 100,
        tolerance: float = 1e-4,
    ) -> list[Center]:
        """Run Lloyd's algorithm.

        Args:
            initialization_strategy: The strategy for initializing centers.
                Defaults to :class:`KMeansPlusPlus`. (Unlike the initialization,
                the ``update_strategy`` is required at construction: its
                canonical choice depends on the space — a graph node update, a
                Karcher mean on a manifold — so there is no universal default.)
            max_iterations: The maximum number of iterations to run.
            tolerance: The tolerance for convergence. If the change in
                energy is less than this value, the algorithm stops.

        Returns:
            A list of the final cluster centers.
        """
        if initialization_strategy is None:
            from .strategies.initialization import KMeansPlusPlus

            initialization_strategy = KMeansPlusPlus()

        # 1. Initialization
        centers = initialization_strategy.initialize_centers(self)

        last_energy = float("inf")

        for i in range(max_iterations):
            # 2. Assignment step
            labels = self.space.assign_clusters(self.points, centers)

            # 3. Update step: always produce exactly k centers, reseeding any
            # cluster that came out empty (or whose update failed), so k never
            # silently shrinks.
            new_centers = []
            for cluster_idx in range(self.k):
                cluster_points = [
                    p for j, p in enumerate(self.points) if labels[j] == cluster_idx
                ]
                new_center = (
                    self.update_strategy.update(cluster_points, self.space)
                    if cluster_points
                    else None
                )
                if new_center is None:
                    # Reseed against the configuration as it stands *now* —
                    # the centers already updated (including earlier reseeds
                    # of this same iteration) plus the not-yet-updated ones.
                    # Reseeding against the pre-iteration centers would hand
                    # every simultaneously-empty cluster the same farthest
                    # point, leaving k shrunk despite the reseeding.
                    reference = new_centers + centers[cluster_idx + 1 :]
                    new_center = self._reseed_center(reference or centers)
                    logger.warning(
                        "Cluster %d is empty; reseeding its center on the point "
                        "farthest from the current centers.",
                        cluster_idx,
                    )
                new_centers.append(new_center)

            centers = new_centers

            # Check for convergence on the empirical objective of this
            # algorithm's own points (the space may be shared with other
            # running algorithms, and may carry unrelated node measures).
            current_energy = self.space.calculate_energy(
                centers, how="empirical", observations=self.points
            )
            if abs(last_energy - current_energy) < tolerance:
                break
            last_energy = current_energy

        return centers

    def _reseed_center(self, centers: list[Center]) -> Center:
        """Center on the point farthest from the current centers.

        The farthest point is the one worst served by the current
        configuration, so seeding there maximally reduces the energy a lone
        empty cluster can recover.
        """
        farthest = max(
            self.points,
            key=lambda p: min(self.space.distance(p, c) for c in centers),
        )
        return self.space.center_from_point(farthest)

observations property

Return the points to be clustered. This is for compatibility with initialization strategies.

__init__(points, k, update_strategy, random_state=None)

Initialize Lloyd's algorithm.

Parameters:

Name Type Description Default
points list[Point]

A list of points to be clustered.

 required
k int

The number of clusters.

 required
update_strategy LloydUpdateStrategy

The strategy for updating cluster centers.

 required
random_state int | Generator | None

Controls randomness for reproducibility.

 None
Source code in kmeanssa_ng/core/lloyd.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    points: list[Point],
    k: int,
    update_strategy: LloydUpdateStrategy,
    random_state: int | np.random.Generator | None = None,
):
    """Initialize Lloyd's algorithm.

    Args:
        points: A list of points to be clustered.
        k: The number of clusters.
        update_strategy: The strategy for updating cluster centers.
        random_state: Controls randomness for reproducibility.
    """
    if not points:
        raise ValueError("Input points list cannot be empty.")
    if k <= 0:
        raise ValueError("Number of clusters k must be positive.")

    self.points = points
    self.k = k
    self.space = points[0].space
    self.update_strategy = update_strategy

    if isinstance(random_state, np.random.Generator):
        self._rng = random_state
    else:
        self._rng = np.random.default_rng(random_state)

run(initialization_strategy=None, max_iterations=100, tolerance=0.0001)

Run Lloyd's algorithm.

Parameters:

Name Type Description Default
initialization_strategy InitializationStrategy | None

The strategy for initializing centers. Defaults to :class:KMeansPlusPlus. (Unlike the initialization, the update_strategy is required at construction: its canonical choice depends on the space — a graph node update, a Karcher mean on a manifold — so there is no universal default.)

 None
max_iterations int

The maximum number of iterations to run.

 100
tolerance float

The tolerance for convergence. If the change in energy is less than this value, the algorithm stops.

 0.0001

Returns:

Type Description
list[Center]

A list of the final cluster centers.
Source code in kmeanssa_ng/core/lloyd.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def run(
    self,
    initialization_strategy: InitializationStrategy | None = None,
    max_iterations: int = 100,
    tolerance: float = 1e-4,
) -> list[Center]:
    """Run Lloyd's algorithm.

    Args:
        initialization_strategy: The strategy for initializing centers.
            Defaults to :class:`KMeansPlusPlus`. (Unlike the initialization,
            the ``update_strategy`` is required at construction: its
            canonical choice depends on the space — a graph node update, a
            Karcher mean on a manifold — so there is no universal default.)
        max_iterations: The maximum number of iterations to run.
        tolerance: The tolerance for convergence. If the change in
            energy is less than this value, the algorithm stops.

    Returns:
        A list of the final cluster centers.
    """
    if initialization_strategy is None:
        from .strategies.initialization import KMeansPlusPlus

        initialization_strategy = KMeansPlusPlus()

    # 1. Initialization
    centers = initialization_strategy.initialize_centers(self)

    last_energy = float("inf")

    for i in range(max_iterations):
        # 2. Assignment step
        labels = self.space.assign_clusters(self.points, centers)

        # 3. Update step: always produce exactly k centers, reseeding any
        # cluster that came out empty (or whose update failed), so k never
        # silently shrinks.
        new_centers = []
        for cluster_idx in range(self.k):
            cluster_points = [
                p for j, p in enumerate(self.points) if labels[j] == cluster_idx
            ]
            new_center = (
                self.update_strategy.update(cluster_points, self.space)
                if cluster_points
                else None
            )
            if new_center is None:
                # Reseed against the configuration as it stands *now* —
                # the centers already updated (including earlier reseeds
                # of this same iteration) plus the not-yet-updated ones.
                # Reseeding against the pre-iteration centers would hand
                # every simultaneously-empty cluster the same farthest
                # point, leaving k shrunk despite the reseeding.
                reference = new_centers + centers[cluster_idx + 1 :]
                new_center = self._reseed_center(reference or centers)
                logger.warning(
                    "Cluster %d is empty; reseeding its center on the point "
                    "farthest from the current centers.",
                    cluster_idx,
                )
            new_centers.append(new_center)

        centers = new_centers

        # Check for convergence on the empirical objective of this
        # algorithm's own points (the space may be shared with other
        # running algorithms, and may carry unrelated node measures).
        current_energy = self.space.calculate_energy(
            centers, how="empirical", observations=self.points
        )
        if abs(last_energy - current_energy) < tolerance:
            break
        last_energy = current_energy

    return centers

Point

Bases: ABC

Abstract base class for points in a metric space.

A point is an element of a metric space with a fixed location. Concrete implementations must define which space the point belongs to.

Source code in kmeanssa_ng/core/abstract.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Point(ABC):
    """Abstract base class for points in a metric space.

    A point is an element of a metric space with a fixed location.
    Concrete implementations must define which space the point belongs to.
    """

    @property
    @abstractmethod
    def space(self) -> Space:
        """The metric space this point belongs to.

        Returns:
            The Space instance containing this point.
        """
        raise NotImplementedError

space abstractmethod property

The metric space this point belongs to.

Returns:

Type Description
Space

The Space instance containing this point.

SimulatedAnnealing

Simulated annealing for offline k-means clustering.

This algorithm solves the k-means problem on arbitrary metric spaces using simulated annealing. Centers perform Brownian motion (exploration) and drift toward observations (exploitation), with temperature controlled by an inhomogeneous Poisson process.

Attributes:

Name Type Description
space Space

The metric space containing the observations.
k int

Number of clusters.
observations list[Point]

List of points to cluster.
centers list[Center]

Current cluster centers.
Example 
from kmeanssa_ng import (
    KMeansPlusPlus,
    MinimizeEnergy,
    SimulatedAnnealing,
    generate_simple_graph,
)
from kmeanssa_ng.quantum_graph.sampling import UniformNodeSampling

# Create a space and sample observations
graph = generate_simple_graph()
points = graph.sample_points(100, strategy=UniformNodeSampling(random_state=0))

# Run simulated annealing with the interleaved algorithm
sa = SimulatedAnnealing(points, k=5, random_state=0)
centers = sa.run(KMeansPlusPlus(), MinimizeEnergy(), robust_prop=0.1)
Source code in kmeanssa_ng/core/simulated_annealing.py 
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
class SimulatedAnnealing:
    """Simulated annealing for offline k-means clustering.

    This algorithm solves the k-means problem on arbitrary metric spaces using
    simulated annealing. Centers perform Brownian motion (exploration) and drift
    toward observations (exploitation), with temperature controlled by an
    inhomogeneous Poisson process.

    Attributes:
        space: The metric space containing the observations.
        k: Number of clusters.
        observations: List of points to cluster.
        centers: Current cluster centers.

    Example:
        ```python
        from kmeanssa_ng import (
            KMeansPlusPlus,
            MinimizeEnergy,
            SimulatedAnnealing,
            generate_simple_graph,
        )
        from kmeanssa_ng.quantum_graph.sampling import UniformNodeSampling

        # Create a space and sample observations
        graph = generate_simple_graph()
        points = graph.sample_points(100, strategy=UniformNodeSampling(random_state=0))

        # Run simulated annealing with the interleaved algorithm
        sa = SimulatedAnnealing(points, k=5, random_state=0)
        centers = sa.run(KMeansPlusPlus(), MinimizeEnergy(), robust_prop=0.1)
        ```
    """

    def __init__(
        self,
        observations: list[Point],
        k: int,
        lambda0: float = 1.0,
        beta0: float = 1.0,
        step_size: float = 0.1,
        energy_mode: str = "uniform",
        random_state: int | np.random.Generator | None = None,
    ) -> None:
        """Initialize the simulated annealing algorithm.

        Args:
            observations: List of points to cluster, all in the same metric space.
            k: Number of clusters.
            lambda0: Intensity scale of the Poisson observation clock (must be > 0).

                Mathematical role: the annealing processes one observation per
                arrival of an inhomogeneous Poisson process of intensity
                lambda(t) = lambda0 * (1 + t) (the paper's schedule). It does
                **not** scale the Brownian steps themselves (each micro-step
                has standard deviation sqrt(step_size), independent of
                lambda0).

                Practical effect:
                - Higher values: arrivals come faster, so the same number of
                  observations spans a shorter annealing horizon (less
                  Brownian exploration per observation)
                - Lower values: longer horizon, more exploration between
                  observation events
                - Recommended default: 1.0

                See the companion paper (References) for the derivation of the
                time schedule.

            beta0: Initial drift intensity parameter (must be > 0).
                Controls how strongly centers are pulled toward observations.

                Mathematical role: The drift proportion at time t is computed as
                alpha(t) = min(h * beta0 * log(1 + t), 1) where h is the time
                interval. This controls the strength of attraction toward the
                nearest observation.

                Practical effect:
                - Higher values (2.0-5.0): Stronger drift, faster convergence,
                  more exploitation of current best positions
                - Lower values (0.3-0.8): Weaker drift, more exploration,
                  slower convergence
                - Recommended default: 1.0-2.0 for most cases

                See the companion paper (References) for the derivation of the
                drift schedule.

            step_size: Time discretization step for the SDE solver (must be > 0).
                Controls the temporal resolution of the stochastic process.

                Mathematical role: Euler discretization step Δt for solving the
                stochastic differential equation. Smaller values give more
                accurate simulation at the cost of more computation.

                Practical effect:
                - Smaller values (0.001-0.01): More accurate simulation, slower
                - Larger values (0.05-0.1): Faster but less accurate
                - Recommended default: 0.01 for good accuracy/speed tradeoff
                - Rule of thumb: Use step_size much smaller than the typical
                  time scale of the Poisson process (~ 1/lambda0)

            energy_mode: Which reference measure the k-means energy (mean
                squared distance to the nearest center) is averaged under —
                used by ``MinimizeEnergy`` to select the best visited state
                and by ``record_energy`` diagnostics:
                - "uniform": average over all graph nodes with equal weight,
                  measuring how well the centers cover the whole geometry
                  irrespective of where the observations lie.
                - "empirical": average over this algorithm's own observation
                  points, exactly where they lie (the empirical k-means
                  objective). The only mode supported on Riemannian
                  manifolds.
                - "node_measure": average under the per-node ``obs_weight``
                  measure registered on the graph by the caller (e.g. a
                  population measure); graph spaces only.
                The former "obs" mode was split into "empirical" and
                "node_measure" and now raises.

            random_state: Controls randomness for reproducibility.
                Determines random number generation for all random operations:
                - Shuffling observations
                - Poisson process time generation
                - Brownian motion (via centers)
                - Initialization strategies (KMeansPlusPlus, RandomInit)
                - Space-specific random operations

                All randomness flows through a single numpy Generator
                (``self._rng``), which is propagated to the centers so that
                every stochastic component is driven by the same stream. No
                global random state (``random.seed``/``np.random.seed``) is
                touched, so runs are isolated and fully reproducible.

                Pass an int for a reproducible seed, a Generator instance for
                fine-grained control, or None for non-deterministic behavior
                (default).

                Example:
                    >>> # Reproducible with seed (recommended)
                    >>> sa1 = SimulatedAnnealing(points, k=3, random_state=42)
                    >>> sa2 = SimulatedAnnealing(points, k=3, random_state=42)
                    >>> # sa1 and sa2 produce identical results
                    >>>
                    >>> # Or pass an explicit Generator
                    >>> rng = np.random.default_rng(42)
                    >>> sa = SimulatedAnnealing(points, k=3, random_state=rng)

        Raises:
            ValueError: If observations is empty, k <= 0, points are in different spaces,
                or hyperparameters are invalid.

        References:
            C. Brécheteau, I. Gavra, N. Klutchnikoff. "Online k-means Clustering
            on Metric Graphs and Geodesic Spaces" (preprint). Derives the
            annealing dynamics and its convergence analysis.

        Example:
            >>> # Quick convergence setup
            >>> sa = SimulatedAnnealing(
            ...     points, k=5,
            ...     lambda0=0.5,  # Less exploration
            ...     beta0=3.0,     # Stronger drift
            ...     step_size=0.01
            ... )
            >>>
            >>> # Thorough search setup (avoid local minima)
            >>> sa = SimulatedAnnealing(
            ...     points, k=5,
            ...     lambda0=2.0,   # More exploration
            ...     beta0=1.0,     # Gentler drift
            ...     step_size=0.01
            ... )
        """
        self._validate_constructor_parameters(
            observations, k, lambda0, beta0, step_size
        )
        if energy_mode == "obs":
            raise ValueError(
                "energy_mode 'obs' was split into two explicit modes: use "
                "'empirical' to average over this algorithm's observation "
                "points, or 'node_measure' to average under the per-node "
                "'obs_weight' measure registered on the graph."
            )
        if energy_mode not in ("uniform", "empirical", "node_measure"):
            raise ValueError(
                "energy_mode must be 'uniform', 'empirical' or "
                f"'node_measure', got {energy_mode!r}"
            )
        self._initialize_random_generator(random_state)

        self._space = observations[0].space
        self._observations = observations.copy()
        self._k = k
        self._lambda = float(lambda0)
        self._beta = float(beta0)
        self._step_size = float(step_size)
        self._energy_mode = energy_mode

        # Shuffle through the instance Generator so ordering is reproducible
        # from random_state without touching global random state.
        self._rng.shuffle(self._observations)
        self._centers: list[Center] = []
        self._energy_history: list[float] = []
        self._time_history: list[float] = []

    def _validate_constructor_parameters(
        self,
        observations: list[Point],
        k: int,
        lambda0: float,
        beta0: float,
        step_size: float,
    ) -> None:
        """Validate parameters for the constructor."""
        if not observations:
            raise ValueError("Observations must be a non-empty list of points.")
        if k <= 0:
            raise ValueError("Number of clusters 'k' must be greater than zero.")
        if any(obs.space != observations[0].space for obs in observations):
            raise ValueError("All observations must belong to the same metric space.")

        self._validate_positive_float(lambda0, "lambda0")
        self._validate_positive_float(beta0, "beta0")
        self._validate_positive_float(step_size, "step_size")

    def _validate_positive_float(self, value: float, name: str) -> None:
        """Validate that a value is a positive float."""
        try:
            float_value = float(value)
        except (TypeError, ValueError) as e:
            raise ValueError(
                f"{name} must be a number, got {type(value).__name__}"
            ) from e
        if float_value <= 0:
            raise ValueError(f"{name} must be positive, got {float_value}")

    def _initialize_random_generator(
        self, random_state: int | np.random.Generator | None
    ) -> None:
        """Initialize the random number generator."""
        if isinstance(random_state, np.random.Generator):
            self._rng = random_state
        else:
            self._rng = np.random.default_rng(random_state)

    @property
    def n(self) -> int:
        """Number of observations."""
        return len(self._observations)

    @property
    def observations(self) -> list[Point]:
        """List of observation points."""
        return self._observations

    @property
    def centers(self) -> list[Center]:
        """Current cluster centers."""
        return self._centers

    @property
    def space(self) -> Space:
        """Metric space containing the observations."""
        return self._space

    @property
    def k(self) -> int:
        """Number of clusters."""
        return self._k

    @property
    def energy_history(self) -> np.ndarray:
        """Energy after each observation from the last ``run(record_energy=True)``.

        Empty until such a run; the first entry is the energy of the initial
        centers (time 0), aligned with :attr:`time_history`.
        """
        return np.asarray(self._energy_history)

    @property
    def time_history(self) -> np.ndarray:
        """Annealing time at each recorded energy (see :attr:`energy_history`)."""
        return np.asarray(self._time_history)

    def _clone_centers(self, centers: list[Center]) -> list[Center]:
        """Create independent copies of centers.

        Uses the clone() method if available (much faster than deepcopy),
        otherwise falls back to deepcopy for compatibility.

        Args:
            centers: List of centers to clone.

        Returns:
            List of cloned centers with independent state.
        """
        if hasattr(centers[0], "clone"):
            return [center.clone() for center in centers]
        else:
            # Fallback for custom Center implementations without clone()
            from copy import deepcopy

            return deepcopy(centers)

    def _initialize_times(self, n: int) -> np.ndarray:
        """Generate inhomogeneous Poisson times.

        Args:
            n: Number of time points to generate.

        Returns:
            Array of n+1 time points.
        """
        # Arrival times of an inhomogeneous Poisson process of intensity
        # lambda(t) = lambda0 * (1 + t) (the paper's schedule). Its cumulative
        # intensity is Lambda(t) = lambda0 * (t + t^2/2), and the i-th arrival
        # is T_i = Lambda^{-1}(S_i) where S_i is a sum of unit-rate
        # exponentials. Here poiss_sum accumulates S_i / lambda0 (draws of mean
        # 1/lambda0), so inverting Lambda gives sqrt(2 * poiss_sum + 1) - 1.
        # The factor 2 was previously missing, which realised the schedule
        # lambda(t) = 2 * lambda0 * (1 + t) instead -- a silent doubling of
        # lambda0 relative to the paper.
        T = np.zeros(n + 1)
        poiss_sum = 0.0
        for i in range(n):
            poiss_sum += self._rng.exponential(1.0 / self._lambda)
            T[i + 1] = np.sqrt(2.0 * poiss_sum + 1.0) - 1.0
        return T

    def calculate_energy(self, centers: list[Center]) -> float:
        """Calculate k-means energy for given centers based on the energy mode.

        Delegates to the space. The algorithm's own observations are the data
        of the "empirical" mode only: "uniform" and "node_measure" define
        their reference measure without them (and reject them, so no mode can
        silently shadow another). Acceleration (e.g. the quantum graph's
        numba kernels) is the space's concern, dispatched inside
        ``Space.calculate_energy``.
        """
        observations = self._observations if self._energy_mode == "empirical" else None
        return self.space.calculate_energy(
            centers, how=self._energy_mode, observations=observations
        )

    def run(
        self,
        initialization_strategy: InitializationStrategy | None = None,
        robustification_strategy: RobustificationStrategy | None = None,
        robust_prop: float = 0.0,
        record_energy: bool = False,
    ):
        """Run the simulated annealing algorithm.

        This is the primary method to execute the simulated annealing algorithm.
        It performs an interleaved sequence of Brownian motion (exploration)
        and drift (exploitation) for the cluster centers.

        Args:
            initialization_strategy: How the initial centers are chosen.
                Defaults to :class:`KMeansPlusPlus`, the canonical k-means++
                seeding, so ``sa.run()`` works out of the box.
            robustification_strategy: How the returned centers are selected
                from the trajectory. Defaults to :class:`MinimizeEnergy`, which
                keeps the lowest-energy state seen during the collection window.
            robust_prop: Fraction of the (trailing) observations over which the
                robustification strategy collects candidate states, in [0, 1].
                Left at 0.0 by default, the window is a single point, so the
                default :class:`MinimizeEnergy` only compares the initial and
                final states — pass ``robust_prop`` around 0.1 for a genuine
                best-of-window robustification.
            record_energy: If True, record the energy and annealing time after
                each observation into :attr:`energy_history` and
                :attr:`time_history` (for convergence diagnostics). Off by
                default so the energy is not recomputed when not needed.

        Example:
            >>> # Zero-config quickstart: k-means++ init, energy-minimizing
            >>> # robustification.
            >>> centers = SimulatedAnnealing(points, k=5, random_state=0).run()
        """
        if initialization_strategy is None:
            initialization_strategy = KMeansPlusPlus()
        if robustification_strategy is None:
            robustification_strategy = MinimizeEnergy()

        logger.info(
            "Starting SA: k=%d, n_obs=%d, lambda0=%.3f, beta0=%.3f, "
            "step_size=%.4f, robust_prop=%.2f",
            self._k,
            self.n,
            self._lambda,
            self._beta,
            self._step_size,
            robust_prop,
        )

        if robust_prop < 0 or robust_prop > 1:
            raise ValueError("The proportion must be in [0,1]")

        i0 = int(np.floor((self.n - 1) * (1 - robust_prop)))

        self._centers = initialization_strategy.initialize_centers(self)
        # Seed each center's RNG from the SA's generator so that all stochastic moves
        # (Brownian step size and vertex routing) are reproducible from random_state.
        for _center in self._centers:
            _center.seed_rng(self._rng)

        robustification_strategy.initialize(self)
        strategy = robustification_strategy

        times = self._initialize_times(self.n)
        time = 0.0
        progress_interval = max(1, self.n // 10)

        if record_energy:
            self._energy_history = [self.calculate_energy(self._centers)]
            self._time_history = [time]

        for i, point in enumerate(self._observations):
            # times[0] is the origin of the clock: observation i is processed
            # over the interval (times[i], times[i + 1]].
            T = times[i + 1]

            if i % progress_interval == 0 and i > 0:
                progress = 100 * i / self.n
                logger.info(
                    "Progress: %.1f%% (%d/%d observations processed)",
                    progress,
                    i,
                    self.n,
                )

            logger.debug("Processing observation %d, target time T=%.4f", i, T)

            while time < T:
                h = min(self._step_size, T - time)
                prop = min(h * self._beta * np.log(1 + time), 1)
                logger.debug(
                    "Time step: time=%.4f, h=%.4f, drift_prop=%.4f", time, h, prop
                )
                for center in self._centers:
                    center.brownian_motion(h)
                distances = self.space.distances_from_centers(self._centers, point)
                closest_idx = np.argmin(distances)
                self._centers[closest_idx].drift(point, prop)
                time += h

            if i >= i0:
                strategy.collect(self)
                logger.debug(
                    "Collected centers for robustification at observation %d", i
                )

            if record_energy:
                self._energy_history.append(self.calculate_energy(self._centers))
                self._time_history.append(time)

        result = strategy.get_result()
        logger.info("SA completed successfully")
        return result

centers property

Current cluster centers.

energy_history property

Energy after each observation from the last run(record_energy=True).

Empty until such a run; the first entry is the energy of the initial centers (time 0), aligned with :attr:time_history.

k property

Number of clusters.

n property

Number of observations.

observations property

List of observation points.

space property

Metric space containing the observations.

time_history property

Annealing time at each recorded energy (see :attr:energy_history).

__init__(observations, k, lambda0=1.0, beta0=1.0, step_size=0.1, energy_mode='uniform', random_state=None)

Initialize the simulated annealing algorithm.

Parameters:

Name Type Description Default
observations list[Point]

List of points to cluster, all in the same metric space.

 required
k int

Number of clusters.

 required
lambda0 float

Intensity scale of the Poisson observation clock (must be > 0).

Mathematical role: the annealing processes one observation per arrival of an inhomogeneous Poisson process of intensity lambda(t) = lambda0 * (1 + t) (the paper's schedule). It does not scale the Brownian steps themselves (each micro-step has standard deviation sqrt(step_size), independent of lambda0).

Practical effect: - Higher values: arrivals come faster, so the same number of observations spans a shorter annealing horizon (less Brownian exploration per observation) - Lower values: longer horizon, more exploration between observation events - Recommended default: 1.0

See the companion paper (References) for the derivation of the time schedule.

 1.0
beta0 float

Initial drift intensity parameter (must be > 0). Controls how strongly centers are pulled toward observations.

Mathematical role: The drift proportion at time t is computed as alpha(t) = min(h * beta0 * log(1 + t), 1) where h is the time interval. This controls the strength of attraction toward the nearest observation.

Practical effect: - Higher values (2.0-5.0): Stronger drift, faster convergence, more exploitation of current best positions - Lower values (0.3-0.8): Weaker drift, more exploration, slower convergence - Recommended default: 1.0-2.0 for most cases

See the companion paper (References) for the derivation of the drift schedule.

 1.0
step_size float

Time discretization step for the SDE solver (must be > 0). Controls the temporal resolution of the stochastic process.

Mathematical role: Euler discretization step Δt for solving the stochastic differential equation. Smaller values give more accurate simulation at the cost of more computation.

Practical effect: - Smaller values (0.001-0.01): More accurate simulation, slower - Larger values (0.05-0.1): Faster but less accurate - Recommended default: 0.01 for good accuracy/speed tradeoff - Rule of thumb: Use step_size much smaller than the typical time scale of the Poisson process (~ 1/lambda0)

 0.1
energy_mode str

Which reference measure the k-means energy (mean squared distance to the nearest center) is averaged under — used by MinimizeEnergy to select the best visited state and by record_energy diagnostics: - "uniform": average over all graph nodes with equal weight, measuring how well the centers cover the whole geometry irrespective of where the observations lie. - "empirical": average over this algorithm's own observation points, exactly where they lie (the empirical k-means objective). The only mode supported on Riemannian manifolds. - "node_measure": average under the per-node obs_weight measure registered on the graph by the caller (e.g. a population measure); graph spaces only. The former "obs" mode was split into "empirical" and "node_measure" and now raises.

 'uniform'
random_state int | Generator | None

Controls randomness for reproducibility. Determines random number generation for all random operations: - Shuffling observations - Poisson process time generation - Brownian motion (via centers) - Initialization strategies (KMeansPlusPlus, RandomInit) - Space-specific random operations

All randomness flows through a single numpy Generator (self._rng), which is propagated to the centers so that every stochastic component is driven by the same stream. No global random state (random.seed/np.random.seed) is touched, so runs are isolated and fully reproducible.

Pass an int for a reproducible seed, a Generator instance for fine-grained control, or None for non-deterministic behavior (default).

Example: >>> # Reproducible with seed (recommended) >>> sa1 = SimulatedAnnealing(points, k=3, random_state=42) >>> sa2 = SimulatedAnnealing(points, k=3, random_state=42) >>> # sa1 and sa2 produce identical results >>> >>> # Or pass an explicit Generator >>> rng = np.random.default_rng(42) >>> sa = SimulatedAnnealing(points, k=3, random_state=rng)

 None

Raises:

Type Description
ValueError

If observations is empty, k <= 0, points are in different spaces, or hyperparameters are invalid.
References

C. Brécheteau, I. Gavra, N. Klutchnikoff. "Online k-means Clustering on Metric Graphs and Geodesic Spaces" (preprint). Derives the annealing dynamics and its convergence analysis.

Example
Quick convergence setup

sa = SimulatedAnnealing( ... points, k=5, ... lambda0=0.5, # Less exploration ... beta0=3.0, # Stronger drift ... step_size=0.01 ... )

Thorough search setup (avoid local minima)

sa = SimulatedAnnealing( ... points, k=5, ... lambda0=2.0, # More exploration ... beta0=1.0, # Gentler drift ... step_size=0.01 ... )

Source code in kmeanssa_ng/core/simulated_annealing.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def __init__(
    self,
    observations: list[Point],
    k: int,
    lambda0: float = 1.0,
    beta0: float = 1.0,
    step_size: float = 0.1,
    energy_mode: str = "uniform",
    random_state: int | np.random.Generator | None = None,
) -> None:
    """Initialize the simulated annealing algorithm.

    Args:
        observations: List of points to cluster, all in the same metric space.
        k: Number of clusters.
        lambda0: Intensity scale of the Poisson observation clock (must be > 0).

            Mathematical role: the annealing processes one observation per
            arrival of an inhomogeneous Poisson process of intensity
            lambda(t) = lambda0 * (1 + t) (the paper's schedule). It does
            **not** scale the Brownian steps themselves (each micro-step
            has standard deviation sqrt(step_size), independent of
            lambda0).

            Practical effect:
            - Higher values: arrivals come faster, so the same number of
              observations spans a shorter annealing horizon (less
              Brownian exploration per observation)
            - Lower values: longer horizon, more exploration between
              observation events
            - Recommended default: 1.0

            See the companion paper (References) for the derivation of the
            time schedule.

        beta0: Initial drift intensity parameter (must be > 0).
            Controls how strongly centers are pulled toward observations.

            Mathematical role: The drift proportion at time t is computed as
            alpha(t) = min(h * beta0 * log(1 + t), 1) where h is the time
            interval. This controls the strength of attraction toward the
            nearest observation.

            Practical effect:
            - Higher values (2.0-5.0): Stronger drift, faster convergence,
              more exploitation of current best positions
            - Lower values (0.3-0.8): Weaker drift, more exploration,
              slower convergence
            - Recommended default: 1.0-2.0 for most cases

            See the companion paper (References) for the derivation of the
            drift schedule.

        step_size: Time discretization step for the SDE solver (must be > 0).
            Controls the temporal resolution of the stochastic process.

            Mathematical role: Euler discretization step Δt for solving the
            stochastic differential equation. Smaller values give more
            accurate simulation at the cost of more computation.

            Practical effect:
            - Smaller values (0.001-0.01): More accurate simulation, slower
            - Larger values (0.05-0.1): Faster but less accurate
            - Recommended default: 0.01 for good accuracy/speed tradeoff
            - Rule of thumb: Use step_size much smaller than the typical
              time scale of the Poisson process (~ 1/lambda0)

        energy_mode: Which reference measure the k-means energy (mean
            squared distance to the nearest center) is averaged under —
            used by ``MinimizeEnergy`` to select the best visited state
            and by ``record_energy`` diagnostics:
            - "uniform": average over all graph nodes with equal weight,
              measuring how well the centers cover the whole geometry
              irrespective of where the observations lie.
            - "empirical": average over this algorithm's own observation
              points, exactly where they lie (the empirical k-means
              objective). The only mode supported on Riemannian
              manifolds.
            - "node_measure": average under the per-node ``obs_weight``
              measure registered on the graph by the caller (e.g. a
              population measure); graph spaces only.
            The former "obs" mode was split into "empirical" and
            "node_measure" and now raises.

        random_state: Controls randomness for reproducibility.
            Determines random number generation for all random operations:
            - Shuffling observations
            - Poisson process time generation
            - Brownian motion (via centers)
            - Initialization strategies (KMeansPlusPlus, RandomInit)
            - Space-specific random operations

            All randomness flows through a single numpy Generator
            (``self._rng``), which is propagated to the centers so that
            every stochastic component is driven by the same stream. No
            global random state (``random.seed``/``np.random.seed``) is
            touched, so runs are isolated and fully reproducible.

            Pass an int for a reproducible seed, a Generator instance for
            fine-grained control, or None for non-deterministic behavior
            (default).

            Example:
                >>> # Reproducible with seed (recommended)
                >>> sa1 = SimulatedAnnealing(points, k=3, random_state=42)
                >>> sa2 = SimulatedAnnealing(points, k=3, random_state=42)
                >>> # sa1 and sa2 produce identical results
                >>>
                >>> # Or pass an explicit Generator
                >>> rng = np.random.default_rng(42)
                >>> sa = SimulatedAnnealing(points, k=3, random_state=rng)

    Raises:
        ValueError: If observations is empty, k <= 0, points are in different spaces,
            or hyperparameters are invalid.

    References:
        C. Brécheteau, I. Gavra, N. Klutchnikoff. "Online k-means Clustering
        on Metric Graphs and Geodesic Spaces" (preprint). Derives the
        annealing dynamics and its convergence analysis.

    Example:
        >>> # Quick convergence setup
        >>> sa = SimulatedAnnealing(
        ...     points, k=5,
        ...     lambda0=0.5,  # Less exploration
        ...     beta0=3.0,     # Stronger drift
        ...     step_size=0.01
        ... )
        >>>
        >>> # Thorough search setup (avoid local minima)
        >>> sa = SimulatedAnnealing(
        ...     points, k=5,
        ...     lambda0=2.0,   # More exploration
        ...     beta0=1.0,     # Gentler drift
        ...     step_size=0.01
        ... )
    """
    self._validate_constructor_parameters(
        observations, k, lambda0, beta0, step_size
    )
    if energy_mode == "obs":
        raise ValueError(
            "energy_mode 'obs' was split into two explicit modes: use "
            "'empirical' to average over this algorithm's observation "
            "points, or 'node_measure' to average under the per-node "
            "'obs_weight' measure registered on the graph."
        )
    if energy_mode not in ("uniform", "empirical", "node_measure"):
        raise ValueError(
            "energy_mode must be 'uniform', 'empirical' or "
            f"'node_measure', got {energy_mode!r}"
        )
    self._initialize_random_generator(random_state)

    self._space = observations[0].space
    self._observations = observations.copy()
    self._k = k
    self._lambda = float(lambda0)
    self._beta = float(beta0)
    self._step_size = float(step_size)
    self._energy_mode = energy_mode

    # Shuffle through the instance Generator so ordering is reproducible
    # from random_state without touching global random state.
    self._rng.shuffle(self._observations)
    self._centers: list[Center] = []
    self._energy_history: list[float] = []
    self._time_history: list[float] = []

calculate_energy(centers)

Calculate k-means energy for given centers based on the energy mode.

Delegates to the space. The algorithm's own observations are the data of the "empirical" mode only: "uniform" and "node_measure" define their reference measure without them (and reject them, so no mode can silently shadow another). Acceleration (e.g. the quantum graph's numba kernels) is the space's concern, dispatched inside Space.calculate_energy.

Source code in kmeanssa_ng/core/simulated_annealing.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def calculate_energy(self, centers: list[Center]) -> float:
    """Calculate k-means energy for given centers based on the energy mode.

    Delegates to the space. The algorithm's own observations are the data
    of the "empirical" mode only: "uniform" and "node_measure" define
    their reference measure without them (and reject them, so no mode can
    silently shadow another). Acceleration (e.g. the quantum graph's
    numba kernels) is the space's concern, dispatched inside
    ``Space.calculate_energy``.
    """
    observations = self._observations if self._energy_mode == "empirical" else None
    return self.space.calculate_energy(
        centers, how=self._energy_mode, observations=observations
    )

run(initialization_strategy=None, robustification_strategy=None, robust_prop=0.0, record_energy=False)

Run the simulated annealing algorithm.

This is the primary method to execute the simulated annealing algorithm. It performs an interleaved sequence of Brownian motion (exploration) and drift (exploitation) for the cluster centers.

Parameters:

Name Type Description Default
initialization_strategy InitializationStrategy | None

How the initial centers are chosen. Defaults to :class:KMeansPlusPlus, the canonical k-means++ seeding, so sa.run() works out of the box.

 None
robustification_strategy RobustificationStrategy | None

How the returned centers are selected from the trajectory. Defaults to :class:MinimizeEnergy, which keeps the lowest-energy state seen during the collection window.

 None
robust_prop float

Fraction of the (trailing) observations over which the robustification strategy collects candidate states, in [0, 1]. Left at 0.0 by default, the window is a single point, so the default :class:MinimizeEnergy only compares the initial and final states — pass robust_prop around 0.1 for a genuine best-of-window robustification.

 0.0
record_energy bool

If True, record the energy and annealing time after each observation into :attr:energy_history and :attr:time_history (for convergence diagnostics). Off by default so the energy is not recomputed when not needed.

 False
Example
Zero-config quickstart: k-means++ init, energy-minimizing
robustification.

centers = SimulatedAnnealing(points, k=5, random_state=0).run()

Source code in kmeanssa_ng/core/simulated_annealing.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
def run(
    self,
    initialization_strategy: InitializationStrategy | None = None,
    robustification_strategy: RobustificationStrategy | None = None,
    robust_prop: float = 0.0,
    record_energy: bool = False,
):
    """Run the simulated annealing algorithm.

    This is the primary method to execute the simulated annealing algorithm.
    It performs an interleaved sequence of Brownian motion (exploration)
    and drift (exploitation) for the cluster centers.

    Args:
        initialization_strategy: How the initial centers are chosen.
            Defaults to :class:`KMeansPlusPlus`, the canonical k-means++
            seeding, so ``sa.run()`` works out of the box.
        robustification_strategy: How the returned centers are selected
            from the trajectory. Defaults to :class:`MinimizeEnergy`, which
            keeps the lowest-energy state seen during the collection window.
        robust_prop: Fraction of the (trailing) observations over which the
            robustification strategy collects candidate states, in [0, 1].
            Left at 0.0 by default, the window is a single point, so the
            default :class:`MinimizeEnergy` only compares the initial and
            final states — pass ``robust_prop`` around 0.1 for a genuine
            best-of-window robustification.
        record_energy: If True, record the energy and annealing time after
            each observation into :attr:`energy_history` and
            :attr:`time_history` (for convergence diagnostics). Off by
            default so the energy is not recomputed when not needed.

    Example:
        >>> # Zero-config quickstart: k-means++ init, energy-minimizing
        >>> # robustification.
        >>> centers = SimulatedAnnealing(points, k=5, random_state=0).run()
    """
    if initialization_strategy is None:
        initialization_strategy = KMeansPlusPlus()
    if robustification_strategy is None:
        robustification_strategy = MinimizeEnergy()

    logger.info(
        "Starting SA: k=%d, n_obs=%d, lambda0=%.3f, beta0=%.3f, "
        "step_size=%.4f, robust_prop=%.2f",
        self._k,
        self.n,
        self._lambda,
        self._beta,
        self._step_size,
        robust_prop,
    )

    if robust_prop < 0 or robust_prop > 1:
        raise ValueError("The proportion must be in [0,1]")

    i0 = int(np.floor((self.n - 1) * (1 - robust_prop)))

    self._centers = initialization_strategy.initialize_centers(self)
    # Seed each center's RNG from the SA's generator so that all stochastic moves
    # (Brownian step size and vertex routing) are reproducible from random_state.
    for _center in self._centers:
        _center.seed_rng(self._rng)

    robustification_strategy.initialize(self)
    strategy = robustification_strategy

    times = self._initialize_times(self.n)
    time = 0.0
    progress_interval = max(1, self.n // 10)

    if record_energy:
        self._energy_history = [self.calculate_energy(self._centers)]
        self._time_history = [time]

    for i, point in enumerate(self._observations):
        # times[0] is the origin of the clock: observation i is processed
        # over the interval (times[i], times[i + 1]].
        T = times[i + 1]

        if i % progress_interval == 0 and i > 0:
            progress = 100 * i / self.n
            logger.info(
                "Progress: %.1f%% (%d/%d observations processed)",
                progress,
                i,
                self.n,
            )

        logger.debug("Processing observation %d, target time T=%.4f", i, T)

        while time < T:
            h = min(self._step_size, T - time)
            prop = min(h * self._beta * np.log(1 + time), 1)
            logger.debug(
                "Time step: time=%.4f, h=%.4f, drift_prop=%.4f", time, h, prop
            )
            for center in self._centers:
                center.brownian_motion(h)
            distances = self.space.distances_from_centers(self._centers, point)
            closest_idx = np.argmin(distances)
            self._centers[closest_idx].drift(point, prop)
            time += h

        if i >= i0:
            strategy.collect(self)
            logger.debug(
                "Collected centers for robustification at observation %d", i
            )

        if record_energy:
            self._energy_history.append(self.calculate_energy(self._centers))
            self._time_history.append(time)

    result = strategy.get_result()
    logger.info("SA completed successfully")
    return result

Space

Bases: ABC

Abstract base class for metric spaces.

A metric space provides: - Distance computation between points - Sampling of random points and centers - Cluster computation and energy calculation

Source code in kmeanssa_ng/core/abstract.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
class Space(ABC):
    """Abstract base class for metric spaces.

    A metric space provides:
    - Distance computation between points
    - Sampling of random points and centers
    - Cluster computation and energy calculation
    """

    @abstractmethod
    def distance(self, p1: Point, p2: Point) -> float:
        """Compute the distance between two points.

        Args:
            p1: First point.
            p2: Second point.

        Returns:
            The distance between p1 and p2.
        """
        raise NotImplementedError

    def sample_points(self, n: int, strategy: SamplingStrategy) -> list[Point]:
        """Sample n points using the specified sampling strategy.

        Args:
            n: Number of points to sample
            strategy: Sampling strategy defining the probability distribution.
                     Must be a SamplingStrategy instance specific to the space type.

        Returns:
            List of n sampled points

        Example:
            ```python
            # For quantum graphs
            from kmeanssa_ng.quantum_graph.sampling import UniformNodeSampling
            points = graph.sample_points(100, strategy=UniformNodeSampling())

            # For Riemannian manifolds
            from kmeanssa_ng.riemannian_manifold.sampling import UniformManifoldSampling
            points = manifold.sample_points(100, strategy=UniformManifoldSampling())
            ```

        Note:
            The strategy parameter is REQUIRED to avoid ambiguity about which
            probability distribution to use. Each space type has its own
            specific sampling strategies in space-specific modules.
        """
        return strategy.sample(self, n)

    def assign_clusters(self, points: list[Point], centers: list[Center]) -> list[int]:
        """Assign points to the nearest center and return cluster labels.

        Args:
            points: A list of points to be clustered.
            centers: A list of centers.

        Returns:
            A list of integer labels, where each label is the index of the
            closest center for the corresponding point in the input list.
        """
        labels = []
        for point in points:
            distances = [self.distance(point, center) for center in centers]
            closest_center_idx = distances.index(min(distances))
            labels.append(closest_center_idx)
        return labels

    @abstractmethod
    def calculate_energy(
        self,
        centers: list[Center],
        how: EnergyMode = "uniform",
        observations: list[Point] | None = None,
    ) -> float:
        """Calculate the k-means energy (distortion) for given centers.

        The energy is the **mean** squared distance to the nearest center,
        taken under a reference measure — the same convention for every
        space, so energies are comparable across modes and implementations.

        The mode names describe the *provenance* of the reference measure —
        the only thing the space can guarantee. Which measure plays which
        statistical role is the caller's declaration: for the population
        objective, register the population measure on the space and use
        ``"node_measure"``; for the empirical objective, pass the sample
        with ``"empirical"``; ``"uniform"`` is a geometric reference measure
        independent of any sampling.

        Args:
            centers: List of cluster centers.
            how: Which reference measure to average under. Each mode has an
                unambiguous data source, and every mismatch is an error —
                there is deliberately no silent fallback between modes:
                - "uniform": the space's own uniform measure (e.g. uniform
                  over graph nodes). Rejects ``observations``. Spaces
                  without a uniform measure (continuous manifolds) reject
                  the mode itself.
                - "empirical": the empirical measure of the ``observations``
                  list, which is **required**. Points are used exactly where
                  they lie (no rounding to nodes). The only mode supported
                  on continuous manifolds.
                - "node_measure": the measure carried by the space's nodes
                  (a graph's per-node ``obs_weight``, set explicitly by the
                  caller). Rejects ``observations`` — the caller picks one
                  source, never both. Graph spaces only.
            observations: The algorithm's observation points, for
                ``"empirical"``. Observations belong to the algorithm, never
                to the space: passing them explicitly keeps two algorithms
                sharing one space independent.

        Returns:
            The mean squared distance to the nearest center.

        Raises:
            ValueError: If the mode is unknown (including the retired
                ``"obs"``), unsupported by the space, or its data source is
                missing or over-specified.
        """
        raise NotImplementedError

    @abstractmethod
    def distances_from_centers(
        self, centers: list[Center], target: Point
    ) -> np.ndarray:
        """Compute distances from multiple centers to a single target point.

        This method is used by the simulated annealing algorithm to efficiently
        find the nearest center to a given observation point.

        Args:
            centers: List of k centers to compute distances from.
            target: The target point.

        Returns:
            Array of shape (k,) with distances from each center to target.

        Example:
            ```python
            centers = space.sample_centers(5)
            target = space.sample_points(1)[0]
            distances = space.distances_from_centers(centers, target)
            closest_idx = np.argmin(distances)
            ```
        """
        raise NotImplementedError

    @abstractmethod
    def center_from_point(self, point: Point) -> Center:
        """Create a Center object from a Point object."""
        raise NotImplementedError

    @abstractmethod
    def get_point_type(self) -> type[Point]:
        """Return the type of points in this space."""
        raise NotImplementedError

assign_clusters(points, centers)

Assign points to the nearest center and return cluster labels.

Parameters:

Name Type Description Default
points list[Point]

A list of points to be clustered.

 required
centers list[Center]

A list of centers.

 required

Returns:

Type Description
list[int]

A list of integer labels, where each label is the index of the
list[int]

closest center for the corresponding point in the input list.
Source code in kmeanssa_ng/core/abstract.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def assign_clusters(self, points: list[Point], centers: list[Center]) -> list[int]:
    """Assign points to the nearest center and return cluster labels.

    Args:
        points: A list of points to be clustered.
        centers: A list of centers.

    Returns:
        A list of integer labels, where each label is the index of the
        closest center for the corresponding point in the input list.
    """
    labels = []
    for point in points:
        distances = [self.distance(point, center) for center in centers]
        closest_center_idx = distances.index(min(distances))
        labels.append(closest_center_idx)
    return labels

calculate_energy(centers, how='uniform', observations=None) abstractmethod

Calculate the k-means energy (distortion) for given centers.

The energy is the mean squared distance to the nearest center, taken under a reference measure — the same convention for every space, so energies are comparable across modes and implementations.

The mode names describe the provenance of the reference measure — the only thing the space can guarantee. Which measure plays which statistical role is the caller's declaration: for the population objective, register the population measure on the space and use "node_measure"; for the empirical objective, pass the sample with "empirical"; "uniform" is a geometric reference measure independent of any sampling.

Parameters:

Name Type Description Default
centers list[Center]

List of cluster centers.

 required
how EnergyMode

Which reference measure to average under. Each mode has an unambiguous data source, and every mismatch is an error — there is deliberately no silent fallback between modes: - "uniform": the space's own uniform measure (e.g. uniform over graph nodes). Rejects observations. Spaces without a uniform measure (continuous manifolds) reject the mode itself. - "empirical": the empirical measure of the observations list, which is required. Points are used exactly where they lie (no rounding to nodes). The only mode supported on continuous manifolds. - "node_measure": the measure carried by the space's nodes (a graph's per-node obs_weight, set explicitly by the caller). Rejects observations — the caller picks one source, never both. Graph spaces only.

 'uniform'
observations list[Point] | None

The algorithm's observation points, for "empirical". Observations belong to the algorithm, never to the space: passing them explicitly keeps two algorithms sharing one space independent.

 None

Returns:

Type Description
float

The mean squared distance to the nearest center.

Raises:

Type Description
ValueError

If the mode is unknown (including the retired "obs"), unsupported by the space, or its data source is missing or over-specified.
Source code in kmeanssa_ng/core/abstract.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@abstractmethod
def calculate_energy(
    self,
    centers: list[Center],
    how: EnergyMode = "uniform",
    observations: list[Point] | None = None,
) -> float:
    """Calculate the k-means energy (distortion) for given centers.

    The energy is the **mean** squared distance to the nearest center,
    taken under a reference measure — the same convention for every
    space, so energies are comparable across modes and implementations.

    The mode names describe the *provenance* of the reference measure —
    the only thing the space can guarantee. Which measure plays which
    statistical role is the caller's declaration: for the population
    objective, register the population measure on the space and use
    ``"node_measure"``; for the empirical objective, pass the sample
    with ``"empirical"``; ``"uniform"`` is a geometric reference measure
    independent of any sampling.

    Args:
        centers: List of cluster centers.
        how: Which reference measure to average under. Each mode has an
            unambiguous data source, and every mismatch is an error —
            there is deliberately no silent fallback between modes:
            - "uniform": the space's own uniform measure (e.g. uniform
              over graph nodes). Rejects ``observations``. Spaces
              without a uniform measure (continuous manifolds) reject
              the mode itself.
            - "empirical": the empirical measure of the ``observations``
              list, which is **required**. Points are used exactly where
              they lie (no rounding to nodes). The only mode supported
              on continuous manifolds.
            - "node_measure": the measure carried by the space's nodes
              (a graph's per-node ``obs_weight``, set explicitly by the
              caller). Rejects ``observations`` — the caller picks one
              source, never both. Graph spaces only.
        observations: The algorithm's observation points, for
            ``"empirical"``. Observations belong to the algorithm, never
            to the space: passing them explicitly keeps two algorithms
            sharing one space independent.

    Returns:
        The mean squared distance to the nearest center.

    Raises:
        ValueError: If the mode is unknown (including the retired
            ``"obs"``), unsupported by the space, or its data source is
            missing or over-specified.
    """
    raise NotImplementedError

center_from_point(point) abstractmethod

Create a Center object from a Point object.

Source code in kmeanssa_ng/core/abstract.py 
234
235
236
237
@abstractmethod
def center_from_point(self, point: Point) -> Center:
    """Create a Center object from a Point object."""
    raise NotImplementedError

distance(p1, p2) abstractmethod

Compute the distance between two points.

Parameters:

Name Type Description Default
p1 Point

First point.

 required
p2 Point

Second point.

 required

Returns:

Type Description
float

The distance between p1 and p2.
Source code in kmeanssa_ng/core/abstract.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@abstractmethod
def distance(self, p1: Point, p2: Point) -> float:
    """Compute the distance between two points.

    Args:
        p1: First point.
        p2: Second point.

    Returns:
        The distance between p1 and p2.
    """
    raise NotImplementedError

distances_from_centers(centers, target) abstractmethod

Compute distances from multiple centers to a single target point.

This method is used by the simulated annealing algorithm to efficiently find the nearest center to a given observation point.

Parameters:

Name Type Description Default
centers list[Center]

List of k centers to compute distances from.

 required
target Point

The target point.

 required

Returns:

Type Description
ndarray

Array of shape (k,) with distances from each center to target.
Example 
centers = space.sample_centers(5)
target = space.sample_points(1)[0]
distances = space.distances_from_centers(centers, target)
closest_idx = np.argmin(distances)
Source code in kmeanssa_ng/core/abstract.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
@abstractmethod
def distances_from_centers(
    self, centers: list[Center], target: Point
) -> np.ndarray:
    """Compute distances from multiple centers to a single target point.

    This method is used by the simulated annealing algorithm to efficiently
    find the nearest center to a given observation point.

    Args:
        centers: List of k centers to compute distances from.
        target: The target point.

    Returns:
        Array of shape (k,) with distances from each center to target.

    Example:
        ```python
        centers = space.sample_centers(5)
        target = space.sample_points(1)[0]
        distances = space.distances_from_centers(centers, target)
        closest_idx = np.argmin(distances)
        ```
    """
    raise NotImplementedError

get_point_type() abstractmethod

Return the type of points in this space.

Source code in kmeanssa_ng/core/abstract.py 
239
240
241
242
@abstractmethod
def get_point_type(self) -> type[Point]:
    """Return the type of points in this space."""
    raise NotImplementedError

sample_points(n, strategy)

Sample n points using the specified sampling strategy.

Parameters:

Name Type Description Default
n int

Number of points to sample

 required
strategy SamplingStrategy

Sampling strategy defining the probability distribution. Must be a SamplingStrategy instance specific to the space type.

 required

Returns:

Type Description
list[Point]

List of n sampled points
Example 
# For quantum graphs
from kmeanssa_ng.quantum_graph.sampling import UniformNodeSampling
points = graph.sample_points(100, strategy=UniformNodeSampling())

# For Riemannian manifolds
from kmeanssa_ng.riemannian_manifold.sampling import UniformManifoldSampling
points = manifold.sample_points(100, strategy=UniformManifoldSampling())
Note

The strategy parameter is REQUIRED to avoid ambiguity about which probability distribution to use. Each space type has its own specific sampling strategies in space-specific modules.

Source code in kmeanssa_ng/core/abstract.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def sample_points(self, n: int, strategy: SamplingStrategy) -> list[Point]:
    """Sample n points using the specified sampling strategy.

    Args:
        n: Number of points to sample
        strategy: Sampling strategy defining the probability distribution.
                 Must be a SamplingStrategy instance specific to the space type.

    Returns:
        List of n sampled points

    Example:
        ```python
        # For quantum graphs
        from kmeanssa_ng.quantum_graph.sampling import UniformNodeSampling
        points = graph.sample_points(100, strategy=UniformNodeSampling())

        # For Riemannian manifolds
        from kmeanssa_ng.riemannian_manifold.sampling import UniformManifoldSampling
        points = manifold.sample_points(100, strategy=UniformManifoldSampling())
        ```

    Note:
        The strategy parameter is REQUIRED to avoid ambiguity about which
        probability distribution to use. Each space type has its own
        specific sampling strategies in space-specific modules.
    """
    return strategy.sample(self, n)

run_parallel(space, n_points, k, sampling_strategy, initialization_strategy, robustification_strategy, n_runs=10, lambda0=1, beta0=1.0, step_size=0.1, energy_mode='uniform', robust_prop=0.0, n_jobs=-1, seeds=None, return_all=False, mp_context=None)

Run simulated annealing multiple times in parallel with different seeds.

This function executes n_runs independent simulated annealing runs in parallel, each with a different random seed. Each run samples its own observations, generates its own Poisson process, and initializes differently, ensuring complete independence between runs.

Parameters:

Name Type Description Default
space 'Space'

The metric space to sample points from.

 required
n_points int

Number of points to sample for each run.

 required
k int

Number of clusters.

 required
n_runs int

Number of parallel runs to execute.

 10
lambda_param

Poisson process intensity parameter (must be > 0).

 required
beta

Inverse temperature parameter (must be > 0).

 required
step_size float

Time step for updating centers (must be > 0).

 0.1
sampling_strategy SamplingStrategy

Strategy for sampling points from the space (required).

 required
initialization_strategy InitializationStrategy

Strategy for initializing centers (required).

 required
robustification_strategy RobustificationStrategy

Strategy for robustifying results (required).

 required
robust_prop float

Proportion of final observations to use for robustification (0-1).

 0.0
n_jobs int

Number of parallel jobs. -1 uses all available cores.

 -1
seeds list[int] | None

Optional list of specific seeds to use. If None, generates random seeds.

 None
return_all bool

If True, return all results; if False, return only the best.

 False
mp_context Literal['fork', 'spawn', 'forkserver'] | None

Multiprocessing context to use ('fork', 'spawn', 'forkserver'). If None, uses the system default. Use 'fork' for Jupyter/Quarto compatibility.

 None

Returns:

Type Description
list[Center] | tuple[list[Center], list[tuple[list[Center], float, int]]]

If return_all is False: List of best centers (lowest energy).
list[Center] | tuple[list[Center], list[tuple[list[Center], float, int]]]

If return_all is True: Tuple of (best_centers, all_results) where all_results is a list of (centers, energy, seed) tuples sorted by energy.

Raises:

Type Description
ValueError

If n_runs <= 0 or other parameters are invalid.
Example 
from kmeanssa_ng import run_parallel

# Generate a graph
graph = QuantumGraph(...)

# Run 10 parallel executions, each sampling its own 100 points
best_centers = run_parallel(graph, n_points=100, k=5, n_runs=10)

# Get all results for analysis
best, all_results = run_parallel(graph, n_points=100, k=5, n_runs=10, return_all=True)
for centers, energy, seed in all_results:
    print(f"Seed {seed}: energy = {energy:.4f}")
Source code in kmeanssa_ng/core/parallel.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def run_parallel(
    space: "Space",
    n_points: int,
    k: int,
    sampling_strategy: SamplingStrategy,
    initialization_strategy: InitializationStrategy,
    robustification_strategy: RobustificationStrategy,
    n_runs: int = 10,
    lambda0: float = 1,
    beta0: float = 1.0,
    step_size: float = 0.1,
    energy_mode: str = "uniform",
    robust_prop: float = 0.0,
    n_jobs: int = -1,
    seeds: list[int] | None = None,
    return_all: bool = False,
    mp_context: Literal["fork", "spawn", "forkserver"] | None = None,
) -> list[Center] | tuple[list[Center], list[tuple[list[Center], float, int]]]:
    """Run simulated annealing multiple times in parallel with different seeds.

    This function executes n_runs independent simulated annealing runs in parallel,
    each with a different random seed. Each run samples its own observations,
    generates its own Poisson process, and initializes differently, ensuring
    complete independence between runs.

    Args:
        space: The metric space to sample points from.
        n_points: Number of points to sample for each run.
        k: Number of clusters.
        n_runs: Number of parallel runs to execute.
        lambda_param: Poisson process intensity parameter (must be > 0).
        beta: Inverse temperature parameter (must be > 0).
        step_size: Time step for updating centers (must be > 0).
        sampling_strategy: Strategy for sampling points from the space (required).
        initialization_strategy: Strategy for initializing centers (required).
        robustification_strategy: Strategy for robustifying results (required).
        robust_prop: Proportion of final observations to use for robustification (0-1).
        n_jobs: Number of parallel jobs. -1 uses all available cores.
        seeds: Optional list of specific seeds to use. If None, generates random seeds.
        return_all: If True, return all results; if False, return only the best.
        mp_context: Multiprocessing context to use ('fork', 'spawn', 'forkserver').
            If None, uses the system default. Use 'fork' for Jupyter/Quarto compatibility.

    Returns:
        If return_all is False: List of best centers (lowest energy).
        If return_all is True: Tuple of (best_centers, all_results) where all_results
            is a list of (centers, energy, seed) tuples sorted by energy.

    Raises:
        ValueError: If n_runs <= 0 or other parameters are invalid.

    Example:
        ```python
        from kmeanssa_ng import run_parallel

        # Generate a graph
        graph = QuantumGraph(...)

        # Run 10 parallel executions, each sampling its own 100 points
        best_centers = run_parallel(graph, n_points=100, k=5, n_runs=10)

        # Get all results for analysis
        best, all_results = run_parallel(graph, n_points=100, k=5, n_runs=10, return_all=True)
        for centers, energy, seed in all_results:
            print(f"Seed {seed}: energy = {energy:.4f}")
        ```
    """
    if n_runs <= 0:
        raise ValueError(f"n_runs must be positive, got {n_runs}")

    # Check n_jobs and issue a warning if it's too high
    import os
    import warnings

    cpu_count = os.cpu_count() or 1
    if n_jobs > cpu_count:
        warnings.warn(
            f"n_jobs={n_jobs} is greater than the number of available CPUs ({cpu_count}). "
            "This may lead to performance degradation.",
            UserWarning,
        )

    # Determine number of workers
    if n_jobs == -1:
        n_jobs = cpu_count
    elif n_jobs == -2:
        n_jobs = max(1, cpu_count - 1)

    # Generate seeds if not provided. Draw them *without* replacement: with
    # replacement, two runs could share a seed, and since each run derives its
    # streams from SeedSequence(seed) they would be exact duplicates rather
    # than independent samples.
    if seeds is None:
        rng = np.random.default_rng()
        seeds = rng.choice(2**31, size=n_runs, replace=False).tolist()
    elif len(seeds) != n_runs:
        raise ValueError(f"Length of seeds ({len(seeds)}) must match n_runs ({n_runs})")

    # Run all jobs in parallel
    results: list[tuple[list[Center], float, int]] = []

    # Set up multiprocessing context if specified
    executor_kwargs = {"max_workers": n_jobs}
    if mp_context is not None:
        if mp_context not in mp.get_all_start_methods():
            raise ValueError(
                f"Multiprocessing context '{mp_context}' not available on this platform. "
                f"Available: {mp.get_all_start_methods()}"
            )
        executor_kwargs["mp_context"] = mp.get_context(mp_context)

    with ProcessPoolExecutor(**executor_kwargs) as executor:
        # Submit all jobs
        futures = [
            executor.submit(
                _run_with_seed,
                space,
                n_points,
                k,
                seed,
                lambda0,
                beta0,
                step_size,
                energy_mode,
                robust_prop,
                sampling_strategy,
                initialization_strategy,
                robustification_strategy,
            )
            for seed in seeds
        ]

        # Collect results as they complete
        for future in as_completed(futures):
            centers, energy, seed = future.result()
            results.append((centers, energy, seed))

    # Sort by energy (best first), breaking exact ties by seed so the returned
    # best run is deterministic and does not depend on completion order.
    results.sort(key=lambda x: (x[1], x[2]))

    # Return results
    if return_all:
        return results[0][0], results
    else:
        return results[0][0]

run_parallel_with_callback(space, n_points, k, sampling_strategy, initialization_strategy, robustification_strategy, n_runs=10, lambda0=1.0, beta0=1.0, step_size=0.1, energy_mode='uniform', robust_prop=0.0, n_jobs=-1, seeds=None, callback=None, mp_context=None)

Run parallel simulated annealing with progress callback.

Similar to run_parallel but calls a callback function after each run completes, useful for progress tracking and real-time monitoring. Each run samples its own observations with its specific seed.

Parameters:

Name Type Description Default
space 'Space'

The metric space to sample points from.

 required
n_points int

Number of points to sample for each run.

 required
k int

Number of clusters.

 required
n_runs int

Number of parallel runs to execute.

 10
lambda_param

Poisson process intensity parameter.

 required
beta

Inverse temperature parameter.

 required
step_size float

Time step for updating centers.

 0.1
robust_prop float

Proportion for robustification.

 0.0
n_jobs int

Number of parallel jobs (-1 = all cores).

 -1
seeds list[int] | None

Optional list of specific seeds.

 None
callback Callable[[int, int, float], None] | None

Optional function(run_index, seed, energy) called after each run.

 None
mp_context Literal['fork', 'spawn', 'forkserver'] | None

Multiprocessing context to use ('fork', 'spawn', 'forkserver'). If None, uses the system default. Use 'fork' for Jupyter/Quarto compatibility.

 None

Returns:

Type Description
list[Center]

List of best centers (lowest energy).
Example 
def progress_callback(run_idx, seed, energy):
    print(f"Run {run_idx+1}/{n_runs}: energy = {energy:.4f} (seed={seed})")

graph = QuantumGraph(...)
centers = run_parallel_with_callback(
    graph, n_points=100, k=5, n_runs=10, callback=progress_callback
)
Source code in kmeanssa_ng/core/parallel.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
def run_parallel_with_callback(
    space: "Space",
    n_points: int,
    k: int,
    sampling_strategy: SamplingStrategy,
    initialization_strategy: InitializationStrategy,
    robustification_strategy: RobustificationStrategy,
    n_runs: int = 10,
    lambda0: float = 1.0,
    beta0: float = 1.0,
    step_size: float = 0.1,
    energy_mode: str = "uniform",
    robust_prop: float = 0.0,
    n_jobs: int = -1,
    seeds: list[int] | None = None,
    callback: Callable[[int, int, float], None] | None = None,
    mp_context: Literal["fork", "spawn", "forkserver"] | None = None,
) -> list[Center]:
    """Run parallel simulated annealing with progress callback.

    Similar to run_parallel but calls a callback function after each run completes,
    useful for progress tracking and real-time monitoring. Each run samples its own
    observations with its specific seed.

    Args:
        space: The metric space to sample points from.
        n_points: Number of points to sample for each run.
        k: Number of clusters.
        n_runs: Number of parallel runs to execute.
        lambda_param: Poisson process intensity parameter.
        beta: Inverse temperature parameter.
        step_size: Time step for updating centers.
        robust_prop: Proportion for robustification.
        n_jobs: Number of parallel jobs (-1 = all cores).
        seeds: Optional list of specific seeds.
        callback: Optional function(run_index, seed, energy) called after each run.
        mp_context: Multiprocessing context to use ('fork', 'spawn', 'forkserver').
            If None, uses the system default. Use 'fork' for Jupyter/Quarto compatibility.

    Returns:
        List of best centers (lowest energy).

    Example:
        ```python
        def progress_callback(run_idx, seed, energy):
            print(f"Run {run_idx+1}/{n_runs}: energy = {energy:.4f} (seed={seed})")

        graph = QuantumGraph(...)
        centers = run_parallel_with_callback(
            graph, n_points=100, k=5, n_runs=10, callback=progress_callback
        )
        ```
    """
    if n_runs <= 0:
        raise ValueError(f"n_runs must be positive, got {n_runs}")

    # Check n_jobs and issue a warning if it's too high
    import os
    import warnings

    cpu_count = os.cpu_count() or 1
    if n_jobs > cpu_count:
        warnings.warn(
            f"n_jobs={n_jobs} is greater than the number of available CPUs ({cpu_count}). "
            "This may lead to performance degradation.",
            UserWarning,
        )

    # Determine number of workers
    if n_jobs == -1:
        n_jobs = cpu_count
    elif n_jobs == -2:
        n_jobs = max(1, cpu_count - 1)

    # Generate seeds if not provided. Draw them *without* replacement: with
    # replacement, two runs could share a seed, and since each run derives its
    # streams from SeedSequence(seed) they would be exact duplicates rather
    # than independent samples.
    if seeds is None:
        rng = np.random.default_rng()
        seeds = rng.choice(2**31, size=n_runs, replace=False).tolist()
    elif len(seeds) != n_runs:
        raise ValueError(f"Length of seeds ({len(seeds)}) must match n_runs ({n_runs})")

    # Run all jobs in parallel with progress tracking
    results: list[tuple[list[Center], float, int]] = []
    completed_count = 0

    # Set up multiprocessing context if specified
    executor_kwargs = {"max_workers": n_jobs}
    if mp_context is not None:
        if mp_context not in mp.get_all_start_methods():
            raise ValueError(
                f"Multiprocessing context '{mp_context}' not available on this platform. "
                f"Available: {mp.get_all_start_methods()}"
            )
        executor_kwargs["mp_context"] = mp.get_context(mp_context)

    with ProcessPoolExecutor(**executor_kwargs) as executor:
        # Submit all jobs
        future_to_index = {
            executor.submit(
                _run_with_seed,
                space,
                n_points,
                k,
                seed,
                lambda0,
                beta0,
                step_size,
                energy_mode,
                robust_prop,
                sampling_strategy,
                initialization_strategy,
                robustification_strategy,
            ): (idx, seed)
            for idx, seed in enumerate(seeds)
        }

        # Collect results as they complete
        for future in as_completed(future_to_index):
            idx, seed = future_to_index[future]
            centers, energy, result_seed = future.result()
            results.append((centers, energy, result_seed))
            completed_count += 1

            # Call callback if provided
            if callback is not None:
                callback(idx, result_seed, energy)

    # Sort by energy (best first), breaking exact ties by seed so the returned
    # best run does not depend on completion order.
    results.sort(key=lambda x: (x[1], x[2]))
    return results[0][0]

:::