Add support for custom edge overlays.

- Allow users to upload an edge overlays data json file. Each file can contain a
  list of edge overlays.
- User can specify color, edge width, label size for each overlay.
- User can enable/disable overlays from the dropdown, or delete overlay set.

PiperOrigin-RevId: 684974829

src/ui/src/components/visualizer/common/edge_overlays.ts

@@ -0,0 +1,82 @@
+/**
+ * @license
+ * Copyright 2024 The Model Explorer Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * 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.
+ * ==============================================================================
+ */
+
+import {TaskData, TaskType} from './task';
+
+/** The data for edge overlays. */
+export declare interface EdgeOverlaysData extends TaskData {
+  type: TaskType.EDGE_OVERLAYS;
+
+  /** The name of this set of overlays, for UI display purposes. */
+  name: string;
+
+  /** A list of edge overlays. */
+  overlays: EdgeOverlay[];
+}
+
+/** An edge overlay. */
+export declare interface EdgeOverlay {
+  /** The name displayed in the UI to identify this overlay. */
+  name: string;
+
+  /** The edges that define the overlay. */
+  edges: Edge[];
+
+  /**
+   * The color of the overlay edges.
+   *
+   * They are rendered in this color when any of the nodes in this overlay is
+   * selected.
+   */
+  edgeColor: string;
+
+  /** The width of the overlay edges. Default to 2. */
+  edgeWidth?: number;
+
+  /** The font size of the edge labels. Default to 7.5. */
+  edgeLabelFontSize?: number;
+}
+
+/** An edge in the overlay. */
+export declare interface Edge {
+  /** The id of the source node. Op node only. */
+  sourceNodeId: string;
+
+  /** The id of the target node. Op node only. */
+  targetNodeId: string;
+
+  /** Label shown on the edge. */
+  label?: string;
+}
+
+/** The processed edge overlays data. */
+export declare interface ProcessedEdgeOverlaysData extends EdgeOverlaysData {
+  /** A random id. */
+  id: string;
+
+  processedOverlays: ProcessedEdgeOverlay[];
+}
+
+/** The processed edge overlay. */
+export declare interface ProcessedEdgeOverlay extends EdgeOverlay {
+  /** A random id. */
+  id: string;
+
+  /** The set of node ids that are in this overlay. */
+  nodeIds: Set<string>;
+}

src/ui/src/components/visualizer/common/model_graph.ts

@@ -271,4 +271,9 @@ export declare interface ModelEdge {
 
   // The following are for webgl rendering.
   curvePoints?: Point[];
+
+  // The label of the edge.
+  //
+  // If set, it will be rendered on edge instead of tensor shape.
+  label?: string;
 }

src/ui/src/components/visualizer/common/sync_navigation.ts

@@ -19,7 +19,7 @@
 import {TaskData, TaskType} from './task';
 
 /** The data for navigation syncing. */
