eigenangles/balanced_network/scripts/plot_eigenspectrum.py

import numpy as np
import scipy as sc
import argparse
from scipy.integrate import quad
import math
from types import SimpleNamespace
import itertools
import matplotlib.pyplot as plt
import seaborn as sns
import sys
from pathlib import Path
sys.path.append(str(Path.cwd().parents[0] / 'scripts'))
from eigenangle_test import EigenangleTest

def process_eigenvectors(vecs, lengths):
    sort_idx = np.argsort(lengths)
    vecs = vecs.T[sort_idx]
    for i, eva in enumerate(vecs):
        vecs[i] = eva * np.sign(eva[np.argmax(np.absolute(eva))])
        vecs[i] /= np.linalg.norm(eva)
    return vecs

def load_eigenspectrum(path_list, N):
    matrix_num = len(path_list)
    all_eigenvalues = np.empty(N*matrix_num, dtype=complex)
    all_eigenangles = np.array([], dtype=float)

    for i, path_a in enumerate(path_list):
        matrix = np.load(path_a)
        eigenvalues, vectors_a = sc.linalg.eig(matrix)
        all_eigenvalues[N*i:N*(i+1)] = eigenvalues

        vectors_a = process_eigenvectors(vectors_a, np.real(eigenvalues))

        for j, path_b in enumerate(path_list[i+1:]):
            matrix = np.load(path_b)
            eigenvalues , vectors_b = sc.linalg.eig(matrix)
            vectors_b = process_eigenvectors(vectors_b, np.real(eigenvalues))

            M = np.real(np.dot(vectors_a, np.conjugate(vectors_b).T))
            M = np.real(M)
            M[np.argwhere(M > 1)] = 1.

            all_eigenangles = np.append(all_eigenangles, np.arccos(np.diag(M)))
    return all_eigenvalues, all_eigenangles

def plot_eigenvalue_dist(test, eigenvalues, ax=None, color=None, width=.9,
                         bins=20, scale=1):
    hist, edges = np.histogram(np.real(eigenvalues), bins=bins, density=True)
    dx = np.diff(edges)[0]

    xvalues = np.linspace(-test.critical_radius, test.critical_radius, 200)
    norm = quad(test.eigenvalue_real_dist,
                -test.critical_radius, test.critical_radius)[0]
    yvalues = [np.nan_to_num(test.eigenvalue_real_dist(x))*scale/norm
              for x in xvalues]

    if ax is None:
        fig, ax = plt.subplots()

    ax.bar(edges[:-1]+dx/2, hist*scale, color=color, width=width*dx)
    ax.plot(xvalues, yvalues, 'k--')
    curve_area = np.nansum(yvalues)*np.diff(xvalues)[0]
    return ax

def plot_eigenangle_dist(test, eigenangles, ax=None, color=None,
                         width=.9, bins=20):
    edges = np.linspace(0, np.pi, bins)
    hist, edges = np.histogram(eigenangles, bins=edges, density=True)
    dx = np.diff(edges)[0]
    xvalues = edges[:-1] + dx/2

    if ax is None:
        fig, ax = plt.subplots()

    ax.bar(xvalues, hist, color=color, width=width*dx)

    edges = np.linspace(0, np.pi, bins*10)
    yvalues = [test.angle_dist(x) for x in edges]
    ax.plot(edges, yvalues, 'k--')
    return ax

def plot_eigenspectrum(test, eigenvalues, eigenangles, original=False):
    sns.set(style='ticks', context='talk')
    sns.set_color_codes('colorblind')
    mosaic = """
             AB
             """
    figsize = (15,7)
    fig, ax = plt.subplot_mosaic(mosaic, figsize=figsize,
                                 gridspec_kw=dict(hspace=.2, width_ratios=[1,1]))
    ax = SimpleNamespace(**ax)

    plot_eigenvalue_dist(eigenangle_test, np.real(eigenvalues), ax.A, scale=1,
                         color=sns.color_palette('deep')[0], width=.9, bins=40)
    ax.A.set_yticks([])
    ax.A.set_xlabel(r'real eigenvalue $\lambda_{\mathbb{R}}$')
    sns.despine(ax=ax.A, left=True, right=True, top=True)

    inset_size = .35
    iax = fig.add_axes((.38,.55,inset_size*figsize[1]/figsize[0],inset_size))
    iax.plot(np.real(eigenvalues), np.imag(eigenvalues), linestyle='None',
             color=sns.color_palette('deep')[0], marker='.', markersize=1)
    vmax = max([np.abs(np.real(eigenvalues)).max(),
                np.abs(np.imag(eigenvalues)).max()])
    iax.set_xlim((-vmax, vmax))
    iax.set_ylim((-vmax, vmax))
    iax.set_axis_off()

    bins = 300
    binwidth = np.pi/bins
    plot_eigenangle_dist(eigenangle_test, eigenangles, ax.B,
                         sns.color_palette('rocket')[2], width=.9, bins=bins,)
    ax.B.set_xlim((min(eigenangles)-binwidth, max(eigenangles)+binwidth))

    if original:
        ax.B.set_xlim((np.pi*11/24, np.pi*13/24))
        ax.B.set_xticks([np.pi*15/32, np.pi/2, np.pi*17/32])
        ax.B.set_xticklabels([r'$\frac{15}{32}\pi$', r'$\frac{1}{2}\pi$',
                              r'$\frac{17}{32}\pi$'])
    else:
        ax.B.set_yscale('log')
        ax.B.set_ylim((0.001,30))

    ax.B.set_yticks([])
    ax.B.set_xlabel(f'eigenangle $\phi$')
    sns.despine(ax=ax.B, left=True, right=True, top=True)

<<<<<<< HEAD
    return fig, ax
=======
    return fig
>>>>>>> refs/remotes/origin/synced/master

if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--input",     nargs='?', type=lambda s: s.split(' '))
    CLI.add_argument("--output",    nargs='?', type=Path)
    CLI.add_argument("--N",         nargs='?', type=int)
    CLI.add_argument("--f",         nargs='?', type=float)
    CLI.add_argument("--mu",        nargs='?', type=float)
    CLI.add_argument("--epsilon",   nargs='?', type=float)
    CLI.add_argument("--sigma_ex",  nargs='?', type=float)
    CLI.add_argument("--sigma_in",  nargs='?', type=float)
    args, unknown = CLI.parse_known_args()

    eigenangle_test = EigenangleTest(N=args.N,
                                     f=args.f,
                                     mu=args.mu,
                                     epsilon=args.epsilon,
                                     sigma_ex=args.sigma_ex,
                                     sigma_in=args.sigma_in,
                                     is_connectivity=True)

    eigenvalues, eigenangles = load_eigenspectrum(args.input, N=args.N)

<<<<<<< HEAD
    fig, ax = plot_eigenspectrum(eigenangle_test, eigenvalues, eigenangles,
                                 original=('original' in args.input[0]))

    fig.suptitle(f'{Path(args.output).name.strip(".pdf")}')
    for axi, panel in zip([ax.A, ax.B], ['A', 'B']):
        axi.text(s=panel, x=0, y=1, ha='right', va='center',
                transform=axi.transAxes, fontsize=20, fontweight='bold')

    fig.align_labels()

    plt.savefig(args.output, dpi=300, bbox_inches='tight')
=======
    fig = plot_eigenspectrum(eigenangle_test, eigenvalues, eigenangles,
                             original=('original' in args.input[0]))

    fig.suptitle(f'{Path(args.output).name.strip(".png")}')
    fig.align_labels()

    plt.savefig(args.output, bbox_inches='tight')
>>>>>>> refs/remotes/origin/synced/master