peeksafe 0.2.0 → 0.3.0

package/src/plan.ts

@@ -32,7 +32,7 @@
  * it.
  */
 
-import { type BaselineStat, baselineNullRate, type CaseRef } from "./baseline.js";
+import type { BaselineStat, CaseRef } from "./baseline.js";
 import {
   PeeksafeError,
   requireCounts,
@@ -44,7 +44,6 @@ import {
   requireUniqueCaseIds,
 } from "./errors.js";
 import {
-  betaQuantile,
   ebhSoloThreshold,
   evidenceCeilingLogE,
   mcnemarSamplesForEvidence,
@@ -52,9 +51,8 @@ import {
   pairPhi,
   sampleSizeTwoProportion,
   samplesForEvidence,
-  sprtExpectedN,
 } from "./stats.js";
-import { DEFAULT_STOP_OPTIONS, settleLogE, shouldStop } from "./stop.js";
+import { shouldStop, shouldStopPaired } from "./stop.js";
 
 export interface PlanConfig {
   /** the drop we want to be able to detect, in probability units (0.15 = 15 points) */
@@ -169,6 +167,13 @@ export interface PlanCase {
    * "undetectable" verdict.
    */
   baselineRunsNeeded: number | null;
+  /**
+   * Expected runs this case costs in a pull request in which nothing moved, at
+   * the design the expected bill prices it at: `config.design`, else `design`,
+   * else the suite's `recommendedDesign` (unpaired when `planCase` is called on
+   * its own). The cases of a plan sum to `totals.typicalRuns`; 0 with no baseline.
+   */
+  typicalRuns: number;
   /** what a fixed-N two-proportion design would spend per arm, for comparison */
   fixedNPerArm: number;
   reason: string;
@@ -330,6 +335,9 @@ export function planCase(
   if (b) requireCounts(b.successes, b.trials, `planCase(${c.id}).baseline`);
   requireErrorRates(cfg.alpha, cfg.beta, "planCase");
   requirePrices(cfg, "planCase");
+  // The expected bill walks the batch grid at the stopping rule's level; a zero batch never ends.
+  requireTrialGrid(cfg, "planCase");
+  requireSettleConfidence(cfg, "planCase");
   const barLogE = Math.log(ebhSoloThreshold(m, cfg.fdr));
   const trials = b?.trials ?? 0;
   const successes = b?.successes ?? 0;
@@ -357,6 +365,7 @@ export function planCase(
       detectability: "IMPOSSIBLE",
       design: null,
       baselineRunsNeeded: null,
+      typicalRuns: 0,
       reason: "no baseline: nothing to compare against, record one first",
     };
   }
@@ -370,6 +379,8 @@ export function planCase(
       detectability: "IMPOSSIBLE",
       design: null,
       baselineRunsNeeded: null,
+      // Undecidable, but the gate still runs it until it stops, so the expected bill pays for it.
+      typicalRuns: typicalCaseRuns(b, cfg.design ?? "unpaired", cfg, m),
       reason:
         `baseline rate ${(rate * 100).toFixed(0)}% cannot lose ${(cfg.mde * 100).toFixed(0)} points, ` +
         `the effect is arithmetically impossible, not merely expensive`,
@@ -445,24 +456,40 @@ export function planCase(
     detectability,
     design,
     baselineRunsNeeded,
+    typicalRuns: typicalCaseRuns(b, cfg.design ?? design ?? "unpaired", cfg, m),
     reason,
   };
 }
 
+// Each design is counted under the rule that stops it: an unpaired case in the runs `shouldStop` takes, a paired one in the pairs `shouldStopPaired` takes.
+function typicalCaseRuns(
+  b: BaselineStat | undefined,
+  design: "unpaired" | "paired",
+  cfg: PlanConfig,
+  m: number,
+): number {
+  if (!b || b.trials === 0) return 0;
+  const rate = b.successes / b.trials;
+  return design === "paired"
+    ? 2 * expectedSettlePairs(rate, cfg.pairCoupling, cfg.mde, m, cfg)
+    : expectedStopRuns(b, rate, m, cfg);
+}
+
 /**
- * Expected total samples for the sequential design, priced under Wald's
- * expected sample number at a supplied "true" rate per case, and clamped into
- * the gate's own [minTrials, maxTrials] batch grid.
+ * Expected total samples for the sequential design: per case, the exact
+ * expected runs `shouldStop` takes to stop it when its candidate passes at
+ * `trueRate`, in a suite of `cases.length`, on the gate's own batch grid
+ * (`expectedStopRuns`). At the default rate this is the figure the plan bills
+ * an unpaired case at, `PlanCase.typicalRuns`.
  *
- * `makePlan` no longer bills from it: an unpaired case there is counted in
- * the runs `shouldStop` takes, which this figure misses either way (#105).
  * `samplesForEvidence`, the one behind `PlanCase.unpaired.runs`, answers a
  * different question: how many runs it takes to *certify a regression*. A case
- * that did not regress never gets there; it hits the H₀ wall in a few batches
- * and stops. Pricing every case at the certification cost over-predicts the
+ * that did not regress never gets there; it stops as futile or settled in a
+ * few batches. Pricing every case at the certification cost over-predicts the
  * bill.
  *
- * `trueRate` defaults to the case's own baseline rate, i.e. "nothing moved", * which is what almost every pull request looks like.
+ * `trueRate` defaults to the case's own baseline rate, i.e. "nothing moved",
+ * which is what almost every pull request looks like.
  */
 export function expectedSequentialSamples(
   cases: Array<CaseRef>,
@@ -472,6 +499,10 @@ export function expectedSequentialSamples(
     b.trials > 0 ? b.successes / b.trials : 0,
 ): { samples: number; perCase: Map<string, number> } {
   requireTrialGrid(cfg, "expectedSequentialSamples");
+  // Checked here so a bad value is refused under this name, not under the stopping rule's.
+  requireOpenProbability(cfg.mde, "mde", "expectedSequentialSamples");
+  requireOpenProbability(cfg.fdr, "fdr", "expectedSequentialSamples");
+  requireSettleConfidence(cfg, "expectedSequentialSamples");
   // A repeated id was added to `samples` once per listing but kept once in `perCase`, so the two disagreed.
   requireUniqueCaseIds(cases, "expectedSequentialSamples");
   const perCase = new Map<string, number>();
@@ -479,16 +510,12 @@ export function expectedSequentialSamples(
   for (const c of cases) {
     const b = baseline.get(c.id);
     if (!b || b.trials === 0) continue;
-    const p0 = baselineNullRate(b, cfg.nullQuantile);
-    const p1 = Math.max(0.005, p0 - cfg.mde);
     const p = Math.max(0, Math.min(1, trueRate(c.id, b)));
-    const raw = sprtExpectedN(p, p0, p1, cfg.alpha, cfg.beta);
-    const capped = onBatchGrid(
-      Number.isFinite(raw) && raw > 0 ? Math.ceil(raw) : cfg.maxTrials,
-      cfg,
-    );
-    perCase.set(c.id, capped);
-    total += capped;
+    // Wald's figure fell back to the cap on a NaN rate; the dynamic program would answer NaN.
+    requireProbability(p, "trueRate", "expectedSequentialSamples");
+    const runs = expectedStopRuns(b, p, cases.length, cfg);
+    perCase.set(c.id, runs);
+    total += runs;
   }
   return { samples: total, perCase };
 }
@@ -505,38 +532,101 @@ export function requireSettleConfidence(
   }
 }
 
-// The gate only ever adds whole batches past the minimum, so a prediction is rounded onto the grid it samples on.
-function onBatchGrid(n: number, cfg: TrialGrid): number {
-  const clamped = Math.min(cfg.maxTrials, Math.max(cfg.minTrials, n));
-  const batched =
-    cfg.minTrials + Math.ceil(Math.max(0, clamped - cfg.minTrials) / cfg.batchSize) * cfg.batchSize;
-  return Math.min(cfg.maxTrials, batched);
-}
+// The frontier prices the same paired cell once per baseline and screen, so a sweep would rerun one program dozens of times.
+const settlePairsMemo = new Map<string, number>();
+// Every case of a suite that shares a baseline, and every planCase call makePlan makes, runs the same program.
+const stopRunsMemo = new Map<string, number>();
 
 /**
- * Expected pairs for a paired case that did not move: the first pair count at
- * which `shouldStopPaired`'s settle e-value, on the mean trajectory at coupling
- * `rho` and `cfg.futilityConfidence` (the stopping rule's default when unset),
- * rules out a drop of `mde`, on the same batch grid as
- * `expectedSequentialSamples`, and `maxTrials` when it never does. The
- * unpaired SPRT figure knows nothing of this rule or of the coupling, which
- * sets how often an unmoved pair disagrees.
+ * Expected pairs `shouldStopPaired` takes to stop a paired case that did not
+ * move, at coupling `rho`, checked at `minTrials` and then every `batchSize`
+ * pairs up to `maxTrials`, in a suite of `suiteSize`. Exact, as
+ * `expectedStopRuns` is: the decision depends only on the worse and better
+ * counts at a given pair count, so a dynamic program carries their joint
+ * distribution from one check to the next. Rounded up to a whole pair.
  */
 export function expectedSettlePairs(
   rate: number,
   rho: number,
   mde: number,
-  cfg: TrialGrid & Pick<PlanConfig, "futilityConfidence">,
+  suiteSize: number,
+  cfg: TrialGrid & Pick<PlanConfig, "fdr" | "futilityConfidence">,
 ): number {
+  const key = [
+    rate,
+    rho,
+    mde,
+    suiteSize,
+    cfg.fdr,
+    cfg.futilityConfidence,
+    cfg.minTrials,
+    cfg.maxTrials,
+    cfg.batchSize,
+  ].join("|");
+  const known = settlePairsMemo.get(key);
+  if (known !== undefined) return known;
+  const options = {
+    suiteSize,
+    mde,
+    fdr: cfg.fdr,
+    futilityConfidence: cfg.futilityConfidence,
+  };
   // Unmoved arms disagree each way only when uncoupled: (1 - rho) p (1 - p).
   const each = (1 - rho) * rate * (1 - rate);
-  const confidence = cfg.futilityConfidence ?? DEFAULT_STOP_OPTIONS.futilityConfidence;
-  const target = -Math.log(1 - confidence);
-  for (let n = 1; n <= cfg.maxTrials; n++) {
-    const counts = { bothPass: n * (1 - 2 * each), worse: n * each, better: n * each, bothFail: 0 };
-    if (settleLogE(counts, mde) >= target) return onBatchGrid(n, cfg);
+  const same = 1 - 2 * each;
+  // live[w * (top + 1) + b]: probability of w worse and b better pairs with the case still running.
+  let live = new Float64Array([1]);
+  // Mass under 1e-18 is dropped so the square stays as small as the spread of the counts; it cannot move the rounded answer.
+  let top = 0;
+  let n = 0;
+  let expected = 0;
+  for (let step = cfg.minTrials; ; step = Math.min(cfg.batchSize, cfg.maxTrials - n)) {
+    for (let i = 0; i < step; i++) {
+      const side = top + 2;
+      const next = new Float64Array(side * side);
+      for (let w = 0; w <= top; w++) {
+        for (let b = 0; b <= top; b++) {
+          const q = live[w * (top + 1) + b]!;
+          if (q === 0) continue;
+          next[w * side + b]! += q * same;
+          next[(w + 1) * side + b]! += q * each;
+          next[w * side + b + 1]! += q * each;
+        }
+      }
+      let reach = 0;
+      for (let at = 0; at < next.length; at++) {
+        if (next[at]! < 1e-18) next[at] = 0;
+        else reach = Math.max(reach, Math.floor(at / side), at % side);
+      }
+      live = new Float64Array((reach + 1) * (reach + 1));
+      for (let w = 0; w <= reach; w++) {
+        for (let b = 0; b <= reach; b++) live[w * (reach + 1) + b] = next[w * side + b]!;
+      }
+      top = reach;
+      n++;
+    }
+    let running = 0;
+    for (let w = 0; w <= top; w++) {
+      for (let b = 0; b <= top; b++) {
+        const at = w * (top + 1) + b;
+        const q = live[at]!;
+        if (q === 0) continue;
+        const counts = { bothPass: n - w - b, worse: w, better: b, bothFail: 0 };
+        if (n >= cfg.maxTrials || shouldStopPaired(counts, options).stop) {
+          expected += q * n;
+          live[at] = 0;
+        } else {
+          running += q;
+        }
+      }
+    }
+    if (running === 0) break;
   }
-  return cfg.maxTrials;
+  // Float residue must not push an exact whole number up a pair.
+  const pairs = Math.ceil(expected - 1e-9);
+  if (settlePairsMemo.size >= 1024) settlePairsMemo.clear();
+  settlePairsMemo.set(key, pairs);
+  return pairs;
 }
 
 /**
@@ -547,12 +637,26 @@ export function expectedSettlePairs(
  * program carries the distribution of the success count from one check to the
  * next. Rounded up to a whole run; an expectation is not a grid point.
  */
-function expectedStopRuns(
+export function expectedStopRuns(
   baseline: BaselineStat,
   rate: number,
   suiteSize: number,
   cfg: TrialGrid & Pick<PlanConfig, "mde" | "fdr" | "futilityConfidence">,
 ): number {
+  const key = [
+    baseline.successes,
+    baseline.trials,
+    rate,
+    suiteSize,
+    cfg.mde,
+    cfg.fdr,
+    cfg.futilityConfidence,
+    cfg.minTrials,
+    cfg.maxTrials,
+    cfg.batchSize,
+  ].join("|");
+  const known = stopRunsMemo.get(key);
+  if (known !== undefined) return known;
   const options = {
     suiteSize,
     mde: cfg.mde,
@@ -586,7 +690,10 @@ function expectedStopRuns(
     if (running === 0) break;
   }
   // Float residue must not push an exact whole number up a run.
-  return Math.ceil(expected - 1e-9);
+  const runs = Math.ceil(expected - 1e-9);
+  if (stopRunsMemo.size >= 1024) stopRunsMemo.clear();
+  stopRunsMemo.set(key, runs);
+  return runs;
 }
 
 /**
@@ -669,6 +776,13 @@ export function makePlan(
   // Per case, like `bestRuns`: one suite-wide design billed a mixed suite as if every case ran it.
   const runsPerObs = (p: PlanCase) =>
     (cfg.design ?? p.design ?? recommendedDesign) === "paired" ? 2 : 1;
+  // `planCase` cannot see the suite, so a case with no design was priced unpaired there.
+  for (const p of planned) {
+    if (p.design === null && runsPerObs(p) === 2) {
+      const b = baseline.get(p.caseId);
+      p.typicalRuns = typicalCaseRuns(b, "paired", { ...cfg, pairCoupling: p.pairCoupling }, m);
+    }
+  }
   // `screenRuns` counts observations, so each case's pass costs its own design's runs.
   const screenPassRuns = planned.reduce((a, p) => a + cfg.screenRuns * runsPerObs(p), 0);
 
@@ -694,27 +808,14 @@ export function makePlan(
   // set from.
   // Both counts are *observations*. A paired observation is two runs, the
   // candidate and the reference revision on the same seed, and so is a paired
-  // screening run. Each case is counted under the rule that stops it: a paired
-  // one in pairs to settle, an unpaired one in the runs `shouldStop` takes.
-  // Wald's SPRT figure, used for both before, runs a different rule, ignores
-  // the coupling and the suite size, and mispriced cases either way.
-  const stopRuns = new Map<string, number>();
-  const observations = (p: PlanCase) => {
-    const b = baseline.get(p.caseId);
-    if (!b || b.trials === 0) return 0;
-    if (runsPerObs(p) === 2) {
-      return expectedSettlePairs(p.baseline.rate, p.pairCoupling, cfg.mde, cfg);
-    }
-    const key = `${b.successes}/${b.trials}`;
-    const known = stopRuns.get(key);
-    if (known !== undefined) return known;
-    const runs = expectedStopRuns(b, b.successes / b.trials, m, cfg);
-    stopRuns.set(key, runs);
-    return runs;
-  };
+  // screening run. Each case is counted under the rule that stops it
+  // (`PlanCase.typicalRuns`): a paired one in pairs to settle, an unpaired one
+  // in the runs `shouldStop` takes. Wald's SPRT figure, used for both before,
+  // runs a different rule, ignores the coupling and the suite size, and
+  // mispriced cases either way.
   const priced = planned.filter((p) => (baseline.get(p.caseId)?.trials ?? 0) > 0).length;
-  const typicalObservations = planned.reduce((a, p) => a + observations(p), 0);
-  const typicalRuns = planned.reduce((a, p) => a + observations(p) * runsPerObs(p), 0);
+  const typicalObservations = planned.reduce((a, p) => a + p.typicalRuns / runsPerObs(p), 0);
+  const typicalRuns = planned.reduce((a, p) => a + p.typicalRuns, 0);
   const runsPerObservation =
     typicalObservations > 0
       ? typicalRuns / typicalObservations
@@ -862,9 +963,7 @@ export function affordabilityGrid(
 ): AffordabilityCell[] {
   requireProbability(baselineRate, "baselineRate", "affordabilityGrid");
   requireErrorRates(cfg.alpha, cfg.beta, "affordabilityGrid");
-  // The typical bill below takes this quantile of the baseline posterior
-  // straight from `betaQuantile`, which answers 0 or 1 for a quantile outside
-  // (0,1) and about 0 for NaN instead of refusing. `makePlan` refuses it.
+  // `makePlan` refuses it, so one options object must not price here and fail there.
   requireOpenProbability(cfg.nullQuantile, "nullQuantile", "affordabilityGrid");
   // makePlan refuses both; unchecked, an fdr of 2 priced a grid cheaper than one at 0.05.
   requireOpenProbability(cfg.fdr, "fdr", "affordabilityGrid");
@@ -936,16 +1035,16 @@ export function affordabilityGrid(
       // The ceiling above prices a pull request in which every case regressed.
       // This prices one in which nothing did, which is what a pull request
       // almost always is.
-      const p0 = betaQuantile(1 + s, 1 + baselineTrials - s, cfg.nullQuantile);
-      const p1 = Math.max(0.005, p0 - mde);
-      const rawN = sprtExpectedN(baselineRate, p0, p1, cfg.alpha, cfg.beta);
-      // A paired cell counts the pairs `shouldStopPaired` takes to settle, as `makePlan` does; the SPRT figure ignores the coupling.
+      // Each design is counted under the rule that stops it, as `makePlan` does; Wald's SPRT figure ignores the coupling and the suite size.
       const perCaseTypical =
         runsPerObs === 2
-          ? expectedSettlePairs(baselineRate, cfg.pairCoupling, mde, cfg)
-          : Number.isFinite(rawN) && rawN > 0
-            ? Math.min(cfg.maxTrials, Math.max(cfg.minTrials, Math.ceil(rawN)))
-            : cfg.maxTrials;
+          ? expectedSettlePairs(baselineRate, cfg.pairCoupling, mde, m, cfg)
+          : expectedStopRuns(
+              { caseId: "typical", successes: s, trials: baselineTrials },
+              baselineRate,
+              m,
+              { ...cfg, mde },
+            );
       // Priced for the same design as the ceiling above (`runsPerObs`), or the
       // cell quotes an unpaired price for a plan only pairing can execute.
       const typicalRuns = Math.ceil(

package/src/stats.ts

@@ -729,9 +729,8 @@ export function sprtExpectedN(p: number, p0: number, p1: number, alpha = 0.05, b
  * outside the two hypotheses, so the two powers are evaluated relative to
  * whichever of them is largest. Written as plain `exp`, a tilt of −662, which is
  * p = 0.0025 against p₀ = 0.999 and p₁ = 0.99, overflows `B^h` to Infinity and
- * the ratio comes back NaN; `expectedSequentialSamples` and
- * `typicalObservationsPerCase` both test the result for `> 0`, so a NaN there
- * silently prices the case at `maxTrials`.
+ * the ratio comes back NaN, which a caller testing the result for `> 0` reads
+ * as no estimate at all.
  */
 function sprtAcceptH1(h: number, A: number, B: number, drift: number): number {
   // No tilt exists at p = 0 or p = 1: one outcome is impossible, so the drift

package/src/stop.ts

@@ -449,7 +449,7 @@ export function shouldStopPaired(
 }
 
 // Mean of prod(1 - lambda (D - mde)) over fixed bets, D = +1 worse, -1 better; a supermartingale whenever E[D] >= mde.
-export function settleLogE(c: PairedCounts, mde: number): number {
+function settleLogE(c: PairedCounts, mde: number): number {
   const terms: number[] = [];
   for (let k = 1; k <= 9; k++) {
     const lambda = k / 10 / (1 - mde);