insomni-plot 0.1.0-alpha.0

package/src/stats/regression.d.ts

@@ -0,0 +1,38 @@
+/**
+ * A fitted regression. `predict(x)` returns the mean-response estimate;
+ * `seMean(x)` returns the standard error of that estimate (for confidence
+ * bands). `summary` carries whatever fit-specific scalars are interesting
+ * (slope, intercept, coefs, r²) — exposed for callers that want to render
+ * an equation or stat label.
+ */
+export interface RegressionFit {
+    readonly n: number;
+    predict(x: number): number;
+    /** Standard error of the mean response at `x`. Returns 0 when undefined. */
+    seMean(x: number): number;
+    readonly residualStdError: number;
+    readonly summary: Readonly<Record<string, number | readonly number[]>>;
+}
+export declare function linearFit(xs: readonly number[], ys: readonly number[]): RegressionFit | null;
+export declare function polyFit(xs: readonly number[], ys: readonly number[], degree: number): RegressionFit | null;
+export interface LoessOptions {
+    /** Fraction of the sample used in each local fit. Default `0.5`. */
+    span?: number;
+    /** Degree of local polynomial: `1` (default) or `2`. */
+    degree?: 1 | 2;
+}
+export declare function loessFit(xs: readonly number[], ys: readonly number[], options?: LoessOptions): RegressionFit | null;
+export interface ConfidenceBandPoint {
+    x: number;
+    yhat: number;
+    lo: number;
+    hi: number;
+}
+export interface ConfidenceBandOptions {
+    samples?: number;
+    /** Confidence level. Default `0.95`. */
+    level?: number;
+    /** Override `[xmin, xmax]`. Defaults to data range. */
+    domain?: readonly [number, number];
+}
+export declare function confidenceBand(fit: RegressionFit, xs: readonly number[], options?: ConfidenceBandOptions): ConfidenceBandPoint[];

package/src/stats/regression.test.ts

@@ -0,0 +1,56 @@
+import { describe, expect, test } from "vite-plus/test";
+
+import { confidenceBand, linearFit, loessFit, polyFit } from "./regression.ts";
+
+describe("regression", () => {
+  test("linearFit recovers slope/intercept from a noiseless line", () => {
+    const xs = [0, 1, 2, 3, 4, 5];
+    const ys = xs.map((x) => 2 * x + 1);
+    const fit = linearFit(xs, ys)!;
+    expect(fit).not.toBeNull();
+    expect(fit.summary.slope).toBeCloseTo(2, 6);
+    expect(fit.summary.intercept).toBeCloseTo(1, 6);
+    expect(fit.predict(10)).toBeCloseTo(21, 6);
+    expect(fit.summary.r2).toBeCloseTo(1, 6);
+  });
+
+  test("linearFit returns null for fewer than 2 finite points", () => {
+    expect(linearFit([1], [1])).toBeNull();
+    expect(linearFit([1, NaN], [2, 3])).toBeNull();
+  });
+
+  test("polyFit recovers a quadratic", () => {
+    const xs = [-3, -2, -1, 0, 1, 2, 3];
+    const ys = xs.map((x) => 0.5 * x * x - x + 2);
+    const fit = polyFit(xs, ys, 2)!;
+    expect(fit).not.toBeNull();
+    expect(fit.predict(4)).toBeCloseTo(0.5 * 16 - 4 + 2, 4);
+    expect(fit.summary.r2).toBeCloseTo(1, 6);
+  });
+
+  test("loessFit produces stable predictions across the domain", () => {
+    const xs = Array.from({ length: 50 }, (_, i) => i / 49);
+    const ys = xs.map((x) => Math.sin(x * Math.PI));
+    const fit = loessFit(xs, ys, { span: 0.4 })!;
+    expect(fit).not.toBeNull();
+    // Smoother should track the underlying sine within a generous tolerance.
+    expect(fit.predict(0.25)).toBeCloseTo(Math.sin(0.25 * Math.PI), 1);
+    expect(fit.predict(0.75)).toBeCloseTo(Math.sin(0.75 * Math.PI), 1);
+  });
+
+  test("confidenceBand emits ordered samples whose interval brackets yhat", () => {
+    const xs = Array.from({ length: 30 }, (_, i) => i);
+    const ys = xs.map((x) => x * 0.3 + (x % 3 === 0 ? 0.4 : -0.2));
+    const fit = linearFit(xs, ys)!;
+    const band = confidenceBand(fit, xs, { samples: 10 });
+    expect(band).toHaveLength(10);
+    for (const p of band) {
+      expect(p.lo).toBeLessThanOrEqual(p.yhat);
+      expect(p.hi).toBeGreaterThanOrEqual(p.yhat);
+    }
+    // x is monotonic.
+    for (let i = 1; i < band.length; i++) {
+      expect(band[i]!.x).toBeGreaterThan(band[i - 1]!.x);
+    }
+  });
+});

