Spaces:
Running
Running
| import { useRef, useState, type RefObject } from "react"; | |
| import { usePlotlyClick } from "@elvis/ui"; | |
| import Plot from "react-plotly.js"; | |
| import type { Camera, Data, Layout } from "plotly.js"; | |
| import type { Trajectory } from "./useOptimisation"; | |
| import type { FunctionValues, Range, SurfaceValues } from "./types"; | |
| interface OptimisationPlotProps { | |
| values: FunctionValues | null; | |
| trajectories: Trajectory[]; | |
| xRange: RefObject<Range>; | |
| yRange: RefObject<Range>; | |
| onViewChange: (x: Range, y: Range) => void; | |
| // Called with data coordinates when the plot is clicked; null ignores clicks | |
| onClick: ((x: number, y: number) => void) | null; | |
| // With two variables: a 3D surface instead of a contour map | |
| surface: boolean; | |
| arrows: boolean; | |
| } | |
| const CONFIG = { scrollZoom: true, displayModeBar: true, doubleClick: false, responsive: true } as const; | |
| const LEGEND = { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.8)" } as const; | |
| const INK = "#1B2A4A"; | |
| export default function OptimisationPlot({ values, trajectories, xRange, yRange, onViewChange, onClick, surface, arrows }: OptimisationPlotProps) { | |
| const plotRef = useRef<Plot | null>(null); | |
| // The 3D camera as the student last left it. Plotly resets it on each redraw | |
| // while stepping (uirevision alone doesn't keep it), so it's passed back in. | |
| const [camera, setCamera] = useState<Partial<Camera> | undefined>(undefined); | |
| // While the student drags the surface, the frame they started on; a redraw | |
| // from stepping would otherwise cancel the drag | |
| const [held, setHeld] = useState<{ data: Data[]; layout: Partial<Layout> } | null>(null); | |
| // el exists at runtime only | |
| // eslint-disable-next-line @typescript-eslint/no-explicit-any | |
| usePlotlyClick(() => (plotRef.current as any)?.el, surface ? null : onClick); | |
| const twoVariables = values !== null && "y" in values; | |
| const showLegend = trajectories.length > 1; | |
| if (twoVariables && surface) { | |
| const data = [surfaceTrace(values), ...trajectories.flatMap((t) => pathTraces(t, true))]; | |
| const layout: Partial<Layout> = { | |
| // Keep the camera while stepping | |
| uirevision: "surface", | |
| showlegend: showLegend, | |
| legend: LEGEND, | |
| margin: { t: 20, r: 20, b: 20, l: 20 }, | |
| scene: { | |
| camera, | |
| aspectmode: "cube", | |
| xaxis: { title: { text: "x" }, range: [...xRange.current] }, | |
| yaxis: { title: { text: "y" }, range: [...yRange.current] }, | |
| zaxis: { title: { text: "f(x, y)" } }, | |
| }, | |
| }; | |
| return ( | |
| <div | |
| className="flex-1 min-h-[320px]" | |
| onPointerDown={() => { | |
| setHeld({ data, layout }); | |
| // Catch up once the drag ends, after Plotly has reported the new camera | |
| window.addEventListener("pointerup", () => setTimeout(() => setHeld(null), 50), { once: true }); | |
| }} | |
| > | |
| <Plot | |
| key="surface" | |
| data={held?.data ?? data} | |
| layout={held?.layout ?? layout} | |
| onRelayout={(event) => { | |
| const moved = (event as { "scene.camera"?: Partial<Camera> })["scene.camera"]; | |
| if (moved) { | |
| setCamera(moved); | |
| } | |
| }} | |
| config={CONFIG} | |
| style={{ width: "100%", height: "100%" }} | |
| /> | |
| </div> | |
| ); | |
| } | |
| const data: Data[] = []; | |
| if (values && "y" in values) { | |
| data.push(contourTrace(values)); | |
| if (arrows) { | |
| data.push(...arrowTraces(values, xRange.current, yRange.current)); | |
| } | |
| } else if (values) { | |
| data.push({ | |
| x: values.x, | |
| y: values.f, | |
| type: "scatter", | |
| mode: "lines", | |
| line: { color: INK, width: 2 }, | |
| name: "f(x)", | |
| showlegend: false, | |
| hovertemplate: "x = %{x:.4g}<br>f(x) = %{y:.4g}<extra></extra>", | |
| }); | |
| } | |
| data.push(...trajectories.flatMap((t) => pathTraces(t, false, twoVariables))); | |
| const layout: Partial<Layout> = { | |
| dragmode: "pan", | |
| // Keep the user's pan and zoom while stepping; a new range (e.g. from a preset) still applies | |
| uirevision: "view", | |
| showlegend: showLegend, | |
| legend: LEGEND, | |
| margin: { t: 40, r: 40, b: 40, l: 40 }, | |
| xaxis: { title: { text: "x" }, range: [...xRange.current] }, | |
| yaxis: { title: { text: twoVariables ? "y" : "f(x)" }, range: [...yRange.current] }, | |
| }; | |
| return ( | |
| <div className="flex-1 min-h-[320px]"> | |
| <Plot | |
| key="flat" | |
| ref={plotRef} | |
| data={data} | |
| layout={layout} | |
| onRelayout={(event) => { | |
| // eslint-disable-next-line @typescript-eslint/no-explicit-any | |
| const full = (plotRef.current as any)?.el?._fullLayout; | |
| const x = (full?.xaxis?.range ?? [event["xaxis.range[0]"], event["xaxis.range[1]"]]) as Range; | |
| const y = (full?.yaxis?.range ?? [event["yaxis.range[0]"], event["yaxis.range[1]"]]) as Range; | |
| const changed = x[0] !== xRange.current[0] || x[1] !== xRange.current[1] || y[0] !== yRange.current[0] || y[1] !== yRange.current[1]; | |
| if (changed && x.every(Number.isFinite) && y.every(Number.isFinite)) { | |
| onViewChange([x[0], x[1]], [y[0], y[1]]); | |
| } | |
| }} | |
| config={CONFIG} | |
| style={{ width: "100%", height: "100%", minWidth: 0 }} | |
| /> | |
| </div> | |
| ); | |
| } | |
| function contourTrace(values: SurfaceValues): Data { | |
| return { | |
| x: values.x, | |
| y: values.y, | |
| z: values.f, | |
| type: "contour", | |
| colorscale: "Viridis", | |
| contours: { coloring: "heatmap", showlines: false }, | |
| colorbar: { thickness: 12, title: { text: "f" } }, | |
| hovertemplate: "x = %{x:.3g}, y = %{y:.3g}<br>f = %{z:.4g}<extra></extra>", | |
| } as Data; | |
| } | |
| function surfaceTrace(values: SurfaceValues): Data { | |
| return { | |
| x: values.x, | |
| y: values.y, | |
| z: values.f, | |
| type: "surface", | |
| colorscale: "Viridis", | |
| opacity: 0.85, | |
| colorbar: { thickness: 12, title: { text: "f" } }, | |
| hovertemplate: "x = %{x:.3g}, y = %{y:.3g}<br>f = %{z:.4g}<extra></extra>", | |
| } as Data; | |
| } | |
| // A run's path and, larger, its current point | |
| function pathTraces(t: Trajectory, threeD: boolean, twoVariables = true): Data[] { | |
| const xs = t.path.map((p) => p.position[0]); | |
| const ys = twoVariables ? t.path.map((p) => p.position[1]) : t.path.map((p) => p.value); | |
| const fs = t.path.map((p) => p.value); | |
| const hovertemplate = twoVariables | |
| ? "x = %{x:.4g}, y = %{y:.4g}<br>f = %{customdata:.4g}<extra>%{fullData.name}</extra>" | |
| : "x = %{x:.4g}<br>f(x) = %{y:.4g}<extra>%{fullData.name}</extra>"; | |
| const last = t.path.length - 1; | |
| if (threeD) { | |
| return [ | |
| { x: xs, y: ys, z: fs, type: "scatter3d", mode: "lines+markers", name: t.label, legendgroup: t.label, | |
| line: { color: t.color, width: 5 }, marker: { color: t.color, size: 2.5 }, hoverinfo: "skip" }, | |
| { x: xs.slice(last), y: ys.slice(last), z: fs.slice(last), type: "scatter3d", mode: "markers", name: t.label, | |
| legendgroup: t.label, showlegend: false, marker: { color: t.color, size: 6, line: { color: "white", width: 1 } }, hoverinfo: "skip" }, | |
| ] as Data[]; | |
| } | |
| return [ | |
| { x: xs, y: ys, customdata: fs, type: "scatter", mode: "lines+markers", name: t.label, legendgroup: t.label, | |
| line: { color: t.color, width: 2 }, marker: { color: t.color, size: 7 }, hovertemplate }, | |
| { x: xs.slice(last), y: ys.slice(last), customdata: fs.slice(last), type: "scatter", mode: "markers", name: t.label, | |
| legendgroup: t.label, showlegend: false, marker: { color: t.color, size: 14, line: { color: "white", width: 2 } }, hovertemplate }, | |
| ] as Data[]; | |
| } | |
| // The gradient as arrows (pointing uphill), scaled so a typical arrow spans | |
| // most of the gap between grid points | |
| function arrowTraces({ arrows }: SurfaceValues, [x0, x1]: Range, [y0, y1]: Range): Data[] { | |
| // Arrows per side of the grid | |
| const side = Math.sqrt(arrows.x.length); | |
| // Lengths as fractions of the view, so x and y scale alike | |
| const lengths = arrows.u.map((u, i) => Math.hypot((u ?? NaN) / (x1 - x0), (arrows.v[i] ?? NaN) / (y1 - y0))); | |
| const finite = lengths.filter((l) => Number.isFinite(l) && l > 0).sort((a, b) => a - b); | |
| const typical = finite[Math.floor(0.9 * (finite.length - 1))] ?? 1; | |
| const shaftX: (number | null)[] = []; | |
| const shaftY: (number | null)[] = []; | |
| const headX: number[] = []; | |
| const headY: number[] = []; | |
| const angles: number[] = []; | |
| const hover: number[][] = []; | |
| arrows.x.forEach((x, i) => { | |
| const u = arrows.u[i]; | |
| const v = arrows.v[i]; | |
| if (u === null || v === null || !(lengths[i] > 0)) { | |
| return; | |
| } | |
| // At most 0.9 of a grid cell (1 / side of the view) | |
| const scale = (0.9 * Math.min(1, lengths[i] / typical)) / (lengths[i] * side); | |
| const dx = u * scale; | |
| const dy = v * scale; | |
| const y = arrows.y[i]; | |
| shaftX.push(x, x + dx, null); | |
| shaftY.push(y, y + dy, null); | |
| headX.push(x + dx); | |
| headY.push(y + dy); | |
| // Plotly turns markers clockwise from up, in screen space | |
| angles.push((Math.atan2(dx / (x1 - x0), dy / (y1 - y0)) * 180) / Math.PI); | |
| hover.push([x, y, u, v]); | |
| }); | |
| const color = "rgba(255, 255, 255, 0.85)"; | |
| return [ | |
| { x: shaftX, y: shaftY, type: "scatter", mode: "lines", line: { color, width: 1 }, hoverinfo: "skip", showlegend: false }, | |
| { | |
| x: headX, y: headY, customdata: hover, type: "scatter", mode: "markers", showlegend: false, | |
| marker: { symbol: "triangle-up", size: 6, color, angle: angles, angleref: "up" }, | |
| hovertemplate: "x = %{customdata[0]:.3g}, y = %{customdata[1]:.3g}<br>∇f = (%{customdata[2]:.3g}, %{customdata[3]:.3g})<extra></extra>", | |
| }, | |
| ] as Data[]; | |
| } | |