path: root/tensorflow/tensorboard/components/tf-graph-common/lib/render.ts

/// <reference path="graph.ts" />
/// <reference path="hierarchy.ts" />

/**
 * Package for the Render Hierarchy for TensorFlow graph.
 */

module tf.graph.render {

/**
 * Color parameters for node encoding.
 * @type {Object}
 */
export let MetanodeColors = {
  SATURATION: 0.6,
  LIGHTNESS: 0.85,
  /**
   * Neutral color to use when the node is expanded (used when coloring by
   * compute time, memory and device).
   */
  EXPANDED_COLOR: "#f0f0f0",
  /**
   * Standard hue values for node color palette.
   */
  HUES: [220, 100, 180, 40, 20, 340, 260, 300, 140, 60],
  STRUCTURE_PALETTE: function(id: number, lightened? : boolean) {
    // The code below is a flexible way to computationally create a set
    // of colors that go well together.
    let hues = MetanodeColors.HUES;
    let n = hues.length;
    let hue = hues[id % n];
    let m = Math.sin(hue * Math.PI / 360);
    let sat = lightened ? 30 : 90 - 60 * m;
    let light = lightened ? 95 : 80;
    return d3.hsl(hue, .01 * sat, .01 * light).toString();
  },
  DEVICE_PALETTE: function (index: number): string {
    return MetanodeColors.STRUCTURE_PALETTE(index);
  },
  UNKNOWN: "#eee",
  GRADIENT_OUTLINE: "#888"
};

/**
 * Parameters that affect how the graph is rendered on the screen.
 */
interface RenderGraphParams {
  /**
   * Whether to extract high degree nodes from the core part of the graph.
   */
  enableExtraction: boolean;
  /**
   * Maximum in-degree that a node can have without being considered as
   * high in-degree node.
   */
  maxInDegree: number;
  /**
   * Maximum out-degree that a node can have without being considered as
   * high out-degree node.
   */
  maxOutDegree: number;
  /**
   * Maximum number of control edges a node can have before they aren't
   * displayed.
   */
  maxControlDegree: number;
  /**
   * Types patterns for predefined out-extract nodes, which are
   * sink-like nodes that will be extracted from the main graph.
   */
  outExtractTypes: string[];

  /**
   * Types patterns for predefined in-extract nodes, which are
   * source-like nodes that will be extracted from the main graph.
   */
  inExtractTypes: string[];

  /**
   * When removing edges from a high degree node, remove all of its edges if
   * detachAllEdgesForHighDegree is true.  Otherwise remove all in-edges if
   * the node has high in-degree, or all out-edges if the node has high
   * out-degree.
   */
  detachAllEdgesForHighDegree: boolean;

  /**
   * After extracting high in/out degree nodes and predefined
   * source-like/sink-like, extract isolated nodes to the side
   * if this extractIsolatedNodesWithAnnotationsOnOneSide is true.
   */
  extractIsolatedNodesWithAnnotationsOnOneSide: boolean;

  /**
   * Whether to add bridge nodes and edges to the core when building the
   * subhierarchy of an expanded metanode. See buildSubhierarchy().
   */
  enableBridgegraph: boolean;

  /**
   * 2 colors, for the minimum and maximum value respectively, whenever we
   * have a gradient scale.
   */
  minMaxColors: string[];

  /**
   * Maximum number of annotations to be displayed on a node.
   */
  maxAnnotations: number;
}

/**
 * Stores the rendering information, such as x and y coordinates,
 * for each node in the graph.
 */
export class RenderGraphInformation {
  private hierarchy: hierarchy.Hierarchy;
  private index: {[nodeName: string]: RenderNodeInformation};
  private params: RenderGraphParams;
  private deviceColorMap: d3.scale.Ordinal<string, string>;
  private memoryUsageScale: d3.scale.Linear<string, string>;
  private computeTimeScale: d3.scale.Linear<string, string>;
  // Since the rendering information for each node is constructed lazily,
  // upon node's expansion by the user, we keep a map between the node's name and
  // whether the rendering information was already constructed for that node.
  private hasSubhierarchy: {[nodeName: string]: boolean};
  root: RenderGroupNodeInformation;

  constructor(hierarchy: hierarchy.Hierarchy, params: RenderGraphParams) {
    this.hierarchy = hierarchy;
    this.index = {};
    this.deviceColorMap = d3.scale.ordinal<string>()
        .domain(hierarchy.devices)
        .range(_.map(d3.range(hierarchy.devices.length),
                     MetanodeColors.DEVICE_PALETTE));

    let topLevelGraph = hierarchy.root.metagraph;
    // Find the maximum and minimum memory usage.
    let memoryExtent = d3.extent(topLevelGraph.nodes(),
        (nodeName, index) => {
      let node = topLevelGraph.node(nodeName);
      // Some ops don't have stats at all.
      if (node.stats != null) {
        return node.stats.totalBytes;
      }
    });
    this.memoryUsageScale = d3.scale.linear<string, string>()
        .domain(memoryExtent)
        .range(params.minMaxColors);

    // Find also the minimum and maximum compute time.
    let computeTimeExtent = d3.extent(topLevelGraph.nodes(), (nodeName, index) => {
      let node = topLevelGraph.node(nodeName);
      // Some ops don't have stats at all.
      if (node.stats != null) {
        return node.stats.totalMicros;
      }
    });
    this.computeTimeScale = d3.scale.linear<string, string>()
        .domain(computeTimeExtent)
        .range(params.minMaxColors);

    // Maps node name to whether the rendering hierarchy was already constructed.
    this.hasSubhierarchy = {};
    this.params = params;
    this.root = new RenderGroupNodeInformation(hierarchy.root);
    this.index[hierarchy.root.name] = this.root;
    this.buildSubhierarchy(hierarchy.root.name);
    this.root.expanded = true;
  }

  getRenderNodeByName(nodeName: string): RenderNodeInformation {
    return this.index[nodeName];
  }

