The connection Laplacian and its use for synchronizing a set of vectors on an NxN grid

ConnectionLaplacian.py

"""
Programmer: Chris Tralie (ctralie@alumni.princeton.edu)
Purpose: To implement the connection Laplacian and show its use
for synchronizing a set of vectors on an NxN grid
"""
import numpy as np
import numpy.linalg as linalg
from scipy import sparse 
from scipy.sparse.linalg import lsqr, cg, eigsh
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d

def getConnectionLaplacian(ws, Os, N, k, weighted=True):
    """
    Given a set of weights and corresponding orientation matrices,
    return k eigenvectors of the connection Laplacian.
    Parameters
    ----------
    ws: ndarray(M, 3)
        An array of weights for each included edge
    Os: M-element list of (ndarray(d, d))
        The corresponding orthogonal transformation matrices
    N: int
        Number of vertices in the graph
    k: int
        Number of eigenvectors to compute
    weighted: boolean
        Whether to normalize by the degree
    
    Returns
    -------
    w: ndarray(k)
        Array of k eigenvalues
    v: ndarray(N*d, k)
        Array of the corresponding eigenvectors
    """
    d = Os[0].shape[0]
    W = sparse.coo_matrix((ws[:, 2], (ws[:, 0], ws[:, 1])), shape=(N, N)).tocsr()
    ## Step 1: Create D^-1 matrix
    deg = np.array(W.sum(1)).flatten()
    deg[deg == 0] = 1

    ## Step 2: Create S matrix
    I = []
    J = []
    V = []
    Jsoff, Isoff = np.meshgrid(np.arange(d), np.arange(d))
    Jsoff = Jsoff.flatten()
    Isoff = Isoff.flatten()
    for idx in range(ws.shape[0]):
        [i, j, wij] = ws[idx, :]
        i, j = int(i), int(j)
        wijOij = wij*Os[idx]
        if weighted:
            wijOij /= deg[i]
        wijOij = (wijOij.flatten()).tolist()
        I.append((i*d + Isoff).tolist())
        J.append((j*d + Jsoff).tolist())
        V.append(wijOij)
    I, J, V = np.array(I).flatten(), np.array(J).flatten(), np.array(V).flatten()
    S = sparse.coo_matrix((V, (I, J)), shape=(N*d, N*d)).tocsr()
    return eigsh(S, which='LA', k=k)


def testConnectionLaplacianSquareGrid(N, seed=0, torus = False):
    """
    Randomly rotate vectors on an NxN grid and use the connection
    Laplacian to come up with a consistent orientation for all of them.
    Animate the results and save to frames.
    ----------
    N: int
        Dimension of grid
    seed: int
        Seed for random initialization of vector angles
    torus: boolean
        Whether this grid is thought of as on the torus or not
    """
    np.random.seed(seed)
    thetas = 2*np.pi*np.random.rand(N, N)
    ws = []
    Os = []
    for i in range(N):
        for j in range(N):
            idxi = i*N+j
            for [di, dj] in [[-1, 0], [1, 0], [0, 1], [0, -1]]:
                i2 = i+di
                j2 = j+dj
                if torus:
                    i2, j2 = i2%N, j2%N
                else:
                    if i2 < 0 or j2 < 0 or i2 >= N or j2 >= N:
                        continue
                idxj = i2*N+j2
                ws.append([idxi, idxj, 1.0])
                # Oij moves vectors from j to i
                theta = thetas[i2, j2] - thetas[i, j]
                c = np.cos(theta)
                s = np.sin(theta)
                Oij = np.array([[c, -s], [s, c]])
                Os.append(Oij)
    ws = np.array(ws)
    w, v = getConnectionLaplacian(ws, Os, N**2, 2)
    thetasnew = np.zeros_like(thetas)
    for idx in range(N*N):
        i, j = np.unravel_index(idx, (N, N))
        vidx = v[idx*2:(idx+1)*2, 0]
        # Bring into world coordinates
        c = np.cos(thetas[i, j])
        s = np.sin(thetas[i, j])
        Oij = np.array([[c, -s], [s, c]])
        vidx = Oij.dot(vidx)
        thetasnew[i, j] = np.arctan2(vidx[1], vidx[0])
    
    # Animate the results
    plt.figure(figsize=(8, 8))
    for fidx, t in enumerate(np.linspace(0, 1, 100)):
        plt.clf()
        ax = plt.gca()
        for i in range(N):
            for j in range(N):
                plt.scatter([j+0.5], [i+0.5], 20, 'k')
                theta = (1-t)*thetas[i, j] + t*thetasnew[i, j]
                v = 0.5*np.array([np.cos(theta), np.sin(theta)])
                ax.arrow(j+0.5, i+0.5, v[0], v[1], head_width=0.1, head_length=0.2)
        plt.scatter([-0.2, N, -0.2, N], [-0.2, -0.2, N, N], alpha=0)
        plt.axis('equal')
        plt.axis('off')
        plt.savefig("%i.png"%fidx, bbox_inches='tight')
    for fidx in range(100, 120):
        plt.savefig("%i.png"%fidx, bbox_inches='tight')

if __name__ == '__main__':
    testConnectionLaplacianSquareGrid(10)