@seantalts/stanli 0.9.6 → 0.11.0

package/README.md

@@ -23,11 +23,51 @@ const fit = await sample({
 
 fit.columns["mu"];     // Float64Array, one entry per draw
 fit.names;             // every CSV column CmdStan would write
+fit.generatedStart;    // index where generated-quantity columns begin
 fit.ms;                // {stanc, lower, sample, total} in milliseconds
 ```
 
+NUTS can take its starting point from single-path Pathfinder. The same seed
+controls initialization and sampling, and an empty options object uses the
+defaults:
+
+```js
+const fit = await sample({
+  code,
+  data,
+  seed: 303,
+  pathfinderInit: { numIterations: 500, numElboDraws: 25 },
+});
+```
+
+`historySize` and Pathfinder's own `initRadius` are also supported. This mode
+does not perform PSIS resampling.
+
+For NUTS, `await diagnose(fit)` returns the same text report as the R and
+Python bindings: divergences, maximum-treedepth saturation, E-BFMI,
+rank-normalized R-hat, and bulk/tail ESS. Pass an array of fits from the
+same model and configuration to diagnose all chains together:
+
+```js
+import { compile, diagnose, sample } from "@seantalts/stanli";
+const { mir } = await compile({ code });
+const fits = await Promise.all([1, 2, 3, 4].map((seed) =>
+  sample({ mir, data, seed })));
+console.log(await diagnose(fits));
+```
+
+`fit.samplerStats` contains seven doubles per post-warmup draw, in order:
+`lp__`, `accept_stat__`, `stepsize__`, `treedepth__`, `n_leapfrog__`,
+`divergent__`, `energy__`. `fit.maxDepth` records the sampling limit (10).
+WALNUTS and Pathfinder return `null` for `samplerStats`; `diagnose()` rejects
+these methods rather than treating missing statistics as successful checks.
+The demo displays the report after NUTS runs, including comparisons, and
+explicitly marks WALNUTS sampler diagnostics as unavailable. Pathfinder keeps
+its existing importance-weight k-hat diagnostic.
+
 Columns cover the full CmdStan CSV: constrained parameters, transformed
 parameters, and generated quantities (RNG draws stream from `seed`).
+When there are no generated quantities, `generatedStart === fit.names.length`.
 The heavy work runs in a worker the package owns, so the page never
 blocks; calls queue and run one at a time.
 

package/index.mjs

@@ -95,6 +95,29 @@ export function compile(opts) {
   return request({ cmd: "compile", code: opts.code }, opts);
 }
 
+function pathfinderInitOptions(value) {
+  if (value == null) return null;
+  if (typeof value !== "object" || Array.isArray(value))
+    throw new TypeError("pathfinderInit must be an options object");
+  const defaults = { numIterations: 1000, numElboDraws: 25,
+                     historySize: 5, initRadius: 2 };
+  const unknown = Object.keys(value).filter((key) => !(key in defaults));
+  if (unknown.length)
+    throw new RangeError("unknown pathfinderInit option" +
+                         (unknown.length > 1 ? "s" : "") + ": " +
+                         unknown.join(", "));
+  const out = { ...defaults, ...value };
+  for (const name of ["numIterations", "numElboDraws", "historySize"])
+    if (!Number.isInteger(out[name]) || out[name] <= 0 ||
+        out[name] > 2147483647)
+      throw new RangeError(`pathfinderInit ${name} must be a positive integer`);
+  if (typeof out.initRadius !== "number" || !Number.isFinite(out.initRadius) ||
+      out.initRadius < 0)
+    throw new RangeError(
+        "pathfinderInit initRadius must be finite and nonnegative");
+  return out;
+}
+
 /** Compile (unless `mir` is given) and draw from the posterior.
  *
  * @param {Object} opts
@@ -108,6 +131,10 @@ export function compile(opts) {
  * @param {number} [opts.warmup=1000]
  * @param {number} [opts.samples=1000]
  * @param {number} [opts.delta=0.8]    Adaptation target acceptance (NUTS).
+ * @param {Object} [opts.pathfinderInit]  Generate the NUTS start with
+ *   single-path Pathfinder. `{}` uses defaults; supported keys are
+ *   `numIterations`, `numElboDraws`, `historySize`, and `initRadius`.
+ *   The sampling seed controls both stages. NUTS only; no PSIS resampling.
  * @param {string} [opts.sampler="nuts"]  "nuts", "walnuts" (within-orbit
  *   adaptive step-length NUTS, arXiv:2506.18746), or "pathfinder"
  *   (a normal approximation fitted along an L-BFGS path). Pathfinder
@@ -123,9 +150,11 @@ export function compile(opts) {
  *   a transferred ArrayBuffer of constrained draws, nCon wide.
  *   Pathfinder streams {live: {phase: "path", iter, lp}} instead, one
  *   message per L-BFGS iterate.
- * @returns {Promise<{names: string[], samples: number,
+ * @returns {Promise<{names: string[], samples: number, generatedStart: number,
  *                    columns: Object<string, Float64Array>,
  *                    exactLp: boolean,
+ *                    sampler: string, maxDepth: number|null,
+ *                    samplerStats: Float64Array|null,
  *                    pathfinder?: {path: {iter, lp}[], khat: number,
  *                                  selectedIter: number,
  *                                  selectedElbo: number,
@@ -134,8 +163,16 @@ export function compile(opts) {
  *                         total: number}}>}
  *   One column per CSV column CmdStan would write: constrained
  *   parameters, transformed parameters, and generated quantities.
+ *   NUTS also returns post-warmup samplerStats in draw-major order:
+ *   lp__, accept_stat__, stepsize__, treedepth__, n_leapfrog__,
+ *   divergent__, energy__. Other methods return null for samplerStats.
  */
 export function sample(opts) {
+  const sampler = opts.sampler === "walnuts" || opts.sampler === "pathfinder"
+      ? opts.sampler : "nuts";
+  const pathfinderInit = pathfinderInitOptions(opts.pathfinderInit);
+  if (pathfinderInit && sampler !== "nuts")
+    throw new RangeError("pathfinderInit is available only with NUTS");
   return request({
     cmd: "run",
     code: opts.code,
@@ -148,16 +185,50 @@ export function sample(opts) {
     warmup: opts.warmup == null ? 1000 : opts.warmup,
     samples: opts.samples == null ? 1000 : opts.samples,
     delta: opts.delta == null ? 0.8 : opts.delta,
-    sampler: opts.sampler === "walnuts" || opts.sampler === "pathfinder"
-        ? opts.sampler : "nuts",
+    sampler,
     maxError: opts.maxError == null ? 0 : opts.maxError,
+    pathfinderInit,
   }, opts).then((done) => {
-    const { names, samples, ms, exactLp, pathfinder } = done;
+    const { names, samples, generatedStart, ms, exactLp, pathfinder,
+            sampler, maxDepth } = done;
     const flat = new Float64Array(done.columns);
     const columns = {};
     names.forEach((name, i) => {
       columns[name] = flat.subarray(i * samples, (i + 1) * samples);
     });
-    return { names, samples, columns, ms, exactLp, pathfinder };
+    const samplerStats = done.samplerStats
+        ? new Float64Array(done.samplerStats) : null;
+    return { names, samples, generatedStart, columns, ms, exactLp, pathfinder,
+             sampler, maxDepth, samplerStats };
   });
 }
+
+/** Diagnose one NUTS fit, or an array of chains from the same model and
+ * configuration. Returns the native R/Python diagnostic report as text,
+ * using only post-warmup draws. Inputs are copied, never transferred away.
+ * WALNUTS and Pathfinder do not expose the required sampler statistics.
+ * @returns {Promise<string>} */
+export async function diagnose(fits) {
+  const chains = Array.isArray(fits) ? fits : [fits];
+  const first = chains[0];
+  if (!first || !Number.isInteger(first.samples) || first.samples < 1 ||
+      !Array.isArray(first.names) || !first.names.length)
+    throw new Error("diagnose requires nonempty NUTS draws");
+  for (const fit of chains) {
+    if (!fit || fit.sampler !== "nuts" || fit.samples !== first.samples ||
+        fit.maxDepth !== first.maxDepth || !Number.isInteger(fit.maxDepth) ||
+        fit.maxDepth < 1 || !Array.isArray(fit.names) ||
+        fit.names.length !== first.names.length ||
+        fit.names.some((name, j) => name !== first.names[j]) ||
+        !(fit.samplerStats instanceof Float64Array) ||
+        fit.samplerStats.length !== first.samples * 7 ||
+        !fit.columns || first.names.some((name) =>
+          !(fit.columns[name] instanceof Float64Array) ||
+          fit.columns[name].length !== first.samples))
+      throw new Error("diagnose requires matching NUTS chains with sampler statistics");
+  }
+  return request({ cmd: "diagnose", names: first.names, samples: first.samples,
+                   maxDepth: first.maxDepth,
+                   chains: chains.map((fit) => fit.columns),
+                   stats: chains.map((fit) => fit.samplerStats) }, {});
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@seantalts/stanli",
-  "version": "0.9.6",
+  "version": "0.11.0",
   "description": "Full Stan in the browser: stanc3 compiles the model in JS, a WASM runtime lowers it to an op graph and runs NUTS. No server, no C++ toolchain.",
   "type": "module",
   "main": "index.mjs",