main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import rcParams
from mpl_toolkits.mplot3d import Axes3D

parser = argparse.ArgumentParser()
parser.add_argument('--dim', '-d', type=int, default=2)
parser.add_argument('--num_points', '-p', type=int, default=1000)
parser.add_argument('--trials', '-n', type=int, default=100)
parser.add_argument('--save', '-s', action='store_true')
args = parser.parse_args()

D = args.dim
N = args.num_points
NUM_TRIALS = args.trials


def sphere(n=N, d=D):
    s = np.random.normal(size=(n, d))
    s /= np.linalg.norm(s, axis=1)[:, np.newaxis]
    return s


def m(d=D):
    '''create square matrix of specified dim `d`.'''
    return np.random.random(size=(d, d))


if __name__ == '__main__':

    # make points smaller
    rcParams.update({'lines.markersize': 3})

    for i in range(NUM_TRIALS):

        f = m()
        s = sphere()

        eigvals, eigvecs = np.linalg.eig(f)

        # Should be a non-degenerate ellipse.
        image = np.array([f @ v for v in s])
        fig = plt.figure()

        if D == 2:
            # fig.gca().set_aspect('equal', adjustable='box')
            ax = fig.add_subplot(
                111,
                aspect='equal',
                adjustable='box',
            )

            ax.scatter(
                s[:, 0],
                s[:, 1],
                s=1,
            )

            ax.scatter(
                image[:, 0],
                image[:, 1],
                s=3,
            )

            ax.axhline(0, color='black')
            ax.axvline(0, color='black')

            ax.quiver(
                [0, 0],
                [0, 0],
                eigvecs[0],
                eigvecs[1],
                angles='xy',
                scale_units='xy',
                scale=1,
            )

        elif D == 3:
            ax = fig.add_subplot(
                111,
                adjustable='box',
                aspect='equal',
                projection='3d',
            )

            ax.scatter(
                s[:, 0],
                s[:, 1],
                s[:, 2],
                s=1,
            )

            ax.scatter(
                image[:, 0],
                image[:, 1],
                image[:, 2],
                s=3,
            )

            ax.quiver(
                [0, 0, 0],
                [0, 0, 0],
                [0, 0, 0],
                eigvecs[0],
                eigvecs[1],
                eigvecs[2],
            )

        else:
            print('Only 2D or 3D.', file=sys.stderr)

        if args.save:

            os.makedirs('assets/2d', exist_ok=True)
            os.makedirs('assets/3d', exist_ok=True)

            fig.savefig(
                f'assets/{D}d/{i}.svg',
                bbox_inches='tight',
                dpi=100,
                format='svg',
            )
            np.savetxt(
                f'assets/{D}d/{i}.vecs.txt',
                eigvecs,
            )
            np.savetxt(
                f'assets/{D}d/{i}.vals.txt',
                eigvals,
            )
        else:
            plt.show()

        plt.close('all')