[tensor-widget] Add colormap selection and Ctrl/Alt/Shift+wheel zooming

tensorboard/components/tensor_widget/colormap-test.ts

@@ -17,36 +17,40 @@ limitations under the License.
 
 import {expect} from 'chai';
 
-import {GrayscaleColorMap} from './colormap';
+import {GrayscaleColorMap, JetColorMap} from './colormap';
 
 describe('GrayscaleColorMap', () => {
   it('max < min causes constructor error', () => {
     const min = 3;
     const max = 2;
-    expect(() => new GrayscaleColorMap(min, max)).to.throw(/max.*<.*min/);
+    expect(() => new GrayscaleColorMap({min, max})).to.throw(/max.*<.*min/);
   });
 
   it('NaN or Infinity min or max causes constructor error', () => {
-    expect(() => new GrayscaleColorMap(0, Infinity)).to.throw(
+    expect(() => new GrayscaleColorMap({min: 0, max: Infinity})).to.throw(
       /max.*not finite/
     );
-    expect(() => new GrayscaleColorMap(0, -Infinity)).to.throw(
+    expect(() => new GrayscaleColorMap({min: 0, max: -Infinity})).to.throw(
       /max.*not finite/
     );
-    expect(() => new GrayscaleColorMap(0, NaN)).to.throw(/max.*not finite/);
-    expect(() => new GrayscaleColorMap(Infinity, 0)).to.throw(
+    expect(() => new GrayscaleColorMap({min: 0, max: NaN})).to.throw(
+      /max.*not finite/
+    );
+    expect(() => new GrayscaleColorMap({min: Infinity, max: 0})).to.throw(
       /min.*not finite/
     );
-    expect(() => new GrayscaleColorMap(-Infinity, 0)).to.throw(
+    expect(() => new GrayscaleColorMap({min: -Infinity, max: 0})).to.throw(
+      /min.*not finite/
+    );
+    expect(() => new GrayscaleColorMap({min: NaN, max: 0})).to.throw(
       /min.*not finite/
     );
-    expect(() => new GrayscaleColorMap(NaN, 0)).to.throw(/min.*not finite/);
   });
 
   it('max > min, finite values', () => {
     const min = 0;
     const max = 10;
-    const colormap = new GrayscaleColorMap(min, max);
+    const colormap = new GrayscaleColorMap({min, max});
     expect(colormap.getRGB(0)).to.eql([0, 0, 0]);
     expect(colormap.getRGB(5)).to.eql([127.5, 127.5, 127.5]);
     expect(colormap.getRGB(10)).to.eql([255, 255, 255]);
@@ -58,7 +62,7 @@ describe('GrayscaleColorMap', () => {
   it('max > min, non-finite values', () => {
     const min = 0;
     const max = 10;
-    const colormap = new GrayscaleColorMap(min, max);
+    const colormap = new GrayscaleColorMap({min, max});
     expect(colormap.getRGB(NaN)).to.eql([255, 0, 0]);
     expect(colormap.getRGB(-Infinity)).to.eql([255, 255 / 2, 0]);
     expect(colormap.getRGB(Infinity)).to.eql([0, 0, 255]);
@@ -67,7 +71,7 @@ describe('GrayscaleColorMap', () => {
   it('max === min, non-finite values', () => {
     const min = -3.2;
     const max = -3.2;
-    const colormap = new GrayscaleColorMap(min, max);
+    const colormap = new GrayscaleColorMap({min, max});
     expect(colormap.getRGB(-32)).to.eql([127.5, 127.5, 127.5]);
     expect(colormap.getRGB(-3.2)).to.eql([127.5, 127.5, 127.5]);
     expect(colormap.getRGB(0)).to.eql([127.5, 127.5, 127.5]);
@@ -77,3 +81,61 @@ describe('GrayscaleColorMap', () => {
     expect(colormap.getRGB(Infinity)).to.eql([0, 0, 255]);
   });
 });
+
+describe('JetColormap', () => {
+  it('max < min causes constructor error', () => {
+    const min = 3;
+    const max = 2;
+    expect(() => new JetColorMap({min, max})).to.throw(/max.*<.*min/);
+  });
+
+  it('NaN or Infinity min or max causes constructor error', () => {
+    expect(() => new JetColorMap({min: 0, max: Infinity})).to.throw(
+      /max.*not finite/
+    );
+    expect(() => new JetColorMap({min: 0, max: -Infinity})).to.throw(
+      /max.*not finite/
+    );
+    expect(() => new JetColorMap({min: 0, max: NaN})).to.throw(
+      /max.*not finite/
+    );
+    expect(() => new JetColorMap({min: Infinity, max: 0})).to.throw(
+      /min.*not finite/
+    );
+    expect(() => new JetColorMap({min: -Infinity, max: 0})).to.throw(
+      /min.*not finite/
+    );
+    expect(() => new JetColorMap({min: NaN, max: 0})).to.throw(
+      /min.*not finite/
+    );
+  });
+
+  it('max > min, finite values', () => {
+    const min = 0;
+    const max = 10;
+    const colormap = new JetColorMap({min, max});
+    expect(colormap.getRGB(0)).to.eql([0, 0, 255]);
+    expect(colormap.getRGB(5)).to.eql([127.5, 255, 127.5]);
+    expect(colormap.getRGB(10)).to.eql([255, 0, 0]);
+    // Over-limits.
+    expect(colormap.getRGB(-100)).to.eql([0, 0, 255]);
+    expect(colormap.getRGB(500)).to.eql([255, 0, 0]);
+  });
+
+  it('max > min, non-finite values', () => {
+    const min = 0;
+    const max = 10;
+    const colormap = new JetColorMap({min, max});
+    expect(colormap.getRGB(NaN)).to.eql([255 * 0.25, 255 * 0.25, 255 * 0.25]);
+    expect(colormap.getRGB(-Infinity)).to.eql([
+      255 * 0.5,
+      255 * 0.5,
+      255 * 0.5,
+    ]);
+    expect(colormap.getRGB(Infinity)).to.eql([
+      255 * 0.75,
+      255 * 0.75,
+      255 * 0.75,
+    ]);
+  });
+});

tensorboard/components/tensor_widget/colormap.ts

@@ -17,6 +17,24 @@ limitations under the License.
 
 const MAX_RGB = 255;
 
+/** Configuration for a colormap. */
+export interface ColorMapConfig {
+  /**
+   * Minimum value that the color map can map to without clipping.
+   * Must be a finite value.
+   */
+  min: number;
+
+  /**
+   * Minimum value that the color map can map to without clipping.
+   * Must be a finite value and be `>= min`.
+   *
+   * In the case of `max === min`, all finite values mapped to the
+   * midpoint of the color scale.
+   */
+  max: number;
+}
+
 /**
  * Abstract base class for colormap.
  *
@@ -28,15 +46,15 @@ export abstract class ColorMap {
    * @param min Minimum. Must be a finite value.
    * @param max Maximum. Must be finite and >= `min`.
    */
-  constructor(protected readonly min: number, protected readonly max: number) {
-    if (!isFinite(min)) {
-      throw new Error(`min value (${min}) is not finite`);
+  constructor(protected config: ColorMapConfig) {
+    if (!isFinite(config.min)) {
+      throw new Error(`min value (${config.min}) is not finite`);
     }
-    if (!isFinite(max)) {
-      throw new Error(`max value (${max}) is not finite`);
+    if (!isFinite(config.max)) {
+      throw new Error(`max value (${config.max}) is not finite`);
     }
-    if (max < min) {
-      throw new Error(`max (${max}) is < min (${min})`);
+    if (config.max < config.min) {
+      throw new Error(`max (${config.max}) is < min (${config.min})`);
     }
   }
 
@@ -47,6 +65,67 @@ export abstract class ColorMap {
    *   The range of RGB values is [0, 255].
    */
   abstract getRGB(value: number): [number, number, number];
+
+  /**
+   * Render the colormap as a horizontal scale, while highlighting a specifc
+   * value.
+   *
+   * The value is highlighted if and only if it is `>= this.config.min` and
+   * `<= this.config.max`.
+   *
+   * @param canvas The canvas on which
+   * @param value The number to highlight (optional).
+   */
+  render(canvas: HTMLCanvasElement, value?: number) {
+    if (this.config.min === this.config.max) {
+      return;
+    }
+    const context = canvas.getContext('2d');
+    if (context == null) {
+      return;
+    }
+    const steps = 100;
+    const cellWidth = canvas.width / steps;
+    const height = canvas.height;
+    const verticalMargin = 0.2;
+    const barHeight = height * (1 - 2 * verticalMargin);
+    for (let i = 0; i < steps; ++i) {
+      const value =
+        (this.config.max - this.config.min) * (i / steps) + this.config.min;
+      const x = cellWidth * i;
+      const y = height * verticalMargin;
+      const [r, g, b] = this.getRGB(value);
+      context.beginPath();
+      context.fillStyle = `rgba(${r}, ${g}, ${b}, 1)`;
+      context.fillRect(x, y, cellWidth, barHeight);
+      context.stroke();
+    }
+
+    if (value != null) {
+      const tickWidth = 8;
+      if (value >= this.config.min && value <= this.config.max) {
+        // Highlight the relative position of `value` along the color scale.
+        const tickX =
+          ((value - this.config.min) / (this.config.max - this.config.min)) *
+          canvas.width;
+
+        // Draw the triangle on the top.
+        context.beginPath();
+        context.fillStyle = 'rgba(0, 0, 0, 1)';
+        context.moveTo(tickX, verticalMargin * height);
+        context.lineTo(tickX - tickWidth / 2, 0);
+        context.lineTo(tickX + tickWidth / 2, 0);
+        context.fill();
+
+        // Draw the triangle on the bottom.
+        context.beginPath();
+        context.moveTo(tickX, (1 - verticalMargin) * height);
+        context.lineTo(tickX - tickWidth / 2, height);
+        context.lineTo(tickX + tickWidth / 2, height);
+        context.fill();
+      }
+    }
+  }
 }
 
 /**
@@ -68,13 +147,52 @@ export class GrayscaleColorMap extends ColorMap {
         return [MAX_RGB, MAX_RGB / 2, 0];
       }
     }
-    let relativeValue =
-      this.min === this.max ? 0.5 : (value - this.min) / (this.max - this.min);
-    relativeValue = Math.max(Math.min(relativeValue, 1), 0);
-    return [
-      MAX_RGB * relativeValue,
-      MAX_RGB * relativeValue,
-      MAX_RGB * relativeValue,
-    ];
+    let relValue =
+      this.config.min === this.config.max
+        ? 0.5
+        : (value - this.config.min) / (this.config.max - this.config.min);
+    relValue = Math.max(Math.min(relValue, 1), 0);
+    return [MAX_RGB * relValue, MAX_RGB * relValue, MAX_RGB * relValue];
+  }
+}
+
+export class JetColorMap extends ColorMap {
+  getRGB(value: number): [number, number, number] {
+    if (isNaN(value)) {
+      // NaN.
+      return [MAX_RGB * 0.25, MAX_RGB * 0.25, MAX_RGB * 0.25];
+    } else if (!isFinite(value)) {
+      if (value < 0) {
+        // -Infinity.
+        return [MAX_RGB * 0.5, MAX_RGB * 0.5, MAX_RGB * 0.5];
+      } else {
+        // +Infinity.
+        return [MAX_RGB * 0.75, MAX_RGB * 0.75, MAX_RGB * 0.75];
+      }
+    }
+
+    let relR = 0;
+    let relG = 0;
+    let relB = 0;
+    const lim0 = 0.35;
+    const lim1 = 0.65;
+
+    let relValue =
+      this.config.min === this.config.max
+        ? 0.5
+        : (value - this.config.min) / (this.config.max - this.config.min);
+    relValue = Math.max(Math.min(relValue, 1), 0);
+    if (relValue <= lim0) {
+      relG = relValue / lim0;
+      relB = 1;
+    } else if (relValue > lim0 && relValue <= lim1) {
+      relR = (relValue - lim0) / (lim1 - lim0);
+      relG = 1;
+      relB = (lim1 - relValue) / (lim1 - lim0);
+    } else if (relValue > lim1) {
+      relR = 1;
+      relG = (1 - relValue) / (1 - lim1);
+    }
+    return [relR * MAX_RGB, relG * MAX_RGB, relB * MAX_RGB];
   }
 }

tensorboard/components/tensor_widget/menu.ts

@@ -21,6 +21,13 @@ type EventCallback = (event: Event) => void | Promise<void>;
 export interface MenuItemConfig {
   /** Caption displayed on the menu item. */
   caption: string;
+
+  /**
+   * A function that determines whether menu item is currently enabled.
+   *
+   * If not provided, the menu item will always be enabled.
+   */
+  isEnabled?: () => boolean;
 }
 
 /**
@@ -29,13 +36,6 @@ export interface MenuItemConfig {
 export interface SingleActionMenuItemConfig extends MenuItemConfig {
   /** The callback that gets called when the menu item is clicked. */
   callback: EventCallback;
-
-  /**
-   * A function that determines whether menu item is currently enabled.
-   *
-   * If not provided, the menu item will always be enabled.
-   */
-  isEnabled?: () => boolean;
 }
 
 export interface ChoiceMenuItemConfig extends MenuItemConfig {
@@ -267,12 +267,9 @@ export class Menu {
         // This is a single-command item.
         const singleActionConfig = item as SingleActionMenuItemConfig;
         outerItemConfig.onClick = singleActionConfig.callback;
-        if (
-          singleActionConfig.isEnabled != null &&
-          !singleActionConfig.isEnabled()
-        ) {
-          outerItemConfig.disabled = true;
-        }
+      }
+      if (item.isEnabled != null && !item.isEnabled()) {
+        outerItemConfig.disabled = true;
       }
       outerItemConfigs.push(outerItemConfig);
     });