Last active
August 22, 2017 18:55
-
-
Save s-boardman/cef9675329a951e89e93 to your computer and use it in GitHub Desktop.
Heatmap functions for Pandas dataframes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""Plots a Pandas dataframe as a heatmap""" | |
import matplotlib as mpl | |
import matplotlib.pyplot as plt | |
def heatmap(df, | |
edgecolors='w', | |
cmap=mpl.cm.RdBu, | |
log=False,vmin=0,vmax=500): | |
width = len(df.columns)/4 | |
height = len(df.index)/4 | |
fig, ax = plt.subplots(figsize=(width,height)) | |
heatmap = ax.pcolor(df, | |
edgecolors=edgecolors, # put white lines between squares in heatmap | |
cmap=cmap, | |
vmin=vmin, # defaults to 0 | |
vmax=vmax, # defaults to 500 | |
norm=mpl.colors.LogNorm() if log else None) | |
ax.autoscale(tight=True) # get rid of whitespace in margins of heatmap | |
ax.set_aspect('equal') # ensure heatmap cells are square | |
ax.xaxis.set_ticks_position('top') # put column labels at the top | |
ax.tick_params(bottom='off', top='off', left='off', right='off') # turn off ticks | |
plt.yticks(np.arange(len(df.index)) + 0.5, df.index) | |
plt.xticks(np.arange(len(df.columns)) + 0.5, df.columns, rotation=90) | |
# ugliness from http://matplotlib.org/users/tight_layout_guide.html | |
from mpl_toolkits.axes_grid1 import make_axes_locatable | |
divider = make_axes_locatable(ax) | |
cax = divider.append_axes("right", "3%", pad="1%") | |
plt.colorbar(heatmap, cax=cax) | |
return fig | |
"""Binary Heatmap""" | |
import matplotlib as mpl | |
import matplotlib.pyplot as plt | |
def binary_heatmap(df): | |
df = df[::-1] # reverse df to put first row at top (last row at origin) | |
width = len(df.columns)/5 | |
height = len(df.index)/5 | |
fig, ax = plt.subplots(figsize=(width,height)) | |
heatmap = ax.pcolor(df, | |
edgecolors='k', # put black lines between squares in heatmap | |
cmap=mpl.cm.binary) # black/white colomarp | |
ax.autoscale(tight=True) # get rid of whitespace in margins of heatmap | |
ax.set_aspect('equal') # ensure heatmap cells are square | |
ax.xaxis.set_ticks_position('top') # put column labels at the top | |
ax.tick_params(bottom='off', top='off', left='off', right='off') # turn off ticks | |
plt.yticks(np.arange(len(df.index)) + 0.5, df.index) | |
plt.xticks(np.arange(len(df.columns)) + 0.5, df.columns, rotation=90) | |
plt.tight_layout() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Looks good, but I would change the vmin default to 0.001 or anything > 0 because:
/anaconda/lib/python3.5/site-packages/matplotlib/colors.py in call(self, value, clip)
1002 raise ValueError("minvalue must be less than or equal to maxvalue")
1003 elif vmin <= 0:
-> 1004 raise ValueError("values must all be positive")