Merge pull request #27 from aerubanov/convergence-diagnostic

add trace plot

Author: aerubanov

src/App.css

@@ -16,7 +16,7 @@
 
 .App-main {
   background: var(--color-bg-primary);
-  overflow: hidden;
+  overflow-y: auto;
   position: relative;
 }
 

src/App.jsx

@@ -1,6 +1,7 @@
 import './App.css';
 import Controls from './components/Controls';
 import Visualizer from './components/Visualizer';
+import TracePlots from './components/TracePlots';
 import useSamplingController from './hooks/useSamplingController';
 
 function App() {
@@ -36,6 +37,7 @@ function App() {
     setUseSecondChain,
     setInitialPosition2,
     setSeed2,
+    burnIn,
   } = useSamplingController();
 
   return (
@@ -109,6 +111,16 @@ function App() {
           acceptedSamples2={samples2}
           useSecondChain={useSecondChain}
         />
+        <div className="trace-plots-section">
+          {contourData && (
+            <TracePlots
+              samples={samples}
+              samples2={samples2}
+              burnIn={burnIn}
+              useSecondChain={useSecondChain}
+            />
+          )}
+        </div>
       </div>
     </div>
   );

src/components/TracePlots.css

@@ -0,0 +1,27 @@
+.trace-plots-container {
+  display: flex;
+  flex-direction: column;
+  gap: 20px;
+  width: 100%;
+  margin-top: 20px;
+  padding: 10px;
+  background-color: #ffffff;
+  border-radius: 8px;
+  box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+}
+
+.trace-plot-wrapper {
+  width: 100%;
+  display: flex;
+  flex-direction: column;
+  align-items: center;
+}
+
+.trace-title {
+  margin: 0 0 10px 0;
+  font-size: 14px;
+  color: #1a1a1a;
+  font-weight: 600;
+  align-self: flex-start;
+  padding-left: 10px;
+}

src/components/TracePlots.jsx

@@ -0,0 +1,117 @@
+import './TracePlots.css';
+import Plot from 'react-plotly.js';
+import PropTypes from 'prop-types';
+import { TRACE_PLOT, HMC_SAMPLER } from '../utils/plotConfig.json';
+import { createTracePlotTrace } from '../utils/plotFunctions';
+
+function TracePlots({ samples, samples2, burnIn, useSecondChain }) {
+  const commonLayout = {
+    ...TRACE_PLOT.layout,
+    showlegend: true,
+    legend: {
+      orientation: 'h',
+      y: -0.2, // Move legend below plot
+    },
+  };
+
+  const xConfig = {
+    displayModeBar: false,
+    responsive: true,
+  };
+
+  // Generate traces for X
+  const xTraces = [];
+  if (samples && samples.length > 0) {
+    xTraces.push(
+      ...createTracePlotTrace(
+        samples,
+        'x',
+        burnIn,
+        HMC_SAMPLER.styles.primaryColor,
+        'Chain 1'
+      )
+    );
+  }
+
+  if (useSecondChain && samples2 && samples2.length > 0) {
+    xTraces.push(
+      ...createTracePlotTrace(
+        samples2,
+        'x',
+        burnIn,
+        HMC_SAMPLER.styles.secondaryColor,
+        'Chain 2'
+      )
+    );
+  }
+
+  // Generate traces for Y
+  const yTraces = [];
+  if (samples && samples.length > 0) {
+    yTraces.push(
+      ...createTracePlotTrace(
+        samples,
+        'y',
+        burnIn,
+        HMC_SAMPLER.styles.primaryColor,
+        'Chain 1'
+      )
+    );
+  }
+
+  if (useSecondChain && samples2 && samples2.length > 0) {
+    yTraces.push(
+      ...createTracePlotTrace(
+        samples2,
+        'y',
+        burnIn,
+        HMC_SAMPLER.styles.secondaryColor,
+        'Chain 2'
+      )
+    );
+  }
+
+  return (
+    <div className="trace-plots-container">
+      <div className="trace-plot-wrapper">
+        <h4 className="trace-title">X Trace</h4>
+        <Plot
+          data={xTraces}
+          layout={{ ...commonLayout, title: '' }} // Remove title from Plotly, use HTML headers
+          config={xConfig}
+          style={{ width: '100%', height: '300px' }}
+          useResizeHandler={true}
+        />
+      </div>
+      <div className="trace-plot-wrapper">
+        <h4 className="trace-title">Y Trace</h4>
+        <Plot
+          data={yTraces}
+          layout={{ ...commonLayout, title: '' }}
+          config={xConfig}
+          style={{ width: '100%', height: '300px' }}
+          useResizeHandler={true}
+        />
+      </div>
+    </div>
+  );
+}
+
+TracePlots.propTypes = {
+  samples: PropTypes.arrayOf(
+    PropTypes.shape({
+      x: PropTypes.number.isRequired,
+      y: PropTypes.number.isRequired,
+    })
+  ),
+  samples2: PropTypes.arrayOf(
+    PropTypes.shape({
+      x: PropTypes.number.isRequired,
+      y: PropTypes.number.isRequired,
+    })
+  ),
+  burnIn: PropTypes.number,
+  useSecondChain: PropTypes.bool,
+};
+
+export default TracePlots;

src/hooks/useSamplingController.js

@@ -35,6 +35,9 @@ export default function useSamplingController() {
   const [rejectedCount2, setRejectedCount2] = useState(0);
   const [seed2, setSeed2State] = useState(null);
 
+  // Visualization params
+  const [burnIn] = useState(10);
+
   // Refs to hold instances/values that don't trigger re-renders or need to be accessed in loops
   const logpInstanceRef = useRef(null);
   const currentParticleRef = useRef(null); // { q, p }
@@ -336,5 +339,6 @@ export default function useSamplingController() {
     setUseSecondChain,
     setInitialPosition2,
     setSeed2,
+    burnIn,
   };
 }

src/utils/plotConfig.json

@@ -36,6 +36,34 @@
       }
     }
   },
