PCA example

pca.py

"""
PCA visualisation based on the example here:
https://gist.github.com/anonymous/7d888663c6ec679ea65428715b99bfdd
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

plt.style.use('dark_background')

X = np.random.randn(100,2);
X = np.matmul(X, np.linalg.cholesky([[1, 0.6], [0.6, 0.6]]))
X = X - np.mean(X, axis=0)


"""
Plot the original data
"""
fig, ax = plt.subplots(1, 1, figsize=(8.67, 6))

ax.scatter(X[:,0], X[:,1])
ax.set_aspect('equal', adjustable='box')
ax.set_xlim(-4,4)
ax.set_ylim(-4,4)
ax.set_xlabel('x')
ax.set_ylabel('y')

# plt.show()
plt.savefig('original_scatter.png', dpi=200); plt.close()

"""
plot the eigenvectors
"""
a, eigvecs = np.linalg.eig(np.cov(X.T))

fig, ax = plt.subplots(1, 1, figsize=(8.67, 6))
ax.scatter(X[:,0], X[:,1])

x = np.linspace(-10,10)
m = eigvecs[1,0] / eigvecs[0,0]
plt.plot(x, m * x)

m = eigvecs[1,1] / eigvecs[0,1]
plt.plot(x, m * x)

ax.set_aspect('equal', adjustable='box')
ax.set_xlim(-4,4)
ax.set_ylim(-4,4)

ax.set_xlabel('x')
ax.set_ylabel('y')

# plt.show()
plt.savefig('pca_scatter.png', dpi=200); plt.close()

"""
Plot rotation, as well as variance
"""
angles = np.arange(180) * (np.pi / 180)

variance = [np.var(X @ np.array([np.cos(alpha), np.sin(alpha)]))
    for alpha in angles]

for j, alpha in enumerate(angles):

    # fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10,6))
    fig, axs = plt.subplot_mosaic([['spin', 'var']], width_ratios=(8,4), figsize=(13,6))
    axs['spin'].set_aspect('equal', adjustable='box')
    axs['spin'].set_xlim(-4,4)
    axs['spin'].set_ylim(-4,4)

    w = np.array([np.cos(alpha), np.sin(alpha)])
    z = X @ (w[np.newaxis].T @ w[np.newaxis])

    # plot variance
    axs['var'].plot(angles, variance)
    axs['var'].scatter(alpha, np.var(X @ w))

    for i in np.arange(X.shape[0]):
        axs['spin'].plot([X[i,0], z[i,0]], [X[i,1], z[i,1]], 'C3', lw=1)

    axs['spin'].plot(w[0] * 3.5 * np.array([-1, 1]),
            w[1] * 3.5 * np.array([-1, 1]), color='C1')

    axs['spin'].plot(-w[1] * 2 * np.array([-1, 1]),
            w[0] * 2 * np.array([-1, 1]), color='C0')

    axs['spin'].scatter(z[:,0], z[:,1], color='C3', s=4, zorder=10)
    axs['spin'].scatter(X[:,0], X[:,1], color='C2', s=4, zorder=10)

    a1, a2 = 3.5, 4.5
    axs['spin'].plot(eigvecs[0,0]*np.array([-a2, -a1]),
            eigvecs[1,0]*np.array([-a2, -a1]), color='C1')

    axs['spin'].plot(eigvecs[0,0]*np.array([ a1,  a2]),
            eigvecs[1,0]*np.array([ a1,  a2]), color='C1')

    axs['spin'].set_xlabel('x')
    axs['spin'].set_ylabel('y')

    axs['var'].set_xlabel('$\\theta$')
    axs['var'].set_ylabel('$\mathrm{var}(P)$')

    axs['var'].set_xlim(0, np.pi)

    plt.savefig(f'plots/animation_{j:03}.png', dpi=150, bbox_inches='tight'); plt.close()