[tfjs-vis] Support custom tickLabels with duplicate values in heatmaps

tfjs-vis/demos/api/index.html

@@ -657,6 +657,44 @@ <h3>render.heatmap(data, container, opts)</h3>
           });
         }
       </script>
+
+      <p>
+        Repeated labels are supported
+      </p>
+      <div id='heatmap-cont5'></div>
+
+      <script class='show-script'>
+        {
+          const data = {
+            values: tf.tensor2d([
+              [10, 2, 3, 4, 5, 6],
+              [1, 20, 3, 4, 5, 6],
+              [1, 2, 30, 4, 5, 6],
+              [1, 2, 3, 40, 5, 6],
+              [1, 2, 3, 4, 50, 6],
+              [1, 2, 3, 4, 5, 60],
+            ]),
+            xTickLabels: ['a', 'token', 'in', 'a', 'sequence', 'repeat'],
+            yTickLabels: ['a', 'token', 'in', 'a', 'sequence', 'repeat'],
+          }
+
+          // Render to visor
+          const surface = tfvis.visor().surface({ name: 'Heatmap', tab: 'Tensor Charts' });
+          tfvis.render.heatmap(surface, data, {
+            xLabel: 'Thing',
+            yLabel: 'Property',
+          });
+
+          // Render to page
+          const container = document.getElementById('heatmap-cont5');
+          tfvis.render.heatmap(container, data, {
+            width: 500,
+            height: 500,
+            xLabel: 'Thing',
+            yLabel: 'Property',
+          });
+        }
+      </script>
     </section>
 
   </article>

tfjs-vis/src/render/confusion_matrix.ts

@@ -95,14 +95,15 @@ export async function confusionMatrix(
         values.push({
           label,
           prediction,
-          diagCount: count,
+          count,
           noFill: true,
         });
       } else {
         values.push({
           label,
           prediction,
           count,
+          scaleCount: count,
         });
         // When not shading the diagonal we want to check if there is a non
         // zero value. If all values are zero we will not color them as the
@@ -122,7 +123,7 @@ export async function confusionMatrix(
     for (const val of values) {
       if (val.noFill === true) {
         val.noFill = false;
-        val.count = val.diagCount;
+        val.scaleCount = val.count;
       }
     }
   }
@@ -171,45 +172,56 @@ export async function confusionMatrix(
     'layer': [
       {
         // The matrix
+        'transform': [
+          {'filter': 'datum.noFill != true'},
+        ],
         'mark': {
           'type': 'rect',
         },
         'encoding': {
-          'fill': {
-            'condition': {
-              'test': 'datum["noFill"] == true',
-              'value': 'white',
-            },
-            'field': 'count',
+          'color': {
+            'field': 'scaleCount',
             'type': 'quantitative',
             'scale': {'range': ['#f7fbff', '#4292c6']},
           },
-          'tooltip': {
-            'condition': {
-              'test': 'datum["noFill"] == true',
-              'field': 'diagCount',
-              'type': 'nominal',
-            },
-            'field': 'count',
-            'type': 'nominal',
-          }
+          'tooltip': [
+            {'field': 'label', 'type': 'nominal'},
+            {'field': 'prediction', 'type': 'nominal'},
+            {'field': 'count', 'type': 'quantitative'},
+          ]
         },
-
       },
     ]
   };
 
+  if (options.shadeDiagonal === false) {
+    spec.layer.push(
+        {
+          // render unfilled rects for the diagonal
+          'transform': [
+            {'filter': 'datum.noFill == true'},
+          ],
+          'mark': {
+            'type': 'rect',
+            'fill': 'white',
+          },
+          'encoding': {
+            'tooltip': [
+              {'field': 'label', 'type': 'nominal'},
+              {'field': 'prediction', 'type': 'nominal'},
+              {'field': 'count', 'type': 'quantitative'},
+            ]
+          },
+        },
+    );
+  }
+
   if (options.showTextOverlay) {
     spec.layer.push({
       // The text labels
       'mark': {'type': 'text', 'baseline': 'middle'},
       'encoding': {
         'text': {
-          'condition': {
-            'test': 'datum["noFill"] == true',
-            'field': 'diagCount',
-            'type': 'nominal',
-          },
           'field': 'count',
           'type': 'nominal',
         },
@@ -218,7 +230,6 @@ export async function confusionMatrix(
   }
 
   await embed(drawArea, spec, embedOpts);
-  return Promise.resolve();
 }
 
 const defaultOpts: ConfusionMatrixOptions = {
@@ -235,7 +246,9 @@ const defaultOpts: ConfusionMatrixOptions = {
 interface MatrixEntry {
   label: string;
   prediction: string;
-  count?: number;
-  diagCount?: number;
+  // The displayed count
+  count: number;
+  // The count values used to compute the color scale
+  scaleCount?: number;
   noFill?: boolean;
 }