package/src/stats/regression.ts

@@ -0,0 +1,396 @@
+// ---------------------------------------------------------------------------
+// Regression
+// ---------------------------------------------------------------------------
+// Pure helpers for fitting smoothers and producing confidence bands. Used by
+// the `smooth` geom; exported via `insomni-plot/core` so callers can fit
+// once and feed the results into custom marks.
+
+/**
+ * A fitted regression. `predict(x)` returns the mean-response estimate;
+ * `seMean(x)` returns the standard error of that estimate (for confidence
+ * bands). `summary` carries whatever fit-specific scalars are interesting
+ * (slope, intercept, coefs, r²) — exposed for callers that want to render
+ * an equation or stat label.
+ */
+export interface RegressionFit {
+  readonly n: number;
+  predict(x: number): number;
+  /** Standard error of the mean response at `x`. Returns 0 when undefined. */
+  seMean(x: number): number;
+  readonly residualStdError: number;
+  readonly summary: Readonly<Record<string, number | readonly number[]>>;
+}
+
+interface NumericSample {
+  xs: number[];
+  ys: number[];
+}
+
+function filterFinite(xs: readonly number[], ys: readonly number[]): NumericSample {
+  const xo: number[] = [];
+  const yo: number[] = [];
+  const n = Math.min(xs.length, ys.length);
+  for (let i = 0; i < n; i++) {
+    const x = xs[i]!;
+    const y = ys[i]!;
+    if (Number.isFinite(x) && Number.isFinite(y)) {
+      xo.push(x);
+      yo.push(y);
+    }
+  }
+  return { xs: xo, ys: yo };
+}
+
+// ---------------------------------------------------------------------------
+// linearFit — ordinary least squares y = β₀ + β₁ x
+// ---------------------------------------------------------------------------
+
+export function linearFit(xs: readonly number[], ys: readonly number[]): RegressionFit | null {
+  const sample = filterFinite(xs, ys);
+  const n = sample.xs.length;
+  if (n < 2) return null;
+
+  let sx = 0;
+  let sy = 0;
+  for (let i = 0; i < n; i++) {
+    sx += sample.xs[i]!;
+    sy += sample.ys[i]!;
+  }
+  const xbar = sx / n;
+  const ybar = sy / n;
+
+  let sxx = 0;
+  let sxy = 0;
+  let syy = 0;
+  for (let i = 0; i < n; i++) {
+    const dx = sample.xs[i]! - xbar;
+    const dy = sample.ys[i]! - ybar;
+    sxx += dx * dx;
+    sxy += dx * dy;
+    syy += dy * dy;
+  }
+
+  if (sxx === 0) return null;
+  const slope = sxy / sxx;
+  const intercept = ybar - slope * xbar;
+
+  // Residual standard error (unbiased) — divide by (n - 2) for two-parameter fit.
+  let ssr = 0;
+  for (let i = 0; i < n; i++) {
+    const r = sample.ys[i]! - (intercept + slope * sample.xs[i]!);
+    ssr += r * r;
+  }
+  const dof = Math.max(1, n - 2);
+  const sigma = Math.sqrt(ssr / dof);
+  const r2 = syy > 0 ? 1 - ssr / syy : 0;
+
+  return {
+    n,
+    predict: (x) => intercept + slope * x,
+    seMean: (x) => sigma * Math.sqrt(1 / n + ((x - xbar) * (x - xbar)) / sxx),
+    residualStdError: sigma,
+    summary: { slope, intercept, r2, xbar, sxx },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// polyFit — least-squares polynomial (degree ≥ 1)
+// ---------------------------------------------------------------------------
+
+/** Solve a small symmetric positive-definite system via Gauss-Jordan. */
+function solveLinear(a: number[][], b: number[]): number[] | null {
+  const n = b.length;
+  const m: number[][] = a.map((row, i) => [...row, b[i]!]);
+  for (let i = 0; i < n; i++) {
+    // Partial pivot.
+    let pivot = i;
+    let pmax = Math.abs(m[i]![i]!);
+    for (let k = i + 1; k < n; k++) {
+      const v = Math.abs(m[k]![i]!);
+      if (v > pmax) {
+        pmax = v;
+        pivot = k;
+      }
+    }
+    if (pmax < 1e-12) return null;
+    if (pivot !== i) {
+      const tmp = m[i]!;
+      m[i] = m[pivot]!;
+      m[pivot] = tmp;
+    }
+    const div = m[i]![i]!;
+    for (let j = i; j <= n; j++) m[i]![j]! /= div;
+    for (let k = 0; k < n; k++) {
+      if (k === i) continue;
+      const factor = m[k]![i]!;
+      if (factor === 0) continue;
+      for (let j = i; j <= n; j++) m[k]![j]! -= factor * m[i]![j]!;
+    }
+  }
+  return m.map((row) => row[n]!);
+}
+
+function evalPoly(coefs: readonly number[], x: number): number {
+  // Horner — coefs[0] = β₀, coefs[1] = β₁, …
+  let acc = 0;
+  for (let i = coefs.length - 1; i >= 0; i--) acc = acc * x + coefs[i]!;
+  return acc;
+}
+
+export function polyFit(
+  xs: readonly number[],
+  ys: readonly number[],
+  degree: number,
+): RegressionFit | null {
+  if (degree < 1) return linearFit(xs, ys);
+  const sample = filterFinite(xs, ys);
+  const n = sample.xs.length;
+  const k = degree + 1;
+  if (n < k) return null;
+
+  // Build normal equations X'X β = X'y for design matrix X of column powers.
+  const xtx: number[][] = Array.from({ length: k }, () => Array.from({ length: k }, () => 0));
+  const xty: number[] = Array.from({ length: k }, () => 0);
+  // Cache power sums to skip O(k²·n) when k > 2.
+  const pow = Array.from({ length: 2 * k - 1 }, () => 0);
+  let sy = 0;
+  for (let i = 0; i < n; i++) {
+    const x = sample.xs[i]!;
+    const y = sample.ys[i]!;
+    let p = 1;
+    for (let j = 0; j < 2 * k - 1; j++) {
+      pow[j]! += p;
+      if (j < k) xty[j]! += p * y;
+      p *= x;
+    }
+    sy += y;
+  }
+  for (let i = 0; i < k; i++) {
+    for (let j = 0; j < k; j++) xtx[i]![j]! = pow[i + j]!;
+  }
+  const coefs = solveLinear(xtx, xty);
+  if (!coefs) return null;
+
+  let ssr = 0;
+  let sst = 0;
+  const ybar = sy / n;
+  for (let i = 0; i < n; i++) {
+    const yhat = evalPoly(coefs, sample.xs[i]!);
+    const r = sample.ys[i]! - yhat;
+    const d = sample.ys[i]! - ybar;
+    ssr += r * r;
+    sst += d * d;
+  }
+  const dof = Math.max(1, n - k);
+  const sigma = Math.sqrt(ssr / dof);
+  const r2 = sst > 0 ? 1 - ssr / sst : 0;
+
+  // Approximate seMean: a stable cheap proxy is the OLS leverage proxy
+  // σ·sqrt(k/n). For exact CI a caller can recompute via (X'X)⁻¹ but for
+  // visualization-grade ribbons this is plenty.
+  const seBase = sigma * Math.sqrt(k / n);
+
+  return {
+    n,
+    predict: (x) => evalPoly(coefs, x),
+    seMean: () => seBase,
+    residualStdError: sigma,
+    summary: { coefs, r2, degree },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// loessFit — local regression with tricube weights
+// ---------------------------------------------------------------------------
+
+/** Tricube weight for normalized distance `u ∈ [0, 1]`; 0 outside. */
+function tricube(u: number): number {
+  if (u >= 1) return 0;
+  const v = 1 - u * u * u;
+  return v * v * v;
+}
+
+export interface LoessOptions {
+  /** Fraction of the sample used in each local fit. Default `0.5`. */
+  span?: number;
+  /** Degree of local polynomial: `1` (default) or `2`. */
+  degree?: 1 | 2;
+}
+
+export function loessFit(
+  xs: readonly number[],
+  ys: readonly number[],
+  options: LoessOptions = {},
+): RegressionFit | null {
+  const sample = filterFinite(xs, ys);
+  const n = sample.xs.length;
+  if (n < 2) return null;
+  const span = Math.max(2 / n, Math.min(1, options.span ?? 0.5));
+  const degree = options.degree ?? 1;
+  const window = Math.max(2, Math.ceil(span * n));
+
+  // Sort once by x — neighborhood selection then is just a slice walk.
+  const order = Array.from({ length: n }, (_, i) => i).sort(
+    (a, b) => sample.xs[a]! - sample.xs[b]!,
+  );
+  const sx = order.map((i) => sample.xs[i]!);
+  const sy = order.map((i) => sample.ys[i]!);
+
+  // Residual standard error from a cheap sliding fit — used as a global σ
+  // for the (constant) confidence band rather than per-x bandwidth tracking.
+  let ssr = 0;
+  function localFit(x: number): number {
+    // Find the `window` nearest neighbors to `x` by linear scan from the
+    // bisection point. Adequate for typical n ≤ a few thousand.
+    let lo = 0;
+    let hi = n;
+    while (lo < hi) {
+      const mid = (lo + hi) >>> 1;
+      if (sx[mid]! < x) lo = mid + 1;
+      else hi = mid;
+    }
+    let l = lo;
+    let r = lo;
+    while (r - l < window) {
+      const dl = l > 0 ? x - sx[l - 1]! : Infinity;
+      const dr = r < n ? sx[r]! - x : Infinity;
+      if (dl <= dr) l = Math.max(0, l - 1);
+      else r = Math.min(n, r + 1);
+      if (l === 0 && r === n) break;
+    }
+    const maxDist = Math.max(Math.abs(x - sx[l]!), Math.abs(sx[r - 1]! - x), 1e-12);
+
+    // Weighted least squares for local polynomial.
+    if (degree === 1) {
+      let sw = 0;
+      let swx = 0;
+      let swy = 0;
+      let swxx = 0;
+      let swxy = 0;
+      for (let i = l; i < r; i++) {
+        const w = tricube(Math.abs(sx[i]! - x) / maxDist);
+        if (w === 0) continue;
+        const xi = sx[i]!;
+        const yi = sy[i]!;
+        sw += w;
+        swx += w * xi;
+        swy += w * yi;
+        swxx += w * xi * xi;
+        swxy += w * xi * yi;
+      }
+      const denom = sw * swxx - swx * swx;
+      if (Math.abs(denom) < 1e-12 || sw === 0) return swy / Math.max(sw, 1e-12);
+      const slope = (sw * swxy - swx * swy) / denom;
+      const intercept = (swy - slope * swx) / sw;
+      return intercept + slope * x;
+    }
+
+    // degree === 2 — solve weighted normal equations of size 3.
+    const xtx: number[][] = [
+      [0, 0, 0],
+      [0, 0, 0],
+      [0, 0, 0],
+    ];
+    const xty = [0, 0, 0];
+    for (let i = l; i < r; i++) {
+      const w = tricube(Math.abs(sx[i]! - x) / maxDist);
+      if (w === 0) continue;
+      const xi = sx[i]!;
+      const yi = sy[i]!;
+      const x2 = xi * xi;
+      xtx[0]![0]! += w;
+      xtx[0]![1]! += w * xi;
+      xtx[0]![2]! += w * x2;
+      xtx[1]![1]! += w * x2;
+      xtx[1]![2]! += w * xi * x2;
+      xtx[2]![2]! += w * x2 * x2;
+      xty[0]! += w * yi;
+      xty[1]! += w * xi * yi;
+      xty[2]! += w * x2 * yi;
+    }
+    xtx[1]![0] = xtx[0]![1]!;
+    xtx[2]![0] = xtx[0]![2]!;
+    xtx[2]![1] = xtx[1]![2]!;
+    const coefs = solveLinear(xtx, xty);
+    if (!coefs) return ys.length > 0 ? ys[0]! : 0;
+    return coefs[0]! + coefs[1]! * x + coefs[2]! * x * x;
+  }
+
+  // Pre-fit to compute residual SE for a (rough) constant-band CI.
+  for (let i = 0; i < n; i++) {
+    const r = sy[i]! - localFit(sx[i]!);
+    ssr += r * r;
+  }
+  const sigma = Math.sqrt(ssr / Math.max(1, n - 2));
+
+  return {
+    n,
+    predict: (x) => localFit(x),
+    seMean: () => sigma * Math.sqrt(1 / n),
+    residualStdError: sigma,
+    summary: { span, degree },
+  };
+}
+
+// ---------------------------------------------------------------------------
+// confidenceBand — sample a fit on a regular grid and return ribbon points
+// ---------------------------------------------------------------------------
+
+export interface ConfidenceBandPoint {
+  x: number;
+  yhat: number;
+  lo: number;
+  hi: number;
+}
+
+export interface ConfidenceBandOptions {
+  samples?: number;
+  /** Confidence level. Default `0.95`. */
+  level?: number;
+  /** Override `[xmin, xmax]`. Defaults to data range. */
+  domain?: readonly [number, number];
+}
+
+/** Two-sided z for a normal-approx confidence level (large-n). */
+function zForLevel(level: number): number {
+  // Quick lookup matches typical levels; falls back to ≈1.96.
+  if (level >= 0.99) return 2.576;
+  if (level >= 0.975) return 2.241;
+  if (level >= 0.95) return 1.96;
+  if (level >= 0.9) return 1.645;
+  if (level >= 0.8) return 1.282;
+  return 1.96;
+}
+
+export function confidenceBand(
+  fit: RegressionFit,
+  xs: readonly number[],
+  options: ConfidenceBandOptions = {},
+): ConfidenceBandPoint[] {
+  const samples = Math.max(2, options.samples ?? 64);
+  const z = zForLevel(options.level ?? 0.95);
+  let lo: number;
+  let hi: number;
+  if (options.domain) {
+    [lo, hi] = options.domain;
+  } else {
+    lo = Infinity;
+    hi = -Infinity;
+    for (const v of xs) {
+      if (!Number.isFinite(v)) continue;
+      if (v < lo) lo = v;
+      if (v > hi) hi = v;
+    }
+    if (!Number.isFinite(lo) || !Number.isFinite(hi) || lo === hi) return [];
+  }
+  const out = Array.from<ConfidenceBandPoint>({ length: samples });
+  const step = (hi - lo) / (samples - 1);
+  for (let i = 0; i < samples; i++) {
+    const x = lo + step * i;
+    const yhat = fit.predict(x);
+    const se = fit.seMean(x);
+    const half = z * se;
+    out[i] = { x, yhat, lo: yhat - half, hi: yhat + half };
+  }
+  return out;
+}

package/src/stats/rolling-window.d.ts

@@ -0,0 +1,55 @@
+export type RollingStatisticKind = "mean" | "median" | "min" | "max" | "sum" | "stddev";
+/**
+ * Statistic to compute over each window. Strings cover the common cases;
+ * `{ kind: 'quantile', p }` picks an arbitrary quantile (0..1); a function
+ * receives the windowed `(values, items)` and returns any scalar — use it for
+ * custom stats (mode, trimmed mean, sign of slope, …).
+ */
+export type RollingStatistic<T> = RollingStatisticKind | {
+    kind: "quantile";
+    p: number;
+} | ((values: readonly number[], items: readonly T[]) => number);
+/**
+ * Window size.
+ *
+ * - A bare number is interpreted as `unit: 'count'` (k nearest items by axis
+ *   value, including self). Most predictable across axis types.
+ * - `{ value, unit: 'domain' }` — sliding window of `value` data-axis units,
+ *   centered on the current item. For a time axis pass milliseconds (e.g.
+ *   `7 * 24 * 3_600_000` for a 7-day window). For a numeric axis the value is
+ *   in the axis's own units.
+ * - `{ value, unit: 'count' }` — explicit count-window form.
+ */
+export type RollingWindow = number | {
+    value: number;
+    unit: "domain" | "count";
+};
+export type RollingAxis = "x" | "y";
+export interface RollingWindowOptions<T> {
+    x: (datum: T, index: number) => number | Date;
+    y: (datum: T, index: number) => number;
+    window: RollingWindow;
+    /** Default `"mean"`. */
+    statistic?: RollingStatistic<T>;
+    /**
+     * Points failing the filter are excluded from every window *and* not
+     * emitted on the output series — they vanish from both the aggregation
+     * and the line that connects the results.
+     */
+    filter?: (datum: T, index: number) => boolean;
+    /**
+     * Axis the window slides along. Default `"x"` (slide on x, aggregate y);
+     * `"y"` flips (slide on y, aggregate x). Output is sorted by the slide
+     * axis so a `lineMark` connecting the points draws monotonically.
+     */
+    axis?: RollingAxis;
+}
+export interface RollingPoint {
+    x: number;
+    y: number;
+    /** Index into the original input `data` array. */
+    sourceIndex: number;
+    /** Number of points (after filter) that fell in this window. */
+    count: number;
+}
+export declare function rollingWindow<T>(data: readonly T[], options: RollingWindowOptions<T>): RollingPoint[];