  /**
   * Return the nearest ancestor node, including itself, that is visible
   * in the visualization. This method is used so that we can select
   * (highlight) a node that isn't drawn yet, by selecting (highlighting)
   * its nearest ancestor that has been drawn.
   */
  getNearestVisibleAncestor(name: string): string {
    let path = getHierarchicalPath(name);
    for (let i = 0; i < path.length; i++) {
      let nodeName = path[i];
      // Op nodes have expanded set to false by default.
      if (!this.getRenderNodeByName(nodeName).expanded) {
        return nodeName;
      }
    }
    // Fallthrough. If everything was expanded return the node.
    return name;
  }

  // TODO(jimbo): Delete this an any code it touches (all deprecated).
  setDepth(depth: number): void {
    setGroupNodeDepth(this.root, +depth);
  }

  buildSubhierarchy(nodeName: string): void {
    // Terminate if the rendering hierarchy was already constructed
    // for this node.
    if (nodeName in this.hasSubhierarchy) {
      return;
    }

    let renderNodeInfo = this.index[nodeName];

    // If it is not a meta node or a series node, don't do anything.
    if (renderNodeInfo.node.type !== NodeType.META &&
        renderNodeInfo.node.type !== NodeType.SERIES) {
      return;
    }

    // At this point we know the rendering information is about a group node.
    let renderGroupNodeInfo = <RenderGroupNodeInformation> renderNodeInfo;
    let metagraph = renderGroupNodeInfo.node.metagraph;
    let coreGraph = renderGroupNodeInfo.coreGraph;

    // Create render nodes to represent each child from the metagraph. Although
    // these will initially be added to the coreGraph, they may later be
    // extracted. Also, due to extraction, the coreGraph may contain disjoint
    // groups between which there is no visible path (other than annotations).
    _.each(metagraph.nodes(), childName => {

      let childNode = metagraph.node(childName);
      let childRenderInfo = childNode.isGroupNode ?
          new RenderGroupNodeInformation(<GroupNode>childNode) :
          new RenderNodeInformation(childNode);
      this.index[childName] = childRenderInfo;
      coreGraph.setNode(childName, childRenderInfo);

      if (childRenderInfo.node.stats != null) {
        childRenderInfo.memoryColor =
          this.memoryUsageScale(childRenderInfo.node.stats.totalBytes);
        childRenderInfo.computeTimeColor =
          this.computeTimeScale(childRenderInfo.node.stats.totalMicros);
      }

      if (!childNode.isGroupNode) {
        _.each((<OpNode>childNode).inEmbeddings, embedding => {
          let renderMetaedgeInfo = new RenderMetaedgeInformation(null);
          addInAnnotation(childRenderInfo, embedding, null, renderMetaedgeInfo,
              AnnotationType.CONSTANT, this.params);
          this.index[embedding.name] = new RenderNodeInformation(embedding);
        });
        _.each((<OpNode>childNode).outEmbeddings, embedding => {
          let renderMetaedgeInfo = new RenderMetaedgeInformation(null);
          addOutAnnotation(childRenderInfo, embedding, null, renderMetaedgeInfo,
              AnnotationType.SUMMARY, this.params);
          this.index[embedding.name] = new RenderNodeInformation(embedding);
        });
        let device = (<OpNode>childRenderInfo.node).device;
        if (device != null) {
          childRenderInfo.deviceColors = [{
            color: this.deviceColorMap(device),
            proportion: 1.0
          }];
        }
      } else {
        // Make a list of tuples (device, proportion), where proportion
        // is the fraction of op nodes that have that device.
        let pairs = _.pairs((<GroupNode> childNode).deviceHistogram);
        if (pairs.length > 0) {
          // Compute the total # of devices.
          let numDevices = _.sum(pairs, _.last);
          childRenderInfo.deviceColors = _.map(pairs, pair => {
            return {
              color: this.deviceColorMap(pair[0]),
              // Normalize to a proportion of total # of devices.
              proportion: pair[1] / numDevices
            };
          });
        }
      }
    });

    // Add render metaedge info for edges in the metagraph.
    _.each(metagraph.edges(), edgeObj => {
      let metaedge = metagraph.edge(edgeObj);
      let renderMetaedgeInfo = new RenderMetaedgeInformation(metaedge);
      coreGraph.setEdge(edgeObj.v, edgeObj.w, renderMetaedgeInfo);
    });

    if (this.params.enableExtraction &&
        renderGroupNodeInfo.node.type === NodeType.META) {
      extractHighDegrees(renderGroupNodeInfo, this.params);
    }

    // Record that we constructed the rendering hierarchy for this node, so we
    // don't construct it another time.
    this.hasSubhierarchy[nodeName] = true;

    // Look up the parent node's render information and short circuit if none.
    let parentNode = renderGroupNodeInfo.node.parentNode;
    if (!parentNode) {
      return;
    }
    let parentNodeInfo =
      <RenderGroupNodeInformation> this.index[parentNode.name];

    // Utility function for computing the name of a bridge node.
    let getBridgeNodeName = (inbound, ...rest) =>
      rest.concat([inbound ? "IN" : "OUT"]).join("~~");

    // Build out the bridgegraph.
    let bridgegraph = this.hierarchy.getBridgegraph(nodeName);

    // Look for popular nodes so we can make annotations instead of paths.
    let otherCounts = {
      // Counts of edges coming INTO other nodes by name (outgoing from self).
      in: <{[nodeName: string]: number}> {},
      // Counts of edges going OUT from other nodes by name (coming into self).
      out: <{[nodeName: string]: number}> {},
      // Counts of all control edges involving other nodes by name.
      control: <{[nodeName: string]: number}> {},
    };
    _.each(bridgegraph.edges(), e => {
      // An edge is inbound if its destination node is in the metagraph.
      let inbound = !!metagraph.node(e.w);
      let otherName = inbound ? e.v : e.w;
      let metaedge = bridgegraph.edge(e);
      if (!metaedge.numRegularEdges) {
        otherCounts.control[otherName] =
          (otherCounts.control[otherName] || 0) + 1;
      } else if (inbound) {
        otherCounts.out[otherName] = (otherCounts.out[otherName] || 0) + 1;
      } else {
        otherCounts.in[otherName] = (otherCounts.in[otherName] || 0) + 1;
      }
    });

    // Add annotations and edges for bridgegraph relationships.
    let hierarchyNodeMap = this.hierarchy.getNodeMap();
    _.each(bridgegraph.edges(), bridgeEdgeObj => {
      let bridgeMetaedge = bridgegraph.edge(bridgeEdgeObj);

      // Determine whether this bridge edge is incoming by checking the
      // metagraph for a node that matches the destination end.
      let inbound = !!metagraph.node(bridgeEdgeObj.w);

      // Based on the direction of the edge, one endpoint will be an immediate
      // child of this renderNodeInfo, and the other endpoint will be a sibling
      // of the parent (or an ancestor further up).
      let [childName, otherName] =
        inbound ?
          [bridgeEdgeObj.w, bridgeEdgeObj.v] :
          [bridgeEdgeObj.v, bridgeEdgeObj.w];

      let childRenderInfo = this.index[childName];
      let otherRenderInfo = this.index[otherName];
      let otherNode =
        otherRenderInfo ?
          otherRenderInfo.node :
          hierarchyNodeMap[otherName];

      // Determine whether this edge is a control edge between nodes where
      // either node is high-degree with respect to control edges. This will
      // be a signal to show it as an annotation instead of a bridge edge.
      let isHighDegreeControlEdge = !bridgeMetaedge.numRegularEdges &&
        otherCounts.control[otherName] > this.params.maxControlDegree;

      let [annotations, childAnnotations] =
        inbound ?
          [renderNodeInfo.inAnnotations, childRenderInfo.inAnnotations] :
          [renderNodeInfo.outAnnotations, childRenderInfo.outAnnotations];

      let isOtherHighDegree =
        inbound ?
          otherCounts.out[otherName] > this.params.maxOutDegree :
          otherCounts.in[otherName] > this.params.maxInDegree;

      // The adjoining render metaedge info from the parent's coreGraph, if any.
      // It will either be a Metaedge involving this node directly, if it
      // previously came from a metagraph, or it'll be a Metaedge involving
      // a previously created bridge node standing in for the other node.
      let adjoiningMetaedge = null;

      // We can only hope to render a bridge path if:
      //  - bridgegraph paths are enabled,
      //  - the other node is not too high-degree,
      //  - the child is in the core (not extracted for being high-degree), and
      //  - there's a path (in the traversal sense) between child and other.
      let canDrawBridgePath = false;
      if (this.params.enableBridgegraph &&
          !isOtherHighDegree &&
          !isHighDegreeControlEdge &&
          childRenderInfo.isInCore()) {

        // Utility function for finding an adjoining metaedge.
        let findAdjoiningMetaedge = targetName => {
          let adjoiningEdgeObj: graphlib.EdgeObject =
            inbound ?
              { v: targetName, w: nodeName } :
              { v: nodeName, w: targetName };
          return <RenderMetaedgeInformation>
            parentNodeInfo.coreGraph.edge(adjoiningEdgeObj);
        };

        adjoiningMetaedge = findAdjoiningMetaedge(otherName);
        if (!adjoiningMetaedge) {
          adjoiningMetaedge = findAdjoiningMetaedge(
              getBridgeNodeName(inbound, otherName, parentNode.name));
        }

        canDrawBridgePath = !!adjoiningMetaedge;
      }

      // Although dataflow edges are acyclic, control dependency edges may
      // actually point "backwards" in the graph. If this bridgeMetaedge is
      // a control dependency, we need to determine whether it's backwards
      // pointing so that we render it appropriately.
      //
      // For instance, say we're rendering a graph with nodes named A/B and Z/Y,
      // and we're currently rendering the bridgegraph for A. Further, let's say
      // that there was an original BaseEdge from A/B->Z/Y and a CONTROL EDGE
      // from Z/Y=>A/B.
      //
      //     +----------------+
      //     | A              |
      //     |  +-----+       |         +------+
      //     |  | B   |>----->|>------->| Z    |
      //     |  |     |       |         |      |
      //     |  |     |   *   |         |      |
      //     |  |     |<=====<|<=======<|      |
      //     |  +-----+       |         +------+
      //     +----------------+
      //
      // When we render the subhierarchy for Metanode A, we'll come across a
      // control-only Metaedge in the bridgegraph from Z=>A/B (*). The question
      // is whether this edge is backwards.
      //
      // To answer that question, we follow the chain of adjoining metaedges
      // until we reach the topmost one. In this case, that's the control-only
      // Metaedge Z=>A in the ROOT's metagraph. We determine that this edge
      // is backwards by looking at the topological ordering of ROOT's metagraph
      // (which ignores control edges) and seeing that Z comes AFTER A.
      //
      // The property of being backwards is independent of whether the edge
      // is inbound or outbound. In the preceeding example, if we were building
      // the subhierarchy for Z, we'd find bridge edge Z/Y=>A, walk to its
      // topmost adjoining metaedge Z=>A and discover that it's backwards.
      let backwards = false;
      if (adjoiningMetaedge && !bridgeMetaedge.numRegularEdges) {
        // Find the top-most adjoining render metaedge information, and the
        // GroupNode whose metagraph must contain the associated metaedge.
        let topAdjoiningMetaedge = adjoiningMetaedge;
        let topGroupNode = parentNodeInfo.node;
        while (topAdjoiningMetaedge.adjoiningMetaedge) {
          topAdjoiningMetaedge = topAdjoiningMetaedge.adjoiningMetaedge;
          topGroupNode = <GroupNode>topGroupNode.parentNode;
        }

        // Check against the topological ordering for the top node. The current
        // bridge metaedge we're evaluating is backwards if its source comes
        // after its destination.
        let ordering = this.hierarchy.getTopologicalOrdering(topGroupNode.name);
        let e = topAdjoiningMetaedge.metaedge;
        backwards = ordering[e.v] > ordering[e.w];
      }

      // Render backwards control edges as annotations.
      canDrawBridgePath = canDrawBridgePath && !backwards;

      // If we can't make a bridge path for any reason, then we add an
      // annotation instead.
      if (!canDrawBridgePath) {
        childAnnotations.push(new Annotation(
            otherNode,
            otherRenderInfo,
            new RenderMetaedgeInformation(bridgeMetaedge),
            AnnotationType.SHORTCUT,
            inbound), this.params);
        return;
      }

      // At this point, all conditions have been met for drawing a bridge path.

      // Find or create the IN/OUT node representing otherNode.
      let bridgeContainerName = getBridgeNodeName(inbound, nodeName);
      let bridgeNodeName = getBridgeNodeName(inbound, otherName, nodeName);
      let bridgeNodeRenderInfo = coreGraph.node(bridgeNodeName);
      if (!bridgeNodeRenderInfo) {

        // Find or create the directional container for the bridge node.
        let bridgeContainerInfo = coreGraph.node(bridgeContainerName);
        if (!bridgeContainerInfo) {
          let bridgeContainerNode: BridgeNode = {
            // Important node properties.
            name: bridgeContainerName,
            type: NodeType.BRIDGE,
            // Unused node properties.
            isGroupNode: false,
            cardinality: 0,
            parentNode: null,
            stats: null,
            // BridgeNode properties.
            inbound: inbound,
          };
          bridgeContainerInfo =
            new RenderNodeInformation(bridgeContainerNode);
          this.index[bridgeContainerName] = bridgeContainerInfo;
          coreGraph.setNode(bridgeContainerName, bridgeContainerInfo);
        }

        let bridgeNode: BridgeNode = {
          // Important node properties.
          name: bridgeNodeName,
          type: NodeType.BRIDGE,
          // Unimportant node properties.
          isGroupNode: false,
          cardinality: 1,
          parentNode: null,
          stats: null,
          // BridgeNode properties.
          inbound: inbound,
        };
        bridgeNodeRenderInfo = new RenderNodeInformation(bridgeNode);
        this.index[bridgeNodeName] = bridgeNodeRenderInfo;
        coreGraph.setNode(bridgeNodeName, bridgeNodeRenderInfo);

        // Set bridgeNode to be a graphlib child of the container node.
        coreGraph.setParent(bridgeNodeName, bridgeContainerName);
        bridgeContainerInfo.node.cardinality++;
      }

      // Create and add a bridge render metaedge.
      let bridgeRenderMetaedge =
        new RenderMetaedgeInformation(bridgeMetaedge);
      bridgeRenderMetaedge.adjoiningMetaedge = adjoiningMetaedge;
      inbound ?
        coreGraph.setEdge(bridgeNodeName, childName, bridgeRenderMetaedge) :
        coreGraph.setEdge(childName, bridgeNodeName, bridgeRenderMetaedge);

    }); // End _.each(bridgegraph.edges).

    // For each bridge container (IN and/or OUT), add structural edges between
    // terminal nodes and that container. A terminal node is one which has no
    // non-bridge edges in the direction of the container.
    //
    // For example, consider a Metanode A which contains two child nodes A/B
    // and A/C. Let's say it has one edge in the metagraph from A/B->A/C, and
    // one edge in the bridgegraph from Z->A/C.
    //
    // At this point, we've added a container bridge node IN to house all
    // incoming bridge nodes. We'v alse added a bridge node Z' (with parent IN)
    // to A, and a bridge edge from Z'->C.
    //
    //     +----------------------+
    //     | A          +---+     |
    //     |    +------>| C |     |
    //     |    |       +---+     |
    //     |    |         ^       |
    //     |    |         |       |
    //     |    |    +----|----+  |
    //     |    |    | IN |    |  |
    //     |  +---+  |  +---+  |  |
    //     |  | B |  |  | Z'|  |  |
    //     |  +---+  |  +---+  |  |
    //     |         +---------+  |
    //     +----------------------+
    //
    // With no other help, dagre would lay out B and Z' on the same level,
    // because both of them have no incoming edges. In other words, B is a
    // terminal node in the INCOMING direction.
    //
    // But we want to force dagre to lay out Z' (and everything in IN) lower
    // than all non-bridge nodes, so that there's enough room for the bridge
    // edges after they've been adjusted to meet up with paths coming in from
    // outside.
    //
    // To force Z' (and all other bridge nodes) to be lowest in the graph, we
    // identify terminal nodes like B and give them structural edges to
    // a new structural bridge node S which we add to IN.
    //
    //     +----------------------+
    //     | A          +---+     |
    //     |       +--->| C |     |
    //     |       |    +---+     |
    //     |     +---+    ^       |
    //     |     | B |    |       |
    //     |     +---+    |       |
    //     |       ^      |       |
    //     |       |      |       |
    //     |  +----|------|----+  |
    //     |  |IN  |      |    |  |
    //     |  |  +---+  +---+  |  |
    //     |  |  | S |  | Z'|  |  |
    //     |  |  +---+  +---+  |  |
    //     |  +----------------+  |
    //     +----------------------+
    //
    // This ensures that dagre will lay out the bridge containers strictly at
    // the ends of the graph. The structural edges will never be seen in the
    // visualization except as a debugging aid.
    _.each([true, false], inbound => {
      let bridgeContainerName = getBridgeNodeName(inbound, nodeName);
      let bridgeContainerInfo = coreGraph.node(bridgeContainerName);
      if (!bridgeContainerInfo) {
        return;
      }
      _.each(coreGraph.nodes(), childName => {
        // Short-circuit if this child is a bridge node or it's not a terminal
        // node in the direction we're interested in.
        let childNodeInfo = coreGraph.node(childName);
        if (childNodeInfo.node.type === NodeType.BRIDGE) {
          return;
        }
        let isTerminal = inbound ?
          !coreGraph.predecessors(childName).length :
          !coreGraph.successors(childName).length;
        if (!isTerminal) {
          return;
        }

        // Find or create a bridge node in the container for all structural
        // metaedges. It would have been nice to skip this step and simply
        // set a metaedge between the terminal node and the container node, but
        // in that case, something about the graph upsets dagre.layout()'s
        // longestPath algorithm (was getting errors due to an undefined).
        let structuralNodeName =
          getBridgeNodeName(inbound, nodeName, "STRUCTURAL_TARGET");
        let structuralRenderInfo = coreGraph.node(structuralNodeName);
        if (!structuralRenderInfo) {
          let bridgeNode: BridgeNode = {
            // Important Node properties.
            name: structuralNodeName,
            type: NodeType.BRIDGE,
            // Unimportant Node properties.
            isGroupNode: false,
            cardinality: 1,
            parentNode: null,
            stats: null,
            // BridgeNode properties.
            inbound: inbound,
          };
          structuralRenderInfo = new RenderNodeInformation(bridgeNode);
          structuralRenderInfo.structural = true;
          this.index[structuralNodeName] = structuralRenderInfo;
          coreGraph.setNode(structuralNodeName, structuralRenderInfo);
          bridgeContainerInfo.node.cardinality++;
          coreGraph.setParent(structuralNodeName, bridgeContainerName);
        }

        // Create the structural Metaedge and insert it.
        let structuralMetaedgeInfo = new RenderMetaedgeInformation(null);
        structuralMetaedgeInfo.structural = true;
        structuralMetaedgeInfo.weight--; // Reduce weight for dagre layout.
        inbound ?
          coreGraph.setEdge(
              structuralNodeName, childName, structuralMetaedgeInfo) :
          coreGraph.setEdge(
              childName, structuralNodeName, structuralMetaedgeInfo);
      });
    });
  }
}

/**
 * A class for rendering annotation object which contains label
 * about the node embedded as annotation, type of annotation and the location
 * of both the annotation's node and edge.
 *
 * Annotation objects include embedded constants, embedded summary, and
 * edge shortcuts.
 */
export class Annotation {
  node: Node;
  renderNodeInfo: RenderNodeInformation;
  renderMetaedgeInfo: RenderMetaedgeInformation;
  annotationType: AnnotationType;
  /**
   * Center position of annotation relative to the host
   * node's center x.
   */
  dx: number;
  /**
   * Center position of annotation relative to the host
   * node's center y.
   */
  dy: number;
  width: number;
  height: number;
  /**
   * A flag whether it is an in-annotation (if true) or
   * out-annotation  (if false).
   */
  isIn: boolean;
  /** Label horizontal offset from the end of the node shape */
  labelOffset: number;
  /**
   * Array of points for edges from the annotation to its host
   * node. Each point contains the point location, relative to
   * the host node's center.
   */
  points: {dx: number, dy: number}[];

