more elegantly handle samples, properly use sample data in SampleSets for bands

packages/components/src/components/DistributionChart.tsx

@@ -19,6 +19,7 @@ import { NumberShower } from "./NumberShower";
 import { Plot, parsePlot } from "../lib/plotParser";
 import { flattenResult } from "../lib/utility";
 import { hasMassBelowZero } from "../lib/distributionUtils";
+import { point } from "@quri/squiggle-lang/src/js/distribution";
 
 export type DistributionPlottingSettings = {
   /** Whether to show a summary of means, stdev, percentiles etc */
@@ -65,6 +66,7 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
           // color: x.color, // not supported yet
           continuous: shape.continuous,
           discrete: shape.discrete,
+          samples: [] as point[],
         }))
       )
     );
@@ -77,6 +79,12 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
       );
     }
 
+    // if this is a sample set, include the samples
+    const t = plot?.distributions[0]?.distribution?.t;
+    if (t.tag === "SampleSet") {
+      shapes.value[0].samples = t.value.map((v) => ({ x: v, y: 0 }));
+    }
+    
     const spec = buildVegaSpec(props);
 
     let widthProp = width ? width : size.width;
@@ -106,12 +114,13 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
               discrete: val.discrete.map((p) => {
                 return { dateTime: p.x, y: p.y };
               }),
+              samples: val.samples.map((p) => {
+                return { dateTime: p, y: 0 };
+              }),
             };
           })
         : shapes.value;
 
-    console.log({ data });
-    
     return (
       <div style={{ width: widthProp }}>
         {logX && shapes.value.some(hasMassBelowZero) ? (

packages/components/src/components/SquiggleChart.tsx

@@ -51,8 +51,6 @@ export interface SquiggleChartProps {
   minX?: number;
   /** Specify the upper bound of the x scale */
   maxX?: number;
-  // /** Whether to include the sample band below the chart */
-  sample?: boolean;
   /** Whether the x-axis should be dates or numbers */
   xAxis?: "number" | "dateTime";
   /** Whether to show vega actions to the user, so they can copy the chart spec */
@@ -83,7 +81,6 @@ export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
     maxX,
     color,
     title,
-    sample = false,
     xAxis = "number",
     distributionChartActions,
     enableLocalSettings = false,
@@ -96,7 +93,7 @@ export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
       onChange,
       executionId,
     });
-    console.log({ result });
+    
     const distributionPlotSettings = {
       showSummary,
       logX,
@@ -107,7 +104,6 @@ export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
       color,
       title,
       xAxis,
-      sample,
       actions: distributionChartActions,
     };
 

packages/components/src/lib/distributionSpecBuilder.ts

@@ -81,7 +81,6 @@ export function buildVegaSpec(
     maxX,
     logX,
     expY,
-    sample = false,
     xAxis = "number",
   } = specOptions;
 
@@ -110,7 +109,7 @@ export function buildVegaSpec(
     width: width,
     height: 100,
     padding: 5,
-    data: [{ name: "data" }, { name: "domain" }],
+    data: [{ name: "data" }, { name: "domain" }, { name: "samples" }],
     signals: [
       {
         name: "hover",
@@ -218,7 +217,6 @@ export function buildVegaSpec(
                   },
                 },
               },
-             
             ],
           },
           {
@@ -299,6 +297,35 @@ export function buildVegaSpec(
               },
             ],
           },
+          {
+            name: "sample_distributions",
+            type: "group",
+            from: {
+              facet: {
+                name: "sample_facet",
+                data: "distribution_facet",
+                field: "samples",
+              },
+            },
+            marks: [
+              {
+                name: "samples",
+                type: "rect",
+                from: { data: "sample_facet" },
+                encode: {
+                  enter: {
+                    x: { scale: "xscale", field: dateTime ? "dateTime" : "x" },
+                    width: { value: 0.1 },
+
+                    y: { value: 25, offset: { signal: "height" } },
+                    height: { value: 5 },
+                    fill: { value: "steelblue" },
+                    fillOpacity: { value: 1 },
+                  },
+                },
+              },
+            ],
+          },
         ],
       },
       {
@@ -384,37 +411,5 @@ export function buildVegaSpec(
     }),
   };
 
-  // include the band at the bottom if specified in the React component
-  if (sample) {
-    spec.marks?.push({
-      name: "sample_distributions",
-      type: "group",
-      from: {
-        facet: {
-          name: "distribution_facet",
-          data: "domain",
-          groupby: ["name"],
-        },
-      },
-      marks: [
-        {
-          name: "samples",
-          type: "rect",
-          from: { data: "distribution_facet" },
-          encode: {
-            enter: {
-              x: { scale: "xscale", field: dateTime ? "dateTime" : "x" },
-              width: { value: 0.5 },
-
-              y: { value: 25, offset: { signal: "height" } },
-              height: { value: 5 },
-              fill: { value: "steelblue" },
-              fillOpacity: { value: 0.8 },
-            },
-          },
-        },
-      ],
-    });
-  }
   return spec;
 }

packages/components/src/stories/SquiggleChart.stories.mdx

@@ -59,7 +59,6 @@ could be continuous, discrete or mixed.
     args={{
       code: "SampleSet.fromDist(normal(5,2))",
       width,
-      sample: true,
     }}
   >
     {Template.bind({})}