Jax kmeans implementation, sklearn removal from main deps

[kylejcaron]

Address #459 and #315

Background

This PR removes sklearn from the main dependencies by implementing KMeans in jax.

Performance

K-means++ normally picks each new starting point by drawing one random candidate. This PR instead draws a few candidates at each step (n_local_trials, controlled by n_init) and keeps whichever one lowers the clustering error the most. Better starting points mean we need fewer restarts to get a good result, so the default n_init drops from 10 to 3, with equal or better clustering quality. See the kmeans docstring in dynamax/utils/cluster.py for the full writeup and measured range (2.8x-40x faster depending on how separated the clusters are). Unsure if its worth keeping that detail in the docstring.

I have a few pseudo-reproducible benchmark scripts that measures three things on the same test data (a GaussianHMM with 10 states, 10-dimensional emissions, 100k timesteps): this PR's k-means, the k-means it replaces (one candidate per step, n_init=10 - what's on main's branch history right before this PR), and scikit-learn's KMeans (what main actually uses today).

How long clustering alone takes. This skips the rest of HMM setup, which costs the same amount of time no matter which k-means you use, and is dwarfed by the EM fitting step anyway — see below.

import time
import jax
import jax.random as jr
from dynamax.hidden_markov_model import GaussianHMM
from dynamax.utils.cluster import kmeans  # this PR

NUM_STATES, EMISSION_DIM, NUM_TIMESTEPS = 10, 10, 100_000
true_key, sample_key, fit_key = jr.split(jr.PRNGKey(0), 3)
true_hmm = GaussianHMM(NUM_STATES, EMISSION_DIM)
true_params, _ = true_hmm.initialize(true_key)
_, emissions = true_hmm.sample(true_params, sample_key, NUM_TIMESTEPS)
X = emissions.reshape(-1, EMISSION_DIM)

def bench(fn, repeats=3):
    out = fn(); jax.block_until_ready(out)
    times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        out = fn(); jax.block_until_ready(out)
        times.append((time.perf_counter() - t0) * 1000)
    return min(times), out

# this PR's default (greedy seeding, n_init=3)
t, out = bench(lambda: kmeans(X, NUM_STATES, fit_key))
print(f"this PR:  {t:.1f} ms, inertia={float(out.inertia):.1f}")

# this PR's kmeans with the old defaults it replaces (single-candidate, n_init=10)
t, out = bench(lambda: kmeans(X, NUM_STATES, fit_key, n_init=10, n_local_trials=1))
print(f"old JAX:  {t:.1f} ms, inertia={float(out.inertia):.1f}")

# sklearn, what main uses today (pip install scikit-learn to run this line)
from sklearn.cluster import KMeans
t0 = time.perf_counter(); r = KMeans(NUM_STATES, random_state=0).fit(X); t = (time.perf_counter() - t0) * 1000
print(f"sklearn:  {t:.1f} ms, inertia={r.inertia_:.1f}")

Measured (Apple M3 Pro, CPU only, min of 3 runs after warmup):

backend	time	inertia
main today - sklearn KMeans	39-43 ms	2090.5
JAX k-means this PR replaces (single-candidate, n_init=10)	80.7 ms	2091.6
this PR (greedy, n_init=3)	16.9 ms	2091.6

The k-means implementation this PR replaces was already about 2x slower than scikit-learn's. This PR doesn't just close that gap — it makes our clustering about 2.4x faster than scikit-learn and about 4.8x faster than the k-means already on main, while finding clusters of the same quality.

How long a full model fit takes (clustering + 50 rounds of EM). Included for completeness, but the speedup above doesn't show up here.

def full_fit(hmm, method_kwargs):
    def run():
        params, props = hmm.initialize(fit_key, method="kmeans", emissions=emissions, **method_kwargs)
        _, lps = hmm.fit_em(params, props, emissions, num_iters=50, verbose=False)
        return lps
    return bench(run)

hmm = GaussianHMM(NUM_STATES, EMISSION_DIM)
t, _ = full_fit(hmm, {})
print(f"init+EM fit: {t:.0f} ms")

backend	init+EM (50 iters)
main - sklearn	7314 ms
JAX k-means this PR replaces	7412 ms
this PR	7407 ms

All three take about the same time, because the EM fitting step — not clustering — is over 95% of the total time at this data size. So fitting one model won't feel any faster. The speedup matters when you run clustering many times without a full model fit each time, e.g. comparing different starting points, running many small fits, or trying different hyperparameters.