Added categorization for data_selector.Selection

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