function dagreLayout()

in tensorboard/plugins/graph/tf_graph_common/layout.ts [322:448]

• 102 lines of code
• 20 McCabe index (conditional complexity)

function dagreLayout(
  graph: graphlib.Graph,
  params
): {height: number; width: number} {
  _.extend(graph.graph(), {
    nodesep: params.nodeSep,
    ranksep: params.rankSep,
    edgesep: params.edgeSep,
  });
  let bridgeNodeNames = [];
  let nonBridgeNodeNames = [];
  // Split out nodes into bridge and non-bridge nodes, and calculate the total
  // width we should use for bridge nodes.
  _.each(graph.nodes(), (nodeName) => {
    let nodeInfo = graph.node(nodeName);
    if (nodeInfo.node.type === NodeType.BRIDGE) {
      bridgeNodeNames.push(nodeName);
    } else {
      nonBridgeNodeNames.push(nodeName);
    }
  });
  // If there are no non-bridge nodes, then the graph has zero size.
  if (!nonBridgeNodeNames.length) {
    return {
      width: 0,
      height: 0,
    };
  }
  dagre.layout(graph);
  // Calculate the true bounding box of the graph by iterating over nodes and
  // edges rather than accepting dagre's word for it. In particular, we should
  // ignore the extra-wide bridge nodes and bridge edges, and allow for
  // annotation boxes and labels.
  let minX = Infinity;
  let minY = Infinity;
  let maxX = -Infinity;
  let maxY = -Infinity;
  _.each(nonBridgeNodeNames, (nodeName) => {
    let nodeInfo = graph.node(nodeName);
    let w = 0.5 * nodeInfo.width;
    let x1 = nodeInfo.x - w;
    let x2 = nodeInfo.x + w;
    minX = x1 < minX ? x1 : minX;
    maxX = x2 > maxX ? x2 : maxX;
    // TODO: Account for the height of labels above op nodes here.
    let h = 0.5 * nodeInfo.height;
    let y1 = nodeInfo.y - h;
    let y2 = nodeInfo.y + h;
    minY = y1 < minY ? y1 : minY;
    maxY = y2 > maxY ? y2 : maxY;
  });
  _.each(graph.edges(), (edgeObj) => {
    let edgeInfo = graph.edge(edgeObj);
    if (edgeInfo.structural) {
      return; // Skip structural edges from min/max calculations.
    }
    // Since the node size passed to dagre includes the in and out
    // annotations, the endpoints of the edge produced by dagre may not
    // point to the actual node shape (rectangle, ellipse). We correct the
    // end-points by finding the intersection of a line between the
    // next-to-last (next-to-first) point and the destination (source)
    // rectangle.
    let sourceNode = graph.node(edgeInfo.metaedge.v);
    let destNode = graph.node(edgeInfo.metaedge.w);
    // Straight 3-points edges are special case, since they are curved after
    // our default correction. To keep them straight, we remove the mid point
    // and correct the first and the last point to be the center of the
    // source and destination node respectively.
    if (edgeInfo.points.length === 3 && isStraightLine(edgeInfo.points)) {
      if (sourceNode != null) {
        let cxSource = sourceNode.expanded
          ? sourceNode.x
          : computeCXPositionOfNodeShape(sourceNode as any);
        edgeInfo.points[0].x = cxSource;
      }
      if (destNode != null) {
        let cxDest = destNode.expanded
          ? destNode.x
          : computeCXPositionOfNodeShape(destNode as any);
        edgeInfo.points[2].x = cxDest;
      }
      // Remove the middle point so the edge doesn't curve.
      edgeInfo.points = [edgeInfo.points[0], edgeInfo.points[1]];
    }
    // Correct the destination endpoint of the edge.
    let nextToLastPoint = edgeInfo.points[edgeInfo.points.length - 2];
    // The destination node might be null if this is a bridge edge.
    if (destNode != null) {
      edgeInfo.points[edgeInfo.points.length - 1] = intersectPointAndNode(
        nextToLastPoint,
        destNode as any
      );
    }
    // Correct the source endpoint of the edge.
    let secondPoint = edgeInfo.points[1];
    // The source might be null if this is a bridge edge.
    if (sourceNode != null) {
      edgeInfo.points[0] = intersectPointAndNode(
        secondPoint,
        sourceNode as any
      );
    }
    _.each(edgeInfo.points, (point: render.Point) => {
      minX = point.x < minX ? point.x : minX;
      maxX = point.x > maxX ? point.x : maxX;
      minY = point.y < minY ? point.y : minY;
      maxY = point.y > maxY ? point.y : maxY;
    });
  });
  // Shift all nodes and edge points to account for the left-padding amount,
  // and the invisible bridge nodes.
  _.each(graph.nodes(), (nodeName) => {
    let nodeInfo = graph.node(nodeName);
    nodeInfo.x -= minX;
    nodeInfo.y -= minY;
  });
  _.each(graph.edges(), (edgeObj) => {
    _.each(graph.edge(edgeObj).points, (point: render.Point) => {
      point.x -= minX;
      point.y -= minY;
    });
  });
  return {
    width: maxX - minX,
    height: maxY - minY,
  };
}