add elkjs to position nodes automatically

optuna_dashboard/ts/components/PreferentialGraph.tsx

@@ -1,5 +1,12 @@
 import React, { FC, useState, useCallback, useMemo, useEffect } from "react"
-import { Card, CardContent, CardHeader, useTheme } from "@mui/material"
+import {
+  Card,
+  CardContent,
+  useTheme,
+  Typography,
+  Box,
+  Chip,
+} from "@mui/material"
 import { MarkdownRenderer } from "./Note"
 import ReactFlow, {
   Node,
@@ -12,17 +19,22 @@ import ReactFlow, {
   MiniMap,
   Position,
   Handle,
-  XYPosition,
 } from "reactflow"
 import "reactflow/dist/style.css"
+import ELK from "elkjs/lib/elk.bundled.js"
+import { ElkNode } from "elkjs/lib/elk-api.js"
 
+const elk = new ELK()
 const nodeWidth = 400
 const nodeHeight = 300
-const nodeMargin = 50
+const nodeMargin = 60
+
 type NodeData = {
   trial?: Trial
+  isBest: boolean
 }
 const GraphNode: FC<NodeProps<NodeData>> = ({ data, isConnectable }) => {
+  const theme = useTheme()
   const trial = data.trial
   if (trial === undefined) {
     return null
@@ -39,7 +51,25 @@ const GraphNode: FC<NodeProps<NodeData>> = ({ data, isConnectable }) => {
         overflow: "hidden",
       }}
     >
-      <CardHeader title={`Trial ${trial.number}`} />
+      <Box
+        sx={{
+          display: "flex",
+          displayDirection: "row",
+          margin: theme.spacing(2),
+        }}
+      >
+        <Typography variant="h5">Trial {trial.number}</Typography>
+        {data.isBest && (
+          <Chip
+            label={"Best Trial"}
+            color="secondary"
+            variant="outlined"
+            sx={{
+              marginLeft: "auto",
+            }}
+          />
+        )}
+      </Box>
       <Handle
         type="target"
         position={Position.Top}
@@ -60,85 +90,84 @@ const GraphNode: FC<NodeProps<NodeData>> = ({ data, isConnectable }) => {
 const nodeTypes: NodeTypes = {
   note: GraphNode,
 }
-
-const createNode = (
-  x: number,
-  y: number,
-  trial: Trial,
-  bestGroupPos: XYPosition,
-  isBest: boolean
-): Node => {
-  return {
-    id: `${trial.number}`,
-    type: "note",
-    data: {
-      label: `Trial ${trial.number}`,
-      trial: trial,
-    },
-    position: {
-      x: bestGroupPos.x + nodeMargin + x * (nodeWidth + nodeMargin),
-      y: bestGroupPos.y + nodeMargin + y * (nodeHeight + nodeMargin),
-    },
-    style: {
-      width: nodeWidth,
-      height: nodeHeight,
-      padding: 0,
-    },
-    parentNode: isBest ? "bestGroup" : undefined,
-  }
-}
-const updateNode = (
-  addX: number,
-  addY: number,
-  node: Node,
-  trial: Trial,
-  isBest: boolean
-): Node => {
-  return {
-    ...node,
-    position: {
-      x: node.position.x + addX * (nodeWidth + nodeMargin),
-      y: node.position.y + addY * (nodeHeight + nodeMargin),
-    },
-    data: {
-      ...node.data,
-      trial: trial,
-    },
-    parentNode: isBest ? "bestGroup" : undefined,
-  }
-}
-
-const initNodes: Node[] = [
-  {
-    id: "bestGroup",
-    type: "default",
-    position: {
-      x: 0,
-      y: 0,
-    },
-    data: {
-      label: "Best Trials",
-    },
-    style: {
-      width: 2 * nodeMargin,
-      height: nodeHeight + 2 * nodeMargin,
-      padding: 0,
-      backgroundColor: "rgb(255,0,0,0.1)",
-    },
-  },
-]
-
 const defaultEdgeOptions: DefaultEdgeOptions = {
   animated: true,
 }
 
+function reductionPreference(
+  input_preferences: [number, number][]
+): [number, number][] {
+  const preferences: [number, number][] = []
+  let n = 0
+  for (const [source, target] of input_preferences) {
+    if (
+      preferences.find((p) => p[0] === source && p[1] === target) !== undefined
+    ) {
+      continue
+    }
+    n = Math.max(n - 1, source, target) + 1
+    preferences.push([source, target])
+  }
+  if (n === 0) {
+    return []
+  }
+  const graph: number[][] = Array.from({ length: n }, () => [])
+  const reverseGraph: number[][] = Array.from({ length: n }, () => [])
+  const degree: number[] = Array.from({ length: n }, () => 0)
+  for (const [source, target] of preferences) {
+    graph[source].push(target)
+    reverseGraph[target].push(source)
+    degree[target]++
+  }
+  const topologicalOrder: number[] = []
+  const q: number[] = []
+  for (let i = 0; i < n; i++) {
+    if (degree[i] === 0) {
+      q.push(i)
+    }
+  }
+  while (q.length > 0) {
+    const v = q.pop()
+    if (v === undefined) break
+    topologicalOrder.push(v)
+    graph[v].forEach((u) => {
+      degree[u]--
+      if (degree[u] === 0) {
+        q.push(u)
+      }
+    })
+  }
+  if (topologicalOrder.length !== n) {
+    console.error("cycle detected")
+    return []
+  }
+
+  const response: [number, number][] = []
+  const descendants: Set<number>[] = Array.from(
+    { length: n },
+    () => new Set<number>()
+  )
+  topologicalOrder.reverse().forEach((v) => {
+    const descendant = new Set<number>([v])
+    graph[v].forEach((u) => {
+      descendants[u].forEach((d) => descendant.add(d))
+    })
+    graph[v].forEach((u) => {
+      if (reverseGraph[u].filter((d) => descendant.has(d)).length === 1) {
+        response.push([v, u])
+      }
+    })
+    descendants[v] = descendant
+  })
+  return response
+}
+
 export const PreferentialGraph: FC<{
   studyDetail: StudyDetail | null
 }> = ({ studyDetail }) => {
   const theme = useTheme()
-  const [nodes, setNodes] = useState<Node[]>(initNodes)
+  const [nodes, setNodes] = useState<Node[]>([])
   const [edges, setEdges] = useState<Edge[]>([])
-  const [historyCount, setHistoryCount] = useState(0)
   const onNodesChange: OnNodesChange = useCallback(
     (changes) => setNodes((nds) => applyNodeChanges(changes, nds)),
     [setNodes]
@@ -147,70 +176,67 @@ export const PreferentialGraph: FC<{
 
   useEffect(() => {
     if (studyDetail === null) return
-    const newHistoryCount =
-      (studyDetail.preference_history?.length ?? 0) - historyCount
-    if (newHistoryCount === 0) return
-
-    setNodes((prev) => {
-      const newNodes: Node[] = []
-      const appendIds: string[] = studyDetail.best_trials.map((t) =>
-        t.number.toString()
-      ) // 新しく追加する Node の id, これと "bestGroup" 以外は newHistoryCount だけ下にスライドする
-      studyDetail.preference_history?.slice(historyCount).forEach((history) => {
-        appendIds.push(history.clicked.toString())
-      })
-
-      const bestGroup = prev.find((node) => node.id === "bestGroup")
-      const bestGroupPos = bestGroup?.position ?? { x: 0, y: 0 }
-      console.log(bestGroupPos)
-      if (bestGroup !== undefined) {
-        newNodes.push({
-          ...bestGroup,
-          style: {
-            ...bestGroup.style,
-            width:
-              nodeMargin +
-              studyDetail.best_trials.length * (nodeWidth + nodeMargin),
-            background: "rgb(255,200,200,0.1)",
-          },
-        })
-      }
-
-      prev.forEach((node) => {
-        if (appendIds.includes(node.id)) return
-        if (node.id === "bestGroup") return
-        const trialNum = parseInt(node.id, 10)
-        if (node.id !== `${trialNum}`) {
-          console.error(`node.id is not trual number: ${node.id}`)
-          return
-        }
-        const trial = studyDetail.trials[trialNum]
-        newNodes.push(updateNode(0, newHistoryCount, node, trial, false))
-      })
-      const histories = studyDetail.preference_history?.slice(historyCount)
-      histories?.reverse().forEach((history, i) => {
-        const x = history.candidates.findIndex((c) => c === history.clicked)
-        newNodes.push(
-          createNode(
-            x,
-            i + 1,
-            studyDetail.trials[history.clicked],
-            bestGroupPos,
-            false
-          )
+    if (!studyDetail.is_preferential || studyDetail.preferences === undefined)
+      return
+    const preferences = reductionPreference(studyDetail.preferences)
+    const graph: ElkNode = {
+      id: "root",
+      layoutOptions: {
+        "elk.algorithm": "layered",
+        "elk.direction": "DOWN",
+        "elk.layered.spacing.nodeNodeBetweenLayers": nodeMargin.toString(),
+        "elk.spacing.nodeNode": nodeMargin.toString(),
+      },
+      children: studyDetail.trials.map((trial) => ({
+        id: `${trial.number}`,
+        targetPosition: "top",
+        sourcePosition: "bottom",
+        width: nodeWidth,
+        height: nodeHeight,
+      })),
+      edges: preferences.map(([source, target]) => ({
+        id: `e${source}-${target}`,
+        sources: [`${source}`],
+        targets: [`${target}`],
+        style: { stroke: isDarkMode ? "#fff" : "#000" },
+      })),
+    }
+    elk
+      .layout(graph)
+      .then((layoutedGraph) => {
+        setNodes(
+          layoutedGraph.children?.map((node, index) => {
+            const trial = studyDetail.trials[index]
+            return {
+              id: `${trial.number}`,
+              type: "note",
+              data: {
+                label: `Trial ${trial.number}`,
+                trial: trial,
+                isBest:
+                  studyDetail.best_trials.find(
+                    (t) => t.number === trial.number
+                  ) !== undefined,
+              },
+              position: {
+                x: node.x ?? 0,
+                y: node.y ?? 0,
+              },
+              style: {
+                width: nodeWidth,
+                height: nodeHeight,
+                padding: 0,
+              },
+              deletable: false,
+              connectable: false,
+              draggable: false,
+            }
+          }) ?? []
         )
       })
-      studyDetail.best_trials.forEach((trial, i) => {
-        newNodes.push(createNode(i, 0, trial, bestGroupPos, true))
-      })
-      return newNodes
-    })
-    setHistoryCount(studyDetail.preference_history?.length ?? 0)
-  }, [studyDetail])
-  useEffect(() => {
-    if (studyDetail?.preferences === undefined) return
+      .catch(console.error)
     setEdges(
-      studyDetail?.preferences?.map((p) => {
+      preferences.map((p) => {
         return {
           id: `e${p[0]}-${p[1]}`,
           source: `${p[0]}`,
@@ -219,7 +245,7 @@ export const PreferentialGraph: FC<{
         } as Edge
       }) ?? []
     )
-  }, [studyDetail?.preferences, isDarkMode])
+  }, [studyDetail, isDarkMode])
 
   if (studyDetail === null || !studyDetail.is_preferential) {
     return null
@@ -233,6 +259,12 @@ export const PreferentialGraph: FC<{
       nodeTypes={nodeTypes}
       zoomOnScroll={false}
       panOnScroll={true}
+      minZoom={0.1}
+      defaultViewport={{
+        x: 0,
+        y: 0,
+        zoom: 0.5,
+      }}
     >
       <MiniMap nodeStrokeWidth={1} zoomable pannable />
     </ReactFlow>

package-lock.json

@@ -18,6 +18,7 @@
         "@react-three/fiber": "^8.13.6",
         "@types/three": "^0.154.0",
         "axios": "^1.2.1",
+        "elkjs": "^0.8.2",
         "notistack": "^3.0.1",
         "plotly.js-dist-min": "^2.22.0",
         "react": "^18.2.0",
@@ -6035,6 +6036,11 @@
       "integrity": "sha512-M8WEXFuKXMYMVr45fo8mq0wUrrJHheiKZf6BArTKk9ZBYCKJEOU5H8cdWgDT+qCVZf7Na4lVUaZsA+h6uA9+PA==",
       "dev": true
     },
+    "node_modules/elkjs": {
+      "version": "0.8.2",
+      "resolved": "https://registry.npmjs.org/elkjs/-/elkjs-0.8.2.tgz",
+      "integrity": "sha512-L6uRgvZTH+4OF5NE/MBbzQx/WYpru1xCBE9respNj6qznEewGUIfhzmm7horWWxbNO2M0WckQypGctR8lH79xQ=="
+    },
     "node_modules/emittery": {
       "version": "0.13.1",
       "resolved": "https://registry.npmjs.org/emittery/-/emittery-0.13.1.tgz",
@@ -19690,6 +19696,11 @@
       "integrity": "sha512-M8WEXFuKXMYMVr45fo8mq0wUrrJHheiKZf6BArTKk9ZBYCKJEOU5H8cdWgDT+qCVZf7Na4lVUaZsA+h6uA9+PA==",
       "dev": true
     },
+    "elkjs": {
+      "version": "0.8.2",
+      "resolved": "https://registry.npmjs.org/elkjs/-/elkjs-0.8.2.tgz",
+      "integrity": "sha512-L6uRgvZTH+4OF5NE/MBbzQx/WYpru1xCBE9respNj6qznEewGUIfhzmm7horWWxbNO2M0WckQypGctR8lH79xQ=="
+    },
     "emittery": {
       "version": "0.13.1",
       "resolved": "https://registry.npmjs.org/emittery/-/emittery-0.13.1.tgz",

package.json

@@ -27,6 +27,7 @@
     "@react-three/fiber": "^8.13.6",
     "@types/three": "^0.154.0",
     "axios": "^1.2.1",
+    "elkjs": "^0.8.2",
     "notistack": "^3.0.1",
     "plotly.js-dist-min": "^2.22.0",
     "react": "^18.2.0",