Quick plotting functions

quickplot.py

"""
Quick plotting for testing and preliminary data visualization
"""
from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np

__all__ = [
    'imshow',
]

# Default color map
dflt_cmap = LinearSegmentedColormap.from_list('rwb', ['red', 'white', 'black'])


def imshow(a, **kwargs):
    mm = np.abs(a).max()
    args = {
        'cmap': dflt_cmap,
        'vmin': -mm,
        'vmax': mm,
        'origin': 'lower',
    }
    args.update(kwargs)
    plt.imshow(a, **args)
    plt.colorbar()

    plt.gca().xaxis.get_major_locator().set_params(integer=True)
    plt.gca().yaxis.get_major_locator().set_params(integer=True)

    plt.gca().set_xticks(range(int(plt.gca().get_xlim()[0]),
                               int(plt.gca().get_xlim()[1]) + 1), minor=True)
    plt.gca().set_yticks(range(int(plt.gca().get_ylim()[0]),
                               int(plt.gca().get_ylim()[1]) + 1), minor=True)

Example usage:

import quickplot as qlt

qlt.imshow(np.random.rand(7, 6) * 5)