Added heatmap support