Add linearRegression!

package.json

@@ -29,7 +29,7 @@
     "@observablehq/plot": "./src/index.js"
   },
   "devDependencies": {
-    "d3": "^6.3.1",
+    "d3": "^6.4.0",
     "eslint": "^7.12.1",
     "esm": "^3.2.25",
     "js-beautify": "^1.13.0",

src/index.js

@@ -10,6 +10,7 @@ export {group, groupX, groupY} from "./marks/group.js";
 export {Line, line, lineX, lineY} from "./marks/line.js";
 export {Link, link} from "./marks/link.js";
 export {Rect, rect, rectX, rectY} from "./marks/rect.js";
+export {linearRegression} from "./marks/regression.js";
 export {RuleX, RuleY, ruleX, ruleY} from "./marks/rule.js";
 export {stackAreaX, stackAreaY, stackBarX, stackBarY} from "./marks/stack.js";
 export {Text, text, textX, textY} from "./marks/text.js";

src/mark.js

@@ -32,7 +32,7 @@ export class Mark {
     let index, data;
     if (this.data !== undefined) {
       if (this.transform === identity) { // optimized common case
-        data = this.data, index = facets !== undefined ? facets : range(data);
+        data = this.data, index = facets !== undefined ? facets : range(data.length);
       } else if (this.transform.length === 2) { // facet-aware transform
         ({index, data} = this.transform(this.data, facets));
         data = arrayify(data);
@@ -46,7 +46,7 @@ export class Mark {
         index = [], data = [];
         for (const facet of facets) {
           const facetData = arrayify(this.transform(take(this.data, facet)), Array);
-          const facetIndex = facetData === undefined ? undefined : offsetRange(facetData, k);
+          const facetIndex = facetData === undefined ? undefined : range(k, k + facetData.length);
           k += facetData.length;
           index.push(facetIndex);
           data.push(facetData);
@@ -65,7 +65,7 @@ export class Mark {
         }
       } else { // basic transform, non-faceted
         data = arrayify(this.transform(this.data));
-        index = data === undefined ? undefined : range(data);
+        index = data === undefined ? undefined : range(data.length);
       }
     }
     return {
@@ -192,15 +192,15 @@ export function titleGroup(L) {
       .text(([i]) => L[i]) : () => {};
 }
 
-// Returns a Uint32Array with elements [0, 1, 2, … data.length - 1].
-export function range(data) {
-  return Uint32Array.from(data, indexOf);
-}
-
-// Returns a Uint32Array with elements [k, k + 1, … k + data.length - 1].
-export function offsetRange(data, k) {
-  k = Math.floor(k);
-  return Uint32Array.from(data, (_, i) => i + k);
+// Returns a Uint32Array with elements [start, start + 1, start + 2, … stop - 1].
+export function range(start, stop) {
+  if (stop === undefined) stop = start, start = 0;
+  start = Math.floor(start), stop = Math.floor(stop);
+  if (!(stop >= start)) throw new Error("invalid range");
+  const n = stop - start;
+  const range = new Uint32Array(n);
+  for (let i = 0; i < n; ++i) range[i] = i + start;
+  return range;
 }
 
 // Returns an array [values[index[0]], values[index[1]], …].

src/marks/regression.js

@@ -0,0 +1,59 @@
+import {extent} from "d3-array";
+import {defined} from "../defined.js";
+import {group} from "../group.js";
+import {maybeColor, range, valueof} from "../mark.js";
+import {link} from "./link.js";
+
+export function linearRegression(data, {stroke, x, y, z, ...options}) {
+  let [vstroke, cstroke] = maybeColor(stroke, "currentColor");
+  if (z === undefined && vstroke != null) z = vstroke;
+  const X1 = [], Y1 = [], X2 = [], Y2 = [], S = vstroke ? [] : undefined; // lazily populated
+  return link(data, {
+    ...options,
+    transform: (data, facets) => {
+      x = valueof(data, x);
+      y = valueof(data, y);
+      z = z !== undefined ? valueof(data, z) : undefined;
+      if (vstroke !== undefined) vstroke = valueof(data, vstroke);
+      const [x1, x2] = extent(x);
+      const index = [];
+      let offset = 0;
+      // TODO it’d be nice if faceting didn’t make this complicated
+      for (let facet of facets === undefined ? [range(data.length)] : facets) {
+        facet = facet.filter(i => defined(x[i]) && defined(y[i]));
+        let n = 0;
+        for (const index of z ? group(facet, z) : [facet]) {
+          const f = linearRegressionLine(index, x, y);
+          X1.push(x1), Y1.push(f(x1)), X2.push(x2), Y2.push(f(x2));
+          if (S) S.push(vstroke[index[0]]);
+          ++n;
+        }
+        index.push(range(offset, offset + n));
+        offset += n;
+      }
+      return {index: facets === undefined ? index[0] : index};
+    },
+    x1: X1,
+    y1: Y1,
+    x2: X2,
+    y2: Y2,
+    stroke: cstroke ? cstroke : vstroke ? S : undefined
+  });
+}
+
+function linearRegressionLine(I, X, Y) {
+  const n = I.length;
+  if (n === 1) return () => Y[I[0]];
+  let sx = 0, sy = 0, sxx = 0, sxy = 0;
+  for (const i of I) {
+    const x = X[i];
+    const y = Y[i];
+    sx += x;
+    sy += y;
+    sxx += x * x;
+    sxy += x * y;
+  }
+  const m = (n * sxy - sx * sy) / (n * sxx - sx * sx);
+  const b = (sy - m * sx) / n;
+  return x => b + m * x;
+}

src/transforms/bin.js

@@ -1,5 +1,5 @@
 import {bin as binner, cross} from "d3-array";
-import {valueof, first, second, maybeValue, range, offsetRange} from "../mark.js";
+import {valueof, first, second, maybeValue, range} from "../mark.js";
 
 export function bin1(options = {}) {
   let {value, domain, thresholds, cumulative} = maybeValue(options);
@@ -34,7 +34,7 @@ function binof({value, domain, thresholds}) {
     const bin = binner().value(i => values[i]);
     if (domain !== undefined) bin.domain(domain);
     if (thresholds !== undefined) bin.thresholds(thresholds);
-    return bin(range(data));
+    return bin(range(data.length));
   };
 }
 
@@ -43,7 +43,7 @@ function rebin(bins, facets, subset, cumulative) {
   if (facets === undefined) {
     if (cumulative) bins = accumulate(cumulative < 0 ? bins.reverse() : bins);
     bins = bins.filter(nonempty);
-    return {index: range(bins), data: bins};
+    return {index: range(bins.length), data: bins};
   }
   const index = [];
   const data = [];
@@ -52,7 +52,7 @@ function rebin(bins, facets, subset, cumulative) {
     let b = bins.map(facet);
     if (cumulative) b = accumulate(cumulative < 0 ? b.reverse() : b);
     b = b.filter(nonempty);
-    index.push(offsetRange(b, k));
+    index.push(range(k, k + b.length));
     data.push(b);
     k += b.length;
   }

src/transforms/group.js

@@ -1,12 +1,12 @@
 import {groups} from "d3-array";
 import {defined} from "../defined.js";
-import {valueof, maybeValue, range, offsetRange} from "../mark.js";
+import {valueof, maybeValue, range} from "../mark.js";
 
 export function group1(x) {
   const {value} = maybeValue({value: x});
   return (data, facets) => {
     const values = valueof(data, value);
-    let g = groups(range(data), i => values[i]).filter(defined1);
+    let g = groups(range(data.length), i => values[i]).filter(defined1);
     return regroup(g, facets);
   };
 }
@@ -17,21 +17,21 @@ export function group2(vx, vy) {
   return (data, facets) => {
     const valuesX = valueof(data, x);
     const valuesY = valueof(data, y);
-    let g = groups(range(data), i => valuesX[i], i => valuesY[i]).filter(defined1);
+    let g = groups(range(data.length), i => valuesX[i], i => valuesY[i]).filter(defined1);
     g = g.flatMap(([x, xgroup]) => xgroup.filter(defined1).map(([y, ygroup]) => [x, y, ygroup]));
     return regroup(g, facets);
   };
 }
 
 // When faceting, subdivides the given groups according to the facet indexes.
 function regroup(groups, facets) {
-  if (facets === undefined) return {index: range(groups), data: groups};
+  if (facets === undefined) return {index: range(groups.length), data: groups};
   const index = [];
   const data = [];
   let k = 0;
   for (const facet of facets.map(subset)) {
     let g = groups.map(facet).filter(nonempty1);
-    index.push(offsetRange(g, k));
+    index.push(range(k, k + g.length));
     data.push(g);
     k += g.length;
   }