Find principle axes of 3D collection of data, rotate positions to match those axes.

rotation_station.py

import sys
import os

import numpy as np
import scipy as sp
import scipy.stats
import matplotlib as mpl
import matplotlib.pyplot as plt

# Silence annoying numpy errors
np.seterr(divide='ignore', invalid='ignore');

# Plotting settings
mpl.rc('font', **{'family': 'serif', 'sans-serif': ['Times']})
mpl.rc('lines', solid_capstyle='round')
mpl.rc('mathtext', fontset='cm')
plt.rcParams.update({'grid.alpha': 0.5})


# Create some covariant data that looks kinda like a galaxy (?)
# -----------------------------------------------------------------------------------
# See: https://scipy-cookbook.readthedocs.io/items/CorrelatedRandomSamples.html

NUM = 2000

# covariance matrix
rr = np.array([
        [  3.40, -2.75, -2.00],
        [  0.0 ,  5.50,  1.50],
        [  0.0 ,  0.00,  1.25]
    ])

# Force matrix to be symmetric (using upper-right)
for ii, jj in np.ndindex((3, 3)):
    if jj < ii:
        rr[ii, jj] = rr.T[ii, jj]

# Generate samples from independent random variables
xx = sp.stats.norm.rvs(size=(3, NUM))

# We need a matrix `c` for which `c*c^T = r`
# Compute the eigenvalues and eigenvectors.
evals, evecs = np.linalg.eigh(rr)
# Construct c, so c*c^T = r.
cc = np.dot(evecs, np.diag(np.sqrt(evals)))

# Convert the data to correlated random variables. 
yy = np.dot(cc, xx)
yy /= np.mean(np.linalg.norm(yy, axis=0))



# Construct moment-of-intertia (MIT) matrix
# ---------------------------------------------------
    
mit = np.zeros((3, 3))
# Generate random weights (e.g. masses)
ww = np.random.uniform(0.5, 1.5, yy.shape[-1])
for ii in range(3):
    for jj in range(3):
        if jj < ii:
            mit[ii, jj] = mit.T[ii, jj]
            continue
        
        t1 = (ii == jj) * np.sum(np.square(yy), axis=0)
        t2 = yy[ii, :] * yy[jj, :]
        mit[ii, jj] = np.sum(ww * (t1 - t2))


# Construct Rotation Matrix
# --------------------------------------------

eig_val, eig_vec = np.linalg.eig(mit)

# Sort by eigenvalue
inds = np.argsort(eig_val)
eig_val = eig_val[inds]
eig_vec = eig_vec[:, inds]

# Create diagonal of eigenvalues
diag_eig = np.eye(3) * eig_val
# Solve `R*A = B`   for the rotation matrix `R` to make matrix `A` diagonal (`B`)
rotate = np.matmul(diag_eig, np.linalg.inv(eig_vec))

v1 = eig_vec
v2 = np.matmul(rotate, eig_vec)
amp = 1.0

# Rotate data to principal axes coordinate system
zz = np.matmul(rotate, yy)
# Normalize
zz /= np.mean(np.linalg.norm(zz, axis=0))

names = ['x', 'y', 'z']

fig, axes_all = plt.subplots(figsize=[10, 7], ncols=3, nrows=2)
plt.subplots_adjust(wspace=0.25)

extr = [np.inf, -np.inf]

for mm, (data, axes) in enumerate(zip([yy, zz], axes_all)):
    
    for kk, ax in enumerate(axes):
        ii = (kk+1)%3
        jj = (kk+2)%3
        ax.set(xlabel=names[ii], ylabel=names[jj])

        # Draw all points, find extrema to plot square axes
        aa = data[ii, :]
        bb = data[jj, :]
        ax.scatter(aa, bb, color='0.5', alpha=0.1, rasterized=True)
        extr[0] = np.min([extr[0], aa.min(), bb.min()])
        extr[1] = np.max([extr[1], aa.max(), bb.max()])

        # Draw principal axes
        colors = ['r', 'b', 'g']
        for ll, cc in enumerate(colors):
            # Unrotated and rotated
            for vv, ls in zip([v1, v2], ['-', '--']):
                nn = np.linalg.norm(vv[:, ll])
                aa = [0.0, amp*vv[ll, ii]/nn]
                bb = [0.0, amp*vv[ll, jj]/nn]
                ax.plot(aa, bb, color=cc, ls=ls)


for idx, ax in np.ndenumerate(axes_all):
    ax.set(xlim=extr, ylim=extr)
    
fig.savefig('test.png')
plt.show()