From 98ae0459c96765126e530918ab1e8494bc870fd1 Mon Sep 17 00:00:00 2001
From: Sam Nolan <samnolan555@gmail.com>
Date: Wed, 13 Jul 2022 14:15:07 +1000
Subject: [PATCH] Refactor specification to include discrete

---
 .../src/components/DistributionChart.tsx      |  48 +++-
 .../src/components/FunctionChart1Dist.tsx     |   3 +-
 .../src/lib/distributionSpecBuilder.ts        | 267 ++++++++----------
 3 files changed, 157 insertions(+), 161 deletions(-)

diff --git a/packages/components/src/components/DistributionChart.tsx b/packages/components/src/components/DistributionChart.tsx
index 0b9535b2..f1253629 100644
--- a/packages/components/src/components/DistributionChart.tsx
+++ b/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") {
-          return x.value;
+        if (
+          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) =>
-      shape.continuous.map((point) => ({ ...point, name: i + 1 }))
-    );
-    let discretePoints = shapes.value.flatMap((shape, i) =>
-      shape.discrete.map((point) => ({ ...point, name: i + 1 }))
+    const domain = shapes.value.flatMap((shape) =>
+      shape.discrete.concat(shape.continuous)
     );
+    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 && (
diff --git a/packages/components/src/components/FunctionChart1Dist.tsx b/packages/components/src/components/FunctionChart1Dist.tsx
index 650d2753..3b203de5 100644
--- a/packages/components/src/components/FunctionChart1Dist.tsx
+++ b/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}
diff --git a/packages/components/src/lib/distributionSpecBuilder.ts b/packages/components/src/lib/distributionSpecBuilder.ts
index 515a1b9f..36f298d2 100644
--- a/packages/components/src/lib/distributionSpecBuilder.ts
+++ b/packages/components/src/lib/distributionSpecBuilder.ts
@@ -25,36 +25,14 @@ export let linearXScale: LinearScale = {
   range: "width",
   zero: false,
   nice: false,
-  domain: {
-    fields: [
-      {
-        data: "con",
-        field: "x",
-      },
-      {
-        data: "dis",
-        field: "x",
-      },
-    ],
-  },
+  domain: { data: "domain", field: "x" },
 };
 export let linearYScale: LinearScale = {
   name: "yscale",
   type: "linear",
   range: "height",
   zero: false,
-  domain: {
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
+  domain: { data: "domain", field: "y" },
 };
 
 export let logXScale: LogScale = {
@@ -65,18 +43,7 @@ export let logXScale: LogScale = {
   base: 10,
   nice: false,
   clamp: true,
-  domain: {
-    fields: [
-      {
-        data: "con",
-        field: "x",
-      },
-      {
-        data: "dis",
-        field: "x",
-      },
-    ],
-  },
+  domain: { data: "domain", field: "x" },
 };
 
 export let expYScale: PowScale = {
@@ -86,24 +53,13 @@ export let expYScale: PowScale = {
   range: "height",
   zero: false,
   nice: false,
-  domain: {
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
+  domain: { data: "domain", 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: "con", field: "name" },
-            { data: "dis", field: "name" },
-          ],
+          data: "data",
+          field: "name",
         },
-        range: { scheme: "category20b" },
+        range: { scheme: "category10" },
       },
     ],
     axes: [
@@ -167,109 +121,132 @@ export function buildVegaSpec(
     ],
     marks: [
       {
-        name: "group",
+        name: "all_distributions",
         type: "group",
         from: {
           facet: {
-            name: "faceted_path_main",
-            data: "con",
+            name: "distribution_facet",
+            data: "data",
             groupby: ["name"],
           },
         },
         marks: [
           {
-            name: "distribution_charts",
-            type: "area",
+            name: "continuous_distribution",
+            type: "group",
             from: {
-              data: "faceted_path_main",
-            },
-            encode: {
-              update: {
-                interpolate: { value: "linear" },
-                x: {
-                  scale: "xscale",
-                  field: "x",
-                },
-                y: {
-                  scale: "yscale",
-                  field: "y",
-                },
-                y2: {
-                  scale: "yscale",
-                  value: 0,
-                },
-                fill: {
-                  field: "name",
-                  scale: "color",
-                },
-                fillOpacity: {
-                  value: 1,
-                },
+              facet: {
+                name: "continuous_facet",
+                data: "distribution_facet",
+                field: "continuous",
               },
             },
+            encode: {
+              update: {},
+            },
+            marks: [
+              {
+                name: "continuous_area",
+                type: "area",
+                from: {
+                  data: "continuous_facet",
+                },
+                encode: {
+                  update: {
+                    interpolate: { value: "linear" },
+                    x: {
+                      scale: "xscale",
+                      field: "x",
+                    },
+                    y: {
+                      scale: "yscale",
+                      field: "y",
+                    },
+                    fill: {
+                      scale: "color",
+                      field: { parent: "name" },
+                    },
+                    y2: {
+                      scale: "yscale",
+                      value: 0,
+                    },
+                    fillOpacity: {
+                      value: 1,
+                    },
+                  },
+                },
+              },
+            ],
+          },
+          {
+            name: "discrete_distribution",
+            type: "group",
+            from: {
+              facet: {
+                name: "discrete_facet",
+                data: "distribution_facet",
+                field: "discrete",
+              },
+            },
+            marks: [
+              {
+                type: "rect",
+                from: {
+                  data: "discrete_facet",
+                },
+                encode: {
+                  enter: {
+                    width: {
+                      value: 1,
+                    },
+                  },
+                  update: {
+                    x: {
+                      scale: "xscale",
+                      field: "x",
+                    },
+                    y: {
+                      scale: "yscale",
+                      field: "y",
+                    },
+                    y2: {
+                      scale: "yscale",
+                      value: 0,
+                    },
+                  },
+                },
+              },
+              {
+                type: "symbol",
+                from: {
+                  data: "discrete_facet",
+                },
+                encode: {
+                  enter: {
+                    shape: {
+                      value: "circle",
+                    },
+                    size: [{ value: 100 }],
+                    tooltip: {
+                      signal: "datum.y",
+                    },
+                  },
+                  update: {
+                    x: {
+                      scale: "xscale",
+                      field: "x",
+                    },
+                    y: {
+                      scale: "yscale",
+                      field: "y",
+                    },
+                  },
+                },
+              },
+            ],
           },
         ],
       },
-      {
-        type: "rect",
-        from: {
-          data: "dis",
-        },
-        encode: {
-          enter: {
-            width: {
-              value: 1,
-            },
-          },
-          update: {
-            x: {
-              scale: "xscale",
-              field: "x",
-            },
-            y: {
-              scale: "yscale",
-              field: "y",
-            },
-            y2: {
-              scale: "yscale",
-              value: 0,
-            },
-            fill: {
-              value: "#2f65a7",
-            },
-          },
-        },
-      },
-      {
-        type: "symbol",
-        from: {
-          data: "dis",
-        },
-        encode: {
-          enter: {
-            shape: {
-              value: "circle",
-            },
-            size: [{ value: 100 }],
-            tooltip: {
-              signal: "datum.y",
-            },
-          },
-          update: {
-            x: {
-              scale: "xscale",
-              field: "x",
-            },
-            y: {
-              scale: "yscale",
-              field: "y",
-            },
-            fill: {
-              value: "#1e4577",
-            },
-          },
-        },
-      },
     ],
   };
   if (title) {