  /**
   * Creates a new Annotation.
   *
   * @param node The underlying node this annotation points to.
   * @param renderNodeInfo The render information for the underlying node
   *     this annotation points to. This can be null if the annotation
   *     denotes an embedding (constant, summary), in which case we
   *     use the node property.
   * @param renderMetaedgeInfo The render information for the edge associated
   *     with the annotation.
   * @param type The type of the annotation.
   * @param isIn True if it is an in-annotation. False if it is an
   *     out-annotation.
   */
  constructor(node: Node, renderNodeInfo: RenderNodeInformation,
      renderMetaedgeInfo: RenderMetaedgeInformation, type: AnnotationType,
      isIn: boolean) {
    this.node = node;
    this.renderNodeInfo = renderNodeInfo;
    this.renderMetaedgeInfo = renderMetaedgeInfo;
    this.annotationType = type;
    // Properties specified by layout
    this.dx = 0;
    this.dy = 0;
    this.width = 0;
    this.height = 0;

    this.isIn = isIn;
    this.points = [];
  }
};

export enum AnnotationType {SHORTCUT, CONSTANT, SUMMARY, ELLIPSIS};

/**
 * Manages a list of annotations. Two will be used for each
 * RenderNodeInformation, one for in annotations and one for out annotations.
 */
export class AnnotationList {
  /**
   * List of visually drawable annotations, may include an ellipses annotation
   * if the number added exceeds the number specified by maxAnnotations.
   */
  list: Annotation[];

