eigenangles/scripts/plot_matrix.py

import neo
import quantities as pq
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
import argparse
from pathlib import Path

if __name__ == '__main__':
    CLI = argparse.ArgumentParser()
    CLI.add_argument("--input", nargs='?', type=Path)
    CLI.add_argument("--output", nargs='?', type=Path)
    args, unknown = CLI.parse_known_args()

    matrix = np.load(args.input)
    np.fill_diagonal(matrix, np.nan)

    cmap = sns.color_palette("vlag", as_cmap=True).copy()
    cmap.set_bad(color='white')

    fig, ax = plt.subplots(ncols=2, figsize=(10,10),
                           gridspec_kw={'width_ratios':[1,0.02]})
    sns.heatmap(matrix, cmap=cmap, center=0, square=True, ax=ax[0],
<<<<<<< HEAD
                cbar=True, cbar_ax=ax[1], rasterized=True)
=======
                cbar=True, cbar_ax=ax[1])
>>>>>>> refs/remotes/origin/synced/master
    ax[0].set_axis_off()

    measure = args.input.name.split('.')[0]
    ax[1].set_ylabel(measure)

    plt.savefig(args.output, bbox_inches='tight')