redash/client/app/visualizations/chart/plotly/prepareHeatmapData.js

import { map, max, uniq, sortBy, flatten, find } from 'lodash';
import { createNumberFormatter } from '@/lib/value-format';

const defaultColorScheme = [
  [0, '#356aff'],
  [0.14, '#4a7aff'],
  [0.28, '#5d87ff'],
  [0.42, '#7398ff'],
  [0.56, '#fb8c8c'],
  [0.71, '#ec6463'],
  [0.86, '#ec4949'],
  [1, '#e92827'],
];

function prepareSeries(series, options, additionalOptions) {
  const { colorScheme, formatNumber } = additionalOptions;

  const plotlySeries = {
    x: [],
    y: [],
    z: [],
    type: 'heatmap',
    name: '',
    colorscale: colorScheme,
  };

  plotlySeries.x = uniq(map(series.data, v => v.x));
  plotlySeries.y = uniq(map(series.data, v => v.y));

  if (options.sortX) {
    plotlySeries.x = sortBy(plotlySeries.x);
  }

  if (options.sortY) {
    plotlySeries.y = sortBy(plotlySeries.y);
  }

  if (options.reverseX) {
    plotlySeries.x.reverse();
  }

  if (options.reverseY) {
    plotlySeries.y.reverse();
  }

  const zMax = max(map(series.data, d => d.zVal));

  // Use text trace instead of default annotation for better performance
  const dataLabels = {
    x: [],
    y: [],
    mode: 'text',
    hoverinfo: 'skip',
    showlegend: false,
    text: [],
    textfont: {
      color: [],
    },
  };

  for (let i = 0; i < plotlySeries.y.length; i += 1) {
    const item = [];
    for (let j = 0; j < plotlySeries.x.length; j += 1) {
      const datum = find(
        series.data,
        { x: plotlySeries.x[j], y: plotlySeries.y[i] },
      );

      const zValue = datum && datum.zVal || 0;
      item.push(zValue);

      if (isFinite(zMax) && options.showDataLabels) {
        dataLabels.x.push(plotlySeries.x[j]);
        dataLabels.y.push(plotlySeries.y[i]);
        dataLabels.text.push(formatNumber(zValue));
        if (options.colorScheme && options.colorScheme === 'Custom...') {
          dataLabels.textfont.color.push('white');
        } else {
          dataLabels.textfont.color.push((zValue / zMax) < 0.25 ? 'white' : 'black');
        }
      }
    }
    plotlySeries.z.push(item);
  }

  if (isFinite(zMax) && options.showDataLabels) {
    return [plotlySeries, dataLabels];
  }
  return [plotlySeries];
}

export default function prepareHeatmapData(seriesList, options) {
  let colorScheme = [];

  if (!options.colorScheme) {
    colorScheme = defaultColorScheme;
  } else if (options.colorScheme === 'Custom...') {
    colorScheme = [[0, options.heatMinColor], [1, options.heatMaxColor]];
  } else {
    colorScheme = options.colorScheme;
  }

  const additionalOptions = {
    colorScheme,
    formatNumber: createNumberFormatter(options.numberFormat),
  };

  return flatten(map(seriesList, series => prepareSeries(series, options, additionalOptions)));
}