  /**
   * Set of nodes which have been added as annotations to this list, so we can
   * prevent duplicates.
   */
  nodeNames: { [nodeName: string]: boolean };

  constructor() {
    this.list = [];
    this.nodeNames = {};
  }

  /**
   * Append an annotation to the list, or a stand-in ellipsis annotation instead
   * if this would make it too many.
   */
  push(annotation: Annotation, params: RenderGraphParams): void {
    if (annotation.node.name in this.nodeNames) {
      return; // Skip duplicate annotation.
    }
    this.nodeNames[annotation.node.name] = true;

    if (this.list.length < params.maxAnnotations) {
      this.list.push(annotation);
      return;
    }

    let lastAnnotation = this.list[this.list.length - 1];
    if (lastAnnotation.annotationType === AnnotationType.ELLIPSIS) {
      let ellipsisNode = <EllipsisNode>lastAnnotation.node;
      ellipsisNode.setNumMoreNodes(++ellipsisNode.numMoreNodes);
      return;
    }

    let ellipsisNode = new tf.graph.EllipsisNodeImpl(1);
    this.list.push(new Annotation(ellipsisNode,
        new RenderNodeInformation(ellipsisNode), null,
        AnnotationType.ELLIPSIS, annotation.isIn));
  }
}

/**
 * Contains rendering information about a node in the hierarchical graph.
 */
export class RenderNodeInformation {
  /** Reference to the original underlying Node from the hierarchical graph. */
  node: Node;
  /** Whether the node is expanded or not. */
  expanded: boolean;
  /**
   * List of rendering information about in-annotations like constants and
   * shortcuts to high-degree nodes.
   */
  inAnnotations: AnnotationList;
  /** List of rendering information about out-annotations (e.g. summary nodes)  */
  outAnnotations: AnnotationList;