+  "TRACE_PLOT": {
+    "styles": {
+      "burnInOpacity": 0.3,
+      "lineWidth": 2
+    },
+    "layout": {
+      "font": {
+        "color": "#1a1a1a",
+        "family": "Inter, system-ui, sans-serif",
+        "size": 10
+      },
+      "xaxis": {
+        "title": "Iteration",
+        "showgrid": false,
+        "zeroline": false
+      },
+      "yaxis": {
+        "showgrid": true,
+        "zeroline": false
+      },
+      "margin": {
+        "l": 50,
+        "r": 20,
+        "t": 30,
+        "b": 40
+      }
+    }
+  },
   "GENERAL": {
     "layout": {
       "autosize": true,

src/utils/plotFunctions.js

@@ -1,4 +1,4 @@
-import { CONTOUR, HMC_SAMPLER } from './plotConfig.json';
+import { CONTOUR, HMC_SAMPLER, TRACE_PLOT } from './plotConfig.json';
 
 /**
  * Creates a Plotly contour trace configuration
@@ -142,3 +142,94 @@ export function createSamplesTrace(
     hovertemplate: 'Sample<br>x: %{x:.2f}<br>y: %{y:.2f}<extra></extra>',
   };
 }
+
+/**
+ * Converts a hex color to rgba string
+ * @param {string} hex - Hex color string (e.g., "#ff0000")
+ * @param {number} alpha - Alpha value (0-1)
+ * @returns {string} Rgba color string
+ */
+function hexToRgba(hex, alpha) {
+  const r = parseInt(hex.slice(1, 3), 16);
+  const g = parseInt(hex.slice(3, 5), 16);
+  const b = parseInt(hex.slice(5, 7), 16);
+  return `rgba(${r}, ${g}, ${b}, ${alpha})`;
+}
+
+/**
+ * Creates Plotly traces for trace plots (iteration vs value)
+ * @param {Array<{x: number, y: number}>} samples - Array of accepted sample points
+ * @param {string} axis - 'x' or 'y' to plot
+ * @param {number} burnIn - Number of samples to treat as burn-in
+ * @param {string} [color] - Color for the valid samples
+ * @param {string} [name] - Name for the valid samples trace
+ * @returns {object[]} Array of Plotly trace objects (burn-in and valid)
+ */
+export function createTracePlotTrace(
+  samples,
+  axis,
+  burnIn = 0,
+  color = HMC_SAMPLER.styles.primaryColor,
+  name = 'Trace'
+) {
+  if (!samples || !Array.isArray(samples) || samples.length === 0) {
+    return [];
+  }
+
+  const traces = [];
+  const validOpacity = 1.0;
+  const burnInOpacity = TRACE_PLOT.styles.burnInOpacity;
+  const lineWidth = TRACE_PLOT.styles.lineWidth;
+
+  // Split samples into burn-in and valid
+  let burnInSamples = [];
+  let validSamples = [];
+
+  if (burnIn > 0) {
+    // If we have valid samples after burn-in, include the first one in burn-in set to connect lines
+    const endIndex = samples.length > burnIn ? burnIn + 1 : burnIn;
+    burnInSamples = samples.slice(0, endIndex);
+    validSamples = samples.slice(burnIn);
+  } else {
+    validSamples = samples;
+  }
+
+  // Helper to create a single trace part
+  const createSubTrace = (data, startIndex, opacity, traceName, showLegend) => {
+    const iterations = data.map((_, i) => i + startIndex);
+    const values = data.map((p) => p[axis]);
+
+    // Use RGBA for color to ensure opacity works reliably on lines
+    const traceColor = opacity < 1 ? hexToRgba(color, opacity) : color;
+
+    return {
+      type: 'scatter',
+      mode: 'lines',
+      x: iterations,
+      y: values,
+      line: {
+        color: traceColor,
+        width: lineWidth,
+      },
+      // Remove top-level opacity to rely on rgba color
+      // opacity: opacity,
+      name: traceName,
+      showlegend: showLegend,
+      hovertemplate: `Iter: %{x}<br>${axis}: %{y:.2f}<extra></extra>`,
+    };
+  };
+
+  // Add burn-in trace
+  if (burnInSamples.length > 0) {
+    traces.push(
+      createSubTrace(burnInSamples, 0, burnInOpacity, `${name} (Burn-in)`, true)
+    );
+  }
+
+  // Add valid samples trace
+  if (validSamples.length > 0) {
+    traces.push(createSubTrace(validSamples, burnIn, validOpacity, name, true));
+  }
+
+  return traces;
+}