Fix overview stage ordering and edge label positioning

- Add ELK partition support to force sequential stage ordering in Overview
- Fix edge label placement bug (was using array index instead of visual midpoint)
- Add z-index to edge labels so they render above nodes
- Increase between-layer spacing for better label readability

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: anth-volk

docs/pipeline-diagrams/src/components/edges/ElkEdge.jsx

@@ -30,10 +30,9 @@ export default function ElkEdge({
       .map((p, i) => `${i === 0 ? "M" : "L"} ${p.x} ${p.y}`)
       .join(" ");
 
-    // Label at midpoint of the path
-    const mid = Math.floor(points.length / 2);
-    labelX = points[mid].x;
-    labelY = points[mid].y;
+    // Label at visual midpoint between start and end of the edge
+    labelX = (points[0].x + points[points.length - 1].x) / 2;
+    labelY = (points[0].y + points[points.length - 1].y) / 2;
   } else {
     // Fallback to smoothstep
     const [path, lx, ly] = getSmoothStepPath({
@@ -66,6 +65,7 @@ export default function ElkEdge({
               position: "absolute",
               transform: `translate(-50%, -50%) translate(${labelX}px, ${labelY}px)`,
               pointerEvents: "all",
+              zIndex: 1000,
             }}
             className="nodrag nopan"
           >

docs/pipeline-diagrams/src/hooks/useElkLayout.js

@@ -8,7 +8,7 @@ const ELK_OPTIONS = {
   "elk.direction": "DOWN",
   "elk.edgeRouting": "ORTHOGONAL",
   "elk.spacing.nodeNode": "60",
-  "elk.layered.spacing.nodeNodeBetweenLayers": "80",
+  "elk.layered.spacing.nodeNodeBetweenLayers": "120",
   "elk.layered.spacing.edgeNodeBetweenLayers": "30",
   "elk.spacing.edgeEdge": "20",
   "elk.spacing.edgeNode": "30",
@@ -77,16 +77,33 @@ export default function useElkLayout(initialNodes, initialEdges) {
     hasRun.current = true;
 
     async function runLayout() {
+      // Check if any node specifies a partition for forced ordering
+      const usePartitioning = initialNodes.some(
+        (n) => n.data?.elkPartition != null,
+      );
+
+      const layoutOptions = { ...ELK_OPTIONS };
+      if (usePartitioning) {
+        layoutOptions["elk.partitioning.activate"] = "true";
+      }
+
       const graph = {
         id: "root",
-        layoutOptions: ELK_OPTIONS,
+        layoutOptions,
         children: initialNodes.map((node) => {
           const size = estimateNodeSize(node);
-          return {
+          const child = {
             id: node.id,
             width: size.width,
             height: size.height,
           };
+          if (node.data?.elkPartition != null) {
+            child.layoutOptions = {
+              "elk.partitioning.partition":
+                String(node.data.elkPartition),
+            };
+          }
+          return child;
         }),
         edges: initialEdges.map((edge) => ({
           id: edge.id,

docs/pipeline-diagrams/src/stages/Overview.jsx

@@ -8,61 +8,62 @@ import {
 } from "../constants/edgeStyles";
 
 const nodes = [
-  // Stage nodes (vertical column, x=400, 200px gaps)
+  // Stage nodes — elkPartition forces sequential layer ordering
   {
     id: "stage0",
     type: "process",
-    data: { label: "Stage 0: Raw Data Download" },
+    data: { label: "Stage 0: Raw Data Download", elkPartition: 0 },
   },
   {
     id: "stage1",
     type: "process",
-    data: { label: "Stage 1: Base Dataset Construction" },
+    data: { label: "Stage 1: Base Dataset Construction", elkPartition: 1 },
   },
   {
     id: "stage2",
     type: "process",
-    data: { label: "Stage 2: Extended CPS (PUF Clone)" },
+    data: { label: "Stage 2: Extended CPS (PUF Clone)", elkPartition: 2 },
   },
   {
     id: "stage3",
     type: "process",
-    data: { label: "Stage 3: Stratified CPS" },
+    data: { label: "Stage 3: Stratified CPS", elkPartition: 3 },
   },
   {
     id: "stage4",
     type: "process",
-    data: { label: "Stage 4: Source Imputation" },
+    data: { label: "Stage 4: Source Imputation", elkPartition: 4 },
   },
   {
     id: "stage5",
     type: "process",
-    data: { label: "Stage 5: Matrix Build" },
+    data: { label: "Stage 5: Matrix Build", elkPartition: 5 },
   },
   {
     id: "stage6",
     type: "process",
-    data: { label: "Stage 6: Weight Fitting (L0)" },
+    data: { label: "Stage 6: Weight Fitting (L0)", elkPartition: 6 },
   },
   {
     id: "stage7",
     type: "process",
-    data: { label: "Stage 7: Local Area H5 Build" },
+    data: { label: "Stage 7: Local Area H5 Build", elkPartition: 7 },
   },
   {
     id: "stage8",
     type: "process",
-    data: { label: "Stage 8: Validation & Promotion" },
+    data: { label: "Stage 8: Validation & Promotion", elkPartition: 8 },
   },
 
-  // Annotation nodes
+  // Annotation nodes — partitioned near associated stages
   {
     id: "modal1",
     type: "external",
     data: {
       label: "Modal (CPU)",
       subtitle: "65GB RAM, 8 CPU",
       details: ["Stages 5-6 build phase"],
+      elkPartition: 5,
     },
   },
   {
@@ -72,6 +73,7 @@ const nodes = [
       label: "Modal (GPU)",
       subtitle: "T4/A10/A100/H100",
       details: ["Stage 6 fit phase"],
+      elkPartition: 6,
     },
   },
   {
@@ -81,6 +83,7 @@ const nodes = [
       label: "Modal (Workers)",
       subtitle: "8 containers, 16GB each",
       details: ["Stage 7 parallel build"],
+      elkPartition: 7,
     },
   },
   {
@@ -90,6 +93,7 @@ const nodes = [
       label: "HuggingFace Hub",
       subtitle: "policyengine/policyengine-us-data",
       details: ["Download + Upload at every stage"],
+      elkPartition: 0,
     },
   },
   {
@@ -99,6 +103,7 @@ const nodes = [
       label: "Local Machine",
       subtitle: "make data / make pipeline",
       details: ["Stages 0-3 run locally"],
+      elkPartition: 1,
     },
   },
 ];