  // --- Params specified by layout --- //

  /** Center x position */
  x: number;
  /** Center y position */
  y: number;
  /** Width of the node's shape */
  width: number;
  /** Height of the node's shape. */
  height: number;
  /** Width of the bounding box for all in-annotations. */
  inboxWidth: number;
  /** Width of the bounding box for all out-annotations. */
  outboxWidth: number;
  /**
   * Whether the node should be excluded from the scene.
   * This is only used when there are too many items in a series so we only
   * want to include top N ones.
   */
  // TODO(jimbo): Now that series rendering is non-recursive, remove this and
  // all its uses from the code base.
  excluded: boolean;

  // --- Params used in drawing the bridge paths --- //

  /**
   * All bridge nodes are meant to be invisible, but whereas most represent a
   * relationship from the underlying graph hierarchy, some exist solely for
   * layout reasons. Specifically, those bridge nodes which have only structural
   * rendering metaedges.
   */
  structural: boolean;

  // --- Params for the size of the node box --- //

  /** Label vertical offset from the center of node shape */
  labelOffset: number;
  /** X-space between each extracted node and the core graph. */
  extractXOffset: number;
  /** Rectangle radius (for making rounded rectangle) */
  radius: number;

  // --- Params for expanded node --- //

  /** Label height for expanded node. */
  labelHeight: number;
  // Paddings between inner subscene and the border of the expanded node.
  paddingTop: number;
  paddingLeft: number;
  paddingRight: number;
  paddingBottom: number;

