<jupyter_start><jupyter_text>Installing packages The code is from https://michielstock.github.io/posts/2017/2017-11-5-OptimalTransport/<jupyter_code>import jax
import jax.numpy as jnp
from jax import jit
import numpy as np
import matplotlib.pyplot as plt
from tqdm.notebook import trange

from sklearn.datasets import make_circles
from scipy.spatial import distance_matrix<jupyter_output><empty_output><jupyter_text>JAX implementation of sinkhorn algorithm<jupyter_code>@jit
def scale_cols_and_rows(P):
    P *= (r / P.sum(1)).reshape((-1, 1))
    P *= (c / P.sum(0)).reshape((1, -1))
    return P
def sinkhorn_knopp_jax(M, r, c, lam, niter=100000):
    M = jnp.array(M)
    n, m = M.shape
    P = jnp.exp(-lam * M)
    P /= P.sum()
    # normalize this matrix
    for i in trange(niter):
        P = scale_cols_and_rows(P)
    return P, jnp.sum(P * M)


def sinkhorn_knopp_np(M, r, c, lam, niter=100000):
    n, m = M.shape
    P = np.exp(-lam * M)
    P /= P.sum()
    # normalize this matrix
    for i in trange(niter):
        P *= r / P.sum(1)
        P *= c / P.sum(0)
    return P, np.sum(P * M)
n_samples = 10000
X, y = make_circles(n_samples=n_samples, noise=0.05, factor=0.5, shuffle=False)
X1 = X[y == 0]
X2 = X[y == 1]

n, m = len(X1), len(X2)<jupyter_output><empty_output><jupyter_text>Comparing numpy and JAX<jupyter_code>r = np.ones(n) / n
c = np.ones(m) / m

M = jnp.array(distance_matrix(X1, X2))

P, d = sinkhorn_knopp_jax(M, r, c, lam=30, niter=1000)
P, d = sinkhorn_knopp_np(M, r, c, lam=30, niter=1000)<jupyter_output><empty_output><jupyter_text>Visualisation of sinkorn algorithm**Warning**: Only taking a subset of points to map otherwise the visualisation will look cluttered on matplotlib<jupyter_code>sampling_factor = 100
plt.scatter(X1[:, 0], X1[:, 1], color="blue")
plt.scatter(X2[:, 0], X2[:, 1], color="red")
for i in trange(0, n, sampling_factor):
    for j in range(0, m, sampling_factor):
        plt.plot([X1[i, 0], X2[j, 0]], [X1[i, 1], X2[j, 1]], color="green", alpha=float(P[i, j] * n) * sampling_factor)
plt.title("Optimal matching")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15318744053282587, 'fraction_numerical': 0.0390104662226451, 'mean_word_length': 4.141809290953545, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 1, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '324642', 'license': 'MIT', 'n_tokens_mistral': 913, 'n_tokens_neox': 861, 'n_words': 212, 'path': 'notebooks/misc/sinkhorn_knopp_algorithm.ipynb', 'repo_name': 'karm-patel/pyprobml', 'size': '2102'}