-export interface SyncNavigationData extends TaskData {
+export declare interface SyncNavigationData extends TaskData {
   type: TaskType.SYNC_NAVIGATION;
 
   mapping: SyncNavigationMapping;

src/ui/src/components/visualizer/common/task.ts

@@ -24,4 +24,5 @@ export declare interface TaskData {
 /** The type of a task. */
 export enum TaskType {
   SYNC_NAVIGATION = 'sync_navigation',
+  EDGE_OVERLAYS = 'edge_overlays',
 }

src/ui/src/components/visualizer/common/utils.ts

@@ -48,6 +48,7 @@ import {
   ProcessedNodeQuery,
   ProcessedNodeRegexQuery,
   ProcessedNodeStylerRule,
+  Rect,
   SearchMatch,
   SearchMatchType,
   SearchNodeType,
@@ -974,3 +975,81 @@ export function splitLabel(label: string): string[] {
 export function getMultiLineLabelExtraHeight(label: string): number {
   return (splitLabel(label).length - 1) * NODE_LABEL_LINE_HEIGHT;
 }
+
+/**
+ * Calculates the closest intersection points of a line (L) connecting
+ * the centers of two rectangles (rect1 and rect2) with the sides of these
+ * rectangles.
+ */
+export function getIntersectionPoints(rect1: Rect, rect2: Rect) {
+  // Function to calculate the center of a rectangle
+  function getCenter(rect: Rect) {
+    return {
+      x: rect.x + rect.width / 2,
+      y: rect.y + rect.height / 2,
+    };
+  }
+
+  // Function to calculate intersection between a line and a rectangle
+  function getIntersection(rect: Rect, center1: Point, center2: Point) {
+    // Line parameters
+    const dx = center2.x - center1.x;
+    const dy = center2.y - center1.y;
+
+    // Check for intersection with each of the four sides of the rectangle
+    let tMin = Number.MAX_VALUE;
+    let intersection: Point = {x: 0, y: 0};
+
+    // Left side (x = rect.x)
+    if (dx !== 0) {
+      const t = (rect.x - center1.x) / dx;
+      const y = center1.y + t * dy;
+      if (t >= 0 && y >= rect.y && y <= rect.y + rect.height && t < tMin) {
+        tMin = t;
+        intersection = {x: rect.x, y};
+      }
+    }
+
+    // Right side (x = rect.x + rect.width)
+    if (dx !== 0) {
+      const t = (rect.x + rect.width - center1.x) / dx;
+      const y = center1.y + t * dy;
+      if (t >= 0 && y >= rect.y && y <= rect.y + rect.height && t < tMin) {
+        tMin = t;
+        intersection = {x: rect.x + rect.width, y};
+      }
+    }
+
+    // Top side (y = rect.y)
+    if (dy !== 0) {
+      const t = (rect.y - center1.y) / dy;
+      const x = center1.x + t * dx;
+      if (t >= 0 && x >= rect.x && x <= rect.x + rect.width && t < tMin) {
+        tMin = t;
+        intersection = {x, y: rect.y};
+      }
+    }
+
+    // Bottom side (y = rect.y + rect.height)
+    if (dy !== 0) {
+      const t = (rect.y + rect.height - center1.y) / dy;
+      const x = center1.x + t * dx;
+      if (t >= 0 && x >= rect.x && x <= rect.x + rect.width && t < tMin) {
+        tMin = t;
+        intersection = {x, y: rect.y + rect.height};
+      }
+    }
+
+    return intersection;
+  }
+
+  // Get the centers of the rectangles
+  const center1 = getCenter(rect1);
+  const center2 = getCenter(rect2);
+
+  // Find the closest intersection point of the line with rect1 and rect2
+  const intersection1 = getIntersection(rect1, center1, center2);
+  const intersection2 = getIntersection(rect2, center2, center1);
+
+  return {intersection1, intersection2};
+}

src/ui/src/components/visualizer/common/visualizer_config.ts

@@ -16,6 +16,7 @@
  * ==============================================================================
  */
 
+import {EdgeOverlaysData} from './edge_overlays';
 import {SyncNavigationData} from './sync_navigation';
 import {NodeStylerRule, RendererType} from './types';
 
@@ -63,6 +64,12 @@ export declare interface VisualizerConfig {
   /** The data for navigation syncing. */
   syncNavigationData?: SyncNavigationData;
 
+  /** List of data for edge overlays that will be applied to the left pane. */
+  edgeOverlaysDataListLeftPane?: EdgeOverlaysData[];
+
+  /** List of data for edge overlays that will be applied to the right pane. */
+  edgeOverlaysDataListRightPane?: EdgeOverlaysData[];
+
   /**
    * Default graph renderer.
    *

src/ui/src/components/visualizer/edge_overlays_dropdown.ng.html

@@ -0,0 +1,97 @@
+<!--
+@license
+Copyright 2024 The Model Explorer Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+==============================================================================
+-->
+
+<div class="container"
+    [bubble]="help"
+    [overlaySize]="helpPopupSize"
+    [hoverDelayMs]="10">
+  <div class="mat-icon-container view"
+      [bubbleClick]="edgeOverlaysPopup"
+      [overlaySize]="edgeOverlaysPopupSize"
+      (opened)="opened=true"
+      (closed)="opened=false"
+      (click)="handleClickOnEdgeOverlaysButton()">
+    <mat-icon class="toolbar-icon">polyline</mat-icon>
+  </div>
+</div>
+
+<ng-template #help>
+  <div class="model-explorer-help-popup">
+    Show custom edge overlays on graph
+  </div>
+</ng-template>
+
+<ng-template #edgeOverlaysPopup>
+  <div class="model-explorer-edge-overlays-popup">
+    <div class="label">
+      <div>Edge overlays</div>
+      <div class="icon-container close" bubbleClose>
+        <mat-icon>close</mat-icon>
+      </div>
+    </div>
+
+    <!-- Loaded overlays -->
+    <div class="loaded-overlays-container">
+      @if (overlaysSets().length === 0) {
+        <div class="no-overlays-label">
+          No loaded edge overlays
+        </div>
+      } @else {
+        @for (overlaySet of overlaysSets(); track overlaySet.id) {
+          <div class="overlay-set-container">
+            <div class="overlay-set-label">
+              {{overlaySet.name}}
+              <div class="icon-container delete" (click)="handleDeleteOverlaySet(overlaySet)">
+                <mat-icon>delete</mat-icon>
+              </div>
+            </div>
+            @for (overlay of overlaySet.overlays; track overlay.id) {
+              <div class="overlay-item">
+                <label>
+                  <input type="checkbox" [checked]="overlay.selected"
+                      (change)="toggleOverlaySelection(overlay)"/>
+                  {{overlay.name}}
+                </label>
+                @if (overlay.selected) {
+                  <div class="view-label" (click)="handleClickViewOverlay(overlay)">
+                    View
+                  </div>
+                }
+              </div>
+            }
+          </div>
+        }
+      }
+    </div>
+
+    <!-- Buttons to load the json -->
+    <div class="upload-container">
+      <div class="description">Load from computer</div>
+      <button class="upload-json-file-button upload"
+          mat-flat-button color="primary"
+          (click)="input.click()">
+        Upload
+      </button>
+    </div>
+    <input class="upload-json-file-input"
+        type="file" #input
+        multiple
+        accept=".json"
+        (change)="handleClickUpload(input)">
+  </div>
+</ng-template>