grid_embedding.py

import numpy as np
from scipy.spatial.distance import cdist
from lapjv import lapjv

def grid_embedding(h):
    assert int(np.sqrt(h.shape[0])) ** 2 == h.shape[0], 'Nb of examples must be a square number'
    size = np.sqrt(h.shape[0])
    grid = np.dstack(np.meshgrid(np.linspace(0, 1, size), np.linspace(0, 1, size))).reshape(-1, 2)
    cost_matrix = cdist(grid, h, "sqeuclidean").astype('float32')
    cost_matrix = cost_matrix * (100000 / cost_matrix.max())
    _, rows, cols = lapjv(cost_matrix)
    return rows

if __name__ == '__main__':
    from sklearn.datasets import load_digits
    from sklearn.manifold import TSNE
    from skimage.io import imsave
    data = load_digits()
    X = data['images']
    sne = TSNE()
    X = X[0:40*40]
    h = sne.fit_transform(X.reshape((X.shape[0], -1)))
    rows = grid_embedding(h)
    X = X[rows]
    X = X.reshape((40, 40, 8, 8))
    X = X.transpose((0, 2, 1, 3))
    X = X.reshape((40 * 8, 40 * 8))
    X = X / 16.
    imsave('out.png', X)