  /** Width of the whole node including its shape and its annotations */
  outerWidth: number;
  /** Height of the whole node including its shape and its annotations */
  outerHeight: number;
  /**
   * Whether a node is extracted as source-like (having high out-degree or matching
   * predefined in-extract pattern.)
   */
  isInExtract: boolean;
  /**
   * Whether a node is extracted as sink-like (having high in-degree or matching
   * predefined out-extract pattern.)
   */
  isOutExtract: boolean;

  /**
   * List of (color, proportion) tuples based on the proportion of devices of
   * its children. If this node is an op node, this list will have only one
   * color with proportion 1.0.
   */
  deviceColors: {color: string, proportion: number}[];

  /**
   * Color according to the memory usage of this node.
   */
  memoryColor: string;

  /**
   * Color according to the compute time of this node.
   */
  computeTimeColor: string;

  constructor(node: Node) {
    this.node = node;
    this.expanded = false;
    this.inAnnotations = new AnnotationList();
    this.outAnnotations = new AnnotationList();
    // Params specified by layout
    this.x = 0;
    this.y = 0;
    this.width = 0;
    this.height = 0;
    this.inboxWidth = 0;
    this.outboxWidth = 0;

    this.excluded = false;

    // Params for bridge paths.
    this.structural = false;

    // Params for node box.
    this.labelOffset = 0;
    this.extractXOffset = 0;
    this.radius = 0;

    // Params for expanded node
    this.labelHeight = 0;
    this.paddingTop = 0;
    this.paddingLeft = 0;
    this.paddingRight = 0;
    this.paddingBottom = 0;

    this.outerWidth = 0;
    this.outerHeight = 0;
    this.isInExtract = false;
    this.isOutExtract = false;
  }

  isInCore(): boolean {
    return !this.isInExtract && !this.isOutExtract;
  }
}

/**
 * Contains rendering information about a Metaedge from the underlying
 * hierarchical graph. It may be from either a metagraph or a bridgegraph.
 */
export class RenderMetaedgeInformation {
  /**
   * Reference to the original underlying Metaedge from the hierarchical graph,
   * if any. This will be null for the edges which connect OpNodes to their
   * embeddings, for example.
   */
  metaedge: Metaedge;

  /**
   * Reference to the adjoining RenderMeteaedgeInformation from the parent's
   * coreGraph. This is used during layout to determine the point at which this
   * edge should touch the node's bounding box. This property will be null for
   * edges which terminate at a node on both ends (all non-bridge edges).
   */
  adjoiningMetaedge: RenderMetaedgeInformation;

  /**
   * Most of the time, a RenderMetaedgeInformation object represents a real
   * edge between nodes in the underlying graph structure. But sometimes, an
   * edge only exsts for layout purposes. These structural edges are added
   * during buildSubhierarchy() to force dagre.layout() to put bridge nodes
   * at the ends of the flow.
   * @see buildSubhierarchy()
   */
  structural: boolean;

  /**
   * Weight of the edge, used by dagre when deciding how important an edge is.
   * Edges with higher weight are made shorter and straighter. The default
   * dagre uses is 1.
   */
  weight: number;

  /**
   * X and Y coordinate pairs of the points in the path of the edge.
   * @see tf.graph.node.subsceneAdjustPaths
   */
  points: any[];

  /**
   * D3 selection of the group containing the path that displays this edge.
   */
  edgeGroup: d3.Selection<RenderMetaedgeInformation>;

  constructor(metaedge: Metaedge) {
    this.metaedge = metaedge;
    this.adjoiningMetaedge = null;
    this.structural = false;
    this.weight = 1;
  }
}

function addInAnnotation(node: RenderNodeInformation, predecessor: Node,
    predecessorRenderInfo: RenderNodeInformation, edge: any,
    type: AnnotationType, params: RenderGraphParams): void {
  let annotation = new Annotation(predecessor, predecessorRenderInfo, edge,
      type, true);
  node.inAnnotations.push(annotation, params);
}

function addOutAnnotation(node: RenderNodeInformation, successor: Node,
    successorRenderInfo: RenderNodeInformation, edge: any,
    type: AnnotationType, params: RenderGraphParams): void {
  let annotation = new Annotation(successor, successorRenderInfo, edge,
      type, false);
  node.outAnnotations.push(annotation, params);
}

function setGraphDepth(graph: graphlib.Graph<RenderNodeInformation, any>,
    depth: number) {
  _.each(graph.nodes(), nodeName => {
    let child = graph.node(nodeName);
    child.expanded = depth > 1; // set all child of depth 1 to collapsed
    if (depth > 0) {
      switch (child.node.type) {
        case NodeType.META:
        case NodeType.SERIES:
          setGroupNodeDepth(<RenderGroupNodeInformation>child, depth - 1);
          break;
        // Do nothing for leaf
      }
    }
  });
};

export class RenderGroupNodeInformation extends RenderNodeInformation {
  node: GroupNode;
  /**
   * The core graph is derived from the underlying node's metagraph, minus
   * the extracted source-like and sink-like nodes.
   */
  coreGraph: graphlib.Graph<RenderNodeInformation, RenderMetaedgeInformation>;
  /** Size of the bounding box for a metanode's core graph. */
  coreBox: {
    width: number,
    height: number,
  };
  /** Size of the bounding box for a metanode's isolated in-extract children. */
  inExtractBox: {width: number, height: number};
  /** Size of the bounding box for a metanode's isolated out-extract children. */
  outExtractBox: {width: number, height: number};
  /** Array of isolated in-extract nodes. */
  isolatedInExtract: RenderNodeInformation[];
  /** Array of isolated out-extract nodes. */
  isolatedOutExtract: RenderNodeInformation[];

