Added categorization for data_selector.Selection (#1324)

Added categorization for data_selector.Selection

The utility divides the data, based on selection, into
tag prefix and tags. After those division, it does not make
assumption about how the data is used and returns list of
experiment/run pair that matches the criteria.

The change also populates tags per run per experiment as well as
tag's plugin information. A tag is created for a specific plugin and,
now, categorization utils take the pluginName to filter out
unrelated tags based on information we pipe from the backend.

Author: stephanwlee

tensorboard/components/tf_backend/type.ts

@@ -16,9 +16,18 @@ namespace tf_backend {
 
 export type Experiment = {id: number, name: string, startTime: number};
 
-// `id` is null when data source is logdir.
-export type Run = {id: number | null, name: string, startTime: number};
+export type Run = {
+  id: number | null,
+  name: string,
+  startTime: number,
+  tags: Tag[],
+};
 
-export type Tag = {id: number, name: string, displayName: string};
+export type Tag = {
+  id: number,
+  name: string,
+  displayName: string,
+  pluginName: string,
+};
 
 }  // namespace tf_backend

tensorboard/components/tf_categorization_utils/BUILD

@@ -18,6 +18,7 @@ tf_web_library(
     deps = [
         "//tensorboard/components/tf_backend",
         "//tensorboard/components/tf_dashboard_common",
+        "//tensorboard/components/tf_data_selector:type",
         "//tensorboard/components/tf_imports:lodash",
         "//tensorboard/components/tf_imports:polymer",
         "//tensorboard/components/tf_storage",

tensorboard/components/tf_categorization_utils/categorizationUtils.ts

@@ -28,10 +28,10 @@ export enum CategoryType {
   PREFIX_GROUP,
 }
 export interface PrefixGroupMetadata {
-  type: CategoryType.PREFIX_GROUP;
+  type: CategoryType;
 }
 export interface SearchResultsMetadata {
-  type: CategoryType.SEARCH_RESULTS;
+  type: CategoryType;
   validRegex: boolean;
   universalRegex: boolean;  // is the search query ".*"? ("(?:)" doesn't count)
 }
@@ -45,6 +45,18 @@ export interface Category<T> {
 export type TagCategory = Category<{tag: string, runs: string[]}>;
 export type RunTagCategory = Category<{tag: string, run: string}>;
 
+/**
+ * Organize data by tagPrefix, tag, then list of series which is comprised of
+ * an experiment and a run.
+ */
+export type SeriesCategory = Category<{
+  tag: string,
+  series: Array<{
+    experiment: string,
+    run: string,
+  }>,
+}>;
+
 export type RawCategory = Category<string>;  // Intermediate structure.
 
 /**
@@ -110,25 +122,103 @@ export function categorizeTags(
     query?: string): TagCategory[] {
   const tags = tf_backend.getTags(runToTag);
   const categories = categorize(tags, query);
-  const tagToRuns: {[tag: string]: string[]} = {};
-  tags.forEach(tag => {
-    tagToRuns[tag] = [];
-  });
-  selectedRuns.forEach(run => {
-    (runToTag[run] || []).forEach(tag => {
-      tagToRuns[tag].push(run);
-    });
-  });
+  const tagToRuns = createTagToRuns(_.pick(runToTag, selectedRuns));
+
   return categories.map(({name, metadata, items}) => ({
     name,
     metadata,
     items: items.map(tag => ({
       tag,
-      runs: tagToRuns[tag].slice(),
+      runs: tagToRuns.get(tag).slice(),
     })),
   }));
 }
 
+/**
+ * Creates grouping of the data based on selection from tf-data-selector. It
+ * groups data by prefixes of tag names and by tag names. Each group contains
+ * series, a tuple of experiment name and run name.
+ */
+export function categorizeSelection(
+    selection: tf_data_selector.Selection[], pluginName: string):
+    SeriesCategory[] {
+  const tagToSeries = new Map();
+  const searchTags = new Set();
+
+  selection.forEach(({experiment, runs, tagRegex}) => {
+    const runNames = runs.map(({name}) => name);
+    const selectedRunToTag = createRunToTagForPlugin(runs, pluginName);
+    const tagToSelectedRuns = createTagToRuns(selectedRunToTag);
+    const tags = tf_backend.getTags(selectedRunToTag);
+
+    const searchCategory = categorizeBySearchQuery(tags, tagRegex);
+    // list of matching tags.
+    searchCategory.items.forEach(tag => searchTags.add(tag));
+
+    // list of all tags that has selected runs.
+    tags.forEach(tag => {
+      const series = tagToSeries.get(tag) || [];
+      series.push(...tagToSelectedRuns.get(tag)
+          .map(run => ({experiment: experiment.name, run})));
+      tagToSeries.set(tag, series);
+    });
+  });
+
+  const searchCategory = {
+    name: selection.length == 1 ? selection[0].tagRegex : 'multi',
+    metadata: {
+      type: CategoryType.SEARCH_RESULTS,
+      validRegex: false,
+      universalRegex: false,
+    },
+    items: Array.from(searchTags)
+        .sort(vz_sorting.compareTagNames)
+        .map(tag => ({
+          tag,
+          series: tagToSeries.get(tag),
+        })),
+  };
+
+  // Organize the tag to items by prefix.
+  const prefixCategories = categorizeByPrefix(Array.from(tagToSeries.keys()))
+      .map(({name, metadata, items}) => ({
+        name,
+        metadata,
+        items: items.map(tag => ({
+          tag,
+          series: tagToSeries.get(tag),
+        })),
+      }));
+
+  return [
+    searchCategory,
+    ...prefixCategories,
+  ];
+}
+
+function createTagToRuns(runToTag: RunToTag): Map<string, string[]> {
+  const tagToRun = new Map();
+  Object.keys(runToTag).forEach(run => {
+    runToTag[run].forEach(tag => {
+      const runs = tagToRun.get(tag) || [];
+      runs.push(run);
+      tagToRun.set(tag, runs);
+    });
+  });
+  return tagToRun;
+}
+
+function createRunToTagForPlugin(runs: tf_backend.Run[], pluginName: string):
+    RunToTag {
+  const runToTag = {};
+  runs.forEach((run) => {
+    runToTag[run.name] = run.tags
+        .filter(tag => tag.pluginName == pluginName)
+        .map(({name}) => name);
+  })
+  return runToTag;
+}
+
 function compareTagRun(a, b: {tag: string, run: string}): number {
   const c = vz_sorting.compareTagNames(a.tag, b.tag);
   if (c != 0) {

tensorboard/components/tf_categorization_utils/test/categorizationUtilsTests.ts

@@ -14,7 +14,7 @@ limitations under the License.
 ==============================================================================*/
 namespace tf_categorization_utils {
 
-const assert = chai.assert;
+const {assert, expect} = chai;
 
 describe('categorizationUtils', () => {
   const {CategoryType} = tf_categorization_utils;
@@ -195,6 +195,201 @@ describe('categorizationUtils', () => {
     });
   });
 
+  describe('categorizeSelection', () => {
+    const {categorizeSelection} = tf_categorization_utils;
+
+    beforeEach(function() {
+      const tag1 = {
+        id: 1, pluginName: 'scalar',
+        name: 'tag1', displayName: 'tag1',
+
+      };
+      const tag2_1 = {
+        id: 2, pluginName: 'scalar',
+        name: 'tag2/subtag1', displayName: 'tag2/subtag1',
+
+      };
+      const tag2_2 = {
+        id: 3, pluginName: 'scalar',
+        name: 'tag2/subtag2', displayName: 'tag2/subtag2',
+
+      };
+      const tag3 = {
+        id: 4, pluginName: 'scalar',
+        name: 'tag3', displayName: 'tag3',
+      };
+      const tag4 = {
+        id: 5, pluginName: 'custom_scalar',
+        name: 'tag4', displayName: 'tag4',
+      };
+
+      this.run1 = {id: 1, name: 'run1', startTime: 10, tags: [tag1, tag4]};
+      this.run2 = {id: 2, name: 'run2', startTime: 5, tags: [tag2_1, tag2_2]};
+      this.run3 = {id: 3, name: 'run3', startTime: 0, tags: [tag2_1, tag3]};
+
+      this.experiment1 = {
+        experiment: {id: 1, name: 'exp1', startTime: 0},
+        runs: [this.run1, this.run2],
+        tagRegex: '',
+      };
+      this.experiment2 = {
+        experiment: {id: 2, name: 'exp2', startTime: 0},
+        runs: [this.run2, this.run3],
+        tagRegex: '(subtag1|tag3)',
+      };
+      this.experiment3 = {
+        experiment: {id: 3, name: 'exp3', startTime: 0},
+        runs: [this.run1, this.run2, this.run3],
+        tagRegex: 'junk',
+      };
+    });
+
+    it('merges the results of the query and the prefix groups', function() {
+      const result = categorizeSelection(
+          [this.experiment1], 'scalar');
+
+      expect(result).to.have.lengthOf(3);
+      expect(result[0]).to.have.property('metadata')
+          .that.has.property('type', CategoryType.SEARCH_RESULTS);
+
+      expect(result[1]).to.have.property('metadata')
+          .that.has.property('type', CategoryType.PREFIX_GROUP);
+      expect(result[2]).to.have.property('metadata')
+          .that.has.property('type', CategoryType.PREFIX_GROUP);
+    });
+
+    describe('search group', () => {
+      it('filters groups by tag with a tagRegex', function() {
+        const [searchResult] = categorizeSelection(
+            [this.experiment2], 'scalar');
+
+        // should match 'tag2/subtag1' and 'tag3'.
+        expect(searchResult).to.have.property('items')
+            .that.has.lengthOf(2);
+        expect(searchResult.items[0]).to.have.property('tag', 'tag2/subtag1');
+        expect(searchResult.items[1]).to.have.property('tag', 'tag3');
+      });
+
+      it('combines selection without tagRegex with one', function() {
+        const [searchResult] = categorizeSelection(
+            [this.experiment1, this.experiment2], 'scalar');
+
+        // should match 'tag1', 'tag2/subtag1', 'tag2/subtag2', and 'tag3'.
+        expect(searchResult).to.have.property('items')
+            .that.has.lengthOf(4);
+        expect(searchResult.items[0]).to.have.property('tag', 'tag1');
+        expect(searchResult.items[1]).to.have.property('tag', 'tag2/subtag1');
+        expect(searchResult.items[2]).to.have.property('tag', 'tag2/subtag2');
+        expect(searchResult.items[3]).to.have.property('tag', 'tag3');
+
+        expect(searchResult.items[1]).to.have.property('series')
+            .that.has.lengthOf(3)
+            .and.that.deep.equal([
+              {experiment: 'exp1', run: 'run2'},
+              {experiment: 'exp2', run: 'run2'},
+              {experiment: 'exp2', run: 'run3'},
+            ]);
+      });
+
+      it('sorts the tag by name', function() {
+        const [searchResult] = categorizeSelection(
+            [this.experiment2, this.experiment1], 'scalar');
+
+        // should match 'tag1', 'tag2/subtag1', 'tag2/subtag2', and 'tag3'.
+        expect(searchResult).to.have.property('items')
+            .that.has.lengthOf(4);
+        expect(searchResult.items[0]).to.have.property('tag', 'tag1');
+        expect(searchResult.items[1]).to.have.property('tag', 'tag2/subtag1');
+        expect(searchResult.items[2]).to.have.property('tag', 'tag2/subtag2');
+        expect(searchResult.items[3]).to.have.property('tag', 'tag3');
+      });
+
+      it('returns name `multi` when there are multiple selections', function() {
+        const [searchResult2] = categorizeSelection(
+            [this.experiment2], 'scalar');
+        expect(searchResult2).to.have.property('name', '(subtag1|tag3)');
+
+        const [searchResult1] = categorizeSelection(
+            [this.experiment1, this.experiment2], 'scalar');
+        expect(searchResult1).to.have.property('name', 'multi');
+      });
+
+      it('returns an empty array when tagRegex does not match any', function() {
+        const result = categorizeSelection(
+            [this.experiment3], 'custom_scalar');
+
+        expect(result).to.have.lengthOf(2);
+        expect(result[0]).to.have.property('items')
+            .that.has.lengthOf(0);
+      });
+    });
+
+    describe('prefix group', () => {
+      it('creates a group when a tag misses separator', function() {
+        const result = categorizeSelection(
+            [this.experiment1], 'scalar');
+
+        expect(result[1]).to.have.property('items')
+            .that.has.lengthOf(1);
+
+        expect(result[1]).to.have.property('name', 'tag1');
+        expect(result[1].items[0]).to.have.property('tag', 'tag1');
+        expect(result[1].items[0]).to.have.property('series')
+            .that.has.lengthOf(1);
+      });
+
+      it('creates a grouping when tag has a separator', function() {
+        const result = categorizeSelection(
+            [this.experiment1], 'scalar');
+
+        expect(result[2]).to.have.property('items')
+            .that.has.lengthOf(2);
+
+        expect(result[2]).to.have.property('name', 'tag2');
+        expect(result[2].items[0]).to.have.property('tag', 'tag2/subtag1');
+        expect(result[2].items[1]).to.have.property('tag', 'tag2/subtag2');
+        expect(result[2].items[0]).to.have.property('series')
+            .that.has.lengthOf(1);
+      });
+
+      it('creates a group with items with experiment and run', function() {
+        const result = categorizeSelection(
+            [this.experiment1], 'scalar');
+
+        expect(result[1].items[0]).to.have.property('series')
+            .that.has.lengthOf(1)
+            .and.that.deep.equal([{experiment: 'exp1', run: 'run1'}]);
+      });
+
+      it('creates distinct subitems when tags exactly match', function() {
+        const result = categorizeSelection(
+            [this.experiment2], 'scalar');
+
+        expect(result[1].items[0]).to.have.property('series')
+            .that.has.lengthOf(2)
+            .and.that.deep.equal([
+              {experiment: 'exp2', run: 'run2'},
+              {experiment: 'exp2', run: 'run3'},
+            ]);
+      });
+
+      it('filters out tags of a different pluguin', function() {
+        const result = categorizeSelection(
+            [this.experiment3], 'custom_scalar');
+
+        expect(result).to.have.lengthOf(2);
+        expect(result[1]).to.have.property('name', 'tag4');
+        expect(result[1]).to.have.property('items')
+            .that.has.lengthOf(1);
+        expect(result[1].items[0]).to.have.property('series')
+            .that.has.lengthOf(1)
+            .and.that.deep.equal([
+              {experiment: 'exp3', run: 'run1'},
+            ]);
+      });
+    });
+  });
+
 });
 
 } // namespace tf_categorization_utils