<jupyter_start><jupyter_text>Gromov Wasserstein In this tutorial we show how to use a [regularised approach](http://proceedings.mlr.press/v48/peyre16.pdf) to solve the [Gromov-Wasserstein](https://link.springer.com/article/10.1007/s10208-011-9093-5) problem. The goal of the GW problem is to match points taken within different spaces endowed with their own geometries. At the core of the GW algorithm is the idea of aligning the structures of two geometries, by aligning cost matrices. We illustrate this by calculating the GW distance and the resulting transport matrix between a 2-dimensional and a 3-dimensional shape.<jupyter_code>!pip install ott-jax
from IPython import display
import jax
from jax import numpy as jnp
from jax import random
import numpy as np
from matplotlib import animation
from matplotlib import cm
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3

import ott
from ott.core import gromov_wasserstein as gw
from ott.geometry import pointcloud<jupyter_output><empty_output><jupyter_text>Matching between spaces with different dimensionsWe apply the Gromov Wasserstein algorithm to a spiral in 2 dimensions and a Swiss roll in 3 dimensions.  To do so, we first generate a spiral and a Swiss roll, and plot them in a 3D space.<jupyter_code># Samples spiral
def sample_spiral(
    n, min_radius, max_radius,
    key, min_angle=0, max_angle=10, noise=1.):
  radius = jnp.linspace(min_radius, max_radius, n)
  angles = jnp.linspace(min_angle, max_angle, n)
  data = []
  noise = jax.random.normal(key, (2, n)) * noise
  for i in range(n): 
    x = (radius[i] + noise[0, i]) * jnp.cos(angles[i])
    y = (radius[i] + noise[1, i]) * jnp.sin(angles[i])
    data.append([x, y])
  data = jnp.array(data)
  return data

# Samples Swiss roll
def sample_swiss_roll(
    n, min_radius, max_radius, length, key,
    min_angle=0, max_angle=10, noise=0.1):
  spiral = sample_spiral(n, min_radius, max_radius, key[0],
                         min_angle, max_angle, noise)
  third_axis = jax.random.uniform(key[1], (n, 1)) * length
  swiss_roll = jnp.hstack((spiral[:, 0:1], third_axis, spiral[:, 1:]))
  return swiss_roll

# Plots spiral and Swiss roll
def plot(swiss_roll, spiral, colormap_angles_swiss_roll, colormap_angles_spiral):
  fig = plt.figure(figsize=(11, 5))
  ax = fig.add_subplot(1, 3, 1)
  ax.scatter(spiral[:, 0], spiral[:, 1], c = colormap_angles_spiral)
  ax.set_ylim(-40, 30)
  ax.grid()
  ax = fig.add_subplot(1, 3, (2, 3), projection='3d')
  ax.view_init(7, -80)
  ax.scatter(swiss_roll[:, 0], swiss_roll[:, 1], swiss_roll[:, 2], 
             c = colormap_angles_swiss_roll)
  ax.set_adjustable('box')
  plt.show()
# Data parameters
n_spiral = 400
n_swiss_roll = 500
length = 10
min_radius = 6
max_radius = 30
noise = 3
min_angle = 0
max_angle = 9
angle_shift = 3
# Seed
seed=14
key = jax.random.PRNGKey(seed)
key, *subkey = random.split(key, 4)
spiral = sample_spiral(n_spiral, min_radius, max_radius, key=subkey[0],
                       min_angle=min_angle + angle_shift, 
                       max_angle=max_angle + angle_shift, noise=noise)
swiss_roll = sample_swiss_roll(n_swiss_roll, min_radius, max_radius, 
                               key=subkey[1:], length=length, min_angle=min_angle, 
                               max_angle=max_angle)

plot(swiss_roll, spiral, 'blue', 'green')<jupyter_output><empty_output><jupyter_text>We then run OTT's Gromov-Wasserstein solver to find a matching between the points of each geometry. In this tutorial, we define two pointclouds (`ott.geometry.pointcloud`) but general geometries can be used as well (`ott.geometry.geometry`). The loss between the distance matrices of the two pointclouds is by default the squared Euclidean loss.<jupyter_code># Applies Gromov Wasserstein
geom_x = pointcloud.PointCloud(x=spiral, y=spiral)
geom_y = pointcloud.PointCloud(x=swiss_roll, y=swiss_roll)
out = gw.gromov_wasserstein(geom_x=geom_x, geom_y=geom_y, 
                            epsilon=100., max_iterations=60, 
                            jit=True)
print(f'The last Sinkhorn iteration has converged: {out.converged_sinkhorn[-1]}')
print(f'The final regularised GW cost is: {out.reg_gw_cost:.3f}')<jupyter_output>The last Sinkhorn iteration has converged: True
The final regularised GW cost is: 24645.891<jupyter_text>The resulting transport matrix between the two pointclouds is as follows:<jupyter_code>transport = out.transport
fig = plt.figure(figsize=(8, 6)) 
plt.imshow(transport, cmap='Purples')
plt.xlabel('IDs of samples from the Swiss roll', fontsize=14)  # IDs are ordered from center to outer part
plt.ylabel('ID of samples from the spiral', fontsize=14)  # IDs are ordered from center to outer part
plt.colorbar()
plt.show()<jupyter_output><empty_output><jupyter_text>The larger the regularisation parameter `epsilon` is, the more diffuse the transport matrix becomes, as we can see in the animation below.<jupyter_code># Animates the transport matrix
fig = plt.figure(figsize=(8, 6))
im=plt.imshow(transport, cmap='Purples')
plt.xlabel('IDs of samples from the Swiss roll', fontsize=14)  # IDs are ordered from center to outer part
plt.ylabel('IDs of samples from the spiral', fontsize=14)  # IDs are ordered from center to outer part
plt.colorbar()

# Initialization function
def init():
    im.set_data(np.zeros(transport.shape))
    return [im]

# Animation function 
def animate(i):
    array = im.get_array()
    geom_x = pointcloud.PointCloud(x=spiral, y=spiral)
    geom_y = pointcloud.PointCloud(x=swiss_roll, y=swiss_roll)
    config = {'threshold': 1e-2}
    out = gw.gromov_wasserstein(geom_x=geom_x, geom_y=geom_y, 
                                epsilon=i, max_iterations=40, 
                                jit=True, sinkhorn_kwargs=config)
    im.set_array(out.transport)
    im.set_clim(0, jnp.max(out.transport[:]))
    return [im]


# Call the animator
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=[100., 500., 750., 1000., 2000., 
                                       10000., 50000., 500000., 1000000.], 
                               interval=1500, blit=True)

