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,40 +47,40 @@ interface FunctionChartProps {
   environment: environment;
 }
 
-export const FunctionChart: React.FC<FunctionChartProps> = ({
-  fn,
-  chartSettings,
-  environment,
-}: FunctionChartProps) => {
-  let [mouseOverlay, setMouseOverlay] = React.useState(0);
-  function handleHover(_name: string, value: unknown) {
-    setMouseOverlay(value as number);
-  }
-  function handleOut() {
-    setMouseOverlay(NaN);
-  }
-  const signalListeners = { mousemove: handleHover, mouseout: handleOut };
-  let mouseItem = runForeign(fn, [mouseOverlay], environment);
-  let showChart =
-    mouseItem.tag === "Ok" && mouseItem.value.tag == "distribution" ? (
-      <DistributionChart
-        distribution={mouseItem.value.value}
-        width={400}
-        height={140}
-        showSummary={false}
-      />
-    ) : (
-      <></>
-    );
-  let data1 = _rangeByCount(
+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
   );
-  type point = { x: number; value: result<Distribution, string> };
-  let valueData: point[] = React.useMemo(
-    () =>
-      data1.map((x) => {
+
+  let chartPointsData: point[] = chartPointsToRender.map((x) => {
     let result = runForeign(fn, [x], environment);
     if (result.tag === "Ok") {
       if (result.value.tag == "distribution") {
@@ -99,15 +101,14 @@ export const FunctionChart: React.FC<FunctionChartProps> = ({
         value: { tag: "Error", value: errorValueToString(result.value) },
       };
     }
-      }),
-    [environment, fn]
-  );
+  });
 
   let initialPartition: [
     { x: number; value: Distribution }[],
     { x: number; value: string }[]
   ] = [[], []];
-  let [functionImage, errors] = valueData.reduce((acc, current) => {
+
+  let [functionImage, errors] = chartPointsData.reduce((acc, current) => {
     if (current.value.tag === "Ok") {
       acc[0].push({ x: current.x, value: current.value.value });
     } else {
@@ -116,35 +117,83 @@ export const FunctionChart: React.FC<FunctionChartProps> = ({
     return acc;
   }, initialPartition);
 
-  let percentiles = functionImage.map(({ x, value }) => {
+  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(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)),
+      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)),
     };
   });
 
-  let groupedErrors = _.groupBy(errors, (x) => x.value);
+  return { percentiles, errors: groupedErrors };
+};
+
+export const FunctionChart: React.FC<FunctionChartProps> = ({
+  fn,
+  chartSettings,
+  environment,
+}: FunctionChartProps) => {
+  let [mouseOverlay, setMouseOverlay] = React.useState(0);
+  function handleHover(_name: string, value: unknown) {
+    setMouseOverlay(value as number);
+  }
+  function handleOut() {
+    setMouseOverlay(NaN);
+  }
+  const signalListeners = { mousemove: handleHover, mouseout: handleOut };
+  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
+        distribution={mouseItem.value.value}
+        width={400}
+        height={140}
+        showSummary={false}
+      />
+    ) : (
+      <></>
+    );
+
+  let getPercentilesMemoized = React.useMemo(
+    () => getPercentiles({ chartSettings, fn, environment }),
+    [environment, fn]
+  );
+
   return (
     <>
       <SquigglePercentilesChart
-        data={{ facet: percentiles }}
+        data={{ facet: getPercentilesMemoized.percentiles }}
         actions={false}
         signalListeners={signalListeners}
       />
       {showChart}
-      {_.entries(groupedErrors).map(([errorName, errorPoints]) => (
+      {_.entries(getPercentilesMemoized.errors).map(
+        ([errorName, errorPoints]) => (
           <ErrorBox key={errorName} heading={errorName}>
             Values:{" "}
             {errorPoints
@@ -155,7 +204,8 @@ export const FunctionChart: React.FC<FunctionChartProps> = ({
                 </>
               ))}
           </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>