ENH: better heatmap

Fixes #73