Spaces:
Running
Running
| import dagre from "@dagrejs/dagre"; | |
| import type { Graph, GraphNode, Operation } from "./types"; | |
| const NODE_SIZE = 40; | |
| // How far a node's value and derivative sit from its centre, above and below its shape | |
| export const SLOT_OFFSET = 28; | |
| // Room kept for a node with its value and derivative | |
| const NODE_ROOM_WIDTH = 60; | |
| const NODE_ROOM_HEIGHT = 76; | |
| const CHAR_WIDTH = 7.2; | |
| // Room kept around each edge for its name | |
| const LABEL_WIDTH = 44; | |
| const LABEL_HEIGHT = 20; | |
| // Flatter than this, a label can stay centred on its edge | |
| const LEVEL_SLOPE = 0.08; | |
| // A hidden node before every input, pulling them all into the first column | |
| const SOURCE = "\0source"; | |
| // The box at the end, f | |
| export const OUTPUT = "\0output"; | |
| interface Point { | |
| x: number; | |
| y: number; | |
| } | |
| export interface PlacedShape { | |
| id: string; | |
| kind: GraphNode["kind"] | "output"; | |
| label: string; | |
| // The centre | |
| x: number; | |
| y: number; | |
| width: number; | |
| height: number; | |
| } | |
| export interface PlacedEdge { | |
| id: string; | |
| from: string; | |
| to: string; | |
| points: Point[]; | |
| // The point on the edge that its label straddles | |
| labelAt: Point; | |
| // The edge's direction there, going right: it rises (-1), falls (1) or is level (0) | |
| labelSlope: -1 | 0 | 1; | |
| } | |
| export interface GraphLayout { | |
| shapes: PlacedShape[]; | |
| edges: PlacedEdge[]; | |
| width: number; | |
| height: number; | |
| } | |
| const INFIX: Partial<Record<Operation, string>> = { add: "+", sub: "−", mul: "·", div: "/", pow: "^" }; | |
| const SYMBOLS: Partial<Record<Operation, string>> = { add: "+", sub: "−", mul: "*", div: "/", pow: "^", neg: "−", sqrt: "√" }; | |
| // What a node computes, e.g. "w1 + 2" or "sin(w1)"; an input's is its name. | |
| // Pass show to write its arguments another way, such as their numbers. | |
| export function expressionOf(node: GraphNode, show: (arg: string) => string = (arg) => arg): string { | |
| if (node.kind === "input") { | |
| return node.name; | |
| } | |
| const [a, b] = node.args.map(show); | |
| if (node.op in INFIX) { | |
| return `${a} ${INFIX[node.op]} ${b}`; | |
| } | |
| return node.op === "neg" ? `−${a}` : `${node.op}(${a})`; | |
| } | |
| // What a node's shape shows: its name, or its operator with any constant, e.g. "2 ^" | |
| function labelOf(node: GraphNode, graph: Graph): string { | |
| if (node.kind === "input") { | |
| return node.name; | |
| } | |
| const symbol = SYMBOLS[node.op] ?? node.op; | |
| const constant = (arg: string | undefined) => arg !== undefined && arg in graph.constants; | |
| if (constant(node.args[0])) { | |
| return `${node.args[0]} ${symbol}`; | |
| } | |
| return constant(node.args[1]) ? `${symbol} ${node.args[1]}` : symbol; | |
| } | |
| // A smooth line through the points | |
| export function edgePath({ points }: PlacedEdge): string { | |
| const [first, ...rest] = points; | |
| const last = rest.pop()!; | |
| let path = `M ${first.x} ${first.y}`; | |
| for (let i = 0; i < rest.length; i++) { | |
| const next = rest[i + 1] ?? last; | |
| path += ` Q ${rest[i].x} ${rest[i].y} ${(rest[i].x + next.x) / 2} ${(rest[i].y + next.y) / 2}`; | |
| } | |
| return `${path} L ${last.x} ${last.y}`; | |
| } | |
| // Where the line from a shape's centre towards a point leaves the shape | |
| function onShape(shape: PlacedShape, toward: Point): Point { | |
| const dx = toward.x - shape.x; | |
| const dy = toward.y - shape.y; | |
| const across = Math.abs(dx / (shape.width / 2)); | |
| const down = Math.abs(dy / (shape.height / 2)); | |
| const scale = 1 / (shape.kind === "op" ? Math.hypot(across, down) : Math.max(across, down)); | |
| return { x: shape.x + dx * scale, y: shape.y + dy * scale }; | |
| } | |
| // Where the drawn curve passes a dagre point, and which way it is heading | |
| function labelPlacement(points: Point[], label: Point): Pick<PlacedEdge, "labelAt" | "labelSlope"> { | |
| const distance = (point: Point) => Math.hypot(point.x - label.x, point.y - label.y); | |
| const i = points.reduce((best, point, j) => (distance(point) < distance(points[best]) ? j : best), 0); | |
| const before = points[Math.max(i - 1, 0)]; | |
| const after = points[Math.min(i + 1, points.length - 1)]; | |
| const slope = (after.y - before.y) / (after.x - before.x); | |
| return { | |
| // Edges are smoothed, so they cut the corner at each point | |
| labelAt: { x: (before.x + 6 * points[i].x + after.x) / 8, y: (before.y + 6 * points[i].y + after.y) / 8 }, | |
| labelSlope: Math.abs(slope) < LEVEL_SLOPE ? 0 : slope < 0 ? -1 : 1, | |
| }; | |
| } | |
| // Left to right in evaluation order, with the inputs in a column on the left | |
| export function layoutGraph(graph: Graph): GraphLayout { | |
| const g = new dagre.graphlib.Graph({ multigraph: true }); | |
| g.setGraph({ rankdir: "LR", nodesep: 12, ranksep: 16, edgesep: 28, marginx: 4, marginy: 4 }); | |
| g.setDefaultEdgeLabel(() => ({})); | |
| const size = (label: string) => ({ width: Math.max(NODE_SIZE, label.length * CHAR_WIDTH + 24), height: NODE_SIZE }); | |
| const unplaced: Omit<PlacedShape, "x" | "y">[] = graph.nodes.map((node) => { | |
| const label = labelOf(node, graph); | |
| return { id: node.id, kind: node.kind, label, ...size(label) }; | |
| }); | |
| unplaced.push({ id: OUTPUT, kind: "output", label: "f", ...size("f") }); | |
| for (const { id, width } of unplaced) { | |
| g.setNode(id, { width: Math.max(width, NODE_ROOM_WIDTH), height: NODE_ROOM_HEIGHT }); | |
| } | |
| const nodeIds = new Set(graph.nodes.map((node) => node.id)); | |
| const edges: { from: string; to: string; id: string }[] = [{ from: graph.output, to: OUTPUT, id: `${graph.output}-output` }]; | |
| for (const node of graph.nodes) { | |
| if (node.kind === "op") { | |
| node.args.forEach((from, i) => { | |
| if (nodeIds.has(from)) { | |
| edges.push({ from, to: node.id, id: `${from}-${node.id}-${i}` }); | |
| } | |
| }); | |
| } | |
| } | |
| for (const edge of edges) { | |
| g.setEdge(edge.from, edge.to, { width: LABEL_WIDTH, height: LABEL_HEIGHT, labelpos: "c" }, edge.id); | |
| } | |
| g.setNode(SOURCE, { width: 0, height: 0 }); | |
| for (const name of graph.nodes.filter((node) => node.kind === "input")) { | |
| g.setEdge(SOURCE, name.id, { weight: 100 }, SOURCE); | |
| } | |
| dagre.layout(g); | |
| // dagre doesn't keep the inputs' order, so mirror the drawing if they came out reversed | |
| const inputY = graph.nodes.filter((node) => node.kind === "input").map((node) => g.node(node.id).y); | |
| const mirrored = inputY[0] > inputY[inputY.length - 1]; | |
| const { height = 0 } = g.graph(); | |
| // Measure without the hidden node's column | |
| const left = Math.min(...unplaced.map(({ id }) => g.node(id).x - g.node(id).width / 2)) - 4; | |
| const place = ({ x, y }: Point): Point => ({ x: x - left, y: mirrored ? height - y : y }); | |
| const shapes: PlacedShape[] = unplaced.map((shape) => ({ ...shape, ...place(g.node(shape.id)) })); | |
| const shapeOf = new Map(shapes.map((shape) => [shape.id, shape])); | |
| const width = Math.max(...unplaced.map(({ id }) => g.node(id).x + g.node(id).width / 2)) - left + 4; | |
| return { | |
| shapes, | |
| width, | |
| height, | |
| edges: edges.map(({ from, to, id }) => { | |
| const edge = g.edge(from, to, id); | |
| const points: Point[] = edge.points.map(place); | |
| // dagre ends edges on the room kept for a node, not on its shape | |
| points[0] = onShape(shapeOf.get(from)!, points[1]); | |
| points[points.length - 1] = onShape(shapeOf.get(to)!, points[points.length - 2]); | |
| return { id, from, to, points, ...labelPlacement(points, place(edge)) }; | |
| }), | |
| }; | |
| } | |