Merge pull request #824 from quantified-uncertainty/multiple-charts

Add multiple plotting

packages/components/src/components/DistributionChart.tsx

@@ -4,6 +4,8 @@ import {
   result,
   distributionError,
   distributionErrorToString,
+  squiggleExpression,
+  resultMap,
 } from "@quri/squiggle-lang";
 import { Vega } from "react-vega";
 import { ErrorAlert } from "./Alert";
@@ -14,6 +16,8 @@ import {
   DistributionChartSpecOptions,
 } from "../lib/distributionSpecBuilder";
 import { NumberShower } from "./NumberShower";
+import { Plot, parsePlot } from "../lib/plotParser";
+import { flattenResult } from "../lib/utility";
 import { hasMassBelowZero } from "../lib/distributionUtils";
 
 export type DistributionPlottingSettings = {
@@ -23,26 +27,41 @@ export type DistributionPlottingSettings = {
 } & DistributionChartSpecOptions;
 
 export type DistributionChartProps = {
-  distribution: Distribution;
+  plot: Plot;
   width?: number;
   height: number;
 } & DistributionPlottingSettings;
 
+export function defaultPlot(distribution: Distribution): Plot {
+  return { distributions: [{ name: "default", distribution }] };
+}
+
+export function makePlot(record: {
+  [key: string]: squiggleExpression;
+}): Plot | void {
+  const plotResult = parsePlot(record);
+  if (plotResult.tag === "Ok") {
+    return plotResult.value;
+  }
+}
+
 export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
-  const {
+  const { plot, height, showSummary, width, logX, actions = false } = props;
-    distribution,
-    height,
-    showSummary,
-    width,
-    logX,
-    actions = false,
-  } = props;
-  const shape = distribution.pointSet();
   const [sized] = useSize((size) => {
-    if (shape.tag === "Error") {
+    let shapes = flattenResult(
+      plot.distributions.map((x) =>
+        resultMap(x.distribution.pointSet(), (shape) => ({
+          name: x.name,
+          // color: x.color, // not supported yet
+          continuous: shape.continuous,
+          discrete: shape.discrete,
+        }))
+      )
+    );
+    if (shapes.tag === "Error") {
       return (
         <ErrorAlert heading="Distribution Error">
-          {distributionErrorToString(shape.value)}
+          {distributionErrorToString(shapes.value)}
         </ErrorAlert>
       );
     }
@@ -56,24 +75,29 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
       );
       widthProp = 20;
     }
+    const domain = shapes.value.flatMap((shape) =>
+      shape.discrete.concat(shape.continuous)
+    );
 
     return (
       <div style={{ width: widthProp }}>
-        {logX && hasMassBelowZero(shape.value) ? (
+        {logX && shapes.value.some(hasMassBelowZero) ? (
           <ErrorAlert heading="Log Domain Error">
             Cannot graph distribution with negative values on logarithmic scale.
           </ErrorAlert>
         ) : (
           <Vega
             spec={spec}
-            data={{ con: shape.value.continuous, dis: shape.value.discrete }}
+            data={{ data: shapes.value, domain }}
             width={widthProp - 10}
             height={height}
             actions={actions}
           />
         )}
         <div className="flex justify-center">
-          {showSummary && <SummaryTable distribution={distribution} />}
+          {showSummary && plot.distributions.length === 1 && (
+            <SummaryTable distribution={plot.distributions[0].distribution} />
+          )}
         </div>
       </div>
     );

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";
@@ -179,7 +180,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/components/SquigglePlayground.tsx

@@ -37,10 +37,7 @@ import { InputItem } from "./ui/InputItem";
 import { Text } from "./ui/Text";
 import { ViewSettings, viewSettingsSchema } from "./ViewSettings";
 import { HeadedSection } from "./ui/HeadedSection";
-import {
+import { defaultTickFormat } from "../lib/distributionSpecBuilder";
-  defaultColor,
-  defaultTickFormat,
-} from "../lib/distributionSpecBuilder";
 import { Button } from "./ui/Button";
 
 type PlaygroundProps = SquiggleChartProps & {
@@ -240,7 +237,6 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
   title,
   minX,
   maxX,
-  color = defaultColor,
   tickFormat = defaultTickFormat,
   distributionChartActions,
   code: controlledCode,
@@ -268,7 +264,6 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
       title,
       minX,
       maxX,
-      color,
       tickFormat,
       distributionChartActions,
       showSummary,

packages/components/src/components/SquiggleViewer/ExpressionViewer.tsx

@@ -1,7 +1,7 @@
 import React from "react";
 import { squiggleExpression, declaration } from "@quri/squiggle-lang";
 import { NumberShower } from "../NumberShower";
-import { DistributionChart } from "../DistributionChart";
+import { DistributionChart, defaultPlot, makePlot } from "../DistributionChart";
 import { FunctionChart, FunctionChartSettings } from "../FunctionChart";
 import clsx from "clsx";
 import { VariableBox } from "./VariableBox";
@@ -102,7 +102,7 @@ export const ExpressionViewer: React.FC<Props> = ({
           {(settings) => {
             return (
               <DistributionChart
-                distribution={expression.value}
+                plot={defaultPlot(expression.value)}
                 {...settings.distributionPlotSettings}
                 height={settings.height}
                 width={width}
@@ -241,9 +241,9 @@ export const ExpressionViewer: React.FC<Props> = ({
     case "module": {
       return (
         <VariableList path={path} heading="Module">
-          {(settings) =>
+          {(_) =>
             Object.entries(expression.value)
-              .filter(([key, r]) => !key.match(/^(Math|System)\./))
+              .filter(([key, _]) => !key.match(/^(Math|System)\./))
               .map(([key, r]) => (
                 <ExpressionViewer
                   key={key}
@@ -257,9 +257,45 @@ export const ExpressionViewer: React.FC<Props> = ({
       );
     }
     case "record":
+      const plot = makePlot(expression.value);
+      if (plot) {
+        return (
+          <VariableBox
+            path={path}
+            heading={"Plot"}
+            renderSettingsMenu={({ onChange }) => {
+              let disableLogX = plot.distributions.some((x) => {
+                let pointSet = x.distribution.pointSet();
+                return (
+                  pointSet.tag === "Ok" && hasMassBelowZero(pointSet.value)
+                );
+              });
+              return (
+                <ItemSettingsMenu
+                  path={path}
+                  onChange={onChange}
+                  disableLogX={disableLogX}
+                  withFunctionSettings={false}
+                />
+              );
+            }}
+          >
+            {(settings) => {
+              return (
+                <DistributionChart
+                  plot={plot}
+                  {...settings.distributionPlotSettings}
+                  height={settings.height}
+                  width={width}
+                />
+              );
+            }}
+          </VariableBox>
+        );
+      } else {
         return (
           <VariableList path={path} heading="Record">
-          {(settings) =>
+            {(_) =>
               Object.entries(expression.value).map(([key, r]) => (
                 <ExpressionViewer
                   key={key}
@@ -271,10 +307,11 @@ export const ExpressionViewer: React.FC<Props> = ({
             }
           </VariableList>
         );
+      }
     case "array":
       return (
         <VariableList path={path} heading="Array">
-          {(settings) =>
+          {(_) =>
             expression.value.map((r, i) => (
               <ExpressionViewer
                 key={i}

packages/components/src/components/SquiggleViewer/ItemSettingsMenu.tsx

@@ -6,10 +6,7 @@ import { Modal } from "../ui/Modal";
 import { ViewSettings, viewSettingsSchema } from "../ViewSettings";
 import { Path, pathAsString } from "./utils";
 import { ViewerContext } from "./ViewerContext";
-import {
+import { defaultTickFormat } from "../../lib/distributionSpecBuilder";
-  defaultColor,
-  defaultTickFormat,
-} from "../../lib/distributionSpecBuilder";
 import { PlaygroundContext } from "../SquigglePlayground";
 
 type Props = {
@@ -46,7 +43,6 @@ const ItemSettingsModal: React.FC<
       tickFormat:
         mergedSettings.distributionPlotSettings.format || defaultTickFormat,
       title: mergedSettings.distributionPlotSettings.title,
-      color: mergedSettings.distributionPlotSettings.color || defaultColor,
       minX: mergedSettings.distributionPlotSettings.minX,
       maxX: mergedSettings.distributionPlotSettings.maxX,
       distributionChartActions: mergedSettings.distributionPlotSettings.actions,
@@ -66,7 +62,6 @@ const ItemSettingsModal: React.FC<
           expY: vars.expY,
           format: vars.tickFormat,
           title: vars.title,
-          color: vars.color,
           minX: vars.minX,
           maxX: vars.maxX,
           actions: vars.distributionChartActions,

packages/components/src/components/ViewSettings.tsx

@@ -5,10 +5,7 @@ import { InputItem } from "./ui/InputItem";
 import { Checkbox } from "./ui/Checkbox";
 import { HeadedSection } from "./ui/HeadedSection";
 import { Text } from "./ui/Text";
-import {
+import { defaultTickFormat } from "../lib/distributionSpecBuilder";
-  defaultColor,
-  defaultTickFormat,
-} from "../lib/distributionSpecBuilder";
 
 export const viewSettingsSchema = yup.object({}).shape({
   chartHeight: yup.number().required().positive().integer().default(350),
@@ -18,7 +15,6 @@ export const viewSettingsSchema = yup.object({}).shape({
   expY: yup.boolean().required(),
   tickFormat: yup.string().default(defaultTickFormat),
   title: yup.string(),
-  color: yup.string().default(defaultColor).required(),
   minX: yup.number(),
   maxX: yup.number(),
   distributionChartActions: yup.boolean(),
@@ -114,12 +110,6 @@ export const ViewSettings: React.FC<{
               register={register}
               label="Tick Format"
             />
-            <InputItem
-              name="color"
-              type="color"
-              register={register}
-              label="Color"
-            />
           </div>
         </HeadedSection>
       </div>

packages/components/src/lib/distributionSpecBuilder.ts

@@ -10,8 +10,6 @@ export type DistributionChartSpecOptions = {
   minX?: number;
   /** The maximum x coordinate shown on the chart */
   maxX?: number;
-  /** The color of the chart */
-  color?: string;
   /** The title of the chart */
   title?: string;
   /** The formatting of the ticks */
@@ -25,36 +23,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: true,
-  domain: {
+  domain: { data: "domain", field: "y" },
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
 };
 
 export let logXScale: LogScale = {
@@ -65,18 +41,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,29 +51,16 @@ export let expYScale: PowScale = {
   range: "height",
   zero: true,
   nice: false,
-  domain: {
+  domain: { data: "domain", field: "y" },
-    fields: [
-      {
-        data: "con",
-        field: "y",
-      },
-      {
-        data: "dis",
-        field: "y",
-      },
-    ],
-  },
 };
 
 export const defaultTickFormat = ".9~s";
-export const defaultColor = "#739ECC";
 
 export function buildVegaSpec(
   specOptions: DistributionChartSpecOptions
 ): VisualizationSpec {
-  let {
+  const {
     format = defaultTickFormat,
-    color = defaultColor,
     title,
     minX,
     maxX,
@@ -127,20 +79,32 @@ export function buildVegaSpec(
 
   let spec: VisualizationSpec = {
     $schema: "https://vega.github.io/schema/vega/v5.json",
-    description: "A basic area chart example",
+    description: "Squiggle plot chart",
     width: 500,
     height: 100,
     padding: 5,
     data: [
       {
-        name: "con",
+        name: "data",
       },
       {
-        name: "dis",
+        name: "domain",
       },
     ],
     signals: [],
-    scales: [xScale, expY ? expYScale : linearYScale],
+    scales: [
+      xScale,
+      expY ? expYScale : linearYScale,
+      {
+        name: "color",
+        type: "ordinal",
+        domain: {
+          data: "data",
+          field: "name",
+        },
+        range: { scheme: "blues" },
+      },
+    ],
     axes: [
       {
         orient: "bottom",
@@ -152,13 +116,40 @@ export function buildVegaSpec(
         domainOpacity: 0.0,
         format: format,
         tickCount: 10,
+        labelOverlap: "greedy",
       },
     ],
     marks: [
       {
+        name: "all_distributions",
+        type: "group",
+        from: {
+          facet: {
+            name: "distribution_facet",
+            data: "data",
+            groupby: ["name"],
+          },
+        },
+        marks: [
+          {
+            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: "con",
+                  data: "continuous_facet",
                 },
                 encode: {
                   update: {
@@ -171,23 +162,37 @@ export function buildVegaSpec(
                       scale: "yscale",
                       field: "y",
                     },
+                    fill: {
+                      scale: "color",
+                      field: { parent: "name" },
+                    },
                     y2: {
                       scale: "yscale",
                       value: 0,
                     },
-            fill: {
-              value: color,
-            },
                     fillOpacity: {
                       value: 1,
                     },
                   },
                 },
               },
+            ],
+          },
+          {
+            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: {
@@ -209,7 +214,8 @@ export function buildVegaSpec(
                       value: 0,
                     },
                     fill: {
-              value: "#2f65a7",
+                      scale: "color",
+                      field: { parent: "name" },
                     },
                   },
                 },
@@ -217,7 +223,7 @@ export function buildVegaSpec(
               {
                 type: "symbol",
                 from: {
-          data: "dis",
+                  data: "discrete_facet",
                 },
                 encode: {
                   enter: {
@@ -239,21 +245,49 @@ export function buildVegaSpec(
                       field: "y",
                     },
                     fill: {
-              value: "#1e4577",
+                      scale: "color",
+                      field: { parent: "name" },
                     },
                   },
                 },
               },
             ],
-  };
+          },
-  if (title) {
+        ],
-    spec = {
+      },
-      ...spec,
+    ],
+    legends: [
+      {
+        fill: "color",
+        orient: "top",
+        labelFontSize: 12,
+        encode: {
+          symbols: {
+            update: {
+              fill: [
+                { test: "length(domain('color')) == 1", value: "transparent" },
+                { scale: "color", field: "value" },
+              ],
+            },
+          },
+          labels: {
+            interactive: true,
+            update: {
+              fill: [
+                { test: "length(domain('color')) == 1", value: "transparent" },
+                { value: "black" },
+              ],
+            },
+          },
+        },
+      },
+    ],
+    ...(title && {
       title: {
         text: title,
       },
+    }),
   };
-  }
 
   return spec;
 }

packages/components/src/lib/plotParser.ts

@@ -0,0 +1,72 @@
+import * as yup from "yup";
+import { Distribution, result, squiggleExpression } from "@quri/squiggle-lang";
+
+export type LabeledDistribution = {
+  name: string;
+  distribution: Distribution;
+  color?: string;
+};
+
+export type Plot = {
+  distributions: LabeledDistribution[];
+};
+
+function error<a, b>(err: b): result<a, b> {
+  return { tag: "Error", value: err };
+}
+
+function ok<a, b>(x: a): result<a, b> {
+  return { tag: "Ok", value: x };
+}
+
+const schema = yup
+  .object()
+  .strict()
+  .noUnknown()
+  .shape({
+    distributions: yup.object().shape({
+      tag: yup.mixed().oneOf(["array"]),
+      value: yup
+        .array()
+        .of(
+          yup.object().shape({
+            tag: yup.mixed().oneOf(["record"]),
+            value: yup.object({
+              name: yup.object().shape({
+                tag: yup.mixed().oneOf(["string"]),
+                value: yup.string().required(),
+              }),
+              // color: yup
+              //   .object({
+              //     tag: yup.mixed().oneOf(["string"]),
+              //     value: yup.string().required(),
+              //   })
+              //   .default(undefined),
+              distribution: yup.object({
+                tag: yup.mixed().oneOf(["distribution"]),
+                value: yup.mixed(),
+              }),
+            }),
+          })
+        )
+        .required(),
+    }),
+  });
+
+export function parsePlot(record: {
+  [key: string]: squiggleExpression;
+}): result<Plot, string> {
+  try {
+    const plotRecord = schema.validateSync(record);
+    return ok({
+      distributions: plotRecord.distributions.value.map((x) => ({
+        name: x.value.name.value,
+        // color: x.value.color?.value, // not supported yet
+        distribution: x.value.distribution.value,
+      })),
+    });
+  } catch (e) {
+    const message = e instanceof Error ? e.message : "Unknown error";
+    return error(message);
+  }
+}

packages/components/src/lib/utility.ts

@@ -0,0 +1,37 @@
+import { result } from "@quri/squiggle-lang";
+
+export function flattenResult<a, b>(x: result<a, b>[]): result<a[], b> {
+  if (x.length === 0) {
+    return { tag: "Ok", value: [] };
+  } else {
+    if (x[0].tag === "Error") {
+      return x[0];
+    } else {
+      let rest = flattenResult(x.splice(1));
+      if (rest.tag === "Error") {
+        return rest;
+      } else {
+        return { tag: "Ok", value: [x[0].value].concat(rest.value) };
+      }
+    }
+  }
+}
+
+export function resultBind<a, b, c>(
+  x: result<a, b>,
+  fn: (y: a) => result<c, b>
+): result<c, b> {
+  if (x.tag === "Ok") {
+    return fn(x.value);
+  } else {
+    return x;
+  }
+}
+
+export function all(arr: boolean[]): boolean {
+  return arr.reduce((x, y) => x && y, true);
+}
+
+export function some(arr: boolean[]): boolean {
+  return arr.reduce((x, y) => x || y, false);
+}

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

@@ -93,6 +93,33 @@ could be continuous, discrete or mixed.
   </Story>
 </Canvas>
 
+## Multiple plots
+
+<Canvas>
+  <Story
+    name="Multiple plots"
+    args={{
+      code: `
+{
+  distributions: [
+    {
+      name: "one",
+      distribution: mx(0.5, normal(0,1))
+    },
+    {
+      name: "two",
+      distribution: mx(2, normal(5, 2)),
+    }
+  ]
+}
+`,
+      width,
+    }}
+  >
+    {Template.bind({})}
+  </Story>
+</Canvas>
+
 ## Constants
 
 A constant is a simple number as a result. This has special formatting rules