<commit_before>from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from pylab import *

def _gen_labels(labels, names=None):
    if names is None:
        names = labels.names
    zips = [zip(names, l) for l in labels]
    new_labels = [', '.join(['{1}'.format(*m) for m in z]) for z in zips]
    return new_labels, names

def heatmap(data, xlabels=None, ylabels=None, title=None):
    fig, ax = plt.subplots()
    heatmap = ax.pcolormesh(data.values, cmap=plt.cm.RdYlGn)
    plt.colorbar(heatmap)


    xaxis = data.columns
    yaxis = data.index
    xlabels, xnames = _gen_labels(xaxis)
    ylabels, ynames = _gen_labels(yaxis)<commit_msg>BUG: heatmap works with np.nan, plots as white<commit_after>from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from pylab import *

def _gen_labels(labels, names=None):
    if names is None:
        names = labels.names
    zips = [zip(names, l) for l in labels]
    new_labels = [', '.join(['{1}'.format(*m) for m in z]) for z in zips]
    return new_labels, names

def heatmap(data, xlabels=None, ylabels=None, title=None):
    fig, ax = plt.subplots()

    values = data.values
    cmap = plt.cm.RdYlGn
    # plot np.nan as white
    cmap.set_bad('w',1.)
    masked_array = np.ma.array(values, mask=np.isnan(values))

    heatmap = ax.pcolormesh(masked_array, cmap=cmap)
    plt.colorbar(heatmap)


    xaxis = data.columns
    yaxis = data.index
    xlabels, xnames = _gen_labels(xaxis)
    ylabels, ynames = _gen_labels(yaxis)
{'fraction_non_alphanumeric': 0.11596091205211727, 'fraction_numerical': 0.001954397394136808, 'mean_word_length': 3.504398826979472, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1179651', 'n_tokens_mistral': 566, 'n_tokens_neox': 557, 'n_words': 150}