Merge pull request #1008 from ideopunk/feature/distribution_tweaks

Distribution component tweaks

.gitignore

@@ -7,4 +7,7 @@ yarn-error.log
 **/.sync.ffs_db
 .direnv
 .log
+
+.vscode
+todo.txt
 result

packages/components/src/components/DistributionChart.tsx

@@ -6,6 +6,7 @@ import {
   resultMap,
   SqRecord,
   environment,
+  SqDistributionTag,
 } from "@quri/squiggle-lang";
 import { Vega } from "react-vega";
 import { ErrorAlert } from "./Alert";
@@ -31,6 +32,7 @@ export type DistributionChartProps = {
   environment: environment;
   width?: number;
   height: number;
+  xAxisType?: "number" | "dateTime";
 } & DistributionPlottingSettings;
 
 export function defaultPlot(distribution: SqDistribution): Plot {
@@ -56,14 +58,15 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
   } = props;
   const [sized] = useSize((size) => {
     const shapes = flattenResult(
-      plot.distributions.map((x) => {
-        return resultMap(x.distribution.pointSet(environment), (pointSet) => ({
-          ...pointSet.asShape(),
+      plot.distributions.map((x) =>
+        resultMap(x.distribution.pointSet(environment), (pointSet) => ({
           name: x.name,
           // color: x.color, // not supported yet
-        }));
-      })
+          ...pointSet.asShape(),
+        }))
+      )
     );
+
     if (shapes.tag === "Error") {
       return (
         <ErrorAlert heading="Distribution Error">
@@ -72,6 +75,14 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
       );
     }
 
+    // if this is a sample set, include the samples
+    const samples: number[] = [];
+    for (const { distribution } of plot?.distributions) {
+      if (distribution.tag === SqDistributionTag.SampleSet) {
+        samples.push(...distribution.value());
+      }
+    }
+
     const spec = buildVegaSpec(props);
 
     let widthProp = width ? width : size.width;
@@ -94,7 +105,7 @@ export const DistributionChart: React.FC<DistributionChartProps> = (props) => {
         ) : (
           <Vega
             spec={spec}
-            data={{ data: shapes.value, domain }}
+            data={{ data: shapes.value, domain, samples }}
             width={widthProp - 10}
             height={height}
             actions={actions}

packages/components/src/components/SquiggleChart.tsx

@@ -48,6 +48,8 @@ export interface SquiggleChartProps {
   minX?: number;
   /** Specify the upper bound of the x scale */
   maxX?: number;
+  /** Whether the x-axis should be dates or numbers */
+  xAxisType?: "number" | "dateTime";
   /** Whether to show vega actions to the user, so they can copy the chart spec */
   distributionChartActions?: boolean;
   enableLocalSettings?: boolean;
@@ -76,6 +78,7 @@ export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
     maxX,
     color,
     title,
+    xAxisType = "number",
     distributionChartActions,
     enableLocalSettings = false,
   }) => {
@@ -96,6 +99,7 @@ export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
       maxX,
       color,
       title,
+      xAxisType,
       actions: distributionChartActions,
     };
 

packages/components/src/lib/distributionSpecBuilder.ts

@@ -1,5 +1,5 @@
 import { VisualizationSpec } from "react-vega";
-import type { LogScale, LinearScale, PowScale } from "vega";
+import type { LogScale, LinearScale, PowScale, TimeScale } from "vega";
 
 export type DistributionChartSpecOptions = {
   /** Set the x scale to be logarithmic by deault */
@@ -14,9 +14,12 @@ export type DistributionChartSpecOptions = {
   title?: string;
   /** The formatting of the ticks */
   format?: string;
+  /** Whether the x-axis should be dates or numbers */
+  xAxisType?: "number" | "dateTime";
 };
 
-export let linearXScale: LinearScale = {
+/** X Scales */
+export const linearXScale: LinearScale = {
   name: "xscale",
   clamp: true,
   type: "linear",
@@ -25,15 +28,8 @@ export let linearXScale: LinearScale = {
   nice: false,
   domain: { data: "domain", field: "x" },
 };
-export let linearYScale: LinearScale = {
-  name: "yscale",
-  type: "linear",
-  range: "height",
-  zero: true,
-  domain: { data: "domain", field: "y" },
-};
 
-export let logXScale: LogScale = {
+export const logXScale: LogScale = {
   name: "xscale",
   type: "log",
   range: "width",
@@ -44,7 +40,25 @@ export let logXScale: LogScale = {
   domain: { data: "domain", field: "x" },
 };
 
-export let expYScale: PowScale = {
+export const timeXScale: TimeScale = {
+  name: "xscale",
+  clamp: true,
+  type: "time",
+  range: "width",
+  nice: false,
+  domain: { data: "domain", field: "x" },
+};
+
+/** Y Scales */
+export const linearYScale: LinearScale = {
+  name: "yscale",
+  type: "linear",
+  range: "height",
+  zero: true,
+  domain: { data: "domain", field: "y" },
+};
+
+export const expYScale: PowScale = {
   name: "yscale",
   type: "pow",
   exponent: 0.1,
@@ -55,20 +69,25 @@ export let expYScale: PowScale = {
 };
 
 export const defaultTickFormat = ".9~s";
+export const timeTickFormat = "%b %d, %Y %H:%M";
+const width = 500;
 
 export function buildVegaSpec(
   specOptions: DistributionChartSpecOptions
 ): VisualizationSpec {
-  const {
-    format = defaultTickFormat,
-    title,
-    minX,
-    maxX,
-    logX,
-    expY,
-  } = specOptions;
+  const { title, minX, maxX, logX, expY, xAxisType = "number" } = specOptions;
+
+  const dateTime = xAxisType === "dateTime";
+
+  // some fallbacks
+  const format = specOptions?.format
+    ? specOptions.format
+    : dateTime
+    ? timeTickFormat
+    : defaultTickFormat;
+
+  let xScale = dateTime ? timeXScale : logX ? logXScale : linearXScale;
 
-  let xScale = logX ? logXScale : linearXScale;
   if (minX !== undefined && Number.isFinite(minX)) {
     xScale = { ...xScale, domainMin: minX };
   }
@@ -77,21 +96,36 @@ export function buildVegaSpec(
     xScale = { ...xScale, domainMax: maxX };
   }
 
-  let spec: VisualizationSpec = {
+  const spec: VisualizationSpec = {
     $schema: "https://vega.github.io/schema/vega/v5.json",
     description: "Squiggle plot chart",
-    width: 500,
+    width: width,
     height: 100,
     padding: 5,
-    data: [
+    data: [{ name: "data" }, { name: "domain" }, { name: "samples" }],
+    signals: [
       {
-        name: "data",
+        name: "hover",
+        value: null,
+        on: [
+          { events: "mouseover", update: "datum" },
+          { events: "mouseout", update: "null" },
+        ],
       },
       {
-        name: "domain",
+        name: "position",
+        value: "[0, 0]",
+        on: [
+          { events: "mousemove", update: "xy() " },
+          { events: "mouseout", update: "null" },
+        ],
+      },
+      {
+        name: "position_scaled",
+        value: null,
+        update: "isArray(position) ? invert('xscale', position[0]) : ''",
       },
     ],
-    signals: [],
     scales: [
       xScale,
       expY ? expYScale : linearYScale,
@@ -115,7 +149,7 @@ export function buildVegaSpec(
         domainColor: "#fff",
         domainOpacity: 0.0,
         format: format,
-        tickCount: 10,
+        tickCount: dateTime ? 3 : 10,
         labelOverlap: "greedy",
       },
     ],
@@ -232,13 +266,16 @@ export function buildVegaSpec(
                     },
                     size: [{ value: 100 }],
                     tooltip: {
-                      signal: "{ probability: datum.y, value: datum.x }",
+                      signal: dateTime
+                        ? "{ probability: datum.y, value: datetime(datum.x) }"
+                        : "{ probability: datum.y, value: datum.x }",
                     },
                   },
                   update: {
                     x: {
                       scale: "xscale",
                       field: "x",
+                      offset: 0.5, // if this is not included, the circles are slightly left of center.
                     },
                     y: {
                       scale: "yscale",
@@ -255,6 +292,69 @@ export function buildVegaSpec(
           },
         ],
       },
+
+      {
+        name: "sampleset",
+        type: "rect",
+        from: { data: "samples" },
+        encode: {
+          enter: {
+            x: { scale: "xscale", field: "data" },
+            width: { value: 0.1 },
+
+            y: { value: 25, offset: { signal: "height" } },
+            height: { value: 5 },
+          },
+        },
+      },
+      {
+        type: "text",
+        name: "announcer",
+        interactive: false,
+        encode: {
+          enter: {
+            x: { signal: String(width), offset: 1 }, // vega would prefer its internal ` "width" ` variable, but that breaks the squiggle playground. Just setting it to the same var as used elsewhere in the spec achieves the same result.
+            fill: { value: "black" },
+            fontSize: { value: 20 },
+            align: { value: "right" },
+          },
+          update: {
+            text: {
+              signal: dateTime
+                ? "position_scaled ? utcyear(position_scaled) + '-' + utcmonth(position_scaled) + '-' + utcdate(position_scaled) + 'T' + utchours(position_scaled)+':' +utcminutes(position_scaled) : ''"
+                : "position_scaled ? format(position_scaled, ',.4r')  : ''",
+            },
+          },
+        },
+      },
+      {
+        type: "rule",
+        interactive: false,
+        encode: {
+          enter: {
+            x: { value: 0 },
+            y: { scale: "yscale", value: 0 },
+
+            y2: {
+              signal: "height",
+              offset: 2,
+            },
+            strokeDash: { value: [5, 5] },
+          },
+
+          update: {
+            x: {
+              signal:
+                "position ? position[0] < 0 ? null : position[0] > width ? null : position[0]: null",
+            },
+
+            opacity: {
+              signal:
+                "position ? position[0] < 0 ? 0 : position[0] > width ? 0 : 1 : 0",
+            },
+          },
+        },
+      },
     ],
     legends: [
       {

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

@@ -79,6 +79,22 @@ could be continuous, discrete or mixed.
   </Story>
 </Canvas>
 
+### Date Distribution
+
+<Canvas>
+  <Story
+    name="Date Distribution"
+    args={{
+      code: "mx(1661819770311, 1661829770311, 1661839770311)",
+      width,
+      xAxisType: "dateTime",
+      width,
+    }}
+  >
+    {Template.bind({})}
+  </Story>
+</Canvas>
+
 ## Mixed distributions
 
 <Canvas>

packages/squiggle-lang/src/js/SqDistribution.ts

@@ -80,27 +80,28 @@ abstract class SqAbstractDistribution {
 }
 
 export class SqPointSetDistribution extends SqAbstractDistribution {
-  tag = Tag.PointSet;
+  tag = Tag.PointSet as const;
 
   value() {
-    return this.valueMethod(RSDistribution.getPointSet);
+    return wrapPointSetDist(this.valueMethod(RSDistribution.getPointSet));
   }
 }
 
 export class SqSampleSetDistribution extends SqAbstractDistribution {
-  tag = Tag.SampleSet;
+  tag = Tag.SampleSet as const;
 
-  value() {
+  value(): number[] {
     return this.valueMethod(RSDistribution.getSampleSet);
   }
 }
 
 export class SqSymbolicDistribution extends SqAbstractDistribution {
-  tag = Tag.Symbolic;
+  tag = Tag.Symbolic as const;
 
-  value() {
-    return this.valueMethod(RSDistribution.getSymbolic);
-  }
+  // not wrapped for TypeScript yet
+  // value() {
+  //   return this.valueMethod(RSDistribution.getSymbolic);
+  // }
 }
 
 const tagToClass = {