  constructor(groupNode: GroupNode) {
    super(groupNode);
    let metagraph = groupNode.metagraph;
    let gl = metagraph.graph();
    this.coreGraph =
        createGraph<RenderNodeInformation, RenderMetaedgeInformation>(
            gl.name, GraphType.CORE, { compound: true });
    this.coreBox = {width: 0, height: 0};
    this.inExtractBox = {width: 0, height: 0};
    this.outExtractBox = {width: 0, height: 0};
    this.isolatedInExtract = [];
    this.isolatedOutExtract = [];
  }
}

function setGroupNodeDepth(renderInfo: RenderGroupNodeInformation,
    depth: number): void {
  if (renderInfo.coreGraph) {
    setGraphDepth(renderInfo.coreGraph, depth);
  }
}

/**
 * Remove an edge from the graph and add annotations to both ends of the edge.
 *
 * @param The core graph.
 * @param v Source name.
 * @param w Sink name.
 */
function createShortcut(graph: graphlib.Graph<RenderNodeInformation, {}>,
  v: string, w: string, params: RenderGraphParams) {
  let src = graph.node(v);
  let sink = graph.node(w);
  let edge = graph.edge(v, w);

  // Add each annotation.
  addOutAnnotation(src, sink.node, sink, edge, AnnotationType.SHORTCUT, params);
  addInAnnotation(sink, src.node, src, edge, AnnotationType.SHORTCUT, params);

  // Remove the edge from the core graph.
  graph.removeEdge(v, w);
}

/**
 * Remove edges from a node, and set its isOutExtract property to true,
 * and remove the node and move it to isolatedOutExtract.
 *
 * If detachAllEdgesForHighDegree is true, extract all of its edges.
 * Otherwise, only extract all in-edges.
 */
function makeOutExtract(renderNode: RenderGroupNodeInformation, n: string,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;

  graph.node(n).isOutExtract = true;

  _.each(graph.predecessors(n), (p, index) => {
    createShortcut(graph, p, n, params);
  });

  if (params.detachAllEdgesForHighDegree) {
    _.each(graph.successors(n), (s, index) => {
      createShortcut(graph, n, s, params);
    });
  }

  if (params.detachAllEdgesForHighDegree || graph.neighbors(n).length === 0) {
    renderNode.isolatedOutExtract.push(graph.node(n));
    graph.removeNode(n);
  }
}

/**
 * Remove edges from a node, set its isInExtract property to true,
 * and remove the node and move it to isolatedInExtract.
 * If detachAllEdgesForHighDegree is true, extract all of its edges.
 * Otherwise, only remove all out-edges.
 */
function makeInExtract(renderNode: RenderGroupNodeInformation, n: string,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;
  graph.node(n).isInExtract = true;

  _.each(graph.successors(n), (s, index) => {
    createShortcut(graph, n, s, params);
  });

  if (params.detachAllEdgesForHighDegree) {
    _.each(graph.predecessors(n), (p, index) => {
      createShortcut(graph, p, n, params);
    });
  }

  // Remove the node from the core graph if conditions are met.
  if (params.detachAllEdgesForHighDegree || graph.neighbors(n).length === 0) {
    renderNode.isolatedInExtract.push(graph.node(n));
    graph.removeNode(n);
  }
}

/**
 * Check whether the node's type is a member of the given list of types.
 *
 * @param node Node.
 * @param types List of type to match.
 */
function hasTypeIn(node: Node, types: string[]): boolean {
  if (node.type === NodeType.OP) {
    for (let i = 0; i < types.length; i++) {
      if ((<OpNode>node).op === types[i]) { return true; }
    }
  } else if (node.type === NodeType.META) {
    let rootOpNode = (<Metanode>node).getRootOp();
    if (rootOpNode) {
      for (let i = 0; i < types.length; i++) {
        if (rootOpNode.op === types[i]) { return true; }
      }
    }
  }
  return false;
}

/** Remove edges from pre-defined out-extract patterns */
function extractPredefinedSink(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;
  _.each(graph.nodes(), n => {
    let renderInfo = graph.node(n);
    if (hasTypeIn(renderInfo.node, params.outExtractTypes)) {
      makeOutExtract(renderNode, n, params);
    }
  });
}

/** Remove edges from pre-defined in-extract patterns */
function extractPredefinedSource(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;

  _.each(graph.nodes(), n => {
    let renderInfo = graph.node(n);
    if (hasTypeIn(renderInfo.node, params.inExtractTypes)) {
      makeInExtract(renderNode, n, params);
    }
  });
}

/** Extract from nodes with in-degree > maxInDegree */
function extractHighInDegree(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;
  let maxInDegree = params.maxInDegree;

  // detect first so degrees don't get affected by other removal
  let highInDegreeNames = _.filter(graph.nodes(), n => {
    // Count the in-degree based on only regular edges, unless there are
    // no regular edges, in which case use the number of control edges.
    // This is done so that control edges don't effect if nodes are extracted
    // from the core graph, unless the node is only used for control.
    let numEdgesToCount = _.reduce(graph.predecessors(n), (numEdgesToCount, pred) => {
      let metaedge = graph.edge(pred, n).metaedge;
      return numEdgesToCount + (metaedge.numRegularEdges ? 1 : 0);
    }, 0);
    if (numEdgesToCount === 0 && graph.predecessors(n).length > 0) {
      numEdgesToCount = graph.predecessors(n).length;
    }
    return numEdgesToCount > maxInDegree;
  });

  _.each(highInDegreeNames, n => {
    makeOutExtract(renderNode, n, params);
  });
}

/** Extract nodes with out-degree > maxOutDegree */
function extractHighOutDegree(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;
  let maxOutDegree = params.maxOutDegree;

  // detect first so degrees don't get affected by other removal
  let highOutDegreeNames = _.filter(graph.nodes(), n => {
    // Count the out-degree based on only regular edges, unless there are
    // no regular edges, in which case use the number of control edges.
    // This is done so that control edges don't effect if nodes are extracted
    // from the core graph, unless the node is only used for control.
    let numEdgesToCount = _.reduce(graph.successors(n), (numEdgesToCount, succ) => {
      let metaedge = graph.edge(n, succ).metaedge;
      return numEdgesToCount + (metaedge.numRegularEdges ? 1 : 0);
    }, 0);
    if (numEdgesToCount === 0 && graph.successors(n).length > 0) {
      numEdgesToCount = graph.successors(n).length;
    }
    return numEdgesToCount > maxOutDegree;
  });

  _.each(highOutDegreeNames, n => {
    makeInExtract(renderNode, n, params);
  });
}

/** Remove control edges from nodes that have too many control edges */
function removeControlEdges(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  let graph = renderNode.coreGraph;

  // Collect control edges into a map by node name.
  let map = <{[nodeName: string]: graphlib.EdgeObject[]}>{};
  _.each(graph.edges(), e => {
    if (!graph.edge(e).metaedge.numRegularEdges) {
      (map[e.v] = map[e.v] || []).push(e);
      (map[e.w] = map[e.w] || []).push(e);
    }
  });

  // For each node with too many control edges, turn them into annotations.
  _.each(map, (edges, nodeName) => {
    if (edges.length > params.maxControlDegree) {
      _.each(edges, e => createShortcut(graph, e.v, e.w, params));
    }
  });
}

/**
 * Given an integer, picks a hue that is far apart from other colors.
 * The formula for picking color that avoid collision is:
 *     hue = (color range * golden ratio * index) % color range
 */
export function mapIndexToHue(id: number): number {
  let GOLDEN_RATIO = 1.61803398875;
  // Hue of 0 is reserved for the gray nodes.
  let MIN_HUE = 1;
  let MAX_HUE = 359;
  let COLOR_RANGE = MAX_HUE - MIN_HUE;
  return MIN_HUE + ((COLOR_RANGE * GOLDEN_RATIO * id) % COLOR_RANGE);
};

/**
 * Remove edges and add to annotation instead.
 *
 * For root node, consider predefined types for source and sink.
 * We do not extract predefined type from non-root so that Variables and the
 * sgd node (op type = "NoOp") do not get extract from inside own group.
 *
 * The order of extraction is important here as swapping the order can totally
 * screw up the graph layout.
 *
 * @param {Render.Node} renderNode Node to manipulate.
 * @param {Object} params render Graph construction parameters. See
 *     <tf-graph-params>'s output
 */
function extractHighDegrees(renderNode: RenderGroupNodeInformation,
    params: RenderGraphParams) {
  if (params.outExtractTypes) {
    extractPredefinedSink(renderNode, params);
  }

  // This has to come before extract high in-degree to protect the core part
  // that takes many variables.
  if (params.inExtractTypes) {
    extractPredefinedSource(renderNode, params);
  }

  // This has to come before extract high out-degree to protect the core part
  // that output to many places as there are more high-degree sinks than
  // sources.

  if (params.maxInDegree) {
    extractHighInDegree(renderNode, params);
  }

  if (params.maxOutDegree) {
    extractHighOutDegree(renderNode, params);
  }

  if (params.maxControlDegree) {
    removeControlEdges(renderNode, params);
  }

  // Extract isolated nodes, which can be
  // (1) source-like and sink-like nodes that are not originally isolated but
  //     become isolated after further removal.
  // (2) isolated nodes with annotations on one-side.  These might be either
  //     - nodes that originally have high out-degree but because we remove
  //       high in-degree nodes first, they no longer have high in-degree when
  //       we check.  (Detecting all high-degree before removing also leads to
  //       another problem.)
  //     - nodes that do not have high degree, but their neighbors are all
  //       extracted, so it might make sense to extract them too.

  let graph = renderNode.coreGraph;
  _.each(graph.nodes(), n => {
    let child = graph.node(n);
    let degree = graph.neighbors(n).length;

    if (degree === 0) {
      let hasOutAnnotations = child.outAnnotations.list.length > 0;
      let hasInAnnotations = child.inAnnotations.list.length > 0;

      if (child.isInExtract) { // Is source-like.
        // This case only happens if detachAllEdgesForHighDegree is false.
        // (Otherwise all source-like nodes are all isolated already.)
        renderNode.isolatedInExtract.push(child);
        graph.removeNode(n);
      } else if (child.isOutExtract) { // Is sink-like.
        // This case only happens if detachAllEdgesForHighDegree is false.
        // // (Otherwise all sink-like nodes are all isolated already.)
        renderNode.isolatedOutExtract.push(child);
        graph.removeNode(n);
      } else if (params.extractIsolatedNodesWithAnnotationsOnOneSide) {
        if (hasOutAnnotations && !hasInAnnotations) {
          child.isInExtract = true; // for ones with high out-annotations
          renderNode.isolatedInExtract.push(child);
          graph.removeNode(n);
        } else if (hasInAnnotations && !hasOutAnnotations) {
          child.isOutExtract = true; // for ones with high in-annotations
          renderNode.isolatedOutExtract.push(child);
          graph.removeNode(n);
        } else {
          // if a low degree node has both in- & out- annotations, do nothing
          // because it is unclear which side it should go to.
        }
      }
    }
  });
}
} // close module tf.graph.render