<jupyter_start><jupyter_text>So our unoptimized kmeans is already 2 times faster than sklearn... not bad. Let's test it for multiple frames:<jupyter_code>n_tests = 1000
key = random.PRNGKey(42)
random_frames = random.choice(key, df.frame.unique(), shape=(n_tests, )).squeeze()
sk_results = []
jax_results = []
for idx in random_frames:
    image = load(original_path, frame=idx)
    
    threshold = 120
    threshold_image = (image < threshold).astype("uint8")
    fly_pixels = cv.findNonZero(threshold_image).squeeze()
    init = df.query(f'frame == {idx-1}')[["x", "y"]].to_numpy().astype(np.float32)
    n_flies = init.shape[0]
    
    #SKlearn
    reg = KMeans(n_clusters=n_flies, n_init=1, init=init, tol=1e-4, algorithm='full', )
    sk_results.append(reg.fit(fly_pixels).cluster_centers_)
    
    # Jax
    jax_results.append(kmeans(fly_pixels, init)[0])
diff = np.array([np.linalg.norm(jax-sk) for jax, sk in zip(sk_results, jax_results)])
plt.semilogy(diff, 'o')
idx_highest_diff = np.argmax(diff)
plt.scatter(*jax_results[idx_highest_diff].T, marker='x', c='r')
plt.scatter(*sk_results[idx_highest_diff].T, marker='x', c='k')

#plt.xlim([300, 500])
#plt.ylim([200, 500])
image = load(original_path, frame=random_frames[idx_highest_diff])
plt.figure(figsize=(25, 15))
plt.imshow(image, cmap='gray')
plt.scatter(*jax_results[idx_highest_diff].T, marker='x', c='b')
plt.scatter(*sk_results[idx_highest_diff].T, marker='x', c='r')
plt.grid(False)
#plt.xlim([300, 500])
plt.ylim([200, 500])
image = load(original_path, frame=random_frames[idx_highest_diff])
plt.figure(figsize=(25, 15))
plt.imshow(image, cmap='gray')
plt.scatter(*df.query(f"frame == {random_frames[idx_highest_diff]}")[["x", "y"]].to_numpy().T, marker='x', c='r')
plt.grid(False)
jax.__version__
from flytracker.kmeans import kmeans
kmeans()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17920952896589062, 'fraction_numerical': 0.027612344342176503, 'mean_word_length': 6.247058823529412, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '471098', 'license': 'MIT', 'n_tokens_mistral': 797, 'n_tokens_neox': 718, 'n_words': 119, 'path': 'dev/kmeans_GPU/implementing_kmeans.ipynb', 'repo_name': 'nicolaseberle/flyTracker', 'size': '1847'}