<jupyter_start><jupyter_text>BasicsRun this first to get the data<jupyter_code># Optional: magic for reloading packages at every cell evaluation, great if you're changing code
#%load_ext autoreload
#%autoreload 2
## Imports

# Basic
import numpy as np
import matplotlib.pyplot as plt
import random

# My code, from bary.py
import bary

# Data: https://www.tensorflow.org/datasets
import tensorflow_datasets as tfds
## Load MNIST data and mass transport cost matrix M with helper function
# M: (n x n) giving the cost of moving from one location in the 28x28 image to another; n = 28**2
# X: (n x N) with each column the histogram for a single digit
# y: length N vector of labels
N = 1000
M, X, y = bary.load_MNIST(N)
## Plot a digit
c = X[:,np.random.randint(0,N)]
img = np.reshape(c, (28,28))
plt.imshow(img)
plt.colorbar()
plt.show()<jupyter_output><empty_output><jupyter_text>Average-case digitPicks out all copies of digit $d$ and averages them together into a single digitThe code below should run in about 30 seconds, 5 seconds per gradient descent step for the first few iterations; you can do fewer steps to save time, and the results are comparableRun things in Jupyter Lab for progress bars<jupyter_code># Which digit to use
d = 7

# Get the data
Xp = X[:,y == d]
m = np.shape(Xp)[1]

# Run the descent
r, _, _ = bary.sinkhorn_barycenter(M, Xp, noise = 0.01, iterations = 30, steps = 10)

# Plot the averaged digit
bary.plot_digits(r, width=1)
plt.title('Averaged together')
plt.show()

# Plot some of the reference digits
Xpp = Xp[:,np.random.choice(m,5)]
bary.plot_digits(Xpp, width=5)
plt.suptitle('Some reference digits (%d total)' % m)
plt.show()<jupyter_output>Gradient descent progress: 100%|███████████████████████████████████████████████████████| 10/10 [00:28<00:00,  2.82s/it]<jupyter_text>ClusteringPicks out all copies of digits $d_1,d_2$ and tries to figure out what two digits are present, creating two average-case digitsDoes this by running k-means with Sinkhorn barycenters for finding cluster centersRuns in about 4 minutesAs-is, probably not applicable: there are better ways to do this. But it's really cool and a good proof-of-concept for Sinkhorn in applications.<jupyter_code># Digits
d1 = 0
d2 = 1

# Get the data
bools = [(label == d1) or (label == d2) for label in y]
Xp = X[:,bools]
m = np.shape(Xp)[1]

# Cluster!
k = 2
_, r, R = bary.k_means_sinkhorn_barycenter(M, Xp, k, noise = 0.01, eta = 0.5, l = 10, iter_sink = 20, iter_Dsink = 20, iter_grad = 6, iter_lloyd = 4)

# Plot the cluster centers
bary.plot_digits(r, width=2)
plt.suptitle('Clusters found')
plt.show()

# Plot some of the reference digits
Xpp = Xp[:,np.random.choice(np.shape(Xp)[1],10)]
bary.plot_digits(Xpp, width=5)
plt.suptitle('Some reference digits (%d total)' % m)
plt.show()

# Plot the Lloyd's algorithm initialization
bary.plot_digits(R[:,:,0], width=2)
plt.suptitle('Lloyd initialization')
plt.show()<jupyter_output>Initializing...
Done initializing
{'fraction_non_alphanumeric': 0.12879298718813217, 'fraction_numerical': 0.027309507754551585, 'mean_word_length': 5.055102040816327, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 11, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '552252', 'license': 'MIT', 'n_tokens_mistral': 1161, 'n_tokens_neox': 1044, 'n_words': 368, 'path': 'Demo digits.ipynb', 'repo_name': 'nicolasbolle/barycenters', 'size': '2966'}