Spaces:
Running
Running
| import { useMemo, useRef } from "react"; | |
| import Plot from "react-plotly.js"; | |
| import type { BivariatePlotView, ComparisonEntry, TrajectoryValues, PlotData } from "./types"; | |
| import { Card } from "@elvis/ui"; | |
| interface OptimizationPlotProps { | |
| comparisons?: (ComparisonEntry & { values?: TrajectoryValues })[]; | |
| data: PlotData; | |
| isLoading: boolean; | |
| xlim: [number, number]; | |
| ylim: [number, number]; | |
| setAxisLimits: (xlim: [number, number], ylim: [number, number]) => void; | |
| currentIteration: number; | |
| showGradientField: boolean; | |
| bivariatePlotView: BivariatePlotView; | |
| } | |
| export default function OptimizationPlot({ comparisons, data, isLoading, xlim, ylim, setAxisLimits, currentIteration, showGradientField, bivariatePlotView }: OptimizationPlotProps) { | |
| const isSurface = bivariatePlotView === "Function surface" && Boolean(data.functionValues?.z); | |
| const comparisonTraces = (comparisons ?? []).flatMap((entry, index) => entry.visible ? [ | |
| { x: entry.values?.x ?? [], y: entry.values?.y ?? [], z: isSurface ? entry.values?.z ?? [] : undefined, | |
| type: isSurface ? "scatter3d" as const : "scatter" as const, mode: "lines+markers" as const, | |
| name: `${index + 1}. ${entry.settings.algorithm}`, legendgroup: String(entry.id), showlegend: true, | |
| line: { color: entry.color, width: isSurface ? 6 : 2 }, | |
| marker: { color: entry.color, size: isSurface ? 2.5 : 6 } }, | |
| { x: entry.values?.x.slice(-1) ?? [], y: entry.values?.y.slice(-1) ?? [], z: isSurface ? entry.values?.z?.slice(-1) ?? [] : undefined, | |
| type: isSurface ? "scatter3d" as const : "scatter" as const, mode: "markers" as const, | |
| name: `${index + 1}. ${entry.settings.algorithm}`, legendgroup: String(entry.id), showlegend: false, | |
| marker: { color: entry.color, size: isSurface ? 4 : 12, line: { color: "white", width: isSurface ? 1 : 2 } } } | |
| ] : []); | |
| // data | |
| let x: number[] = data.functionValues ? data.functionValues.x : []; | |
| let y: number[] = data.functionValues ? data.functionValues.y : []; | |
| let z: number[][] = data.functionValues && data.functionValues.z ? data.functionValues.z : []; | |
| let gradientX: number[][] = data.functionValues?.gradientX ?? []; | |
| let gradientY: number[][] = data.functionValues?.gradientY ?? []; | |
| console.log('gradientX', gradientX) | |
| console.log('gradientY', gradientY) | |
| let trajX: number[] = data.trajectoryValues ? data.trajectoryValues.x : []; | |
| let trajY: number[] = data.trajectoryValues ? data.trajectoryValues.y : []; | |
| let trajZ: number[] = data.trajectoryValues && data.trajectoryValues.z ? data.trajectoryValues.z : []; | |
| const colorScaleRange = useMemo<[number, number] | null>(() => { | |
| if (z.length === 0) return null; | |
| let zMin = Infinity; | |
| let zMax = -Infinity; | |
| for (const row of z) { | |
| for (const value of row) { | |
| if (!Number.isFinite(value)) continue; | |
| zMin = Math.min(zMin, value); | |
| zMax = Math.max(zMax, value); | |
| } | |
| } | |
| if (!Number.isFinite(zMin) || !Number.isFinite(zMax)) return null; | |
| const padding = zMax === zMin ? Math.max(Math.abs(zMin) * 0.1, 1) : (zMax - zMin) * 0.1; | |
| return [zMin - padding, zMax + padding]; | |
| }, [z]); | |
| const gradientField = useMemo(() => { | |
| const emptyField = { | |
| shaftX: [] as (number | null)[], | |
| shaftY: [] as (number | null)[], | |
| headX: [] as number[], | |
| headY: [] as number[], | |
| headAngles: [] as number[], | |
| hoverData: [] as number[][], | |
| }; | |
| if ( | |
| !showGradientField | |
| || gradientX.length !== y.length | |
| || gradientY.length !== y.length | |
| || x.length === 0 | |
| || y.length === 0 | |
| ) return emptyField; | |
| const rowStep = 5 //Math.max(1, Math.ceil(y.length / 10)); | |
| const columnStep = 5 //Math.max(1, Math.ceil(x.length / 10)); | |
| const shaftX: (number | null)[] = []; | |
| const shaftY: (number | null)[] = []; | |
| const headX: number[] = []; | |
| const headY: number[] = []; | |
| const headAngles: number[] = []; | |
| const hoverData: number[][] = []; | |
| const gradientLengths: number[] = []; | |
| for (let row = 0; row < gradientX.length; row++) { | |
| for (let column = 0; column < gradientX[row].length; column++) { | |
| const gx = gradientX[row][column]; | |
| const gy = gradientY[row][column]; | |
| const length = Math.hypot(gx, gy); | |
| if (Number.isFinite(length)) { | |
| gradientLengths.push(length); | |
| } | |
| } | |
| } | |
| const averageGradientLength = gradientLengths.length > 0 ? | |
| gradientLengths.reduce((sum, length) => sum + length, 0) / gradientLengths.length | |
| : 0; | |
| console.log('average_grad_length', averageGradientLength); | |
| console.log('cell diagonal', Math.hypot(x[columnStep]-x[0], y[rowStep]-y[0])); | |
| const grad_scale = 1.8*averageGradientLength / Math.hypot(x[columnStep]-x[0], y[rowStep]-y[0]); | |
| console.log('grad_scale', grad_scale); | |
| for (let row = 0; row < y.length; row += rowStep) { | |
| if (!gradientX[row] || !gradientY[row]) continue; | |
| for (let column = 0; column < x.length; column += columnStep) { | |
| const gx = gradientX[row][column]; | |
| const gy = gradientY[row][column]; | |
| const magnitude = Math.hypot(gx, gy); | |
| if (!Number.isFinite(magnitude) || magnitude === 0) continue; | |
| const vx = gx / grad_scale; | |
| const vy = gy / grad_scale; | |
| const arrowX = x[column] + vx; | |
| const arrowY = y[row] + vy; | |
| shaftX.push(x[column], arrowX, null); | |
| shaftY.push(y[row], arrowY, null); | |
| headX.push(arrowX); | |
| headY.push(arrowY); | |
| // Plotly marker angles rotate clockwise from the upward direction. | |
| headAngles.push(Math.atan2(vx, vy) * 180 / Math.PI); | |
| hoverData.push([x[column], y[row], gx, gy]); | |
| } | |
| } | |
| return { shaftX, shaftY, headX, headY, headAngles, hoverData }; | |
| }, [showGradientField, x, y, gradientX, gradientY]); | |
| const plotRef = useRef<any>(null); | |
| const layoutRef = useRef<any>({ | |
| dragmode: 'pan', | |
| showlegend: false, | |
| xaxis: { | |
| title: { text: 'x' }, | |
| range: xlim, | |
| }, | |
| yaxis: { | |
| title: { text: 'y' }, | |
| range: ylim, | |
| }, | |
| margin: { t: 40, r: 40, b: 40, l: 40 } | |
| }) | |
| const hasNoData = !data.functionValues && !data.trajectoryValues; | |
| function zoomSurfaceAxes(event: React.WheelEvent) { | |
| event.preventDefault(); | |
| const factor = event.deltaY > 0 ? 1.2 : 1 / 1.2; | |
| const resizeRange = (range: [number, number]): [number, number] => { | |
| const center = (range[0] + range[1]) / 2; | |
| const halfWidth = (range[1] - range[0]) * factor / 2; | |
| return [center - halfWidth, center + halfWidth]; | |
| }; | |
| setAxisLimits(resizeRange(xlim), resizeRange(ylim)); | |
| } | |
| if (isLoading && hasNoData) { | |
| return ( | |
| <Card className="min-h-[320px] flex flex-1 items-center justify-center p-6"> | |
| <div className="text-center"> | |
| <div className="text-lg font-medium">Loading...</div> | |
| </div> | |
| </Card> | |
| ); | |
| } | |
| if (z.length === 0) { | |
| return ( | |
| <Card className="min-h-[320px] flex-1"> | |
| <Plot | |
| ref={plotRef} | |
| onRelayout={(event) => { | |
| const x0 = event['xaxis.range[0]']; | |
| const x1 = event['xaxis.range[1]']; | |
| const y0 = event['yaxis.range[0]']; | |
| const y1 = event['yaxis.range[1]']; | |
| if ( | |
| typeof x0 === "number" | |
| && typeof x1 === "number" | |
| && typeof y0 === "number" | |
| && typeof y1 === "number" | |
| ) { | |
| setAxisLimits([x0, x1], [y0, y1]); | |
| } | |
| }} | |
| data={[ | |
| { | |
| x: x, | |
| y: y, | |
| type: 'scatter', | |
| mode: 'lines', | |
| line: { color: '#1f77b4', width: 2 }, | |
| hoverinfo: "skip", | |
| showlegend: false, | |
| }, | |
| { | |
| x: trajX, | |
| y: trajY, | |
| type: 'scatter', | |
| mode: 'lines+markers', | |
| line: { color: '#d97871', width: 2 }, | |
| marker: { color: '#d97871', size: 10 }, | |
| hoverinfo: "skip", | |
| showlegend: false, | |
| }, | |
| { | |
| x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [], | |
| y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [], | |
| type: 'scatter', | |
| mode: 'markers', | |
| marker: { color: 'red', size: 12 }, | |
| hoverinfo: "skip", | |
| showlegend: false, | |
| } | |
| ,...comparisonTraces | |
| ]} | |
| layout={{ ...layoutRef.current, showlegend: comparisons !== undefined, legend: { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.75)" } }} | |
| style={{ width: '100%', height: '100%' }} | |
| config={{ | |
| responsive: true, | |
| displayModeBar: true, | |
| scrollZoom: true, | |
| }} | |
| /> | |
| </Card> | |
| ); | |
| } else { | |
| if (isSurface) { | |
| return ( | |
| <div className="min-h-[420px] flex-1" onWheel={zoomSurfaceAxes}> | |
| <Card className="h-full"> | |
| <Plot | |
| ref={plotRef} | |
| data={[ | |
| { | |
| x, y, z, | |
| type: 'surface', | |
| colorscale: 'Viridis', | |
| opacity: 0.82, | |
| showscale: true, | |
| hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>z: %{z:.5g}<extra></extra>', | |
| }, | |
| { | |
| x: trajX, y: trajY, z: trajZ, | |
| type: 'scatter3d', mode: 'lines+markers', | |
| line: { color: '#d97871', width: 6 }, | |
| marker: { color: '#d97871', size: 2.5 }, | |
| hoverinfo: 'skip', showlegend: false, | |
| }, | |
| { | |
| x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [], | |
| y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [], | |
| z: trajZ[currentIteration] !== undefined ? [trajZ[currentIteration]] : [], | |
| type: 'scatter3d', mode: 'markers', | |
| marker: { color: 'red', size: 7 }, | |
| hoverinfo: 'skip', showlegend: false, | |
| }, | |
| ...comparisonTraces, | |
| ]} | |
| layout={{ | |
| // Keep the user's camera angle while animation updates the trajectory data. | |
| uirevision: 'bivariate-surface', | |
| dragmode: 'orbit', | |
| showlegend: comparisons !== undefined, | |
| legend: { x: 0.02, y: 0.98, xanchor: 'left', yanchor: 'top', bgcolor: 'rgba(255,255,255,0.75)' }, | |
| scene: { | |
| aspectmode: 'cube', | |
| xaxis: { title: { text: 'x' }, range: xlim, autorange: false }, | |
| yaxis: { title: { text: 'y' }, range: ylim, autorange: false }, | |
| zaxis: { title: { text: 'z' }, range: colorScaleRange ?? undefined, autorange: false }, | |
| }, | |
| margin: { t: 30, r: 20, b: 20, l: 20 }, | |
| }} | |
| style={{ width: '100%', height: '100%' }} | |
| config={{ responsive: true, displayModeBar: true, scrollZoom: false }} | |
| /> | |
| </Card> | |
| </div> | |
| ); | |
| } | |
| return ( | |
| <Card className="min-h-[420px] flex-1"> | |
| <Plot | |
| ref={plotRef} | |
| onRelayout={(event) => { | |
| const x0 = event['xaxis.range[0]']; | |
| const x1 = event['xaxis.range[1]']; | |
| const y0 = event['yaxis.range[0]']; | |
| const y1 = event['yaxis.range[1]']; | |
| if ( | |
| typeof x0 === "number" | |
| && typeof x1 === "number" | |
| && typeof y0 === "number" | |
| && typeof y1 === "number" | |
| ) { | |
| setAxisLimits([x0, x1], [y0, y1]); | |
| } | |
| }} | |
| data={[ | |
| { | |
| x: x, | |
| y: y, | |
| z: z, | |
| zmin: colorScaleRange?.[0], | |
| zmax: colorScaleRange?.[1], | |
| type: 'contour', | |
| colorscale: 'Viridis', | |
| hoverinfo: "skip", | |
| showlegend: false, | |
| contours: { | |
| coloring: "heatmap", | |
| showlines: false, | |
| } | |
| }, | |
| { | |
| x: gradientField.shaftX, | |
| y: gradientField.shaftY, | |
| type: 'scatter', | |
| mode: 'lines', | |
| line: { color: 'rgba(255, 255, 255, 0.8)', width: 1 }, | |
| hoverinfo: 'skip', | |
| showlegend: false, | |
| cliponaxis: true, | |
| }, | |
| { | |
| x: gradientField.headX, | |
| y: gradientField.headY, | |
| customdata: gradientField.hoverData, | |
| type: 'scatter', | |
| mode: 'markers', | |
| marker: { | |
| symbol: 'triangle-up', | |
| size: 5, | |
| color: 'rgba(255, 255, 255, 0.8)', | |
| angle: gradientField.headAngles, | |
| angleref: 'up', | |
| } as any, | |
| hovertemplate: | |
| 'x: %{customdata[0]:.5g}<br>' | |
| + 'y: %{customdata[1]:.5g}<br>' | |
| + 'gx: %{customdata[2]:.5g}<br>' | |
| + 'gy: %{customdata[3]:.5g}<extra></extra>', | |
| showlegend: false, | |
| cliponaxis: true, | |
| }, | |
| { | |
| x: trajX, | |
| y: trajY, | |
| customdata: trajZ, | |
| type: 'scatter', | |
| mode: 'lines+markers', | |
| line: { color: '#d97871', width: 2 }, | |
| marker: { color: '#d97871', size: 10 }, | |
| hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>f(x, y): %{customdata:.5g}<extra></extra>', | |
| showlegend: false, | |
| }, | |
| { | |
| x: trajX[currentIteration] !== undefined ? [trajX[currentIteration]] : [], | |
| y: trajY[currentIteration] !== undefined ? [trajY[currentIteration]] : [], | |
| customdata: trajZ[currentIteration] !== undefined ? [trajZ[currentIteration]] : [], | |
| type: 'scatter', | |
| mode: 'markers', | |
| marker: { color: 'red', size: 12 }, | |
| hovertemplate: 'x: %{x:.5g}<br>y: %{y:.5g}<br>f(x, y): %{customdata:.5g}<extra></extra>', | |
| showlegend: false, | |
| } | |
| ,...comparisonTraces | |
| ]} | |
| layout={{ ...layoutRef.current, annotations: [], showlegend: comparisons !== undefined, legend: { x: 0.02, y: 0.98, xanchor: "left", yanchor: "top", bgcolor: "rgba(255,255,255,0.75)" } }} | |
| style={{ width: '100%', height: '100%' }} | |
| config={{ | |
| responsive: true, | |
| displayModeBar: true, | |
| scrollZoom: true, | |
| }} | |
| /> | |
| </Card> | |
| ); | |
| } | |
| } | |