hparams: fix the parallel coordinate layout (#4988)

Bug
---
Previously, we were setting the width of the container based on count of
all hparams and metrics without considering the visibility. This had
caused the UI to be in very awkward state where each visible hparams and
metrics are spaced a lot even if the selection count is quite low.

Fix
---
We layout the parallel coordinates based on number of visible
{hparam,metric}s. The change is a bit larger because:

- I had to add type information and refactor it a little for
  understanding the code.
- Undo the performance optimization where we do not mutate the layout
  unless important piece of information in the configuration has
  changed. To detect the old configuration vs. new configuration, we
  have to now remember an additional state in the component.

Fixes #3414.

Author: stephanwlee

tensorboard/plugins/hparams/tf_hparams_parallel_coords_plot/BUILD

@@ -17,8 +17,8 @@ tf_ts_library(
     deps = [
         "//tensorboard/components/polymer:legacy_element_mixin",
         "//tensorboard/plugins/hparams:types",
-        "//tensorboard/plugins/hparams/tf_hparams_query_pane:schema_d_ts",
         "//tensorboard/plugins/hparams/tf_hparams_session_group_values",
+        "//tensorboard/plugins/hparams/tf_hparams_types",
         "//tensorboard/plugins/hparams/tf_hparams_utils",
         "@npm//@polymer/decorators",
         "@npm//@polymer/polymer",

tensorboard/plugins/hparams/tf_hparams_parallel_coords_plot/axes.ts

@@ -18,7 +18,7 @@ import * as d3 from 'd3';
 import * as tf_hparams_parallel_coords_plot_interaction_manager from './interaction_manager';
 import * as tf_hparams_utils from '../tf_hparams_utils/tf-hparams-utils';
 import * as tf_hparams_parallel_coords_plot_utils from './utils';
-import * as tf_hparams_query_pane from '../tf_hparams_query_pane/schema.d';
+import {Schema} from '../tf_hparams_types/types';
 import * as tf_hparams_api from '../types';
 
 export enum ScaleType {
@@ -120,7 +120,7 @@ export class Axis {
    */
   public constructor(
     svgProps: tf_hparams_parallel_coords_plot_interaction_manager.SVGProperties,
-    schema: tf_hparams_query_pane.Schema,
+    schema: Schema,
     interactionManager: tf_hparams_parallel_coords_plot_interaction_manager.InteractionManager,
     colIndex: number
   ) {
@@ -362,7 +362,7 @@ export class Axis {
     return new AlwaysPassingBrushFilter();
   }
   private readonly _svgProps: tf_hparams_parallel_coords_plot_interaction_manager.SVGProperties;
-  private readonly _schema: tf_hparams_query_pane.Schema;
+  private readonly _schema: Schema;
   private readonly _interactionManager: tf_hparams_parallel_coords_plot_interaction_manager.InteractionManager;
   private readonly _colIndex: number;
   private _isDisplayed: boolean;
@@ -379,7 +379,7 @@ export class Axis {
 export class AxesCollection {
   public constructor(
     svgProps: tf_hparams_parallel_coords_plot_interaction_manager.SVGProperties,
-    schema: tf_hparams_query_pane.Schema,
+    schema: Schema,
     interactionManager: tf_hparams_parallel_coords_plot_interaction_manager.InteractionManager
   ) {
     this._svgProps = svgProps;
@@ -556,7 +556,7 @@ export class AxesCollection {
       );
   }
   private _svgProps: tf_hparams_parallel_coords_plot_interaction_manager.SVGProperties;
-  private _schema: tf_hparams_query_pane.Schema;
+  private _schema: Schema;
   private _axes: Axis[];
   /**
    * The current assignment of stationary positions to axes.

tensorboard/plugins/hparams/tf_hparams_parallel_coords_plot/interaction_manager.ts

@@ -20,7 +20,7 @@ import * as d3 from 'd3';
 import * as tf_hparams_utils from '../tf_hparams_utils/tf-hparams-utils';
 import {AxesCollection} from './axes';
 import {LinesCollection, LineType, SessionGroupHandle} from './lines';
-import * as tf_hparams_query_pane from '../tf_hparams_query_pane/schema.d';
+import * as tf_hparams_query_pane from '../tf_hparams_types/types';
 
 import * as tf_hparams_api from '../types';
 

tensorboard/plugins/hparams/tf_hparams_parallel_coords_plot/lines.ts

@@ -22,7 +22,7 @@ import * as tf_hparams_parallel_coords_plot_utils from './utils';
 import {AxesCollection} from './axes';
 import * as tf_hparams_parallel_coords_plot_interaction_manager from './interaction_manager';
 
-import * as tf_hparams_query_pane from '../tf_hparams_query_pane/schema.d';
+import * as tf_hparams_query_pane from '../tf_hparams_types/types';
 
 import * as tf_hparams_api from '../types';
 

tensorboard/plugins/hparams/tf_hparams_parallel_coords_plot/tf-hparams-parallel-coords-plot.ts

@@ -60,16 +60,27 @@ limitations under the License.
  *
  * See the individual class comments in the respective files for more details.
  */
-import {PolymerElement, html} from '@polymer/polymer';
 import {customElement, observe, property} from '@polymer/decorators';
-import * as _ from 'lodash';
+import {html, PolymerElement} from '@polymer/polymer';
 import * as d3 from 'd3';
+import * as _ from 'lodash';
 
-import * as tf_hparams_utils from '../tf_hparams_utils/tf-hparams-utils';
+import {LegacyElementMixin} from '../../../components/polymer/legacy_element_mixin';
 import '../tf_hparams_session_group_values/tf-hparams-session-group-values';
+import {HparamInfo, MetricInfo, Schema} from '../tf_hparams_types/types';
+import * as tf_hparams_utils from '../tf_hparams_utils/tf-hparams-utils';
 import * as tf_hparams_parallel_coords_plot_interaction_manager from './interaction_manager';
 
-import {LegacyElementMixin} from '../../../components/polymer/legacy_element_mixin';
+interface Option {
+  configuration: {
+    columnsVisibility: boolean[];
+    schema: Schema;
+    visibleSchema: {
+      hparamInfos: HparamInfo[];
+      metricInfos: MetricInfo[];
+    };
+  };
+}
 
 @customElement('tf-hparams-parallel-coords-plot')
 class TfHparamsParallelCoordsPlot extends LegacyElementMixin(PolymerElement) {
@@ -136,9 +147,13 @@ class TfHparamsParallelCoordsPlot extends LegacyElementMixin(PolymerElement) {
   // See the property description in tf-hparams-query-pane.html
   @property({type: Array})
   sessionGroups: any[];
+
   // See the description in tf-hparams-scale-and-color-controls.html
   @property({type: Object})
-  options: any;
+  options: Option;
+
+  private _prevOptions?: Option;
+
   // The last session group that was clicked on or null if no
   // session group was clicked on yet.
   /**
@@ -179,33 +194,37 @@ class TfHparamsParallelCoordsPlot extends LegacyElementMixin(PolymerElement) {
   // element. Defined in tf-hparams-parallel-coords-plot.ts.
   @property({type: Object})
   _interactionManager: any;
+
   @observe('options.*', 'sessionGroups.*')
   _optionsOrSessionGroupsChanged() {
     if (!this.options) {
       return;
     }
-    const configuration = this.options.configuration;
+
+    const {configuration: prevConfig} = this._prevOptions ?? {};
+    const {configuration: nextConfig} = this.options;
     // See if we need to redraw from scratch. We redraw from scratch if
     // this is initialization or if configuration.schema has changed.
     if (
       this._interactionManager === undefined ||
-      !_.isEqual(this._interactionManager.schema(), configuration.schema)
+      !_.isEqual(prevConfig.schema, nextConfig.schema) ||
+      !_.isEqual(prevConfig.columnsVisibility, nextConfig.columnsVisibility)
     ) {
       // Remove any pre-existing DOM children of our SVG.
       d3.select(this.$.svg as SVGElement)
         .selectAll('*')
         .remove();
       const svgProps = new tf_hparams_parallel_coords_plot_interaction_manager.SVGProperties(
         this.$.svg as HTMLElement,
-        tf_hparams_utils.numColumns(configuration.schema)
+        nextConfig.columnsVisibility.filter(Boolean).length
       );
       // Listen to DOM changes underneath this.$.svg, and apply local CSS
       // scoping rules so that our rules in the <style> section above
       // would apply.
       this.scopeSubtree(this.$.svg as SVGElement, true);
       this._interactionManager = new tf_hparams_parallel_coords_plot_interaction_manager.InteractionManager(
         svgProps,
-        configuration.schema,
+        nextConfig.schema,
         (sessionGroup) => this.closestSessionGroupChanged(sessionGroup),
         (sessionGroup) => this.selectedSessionGroupChanged(sessionGroup)
       );
@@ -216,6 +235,7 @@ class TfHparamsParallelCoordsPlot extends LegacyElementMixin(PolymerElement) {
       this._validSessionGroups
     );
     this.redrawCount++;
+    this._prevOptions = this.options;
   }
   closestSessionGroupChanged(sessionGroup) {
     this.closestSessionGroup = sessionGroup;

tensorboard/plugins/hparams/tf_hparams_query_pane/BUILD

@@ -5,13 +5,6 @@ package(default_visibility =
 
 licenses(["notice"])
 
-tf_ts_library(
-    name = "schema_d_ts",
-    srcs = [
-        "schema.d.ts",
-    ],
-)
-
 tf_ts_library(
     name = "tf_hparams_query_pane",
     srcs = [

tensorboard/plugins/hparams/tf_hparams_types/BUILD

@@ -0,0 +1,12 @@
+load("//tensorboard/defs:defs.bzl", "tf_ts_library")
+
+package(default_visibility = ["//tensorboard/plugins/hparams:__subpackages__"])
+
+licenses(["notice"])
+
+tf_ts_library(
+    name = "tf_hparams_types",
+    srcs = [
+        "types.d.ts",
+    ],
+)

…params/tf_hparams_query_pane/schema.d.ts …gins/hparams/tf_hparams_types/types.d.tstensorboard/plugins/hparams/tf_hparams_query_pane/schema.d.ts renamed to tensorboard/plugins/hparams/tf_hparams_types/types.d.ts tensorboard/plugins/hparams/tf_hparams_query_pane/schema.d.ts renamed to tensorboard/plugins/hparams/tf_hparams_types/types.d.ts

@@ -12,11 +12,30 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
+
+export interface HparamInfo {
+  description: string;
+  displayName: string;
+  name: string;
+  type:
+    | 'DATA_TYPE_UNSET'
+    | 'DATA_TYPE_STRING'
+    | 'DATA_TYPE_BOOL'
+    | 'DATA_TYPE_FLOAT64';
+}
+
+export interface MetricInfo {
+  datasetType: 'DATASET_UNKNOWN' | 'DATASET_TRAINING' | 'DATASET_VALIDATION';
+  description: string;
+  displayName: string;
+  name: {tag: string; group: string};
+}
+
 export interface Schema {
   hparamColumn: Array<{
-    hparamInfo: Object;
+    hparamInfo: HparamInfo;
   }>;
   metricColumn: Array<{
-    metricInfo: Object;
+    metricInfo: MetricInfo;
   }>;
 }