html = display.HTML(anim.to_jshtml())
display.display(html)
plt.close()<jupyter_output><empty_output><jupyter_text>To better understand the correspondance found by the GW solver, we plot in the same color, for each green, blue or red point in the spiral, the point in the Swiss roll it is the most coupled to.<jupyter_code># For each sample from the spiral, we get the most coupled point from the Swiss roll
indices_swiss_roll = jnp.array(np.argmax(transport, axis=1))
# Sets colors to visualise matching of some areas between each shape
# IDs of spiral and Swiss roll are ordered from center to outer part
colors_input_spiral = ['b']*40 + ['silver']*160 + ['g']*40 + ['silver']*90 + ['r']*40 + ['silver']*30
colors_swiss_roll = np.array(['silver']*500)
colors_swiss_roll[indices_swiss_roll[:40]] = 'b'
colors_swiss_roll[indices_swiss_roll[200:240]] = 'g'
colors_swiss_roll[indices_swiss_roll[330:370]] = 'r'

plot(swiss_roll, spiral, colors_swiss_roll, colors_input_spiral)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10491757474154792, 'fraction_numerical': 0.027801061749091926, 'mean_word_length': 4.109921484653818, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '342649', 'license': 'Apache-2.0', 'n_tokens_mistral': 2581, 'n_tokens_neox': 2380, 'n_words': 700, 'path': 'docs/notebooks/gromov_wasserstein.ipynb', 'repo_name': 'theouscidda6/ott', 'size': '7158'}