Merge pull request #528 from quantified-uncertainty/function-chart-fix

Function chart fixes

packages/components/src/components/FunctionChart.tsx

@@ -7,6 +7,8 @@ import {
   lambdaValue,
   environment,
   runForeign,
+  squiggleExpression,
+  errorValue,
   errorValueToString,
 } from "@quri/squiggle-lang";
 import { createClassFromSpec } from "react-vega";
@@ -45,6 +47,104 @@ interface FunctionChartProps {
   environment: environment;
 }
 
+type percentiles = {
+  x: number;
+  p1: number;
+  p5: number;
+  p10: number;
+  p20: number;
+  p30: number;
+  p40: number;
+  p50: number;
+  p60: number;
+  p70: number;
+  p80: number;
+  p90: number;
+  p95: number;
+  p99: number;
+}[];
+
+type errors = _.Dictionary<
+  {
+    x: number;
+    value: string;
+  }[]
+>;
+
+type point = { x: number; value: result<Distribution, string> };
+
+let getPercentiles = ({ chartSettings, fn, environment }) => {
+  let chartPointsToRender = _rangeByCount(
+    chartSettings.start,
+    chartSettings.stop,
+    chartSettings.count
+  );
+
+  let chartPointsData: point[] = chartPointsToRender.map((x) => {
+    let result = runForeign(fn, [x], environment);
+    if (result.tag === "Ok") {
+      if (result.value.tag == "distribution") {
+        return { x, value: { tag: "Ok", value: result.value.value } };
+      } else {
+        return {
+          x,
+          value: {
+            tag: "Error",
+            value:
+              "Cannot currently render functions that don't return distributions",
+          },
+        };
+      }
+    } else {
+      return {
+        x,
+        value: { tag: "Error", value: errorValueToString(result.value) },
+      };
+    }
+  });
+
+  let initialPartition: [
+    { x: number; value: Distribution }[],
+    { x: number; value: string }[]
+  ] = [[], []];
+
+  let [functionImage, errors] = chartPointsData.reduce((acc, current) => {
+    if (current.value.tag === "Ok") {
+      acc[0].push({ x: current.x, value: current.value.value });
+    } else {
+      acc[1].push({ x: current.x, value: current.value.value });
+    }
+    return acc;
+  }, initialPartition);
+
+  let groupedErrors: errors = _.groupBy(errors, (x) => x.value);
+
+  let percentiles: percentiles = functionImage.map(({ x, value }) => {
+    // We convert it to to a pointSet distribution first, so that in case its a sample set
+    // distribution, it doesn't internally convert it to a pointSet distribution for every
+    // single inv() call.
+    let toPointSet: Distribution = unwrap(value.toPointSet());
+    return {
+      x: x,
+      p1: unwrap(toPointSet.inv(0.01)),
+      p5: unwrap(toPointSet.inv(0.05)),
+      p10: unwrap(toPointSet.inv(0.1)),
+      p20: unwrap(toPointSet.inv(0.2)),
+      p30: unwrap(toPointSet.inv(0.3)),
+      p40: unwrap(toPointSet.inv(0.4)),
+      p50: unwrap(toPointSet.inv(0.5)),
+      p60: unwrap(toPointSet.inv(0.6)),
+      p70: unwrap(toPointSet.inv(0.7)),
+      p80: unwrap(toPointSet.inv(0.8)),
+      p90: unwrap(toPointSet.inv(0.9)),
+      p95: unwrap(toPointSet.inv(0.95)),
+      p99: unwrap(toPointSet.inv(0.99)),
+    };
+  });
+
+  return { percentiles, errors: groupedErrors };
+};
+
 export const FunctionChart: React.FC<FunctionChartProps> = ({
   fn,
   chartSettings,
@@ -58,7 +158,15 @@ export const FunctionChart: React.FC<FunctionChartProps> = ({
     setMouseOverlay(NaN);
   }
   const signalListeners = { mousemove: handleHover, mouseout: handleOut };
-  let mouseItem = runForeign(fn, [mouseOverlay], environment);
+  let mouseItem: result<squiggleExpression, errorValue> = !!mouseOverlay
+    ? runForeign(fn, [mouseOverlay], environment)
+    : {
+        tag: "Error",
+        value: {
+          tag: "REExpectedType",
+          value: "Hover x-coordinate returned NaN. Expected a number.",
+        },
+      };
   let showChart =
     mouseItem.tag === "Ok" && mouseItem.value.tag == "distribution" ? (
       <DistributionChart
@@ -70,92 +178,34 @@ export const FunctionChart: React.FC<FunctionChartProps> = ({
     ) : (
       <></>
     );
-  let data1 = _rangeByCount(
-    chartSettings.start,
-    chartSettings.stop,
-    chartSettings.count
-  );
-  type point = { x: number; value: result<Distribution, string> };
-  let valueData: point[] = React.useMemo(
-    () =>
-      data1.map((x) => {
-        let result = runForeign(fn, [x], environment);
-        if (result.tag === "Ok") {
-          if (result.value.tag == "distribution") {
-            return { x, value: { tag: "Ok", value: result.value.value } };
-          } else {
-            return {
-              x,
-              value: {
-                tag: "Error",
-                value:
-                  "Cannot currently render functions that don't return distributions",
-              },
-            };
-          }
-        } else {
-          return {
-            x,
-            value: { tag: "Error", value: errorValueToString(result.value) },
-          };
-        }
-      }),
+
+  let getPercentilesMemoized = React.useMemo(
+    () => getPercentiles({ chartSettings, fn, environment }),
     [environment, fn]
   );
 
-  let initialPartition: [
-    { x: number; value: Distribution }[],
-    { x: number; value: string }[]
-  ] = [[], []];
-  let [functionImage, errors] = valueData.reduce((acc, current) => {
-    if (current.value.tag === "Ok") {
-      acc[0].push({ x: current.x, value: current.value.value });
-    } else {
-      acc[1].push({ x: current.x, value: current.value.value });
-    }
-    return acc;
-  }, initialPartition);
-
-  let percentiles = functionImage.map(({ x, value }) => {
-    return {
-      x: x,
-      p1: unwrap(value.inv(0.01)),
-      p5: unwrap(value.inv(0.05)),
-      p10: unwrap(value.inv(0.12)),
-      p20: unwrap(value.inv(0.2)),
-      p30: unwrap(value.inv(0.3)),
-      p40: unwrap(value.inv(0.4)),
-      p50: unwrap(value.inv(0.5)),
-      p60: unwrap(value.inv(0.6)),
-      p70: unwrap(value.inv(0.7)),
-      p80: unwrap(value.inv(0.8)),
-      p90: unwrap(value.inv(0.9)),
-      p95: unwrap(value.inv(0.95)),
-      p99: unwrap(value.inv(0.99)),
-    };
-  });
-
-  let groupedErrors = _.groupBy(errors, (x) => x.value);
   return (
     <>
       <SquigglePercentilesChart
-        data={{ facet: percentiles }}
+        data={{ facet: getPercentilesMemoized.percentiles }}
         actions={false}
         signalListeners={signalListeners}
       />
       {showChart}
-      {_.entries(groupedErrors).map(([errorName, errorPoints]) => (
-        <ErrorBox key={errorName} heading={errorName}>
-          Values:{" "}
-          {errorPoints
-            .map((r, i) => <NumberShower key={i} number={r.x} />)
-            .reduce((a, b) => (
-              <>
-                {a}, {b}
-              </>
-            ))}
-        </ErrorBox>
-      ))}
+      {_.entries(getPercentilesMemoized.errors).map(
+        ([errorName, errorPoints]) => (
+          <ErrorBox key={errorName} heading={errorName}>
+            Values:{" "}
+            {errorPoints
+              .map((r, i) => <NumberShower key={i} number={r.x} />)
+              .reduce((a, b) => (
+                <>
+                  {a}, {b}
+                </>
+              ))}
+          </ErrorBox>
+        )
+      )}
     </>
   );
 };

packages/components/src/components/SquiggleChart.tsx

@@ -194,7 +194,10 @@ const SquiggleItem: React.FC<SquiggleItemProps> = ({
         <FunctionChart
           fn={expression.value}
           chartSettings={chartSettings}
-          environment={environment}
+          environment={{
+            sampleCount: environment.sampleCount / 10,
+            xyPointLength: environment.xyPointLength / 10,
+          }}
         />
       );
   }
@@ -232,7 +235,8 @@ const ChartWrapper = styled.div`
     "Segoe UI Emoji", "Segoe UI Symbol", "Noto Color Emoji";
 `;
 
-let defaultChartSettings = { start: 0, stop: 10, count: 100 };
+let defaultChartSettings = { start: 0, stop: 10, count: 20 };
+
 export const SquiggleChart: React.FC<SquiggleChartProps> = ({
   squiggleString = "",
   environment,

packages/components/src/components/SquiggleEditor.tsx

@@ -56,7 +56,7 @@ export let SquiggleEditor: React.FC<SquiggleEditorProps> = ({
   environment,
   diagramStart = 0,
   diagramStop = 10,
-  diagramCount = 100,
+  diagramCount = 20,
   onChange,
   bindings = defaultBindings,
   jsImports = defaultImports,

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

@@ -153,6 +153,20 @@ to allow large and small numbers being printed cleanly.
   </Story>
 </Canvas>
 
+## Functions
+
+<Canvas>
+  <Story
+    name="Function"
+    args={{
+      squiggleString: "foo(t) = normal(t,2)*normal(5,3); foo",
+      width,
+    }}
+  >
+    {Template.bind({})}
+  </Story>
+</Canvas>
+
 ## Records
 
 <Canvas>