Refactor specification to include discrete

packages/components/src/components/DistributionChart.tsx

@@ -5,6 +5,7 @@ import {
   distributionError,
   distributionErrorToString,
   squiggleExpression,
+  resultMap,
 } from "@quri/squiggle-lang";
 import { Vega } from "react-vega";
 import { ErrorAlert } from "./Alert";
@@ -24,8 +25,10 @@ export type DistributionPlottingSettings = {
   showControls: boolean;
 } & DistributionChartSpecOptions;
 
+export type LabeledDistribution = { name: string; distribution: Distribution };
+
 export type Plot = {
-  distributions: Distribution[];
+  distributions: LabeledDistribution[];
 };
 
 export type DistributionChartProps = {
@@ -36,19 +39,29 @@ export type DistributionChartProps = {
 } & DistributionPlottingSettings;
 
 export function defaultPlot(distribution: Distribution): Plot {
-  return { distributions: [distribution] };
+  return { distributions: [{ name: "default", distribution }] };
 }
+
 export function makePlot(expression: {
   [key: string]: squiggleExpression;
 }): Plot | void {
   if (expression["distributions"].tag === "array") {
-    let distributions: Distribution[] = expression["distributions"].value
+    let distributions: LabeledDistribution[] = expression["distributions"].value
       .map((x) => {
-        if (x.tag === "distribution") {
+        if (
-          return x.value;
+          x.tag === "record" &&
+          x.value["name"] &&
+          x.value["name"].tag === "string" &&
+          x.value["distribution"] &&
+          x.value["distribution"].tag === "distribution"
+        ) {
+          return {
+            name: x.value["name"].value,
+            distribution: x.value["distribution"].value,
+          };
         }
       })
-      .filter((x): x is Distribution => x !== undefined);
+      .filter((x): x is LabeledDistribution => x !== undefined);
     return { distributions };
   }
 }
@@ -91,7 +104,15 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
   React.useEffect(() => setExpY(expY), [expY]);
 
   const [sized] = useSize((size) => {
-    let shapes = flattenResult(plot.distributions.map((x) => x.pointSet()));
+    let shapes = flattenResult(
+      plot.distributions.map((x) =>
+        resultMap(x.distribution.pointSet(), (shape) => ({
+          name: x.name,
+          continuous: shape.continuous,
+          discrete: shape.discrete,
+        }))
+      )
+    );
     if (shapes.tag === "Error") {
       return (
         <ErrorAlert heading="Distribution Error">
@@ -116,20 +137,17 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
       );
       widthProp = 20;
     }
-    let continuousPoints = shapes.value.flatMap((shape, i) =>
+    const domain = shapes.value.flatMap((shape) =>
-      shape.continuous.map((point) => ({ ...point, name: i + 1 }))
+      shape.discrete.concat(shape.continuous)
-    );
-    let discretePoints = shapes.value.flatMap((shape, i) =>
-      shape.discrete.map((point) => ({ ...point, name: i + 1 }))
     );
+    console.log(shapes.value);
 
-    console.log(continuousPoints);
     return (
       <div style={{ width: widthProp }}>
         {!(isLogX && massBelow0) ? (
           <Vega
             spec={spec}
-            data={{ con: continuousPoints, dis: discretePoints }}
+            data={{ data: shapes.value, domain }}
             width={widthProp - 10}
             height={height}
             actions={actions}
@@ -141,7 +159,7 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
         )}
         <div className="flex justify-center">
           {showSummary && plot.distributions.length == 1 && (
-            <SummaryTable distribution={plot.distributions[0]} />
+            <SummaryTable distribution={plot.distributions[0].distribution} />
           )}
         </div>
         {showControls && (

packages/components/src/components/FunctionChart1Dist.tsx

@@ -16,6 +16,7 @@ import * as percentilesSpec from "../vega-specs/spec-percentiles.json";
 import {
   DistributionChart,
   DistributionPlottingSettings,
+  defaultPlot,
 } from "./DistributionChart";
 import { NumberShower } from "./NumberShower";
 import { ErrorAlert } from "./Alert";
@@ -177,7 +178,7 @@ export const FunctionChart1Dist: React.FC<FunctionChart1DistProps> = ({
   let showChart =
     mouseItem.tag === "Ok" && mouseItem.value.tag === "distribution" ? (
       <DistributionChart
-        distribution={mouseItem.value.value}
+        plot={defaultPlot(mouseItem.value.value)}
         width={400}
         height={50}
         {...distributionPlotSettings}

packages/components/src/lib/distributionSpecBuilder.ts

@@ -25,36 +25,14 @@ export let linearXScale: LinearScale = {
   range: "width",
   zero: false,
   nice: false,
-  domain: {
+  domain: { data: "domain", field: "x" },
-    fields: [
-      {
-        data: "con",
-        field: "x",
-      },
-      {
-        data: "dis",
-        field: "x",
-      },
-    ],
-  },
 };
 export let linearYScale: LinearScale = {
   name: "yscale",
   type: "linear",
   range: "height",
   zero: false,
-  domain: {
+  domain: { data: "domain", field: "y" },
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
 };
 
 export let logXScale: LogScale = {
@@ -65,18 +43,7 @@ export let logXScale: LogScale = {
   base: 10,
   nice: false,
   clamp: true,
-  domain: {
+  domain: { data: "domain", field: "x" },
-    fields: [
-      {
-        data: "con",
-        field: "x",
-      },
-      {
-        data: "dis",
-        field: "x",
-      },
-    ],
-  },
 };
 
 export let expYScale: PowScale = {
@@ -86,24 +53,13 @@ export let expYScale: PowScale = {
   range: "height",
   zero: false,
   nice: false,
-  domain: {
+  domain: { data: "domain", field: "y" },
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
 };
 
 export function buildVegaSpec(
   specOptions: DistributionChartSpecOptions
 ): VisualizationSpec {
-  let {
+  const {
     format = ".9~s",
     color = "#739ECC",
     title,
@@ -130,10 +86,10 @@ export function buildVegaSpec(
     padding: 5,
     data: [
       {
-        name: "con",
+        name: "data",
       },
       {
-        name: "dis",
+        name: "domain",
       },
     ],
     signals: [],
@@ -144,12 +100,10 @@ export function buildVegaSpec(
         name: "color",
         type: "ordinal",
         domain: {
-          fields: [
+          data: "data",
-            { data: "con", field: "name" },
+          field: "name",
-            { data: "dis", field: "name" },
-          ],
         },
-        range: { scheme: "category20b" },
+        range: { scheme: "category10" },
       },
     ],
     axes: [
@@ -167,21 +121,35 @@ export function buildVegaSpec(
     ],
     marks: [
       {
-        name: "group",
+        name: "all_distributions",
         type: "group",
         from: {
           facet: {
-            name: "faceted_path_main",
+            name: "distribution_facet",
-            data: "con",
+            data: "data",
             groupby: ["name"],
           },
         },
         marks: [
           {
-            name: "distribution_charts",
+            name: "continuous_distribution",
+            type: "group",
+            from: {
+              facet: {
+                name: "continuous_facet",
+                data: "distribution_facet",
+                field: "continuous",
+              },
+            },
+            encode: {
+              update: {},
+            },
+            marks: [
+              {
+                name: "continuous_area",
                 type: "area",
                 from: {
-              data: "faceted_path_main",
+                  data: "continuous_facet",
                 },
                 encode: {
                   update: {
@@ -194,14 +162,14 @@ export function buildVegaSpec(
                       scale: "yscale",
                       field: "y",
                     },
+                    fill: {
+                      scale: "color",
+                      field: { parent: "name" },
+                    },
                     y2: {
                       scale: "yscale",
                       value: 0,
                     },
-                fill: {
-                  field: "name",
-                  scale: "color",
-                },
                     fillOpacity: {
                       value: 1,
                     },
@@ -210,10 +178,21 @@ export function buildVegaSpec(
               },
             ],
           },
+          {
+            name: "discrete_distribution",
+            type: "group",
+            from: {
+              facet: {
+                name: "discrete_facet",
+                data: "distribution_facet",
+                field: "discrete",
+              },
+            },
+            marks: [
               {
                 type: "rect",
                 from: {
-          data: "dis",
+                  data: "discrete_facet",
                 },
                 encode: {
                   enter: {
@@ -234,16 +213,13 @@ export function buildVegaSpec(
                       scale: "yscale",
                       value: 0,
                     },
-            fill: {
-              value: "#2f65a7",
-            },
                   },
                 },
               },
               {
                 type: "symbol",
                 from: {
-          data: "dis",
+                  data: "discrete_facet",
                 },
                 encode: {
                   enter: {
@@ -264,11 +240,12 @@ export function buildVegaSpec(
                       scale: "yscale",
                       field: "y",
                     },
-            fill: {
-              value: "#1e4577",
                   },
                 },
               },
+            ],
+          },
+        ],
       },
     ],
   };