diff --git a/docs/math-objects.md b/docs/math-objects.md index b2f927a..8af903c 100644 --- a/docs/math-objects.md +++ b/docs/math-objects.md @@ -59,8 +59,10 @@ object must respect): | `a = 2`, `b = a² + t` | constant (slider / computed) | widget; uniform | | `f(x) = …` | function | inlined | | definition using x/y (`r = sqrt(x²+y²)`) | coordinate field | grid family (level sets) | -| `X ~ Normal(m, s)` | random variable | its density curve | -| `P(X < b)` | probability | shaded area + numeric readout | +| `X ~ Normal(m, s)` (also Uniform, Exponential) | random variable | its exact density curve | +| `Y = X^2`, `S = X1 + X2`, bare `X + Y` (X random) | derived random variable | affine-in-normals: exact pdf (shader); otherwise sampled density estimate (KDE polyline); μ/σ readout | +| `P(X < b)` | probability | shaded area + exact numeric readout | +| `P(Y > 0.5)`, `P(Y > X)` (derived / joint) | probability | Monte Carlo readout (+ shaded density area when one-variable) | Two consequences of invariant 4 worth naming because they already answer part of "what should render as what" in custom coordinates, and should stay: diff --git a/lib/defs.ts b/lib/defs.ts index e5a08e0..37726d5 100644 --- a/lib/defs.ts +++ b/lib/defs.ts @@ -15,9 +15,14 @@ * - `sum(n=1..N, …)` / `prod(…)` (also Σ/Π, and `sum[n=1..N] …` binding the * trailing product like d/dx) expand symbolically at resolve time, so the * bounds must be numbers or already-known constants. + * - `int(f dx)` / `int[a..b] f dx` (also ∫) integrate at resolve time: + * symbolically when integrate.ts finds a verified antiderivative, and + * otherwise by expanding a fixed Gauss–Legendre sum the same way Σ + * expands — so every downstream consumer still sees ordinary expressions. */ import { add, diff, div, mul, neg, pow, sub } from './diff.ts'; import { FUNCTIONS, type Expr, evaluate, freeVars, parseExpr, substVars } from './expr.ts'; +import { QUAD_TERMS, antiderivative, improperSum, quadratureSum, verifyDefinite } from './integrate.ts'; import { lowerGeom, pointComps, vecStateComps } from './geom.ts'; import { type Mat, matrixFromList } from './mat.ts'; @@ -218,12 +223,19 @@ type SumCall = Expr & { kind: 'call'; name: 'sum' | 'prod' }; const isSumHeader = (e: Expr): e is SumCall => e.kind === 'call' && (e.name === 'sum' || e.name === 'prod') && e.args.length === 3; +/** An ∫ header awaiting its body: `int[a..b]` (bounds) or a bare `int`/∫. */ +const isIntHeader = (e: Expr): boolean => + (e.kind === 'call' && e.name === 'int' && e.args.length === 2) + || (e.kind === 'var' && e.name === 'int'); + /** - * A header in a product chain sums its trailing factors: in + * A header in a product chain binds its trailing factors: in * `2 sum[n=1..N] sin(n x)/n` the body is sin(n x)/n and the 2 stays outside - * (like d/dx, which also binds the rest of its product chain). + * (like d/dx, which also binds the rest of its product chain). ∫ headers + * (`int[a..b] f dx`, bare `∫ f dx`) bind the same way; the LEFTMOST header + * wins, so `int[0..1] sum[n=1..2] x^n dx` nests the Σ inside the ∫'s body. */ -function splitSumChain(e: Expr): { coeff: Expr | null; op: '*' | '/'; header: SumCall; body: Expr } | null { +function splitSumChain(e: Expr): { coeff: Expr | null; op: '*' | '/'; header: Expr; body: Expr } | null { const factors: Array<{ e: Expr; op: '*' | '/' }> = []; let node: Expr = e; while (node.kind === 'bin' && (node.op === '*' || node.op === '/')) { @@ -231,7 +243,7 @@ function splitSumChain(e: Expr): { coeff: Expr | null; op: '*' | '/'; header: Su node = node.a; } factors.unshift({ e: node, op: '*' }); - const at = factors.findIndex(f => isSumHeader(f.e)); + const at = factors.findIndex(f => isSumHeader(f.e) || isIntHeader(f.e)); if (at < 0) return null; const rest = factors.slice(at + 1); if (!rest.length) return null; // bodyless header: the call case reports it @@ -241,7 +253,80 @@ function splitSumChain(e: Expr): { coeff: Expr | null; op: '*' | '/'; header: Su for (let k = 0; k < at; k++) { coeff = coeff === null ? factors[k].e : { kind: 'bin', op: factors[k].op, a: coeff, b: factors[k].e }; } - return { coeff, op: factors[at].op, header: factors[at].e as SumCall, body }; + return { coeff, op: factors[at].op, header: factors[at].e, body }; +} + +/** The bounds of an ∫ header, or null for the bare indefinite form. */ +const intBounds = (header: Expr): [Expr, Expr] | null => + header.kind === 'call' && header.args.length === 2 ? [header.args[0], header.args[1]] : null; + +/** Canonical call node for an ∫ header + its chain-bound body. */ +const intCallOf = (header: Expr, body: Expr): Expr => { + const b = intBounds(header); + return { kind: 'call', name: 'int', args: b ? [b[0], b[1], body] : [body] }; +}; + +interface StripDx { + v: string; + integrand: Expr; + /** Factors after this integral's measure that carry ANOTHER d-var: they + * belong to an enclosing ∫ (`int[0..1] int[0..y] x dx dy` pairs inside + * out), so the expansion multiplies them back on for the outer level. */ + residual: Expr | null; +} + +/** + * Split the integration variable off a body: the first d factor in + * its multiplicative structure (`x^2 dx` → v = x, integrand x^2). Implicit + * multiplication binds tighter than '/', so in `sin(t)/t dt` the dt sits + * inside the denominator product — the measure is recognized on either side + * and the rest of that denominator stays a true denominator. A tail after + * the measure folds into the integrand (`∫ dx/(1+x^2)`) unless it carries + * its own d-var, in which case it is the enclosing integral's (residual). + * Sums integrate termwise, so every term must end in the same dx. + */ +function stripDx(body: Expr): StripDx | null { + if (body.kind === 'neg') { + const m = stripDx(body.a); + return m && { ...m, integrand: neg(m.integrand) }; + } + if (body.kind === 'bin' && (body.op === '+' || body.op === '-')) { + const a = stripDx(body.a); + const b = stripDx(body.b); + if (!a && !b) return null; + if (!a || !b || a.v !== b.v || a.residual || b.residual) { + throw new Error(`Every term under one ∫ must end in the same d${a?.v ?? b?.v ?? 'x'}.`); + } + return { v: a.v, integrand: { kind: 'bin', op: body.op, a: a.integrand, b: b.integrand }, residual: null }; + } + const factors: Array<{ e: Expr; inv: boolean }> = []; + const walk = (e: Expr, inv: boolean): void => { + if (e.kind === 'bin' && (e.op === '*' || e.op === '/')) { + walk(e.a, inv); + walk(e.b, e.op === '/' ? !inv : inv); + return; + } + factors.push({ e, inv }); + }; + walk(body, false); + const at = factors.findIndex(f => dVarName(f.e) !== null); + if (at < 0) return null; + const v = dVarName(factors[at].e)!; + const tail = factors.slice(at + 1); + const tailHasDx = tail.some(f => dVarName(f.e) !== null); + const inside = tailHasDx ? factors.slice(0, at) : factors.filter((_, i) => i !== at); + const product = (fs: Array<{ e: Expr; inv: boolean }>): Expr => { + let numr: Expr | null = null; + let den: Expr | null = null; + for (const f of fs) { + if (f.inv) den = den === null ? f.e : { kind: 'bin', op: '*', a: den, b: f.e }; + else numr = numr === null ? f.e : { kind: 'bin', op: '*', a: numr, b: f.e }; + } + let out: Expr = numr ?? num(1); + if (den) out = { kind: 'bin', op: '/', a: out, b: den }; + return out; + }; + return { v, integrand: product(inside), residual: tailHasDx ? product(tail) : null }; } /** substVars for a Σ/Π index, stopping at nested Σ/Π that rebind the same name. */ @@ -255,14 +340,13 @@ function substIdx(e: Expr, idx: string, val: Expr): Expr { const m = splitSumChain(e); if (m) { // A bodyless header binds the rest of its product chain as its body - // (`sum[n=1..2] sum[n=1..3] n`). Canonicalize to the 4-arg call so - // the rebinding guard below sees that body as the inner Σ's own, - // not as a sibling factor ours may substitute into. - const call = substIdx( - { kind: 'call', name: m.header.name, args: [...m.header.args, m.body] }, - idx, - val, - ); + // (`sum[n=1..2] sum[n=1..3] n`). Canonicalize to the full call so + // the rebinding guard below sees that body as the inner binder's + // own, not as a sibling factor ours may substitute into. + const canonical: Expr = isSumHeader(m.header) + ? { kind: 'call', name: m.header.name, args: [...m.header.args, m.body] } + : intCallOf(m.header, m.body); + const call = substIdx(canonical, idx, val); return m.coeff ? { kind: 'bin', op: m.op, a: substIdx(m.coeff, idx, val), b: call } : call; } } @@ -274,6 +358,18 @@ function substIdx(e: Expr, idx: string, val: Expr): Expr { const args = e.args.map((a, k) => (k === 0 || k === 3 ? a : substIdx(a, idx, val))); return { kind: 'call', name: e.name, args }; } + if (e.name === 'int' && (e.args.length === 1 || e.args.length === 3)) { + // An ∫ whose dx names this index rebinds it: bounds only. + const bodyAt = e.args.length - 1; + let dx: { v: string } | null = null; + try { + dx = stripDx(e.args[bodyAt]); + } catch { /* multiple dx factors: expansion will report it */ } + if (dx && dx.v === idx) { + const args = e.args.map((a, k) => (k === bodyAt ? a : substIdx(a, idx, val))); + return { kind: 'call', name: 'int', args }; + } + } return { kind: 'call', name: e.name, args: e.args.map(a => substIdx(a, idx, val)) }; } case 'eq': return { kind: 'eq', l: substIdx(e.l, idx, val), r: substIdx(e.r, idx, val) }; @@ -362,9 +458,118 @@ function expandSum(header: SumCall, body: Expr, ctx: Ctx): Expr { return acc ?? num(header.name === 'sum' ? 0 : 1); } +/** + * Cache of expanded integrals. The expansion is a pure function of the + * resolved integrand and bounds — slider VALUES stay symbolic inside it — so + * a slider drag's per-keystroke recompiles hit this instead of re-running + * the symbolic engine and its verification quadratures. + */ +const intMemo = new Map(); + +/** + * Expand an ∫: a verified antiderivative when integrate.ts finds one + * (definite values additionally checked against adaptive quadrature — the + * fundamental theorem lies across a non-integrable singularity), otherwise + * a fixed Gauss–Legendre sum in ordinary Expr form. + */ +/** ±1 for a bound written as ±inf/±∞, 0 for a finite (or absent) bound. */ +const infOf = (e: Expr | null): 1 | -1 | 0 => { + if (!e) return 0; + if (e.kind === 'var' && e.name === 'inf') return 1; + if (e.kind === 'neg' && e.a.kind === 'var' && e.a.name === 'inf') return -1; + return 0; +}; + +/** + * Stand-in argument for an antiderivative's limit at ±∞. Far beyond any plot + * range, so F(±BIG) matches the true limit wherever F converges by then — + * and verifyDefinite runs the REAL improper quadrature, so a + * not-yet-converged (or divergent) limit is rejected, never reported. + */ +const INT_BIG = 1e8; + +function expandInt(bounds: [Expr, Expr] | null, rawBody: Expr, ctx: Ctx): Expr { + // Resolve the body FIRST: a d/dt inside consumes its own dt, user + // functions inline, and nested (parenthesized) integrals expand — only + // then is the surviving d factor unambiguous. + const m = stripDx(rx(rawBody, ctx)); + if (!m) throw new Error('∫ needs its variable as a dx factor: int(x^2 dx) or int[0..2] x^2 dx.'); + const v = m.v; + const integrand = m.integrand; + let lo = bounds && rx(bounds[0], ctx); + let hi = bounds && rx(bounds[1], ctx); + let loI = infOf(lo); + let hiI = infOf(hi); + // Normalize a downhill infinite range (int[inf..0]) to the negated uphill one. + let flip = false; + if (loI === 1 || hiI === -1) { + [lo, hi] = [hi, lo]; + [loI, hiI] = [hiI, loI]; + flip = true; + if (loI === 1 || hiI === -1) return num(0); // int[inf..inf]: equal bounds + } + const memoKey = JSON.stringify([v, integrand, lo, hi]); + const done = (out: Expr): Expr => { + const signed = flip ? neg(out) : out; + // An enclosing integral's measure rides along: (∫ inner) · residual. + return m.residual ? { kind: 'bin', op: '*', a: signed, b: m.residual } : signed; + }; + const hit = intMemo.get(memoKey); + if (hit) return done(hit); + + const F = antiderivative(integrand, v); + let out: Expr | null = null; + if (!bounds) { + out = F; + } else if (F) { + const loSub = loI ? num(loI * INT_BIG) : lo!; + const hiSub = hiI ? num(hiI * INT_BIG) : hi!; + const val = foldNums(sub(substVars(F, { [v]: hiSub }), substVars(F, { [v]: loSub }))); + const loChk = loI ? num(loI * Infinity) : lo!; + const hiChk = hiI ? num(hiI * Infinity) : hi!; + if (verifyDefinite(val, integrand, v, loChk, hiChk)) out = val; + } + if (!out) { + // Numeric fallback. An indefinite ∫ anchors at 0: F(x) = ∫₀ˣ; infinite + // ranges transform onto a finite interval first (improperSum). + ctx.terms += QUAD_TERMS; + if (ctx.terms > SUM_MAX_TOTAL) { + throw new Error(`Nested Σ/∫ expand to too many terms (limit ${SUM_MAX_TOTAL} total).`); + } + out = foldNums(loI || hiI + ? improperSum(integrand, v, loI ? null : lo, hiI ? null : hi) + : quadratureSum(integrand, v, lo ?? num(0), hi ?? { kind: 'var', name: v })); + if (JSON.stringify(out).length > 400_000) { + throw new Error('∫ has no closed form here and its numeric expansion is too large.'); + } + } + if (intMemo.size > 500) intMemo.clear(); + intMemo.set(memoKey, out); + return done(out); +} + +/** Whether a parsed (unresolved) expression uses ∫ anywhere — after + * resolution the integral is gone, so row readouts test the parse. */ +export function usesIntegral(e: Expr): boolean { + switch (e.kind) { + case 'num': return false; + case 'var': return e.name === 'int'; + case 'neg': return usesIntegral(e.a); + case 'bin': return usesIntegral(e.a) || usesIntegral(e.b); + case 'call': return e.name === 'int' || e.args.some(usesIntegral); + case 'eq': return usesIntegral(e.l) || usesIntegral(e.r); + case 'ineq': return usesIntegral(e.l) || usesIntegral(e.r); + case 'vec': return e.items.some(usesIntegral); + case 'list': return e.items.some(usesIntegral); + case 'piecewise': + return e.cases.some(c => usesIntegral(c.cond) || usesIntegral(c.value)) + || (e.otherwise ? usesIntegral(e.otherwise) : false); + } +} + /** * Inline user-function calls, resolve d/dx derivative notation, and expand - * Σ/Π sums (post-order). + * Σ/Π sums and ∫ integrals (post-order). */ export function resolveExpr(e: Expr, getFn: GetFn, opts: ResolveOpts = {}): Expr { return rx(e, { getFn, opts, terms: 0 }); @@ -379,11 +584,14 @@ function rx(e: Expr, ctx: Ctx): Expr { case 'neg': return { kind: 'neg', a: rx(e.a, ctx) }; case 'bin': { if (e.op === '*' || e.op === '/') { - // Σ headers capture their trailing product chain before it resolves, - // so `sum[n=1..N] sin(n x)/n` divides each term, not the whole sum. + // Σ/∫ headers capture their trailing product chain before it + // resolves, so `sum[n=1..N] sin(n x)/n` divides each term, not the + // whole sum, and `int[0..1] x^2 dx` binds through to its dx. const m = splitSumChain(e); if (m) { - const body = expandSum(m.header, m.body, ctx); + const body = isSumHeader(m.header) + ? expandSum(m.header, m.body, ctx) + : expandInt(intBounds(m.header), m.body, ctx); return m.coeff ? { kind: 'bin', op: m.op, a: rx(m.coeff, ctx), b: body } : body; } } @@ -408,7 +616,12 @@ function rx(e: Expr, ctx: Ctx): Expr { } return expandSum(e as SumCall, e.args[3], ctx); } - if (e.name === '[range]') throw new Error("'..' ranges only appear in sum(n=1..N, …) or prod(…)."); + if (e.name === 'int') { + if (e.args.length === 2) throw new Error('∫ needs a body: write int[a..b] f(x) dx.'); + const body = e.args[e.args.length - 1]; + return expandInt(e.args.length === 3 ? [e.args[0], e.args[1]] : null, body, ctx); + } + if (e.name === '[range]') throw new Error("'..' ranges only appear in sum(n=1..N, …), prod(…) or int[a..b]."); const args = e.args.map(x => rx(x, ctx)); const fn = getFn(e.name); if (fn) { @@ -744,7 +957,9 @@ export function buildDefs(raw: Definition[]): BuiltDefs { for (const [name, e] of defs.consts) { for (const fv of freeVars(e)) { if (fv !== 't' && !constNames.has(fv) && !stateNames.has(fv)) { - errors.set(name, `${name} can only depend on other constants and t (found ${fv}).`); + errors.set(name, fv === 'inf' + ? `inf only works as an ∫ bound — write it inline: int[-inf..x] f(x) dx.` + : `${name} can only depend on other constants and t (found ${fv}).`); defs.consts.delete(name); break; } diff --git a/lib/diff.ts b/lib/diff.ts index 7c28684..6305495 100644 --- a/lib/diff.ts +++ b/lib/diff.ts @@ -92,6 +92,16 @@ export function diff(e: Expr, v: string): Expr { const n = sub(mul(diff(y, v), x), mul(y, diff(x, v))); return div(n, add(pow(x, num(2)), pow(y, num(2)))); } + if ((e.name === 'normalpdf' || e.name === 'normalcdf') && e.args.length === 3) { + // Full chain rule in all three arguments (x, mean, sd may all move). + const [x, m, s] = e.args; + const z = div(sub(x, m), s); + const dz = div(sub(sub(diff(x, v), diff(m, v)), mul(z, diff(s, v))), s); + const pdf = call('normalpdf', x, m, s); + if (e.name === 'normalcdf') return mul(mul(pdf, s), dz); + // φ′ = φ·(−z·z′ − s′/s): the −z z′ from the exponent, −s′/s from 1/s. + return mul(pdf, sub(mul(neg(z), dz), div(diff(s, v), s))); + } if (e.args.length !== 1) throw new Error(`Cannot differentiate ${e.name}.`); const a = e.args[0]; const da = diff(a, v); diff --git a/lib/dist.test.ts b/lib/dist.test.ts index 7f18b53..bf02ed2 100644 --- a/lib/dist.test.ts +++ b/lib/dist.test.ts @@ -1,25 +1,47 @@ import { describe, expect, it } from 'vitest'; import { - type DistDef, + type BaseDist, + RVSystem, + buildRVSystem, + checkDerived, + densityAt, densityExpr, + matchExpectation, matchProbability, parseDistribution, probabilityValue, regionExpr, scanDistribution, + scanRandomRows, + shadePolygon, + toExpectation, toProbability, } from './dist.ts'; -import { evaluate, parseExpr } from './expr.ts'; +import { evaluate, normalcdf, normalpdf, parseExpr } from './expr.ts'; import { classify } from './plot.ts'; const none = new Set(); -const dist = (rhs: string, name = 'X'): DistDef => parseDistribution(name, rhs, none); +const dist = (rhs: string): BaseDist => parseDistribution(rhs, none); -const dists = (...ds: DistDef[]) => new Map(ds.map(d => [d.name, d])); +const names = (...ns: string[]) => new Set(ns); -const prob = (inner: string, ds = dists(dist('Normal(0, 1)'))) => - toProbability(parseExpr(inner), ds); +const prob = (inner: string, ns = names('X')) => toProbability(parseExpr(inner), ns); + +/** Build a system from document rows, as the app and worker do. */ +const build = (rows: string[], constNames = names('a', 'b', 'm', 's')) => { + const sys = new RVSystem(); + const built = buildRVSystem(sys, scanRandomRows(rows), { + fnNames: none, + getFn: () => undefined, + constNames, + taken: () => false, + }); + return { sys, built }; +}; + +const P = (sys: RVSystem, body: string, env: Record = {}) => + sys.probability(parseExpr(body), env); describe('erf / normal built-ins', () => { it('evaluates erf accurately', () => { @@ -30,34 +52,37 @@ describe('erf / normal built-ins', () => { it('evaluates normalpdf and normalcdf', () => { expect(evaluate(parseExpr('normalpdf(0, 0, 1)'), {})).toBeCloseTo(0.3989423, 6); - expect(evaluate(parseExpr('normalpdf(3, 1, 2)'), {})).toBeCloseTo(0.1209854, 6); - expect(evaluate(parseExpr('normalcdf(0, 0, 1)'), {})).toBeCloseTo(0.5, 6); expect(evaluate(parseExpr('normalcdf(1.959964, 0, 1)'), {})).toBeCloseTo(0.975, 5); }); - - it('plots y = normalcdf(x, 0, 1) as a curve', () => { - expect(classify(parseExpr('y = normalcdf(x, 0, 1)')).plot.type).toBe('implicit2d'); - }); }); describe('scanDistribution / parseDistribution', () => { it('detects name ~ rhs rows', () => { expect(scanDistribution('X ~ Normal(0, 1)')).toEqual({ name: 'X', rhs: 'Normal(0, 1)' }); expect(scanDistribution('y = x^2')).toBeNull(); - expect(scanDistribution('a = 2')).toBeNull(); }); it('parses Normal with symbolic parameters', () => { const d = dist('Normal(0, a)'); - expect(d.mean).toEqual({ kind: 'num', value: 0 }); - expect(d.sd).toEqual({ kind: 'var', name: 'a' }); - expect(dist('normal(1, 2)').mean).toEqual({ kind: 'num', value: 1 }); + expect(d.kind).toBe('normal'); + expect(d.args[0]).toEqual({ kind: 'num', value: 0 }); + expect(d.args[1]).toEqual({ kind: 'var', name: 'a' }); + }); + + it('parses the whole zoo, with aliases and standard defaults', () => { + expect(dist('N').kind).toBe('normal'); + expect(dist('N').args.map(a => evaluate(a, {}))).toEqual([0, 1]); + expect(dist('U(2, 3)').kind).toBe('uniform'); + expect(dist('uniform(2, 3)').kind).toBe('uniform'); + expect(dist('Exp(2)').kind).toBe('exponential'); + expect(dist('Exponential(2)').args).toEqual([{ kind: 'num', value: 2 }]); }); it('rejects unknown distributions and wrong arity', () => { expect(() => dist('Poisson(3)')).toThrow(/Unknown distribution/); expect(() => dist('Normal(1)')).toThrow(/2 arguments/); expect(() => dist('Normal(1, 2, 3)')).toThrow(/2 arguments/); + expect(() => dist('Exponential(1, 2)')).toThrow(/1 argument/); expect(() => dist('2x + 1')).toThrow(/Expected a distribution/); }); }); @@ -72,59 +97,97 @@ describe('densityExpr', () => { expect(field).toContain('u_a'); }); - it('evaluates to the density (y - pdf = 0 on the curve)', () => { - const e = densityExpr(dist('Normal(0, 1)')); - expect(evaluate(e, { x: 0, y: 0.3989423 })).toBeCloseTo(0, 6); + it('evaluates the uniform density as a piecewise box', () => { + const e = densityExpr(dist('Uniform(1, 3)')); + expect(evaluate(e, { x: 2, y: 0.5 })).toBeCloseTo(0, 9); // on the curve + expect(evaluate(e, { x: 0, y: 0 })).toBeCloseTo(0, 9); // outside the support + }); + + it('classifies the uniform and exponential densities as curves', () => { + expect(classify(densityExpr(dist('Uniform(0, 1)'))).plot.type).toBe('implicit2d'); + expect(classify(densityExpr(dist('Exponential(2)'))).plot.type).toBe('implicit2d'); + }); + + it('degrades invalid parameters to a flat 0 instead of a negative density', () => { + const uni = densityExpr(dist('Uniform(3, 1)')); + expect(evaluate(uni, { x: 2, y: 0 })).toBeCloseTo(0, 9); + const exp = densityExpr(dist('Exponential(-1)')); + expect(evaluate(exp, { x: 2, y: 0 })).toBeCloseTo(0, 9); }); }); -describe('toProbability', () => { - it('reads X < b and b < X', () => { - const upper = prob('X < 2'); - expect(upper.lo).toBeUndefined(); - expect(upper.hi).toEqual({ kind: 'num', value: 2 }); - const lower = prob('-1 < X'); - expect(lower.lo).toEqual({ kind: 'neg', a: { kind: 'num', value: 1 } }); - expect(lower.hi).toBeUndefined(); +describe('scanRandomRows', () => { + it('finds base rows and follows derived rows transitively', () => { + const scan = scanRandomRows(['X ~ Normal(0, 1)', 'Z = Y + 1', 'Y = X^2', 'a = 2']); + expect([...scan.base.keys()]).toEqual([0]); + expect(new Set(scan.derived.keys())).toEqual(new Set([1, 2])); + expect(scan.derived.get(2)).toEqual({ name: 'Y', rhs: ' X^2' }); + }); + + it('leaves plain constants and claimed rows alone', () => { + const scan = scanRandomRows(['X ~ N', 'a = 2', null, 'b = a + 1']); + expect(scan.derived.size).toBe(0); + }); + + it('matches identifier tokens, not word boundaries', () => { + expect(scanRandomRows(['X ~ N', 'Y = X_1 + 2']).derived.size).toBe(0); // X_1 is not X + expect(scanRandomRows(['X ~ N', 'Y = aX + 2']).derived.size).toBe(0); // aX is one name + expect(scanRandomRows(['X ~ N', 'Y = 2X']).derived.size).toBe(1); // 2X is 2·X }); - it('reads > by flipping', () => { - const p = prob('X > 2'); - expect(p.lo).toEqual({ kind: 'num', value: 2 }); - expect(p.hi).toBeUndefined(); + it('never claims reserved names', () => { + const scan = scanRandomRows(['X ~ N', 'e = X']); + expect(scan.derived.size).toBe(0); // `e = X` stays an equation row }); +}); - it('reads two-sided chains in either direction', () => { - const asc = prob('-1 < X <= 2'); +describe('toProbability', () => { + it('reads bounds around one variable, both directions', () => { + expect(prob('X < 2').single).toEqual({ rv: 'X', lo: undefined, hi: { kind: 'num', value: 2 } }); + expect(prob('X > 2').single).toEqual({ rv: 'X', lo: { kind: 'num', value: 2 }, hi: undefined }); + const asc = prob('-1 < X <= 2').single!; expect(asc.lo).toEqual({ kind: 'neg', a: { kind: 'num', value: 1 } }); expect(asc.hi).toEqual({ kind: 'num', value: 2 }); - const desc = prob('2 > X > -1'); + const desc = prob('2 > X > -1').single!; expect(desc.lo).toEqual({ kind: 'neg', a: { kind: 'num', value: 1 } }); expect(desc.hi).toEqual({ kind: 'num', value: 2 }); }); - it('picks the referenced variable among several', () => { - const ds = dists(dist('Normal(0, 1)'), dist('Normal(5, 2)', 'Y')); - expect(prob('Y < 4', ds).dist.name).toBe('Y'); + it('reports every referenced variable', () => { + expect(prob('Y < 4', names('X', 'Y')).rvs).toEqual(['Y']); + expect(new Set(prob('X < Y', names('X', 'Y')).rvs)).toEqual(new Set(['X', 'Y'])); + }); + + it('captures bounds around one inline expression', () => { + const sum = prob('0.5 < X + Y < 1.5', names('X', 'Y')); + expect(sum.single).toBeUndefined(); + expect(sum.inline!.e).toEqual(parseExpr('X + Y')); + expect(sum.inline!.lo).toEqual({ kind: 'num', value: 0.5 }); + expect(prob('X^2 < 1').inline!.e).toEqual(parseExpr('X^2')); + }); + + it('handles bodies with no single-variable shape', () => { + expect(prob('X < Y', names('X', 'Y')).single).toBeUndefined(); + expect(prob('X < Y', names('X', 'Y')).inline).toBeUndefined(); // two terms carry variables + expect(prob('X < a < b').single).toBeUndefined(); // extra constraint beyond the bounds + expect(prob('a < b < X').single).toBeUndefined(); // ditto, from the left + expect(prob('X > X').single).toBeUndefined(); }); it('rejects malformed bodies', () => { expect(() => prob('X + 1')).toThrow(/expects an inequality/); expect(() => prob('a < b')).toThrow(/must reference a random variable/); expect(() => prob('-1 < X > 2')).toThrow(/same way/); - expect(() => prob('X < a < b')).toThrow(/at most two bounds/); - expect(() => prob('X < x')).toThrow(/cannot use x/); - const two = dists(dist('Normal(0, 1)'), dist('Normal(0, 1)', 'Y')); - expect(() => prob('X < Y', two)).toThrow(/Only one random variable/); + expect(() => prob('X < x')).toThrow(/plot coordinate x/); }); }); describe('regionExpr', () => { it('classifies as a shaded region with an outline', () => { - const c = classify(regionExpr(prob('X < b')), new Set(['b'])); + const p = prob('X < b').single!; + const c = classify(regionExpr(dist('Normal(0, 1)'), p.lo, p.hi), new Set(['b'])); expect(c.plot.type).toBe('ineq2d'); const plot = c.plot as { field: string; edges: string[] }; - expect(plot.field).toContain('max('); expect(plot.field).toContain('eq_normalpdf'); expect(plot.field).toContain('u_b'); expect(plot.edges).toHaveLength(1); @@ -132,7 +195,8 @@ describe('regionExpr', () => { }); it('is negative inside the area and positive outside', () => { - const region = regionExpr(prob('-1 < X < 1')); + const p = prob('-1 < X < 1').single!; + const region = regionExpr(dist('Normal(0, 1)'), p.lo, p.hi); if (region.kind !== 'ineq') throw new Error('expected ineq'); const f = (x: number, y: number) => evaluate(region.l, { x, y }); expect(f(0, 0.2)).toBeLessThan(0); // under the peak @@ -142,23 +206,32 @@ describe('regionExpr', () => { }); }); -describe('probabilityValue', () => { - it('computes one- and two-sided probabilities', () => { - expect(probabilityValue(prob('X < 0'), {})).toBeCloseTo(0.5, 6); - expect(probabilityValue(prob('X < 1.959964'), {})).toBeCloseTo(0.975, 5); - expect(probabilityValue(prob('X > 1'), {})).toBeCloseTo(0.1586553, 5); - expect(probabilityValue(prob('-1 < X < 1'), {})).toBeCloseTo(0.6826895, 5); +describe('probabilityValue (exact)', () => { + const value = (d: string, body: string, env: Record = {}) => { + const p = prob(body).single!; + return probabilityValue(dist(d), p.lo, p.hi, env); + }; + + it('computes normal probabilities', () => { + expect(value('Normal(0, 1)', 'X < 0')).toBeCloseTo(0.5, 6); + expect(value('Normal(0, 1)', 'X < 1.959964')).toBeCloseTo(0.975, 5); + expect(value('Normal(0, 1)', 'X > 1')).toBeCloseTo(0.1586553, 5); + expect(value('Normal(0, 1)', '-1 < X < 1')).toBeCloseTo(0.6826895, 5); + expect(value('Normal(m, s)', 'X < b', { m: 1, s: 2, b: 1 })).toBeCloseTo(0.5, 6); }); - it('uses the constant environment for parameters and bounds', () => { - const ds = dists(dist('Normal(m, s)')); - const p = toProbability(parseExpr('X < b'), ds); - expect(probabilityValue(p, { m: 1, s: 2, b: 1 })).toBeCloseTo(0.5, 6); + it('computes uniform and exponential probabilities', () => { + expect(value('Uniform(0, 2)', 'X < 0.5')).toBeCloseTo(0.25, 9); + expect(value('Uniform(0, 2)', 'X > 3')).toBe(0); + expect(value('Uniform(0, 2)', '-1 < X < 5')).toBeCloseTo(1, 9); + expect(value('Exponential(2)', 'X < 1')).toBeCloseTo(1 - Math.exp(-2), 9); + expect(value('Exponential(2)', 'X < -1')).toBe(0); }); - it('is NaN for a non-positive sd', () => { - const ds = dists(dist('Normal(0, s)')); - expect(probabilityValue(toProbability(parseExpr('X < 1'), ds), { s: 0 })).toBeNaN(); + it('is NaN while parameters are invalid', () => { + expect(value('Normal(0, s)', 'X < 1', { s: 0 })).toBeNaN(); + expect(value('Uniform(3, 1)', 'X < 1')).toBeNaN(); + expect(value('Exponential(-2)', 'X < 1')).toBeNaN(); }); }); @@ -170,3 +243,480 @@ describe('P(…) row matching', () => { expect(matchProbability('Q(X < 2)')).toBeNull(); }); }); + +describe('E(…) rows', () => { + const expectation = (inner: string, ns = names('X')) => toExpectation(parseExpr(inner), ns); + + it('matches only whole E(...) rows', () => { + expect(matchExpectation('E(X)')).toBe('X'); + expect(matchExpectation(' E ( X^2 + Y ) ')).toBe(' X^2 + Y '); + expect(matchExpectation('y = E(X)')).toBeNull(); + expect(matchExpectation('F(X)')).toBeNull(); + }); + + it('validates the body against the declared variables', () => { + expect(expectation('X + 1').rvs).toEqual(['X']); + expect(() => expectation('X < 2')).toThrow(/P\(…\)/); + expect(() => expectation('(X, 1)')).toThrow(/single value/); + expect(() => expectation('a + 1')).toThrow(/must reference a random variable/); + expect(() => expectation('X + x')).toThrow(/plot coordinate x/); + }); + + it('is exact under a derivable law', () => { + const { sys } = build(['X ~ Normal(2, 3)', 'Y = 2X + 1']); + expect(sys.mean('X', {})).toBe(2); + expect(sys.mean('Y', {})).toBe(5); + const { sys: u } = build(['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 3)', 'S = X1 + X2']); + expect(u.mean('S', {})).toBe(2); + const { sys: e } = build(['X ~ Exponential(4)']); + expect(e.mean('X', {})).toBe(0.25); + }); + + it('integrates one-variable transforms against the base pdf (quadrature)', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y = X^2']); + expect(sys.exactMoments('Y', {})).toBeNull(); + const qm = sys.quadMoments('Y', {})!; + expect(qm.mean).toBeCloseTo(1, 8); // E[X²] = Var(X) = 1, to quadrature digits + expect(qm.sd).toBeCloseTo(Math.SQRT2, 7); // Var(X²) = 2 + expect(sys.mean('Y', {})).toBeCloseTo(1, 8); + const { sys: r } = build(['X ~ Normal(0, 1)', 'R = sqrt(X)']); + // Partial support averages where defined: E[√X | X > 0] = + // 2^(-1/4)·Γ(3/4)/√(2π) / P(X > 0) ≈ 0.8222. + const rq = r.quadMoments('R', {})!; + expect(rq.mass).toBeCloseTo(0.5, 5); + expect(rq.mean).toBeCloseTo(0.8222, 3); + }); + + it('quadrature moments follow chains and every base family', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y = X^2', 'Z = Y + 1']); + expect(sys.quadMoments('Y', {})!.mean).toBeCloseTo(1 / 3, 9); + expect(sys.quadMoments('Y', {})!.sd).toBeCloseTo(Math.sqrt(4 / 45), 8); + expect(sys.mean('Z', {})).toBeCloseTo(4 / 3, 8); // grounded through Y + const { sys: e } = build(['X ~ Exponential(2)', 'Y = X^2']); + expect(e.mean('Y', {})).toBeCloseTo(0.5, 7); // E[X²] = 2/λ² + const { sys: s } = build(['X ~ Normal(m, s)', 'Y = X^2']); + expect(s.mean('Y', { m: 2, s: 1 })).toBeCloseTo(5, 7); // μ² + σ² + }); + + it('leaves joint dependence to the sampler', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'W ~ Normal(0, 1)', 'M = X W']); + expect(sys.quadMoments('M', {})).toBeNull(); + expect(sys.mean('M', {})).toBeCloseTo(0, 1); // Monte Carlo still answers + }); + + it('responds to slider constants and is NaN when broken', () => { + const { sys } = build(['X ~ Normal(m, s)', 'Y = X^3 + a']); + expect(sys.mean('Y', { m: 0, s: 1, a: 10 })).toBeCloseTo(10, 1); + expect(sys.mean('X', { m: 1, s: -1 })).toBeNaN(); + }); + + it('interpolates the density polyline for the marker height', () => { + const curve = { pts: [0, 0, 1, 2, 2, 0], mean: 1, sd: 0.5, mass: 1 }; + expect(densityAt(curve, 0.5)).toBeCloseTo(1, 9); + expect(densityAt(curve, 1)).toBeCloseTo(2, 9); + expect(densityAt(curve, 5)).toBe(0); + }); +}); + +describe('buildRVSystem', () => { + it('claims base and derived rows and reports row errors', () => { + const { sys, built } = build(['X ~ Normal(0, 1)', 'Y = X^2', 'W ~ Poisson(3)']); + expect(built.rowRV.get(0)).toBe('X'); + expect(built.rowRV.get(1)).toBe('Y'); + expect(sys.get('Y')?.kind).toBe('derived'); + expect(built.errors.get(2)).toMatch(/Unknown distribution/); + }); + + it('rejects name collisions and reserved names', () => { + const { built } = build(['X ~ N', 'X ~ U']); + expect(built.errors.get(1)).toBe('X is already defined.'); + expect(build(['pi ~ N']).built.errors.get(0)).toMatch(/Cannot use pi/); + const taken = buildRVSystem(new RVSystem(), scanRandomRows(['X ~ N']), { + fnNames: none, getFn: () => undefined, constNames: none, taken: () => true, + }); + expect(taken.errors.get(0)).toBe('X is already defined.'); + }); + + it('rejects random variables inside distribution parameters', () => { + const { built } = build(['X ~ Normal(0, 1)', 'Y ~ Normal(X, 1)']); + expect(built.errors.get(1)).toMatch(/cannot depend on a random variable/); + }); + + it('reports cycles and ripples errors to dependents', () => { + const { sys, built } = build(['X ~ N', 'Y = Z + X', 'Z = Y + 1', 'W = Z^2']); + expect(built.errors.get(1)).toMatch(/circular/); + expect(built.errors.get(2)).toMatch(/circular/); + expect(built.errors.get(3)).toMatch(/Z has an error/); + expect(sys.has('X')).toBe(true); + expect(sys.has('W')).toBe(false); + }); + + it('validates derived expressions', () => { + const { built } = build(['X ~ N', 'Y = X + x']); + expect(built.errors.get(1)).toMatch(/plot coordinate x/); + expect(build(['X ~ N', 'Y = X + q']).built.errors.get(1)).toBe('q is not defined.'); + expect(() => checkDerived(parseExpr('(X, 1)'), names('X'), none)).toThrow(/single value/); + }); +}); + +describe('RVSystem sampling', () => { + it('is deterministic and matches the declared moments', () => { + const { sys } = build(['X ~ Normal(2, 3)']); + const { sys: sys2 } = build(['X ~ Normal(2, 3)']); + const c = sys.curve('X', {})!; + expect(c.mean).toBeCloseTo(2, 2); + expect(c.sd).toBeCloseTo(3, 2); + expect(c.mass).toBe(1); + expect(sys.columns('X', {})).toEqual(sys2.columns('X', {})); + }); + + it('keeps distinct names independent: X + Y is the convolution', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y ~ Normal(0, 1)', 'S = X + Y']); + const c = sys.curve('S', {})!; + expect(c.mean).toBeCloseTo(0, 1); + expect(c.sd).toBeCloseTo(Math.SQRT2, 1); + expect(P(sys, 'S < 0')).toBeCloseTo(0.5, 1.5); + // Against the exact normal CDF at a non-symmetric point. + expect(P(sys, 'S < 1')).toBeCloseTo(normalcdf(1, 0, Math.SQRT2), 1.5); + }); + + it('keeps the same name dependent: X + X is 2X, not a convolution', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'D = X + X']); + expect(sys.curve('D', {})!.sd).toBeCloseTo(2, 1); + expect(P(sys, 'X > X')).toBe(0); + }); + + it('sums of uniforms make the CLT triangle', () => { + const { sys } = build(['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 1)', 'S = X1 + X2']); + const c = sys.curve('S', {})!; + expect(c.mean).toBeCloseTo(1, 2); + expect(P(sys, 'S < 1')).toBeCloseTo(0.5, 2); + expect(P(sys, 'S < 0.5')).toBeCloseTo(0.125, 1.5); + }); + + it('estimates product distributions', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y ~ Normal(0, 1)', 'M = X Y']); + expect(P(sys, 'M > 0')).toBeCloseTo(0.5, 1.5); + }); + + it('handles piecewise conditionals: Y = {X > 0: X^2, 1}', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y = {X > 0: X^2, 1}']); + // P(Y > 1/2) = P(X > 1/√2) + P(X <= 0). + const exact = 1 - normalcdf(Math.SQRT1_2, 0, 1) + 0.5; + expect(P(sys, 'Y > 0.5')).toBeCloseTo(exact, 1.5); + expect(P(sys, 'Y >= 0')).toBe(1); + }); + + it('estimates joint probabilities of dependent variables', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'W ~ Normal(0, 1)', 'Y = {X > 0: X^2, 1}']); + expect(P(sys, 'W > X')).toBeCloseTo(0.5, 1.5); + // P(Y > X): on X <= 0, Y = 1 > X always (prob 1/2); on X > 0, X^2 > X iff + // X > 1, so P(X > 1) adds. Exact: 0.5 + (1 - Φ(1)). + const exact = 0.5 + 1 - normalcdf(1, 0, 1); + expect(P(sys, 'Y > X')).toBeCloseTo(exact, 1.5); + }); + + it('responds to slider constants through the environment', () => { + const { sys } = build(['X ~ Normal(m, s)', 'Y = X + a']); + expect(sys.curve('Y', { m: 1, s: 2, a: 10 })!.mean).toBeCloseTo(11, 1); + expect(P(sys, 'Y < 11', { m: 1, s: 2, a: 10 })).toBeCloseTo(0.5, 1.5); + // Re-query under new values: the cache must not serve the old ones. + expect(sys.curve('Y', { m: 5, s: 2, a: 0 })!.mean).toBeCloseTo(5, 1); + }); + + it('caches columns per variable: unrelated constants never resample', () => { + // Slider drags recompile on every input event; this stays cheap only + // because a variable resamples exactly when ITS OWN parameters move. + const { sys } = build(['X ~ Normal(m, 1)', 'Y ~ Uniform(0, 1)']); + const x = sys.columns('X', { m: 0, s: 7 }); + expect(sys.columns('X', { m: 0, s: 8 })).toBe(x); // s is not X's parameter + const y = sys.columns('Y', { m: 0 }); + expect(sys.columns('Y', { m: 1 })).toBe(y); // m is not Y's parameter + expect(sys.columns('X', { m: 1 })).not.toBe(x); + }); + + it('invalidates cached samples when a dependency is redeclared', () => { + const sys = new RVSystem(); + const opts = { fnNames: none, getFn: () => undefined, constNames: none, taken: () => false }; + buildRVSystem(sys, scanRandomRows(['X ~ Normal(0, 1)', 'Y = X + 0']), opts); + const before = sys.curve('Y', {})!; + buildRVSystem(sys, scanRandomRows(['X ~ Uniform(0, 1)', 'Y = X + 0']), opts); + const after = sys.curve('Y', {})!; + expect(before.sd).toBeCloseTo(1, 1); + expect(after.sd).toBeCloseTo(Math.sqrt(1 / 12), 1); + }); + + it('treats partial support honestly', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'R = sqrt(X)']); + const c = sys.curve('R', {})!; + expect(c.mass).toBeCloseTo(0.5, 2); + // The event "R > -1" happens only where R is defined. + expect(P(sys, 'R > -1')).toBeCloseTo(0.5, 2); + }); + + it('is NaN when parameters are broken', () => { + const { sys } = build(['X ~ Normal(0, s)']); + expect(P(sys, 'X < 1', { s: -1 })).toBeNaN(); + expect(sys.curve('X', { s: -1 })).toBeNull(); + }); +}); + +describe('exact normal propagation (affine in normal bases)', () => { + it('recognizes affine combinations of independent normals', () => { + const { sys } = build(['X ~ Normal(1, 0.5)', 'Y ~ Normal(3.35, 0.5)', 'Z = (X + Y)/2']); + const d = sys.exactDist('Z')!; + expect(d.kind).toBe('normal'); + expect(evaluate(d.args[0], {})).toBeCloseTo(2.175, 9); + expect(evaluate(d.args[1], {})).toBeCloseTo(Math.sqrt(0.5) / 2, 9); + }); + + it('accounts for dependence through shared names: X + X is 2X', () => { + const { sys } = build(['X ~ Normal(1, 0.5)', 'D = X + X']); + const d = sys.exactDist('D')!; + expect(evaluate(d.args[0], {})).toBeCloseTo(2, 9); + expect(evaluate(d.args[1], {})).toBeCloseTo(1, 9); // (1+1)·σ, not √2·σ + }); + + it('keeps coefficients symbolic (sliders, chains through derived names)', () => { + const { sys } = build(['X ~ Normal(1, 0.5)', 'V = a X + 1', 'W = V - X']); + const v = sys.exactDist('V')!; + expect(evaluate(v.args[0], { a: 2 })).toBeCloseTo(3, 9); + expect(evaluate(v.args[1], { a: 2 })).toBeCloseTo(1, 9); + const w = sys.exactDist('W')!; // (a−1)·X + 1 + expect(evaluate(w.args[0], { a: 3 })).toBeCloseTo(3, 9); + expect(evaluate(w.args[1], { a: 3 })).toBeCloseTo(1, 9); + }); + + it('declines everything without a closed form', () => { + const { sys } = build([ + 'X ~ Normal(0, 1)', 'Y ~ Normal(0, 1)', 'U1 ~ Uniform(0, 1)', + 'Q = X^2', 'M = X Y', 'C = {X > 0: X^2, 1}', 'S = X + U1', + ]); + for (const name of ['Q', 'M', 'C', 'S']) expect(sys.exactDist(name)).toBeNull(); + expect(sys.exactDist('U1')!.kind).toBe('uniform'); // bases pass through + }); + + it('agrees with the sampled estimate', () => { + const { sys } = build(['X ~ Normal(1, 0.5)', 'Y ~ Normal(3.35, 0.5)', 'Z = (X + Y)/2']); + const d = sys.exactDist('Z')!; + const c = sys.curve('Z', {})!; + expect(c.mean).toBeCloseTo(evaluate(d.args[0], {}), 2); + expect(c.sd).toBeCloseTo(evaluate(d.args[1], {}), 2); + }); +}); + +describe('exact laws (law propagation + uniform convolution)', () => { + it('passes a bare or affine variable of a uniform through exactly', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y = X', 'W = 2X + 1', 'V = 1 - X']); + const y = sys.exactDist('Y')!; + expect(y.kind).toBe('uniform'); + expect(y.args.map(e => evaluate(e, {}))).toEqual([0, 1]); + expect(sys.exactDist('W')!.args.map(e => evaluate(e, {}))).toEqual([1, 3]); + expect(sys.exactDist('V')!.args.map(e => evaluate(e, {}))).toEqual([0, 1]); // flipped + expect(sys.exactMoments('W', {})).toEqual({ mean: 2, sd: 2 / Math.sqrt(12) }); + }); + + it('keeps scaled exponentials exponential, and only those', () => { + const { sys } = build(['X ~ Exponential(2)', 'Y = X', 'H = 2X', 'S = X + 1']); + expect(evaluate(sys.exactDist('Y')!.args[0], {})).toBe(2); + expect(evaluate(sys.exactDist('H')!.args[0], {})).toBe(1); // rate λ/c + expect(sys.exactDist('S')).toBeNull(); // a shift leaves the family + }); + + it('collapses a repeated name before choosing a law: X + X is a box', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'D = X + X']); + const d = sys.exactDist('D')!; + expect(d.kind).toBe('uniform'); + expect(d.args.map(e => evaluate(e, {}))).toEqual([0, 2]); + }); + + it('convolves two uniforms into the exact triangle', () => { + const { sys } = build(['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 1)', 'S = X1 + X2']); + const c = sys.curve('S', {})!; + const at = (x: number) => { + for (let i = 0; i + 1 < c.pts.length; i += 2) if (c.pts[i] === x) return c.pts[i + 1]; + return NaN; + }; + expect(at(1)).toBeCloseTo(1, 12); // the apex is a corner, not a KDE shoulder + expect(at(0)).toBeCloseTo(0, 12); + expect(at(2)).toBeCloseTo(0, 12); + expect(at(0.5)).toBeCloseTo(0.5, 12); + expect(sys.exactProbability('S', undefined, parseExpr('1'), {})).toBeCloseTo(0.5, 12); + expect(sys.exactMoments('S', {})!.sd).toBeCloseTo(Math.sqrt(1 / 6), 12); + expect(c.mass).toBe(1); + }); + + it('matches Irwin–Hall for the four-fold sum', () => { + const rows = ['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 1)', 'X3 ~ Uniform(0, 1)', 'X4 ~ Uniform(0, 1)', + 'S = X1 + X2 + X3 + X4']; + const { sys } = build(rows); + const c = sys.curve('S', {})!; + const mid = c.pts.findIndex((v, i) => i % 2 === 0 && v === 2); + expect(c.pts[mid + 1]).toBeCloseTo(2 / 3, 12); // Irwin–Hall density at n/2 + expect(sys.exactProbability('S', parseExpr('3'), undefined, {})).toBeCloseTo(1 / 24, 12); + expect(sys.exactMoments('S', {})).toEqual({ mean: 2, sd: Math.sqrt(4 / 12) }); + }); + + it('handles slider coefficients and differences', () => { + const { sys } = build(['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 1)', 'S = a X1 + X2', 'D = X1 - X2']); + // a = 2: U(0,2) ∗ U(0,1) is a trapezoid on [0, 3]; its CDF at 1.5 is 1/2. + expect(sys.exactProbability('S', undefined, parseExpr('1.5'), { a: 2 })).toBeCloseTo(0.5, 12); + expect(sys.exactMoments('S', { a: 2 })!.mean).toBeCloseTo(1.5, 12); + const d = sys.curve('D', {})!; + expect(d.pts[0]).toBeCloseTo(-1, 12); // support [-1, 1], apex at 0 + expect(Math.max(...d.pts.filter((_, i) => i % 2 === 1))).toBeCloseTo(1, 12); + expect(sys.exactProbability('D', undefined, parseExpr('0'), {})).toBeCloseTo(0.5, 12); + }); + + it('agrees with the sampled estimate', () => { + const { sys } = build(['X1 ~ Uniform(0, 1)', 'X2 ~ Uniform(0, 1)', 'S = X1 + X2']); + const exact = sys.exactProbability('S', undefined, parseExpr('0.75'), {})!; + expect(Math.abs(exact - P(sys, 'S < 0.75'))).toBeLessThan(0.01); + }); + + it('leaves mixed and nonlinear forms to the sampler', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y ~ Normal(0, 1)', 'M = X + Y', 'Q = X^2']); + expect(sys.exactLaw('M')).toBeNull(); + expect(sys.exactLaw('Q')).toBeNull(); + expect(sys.curve('M', {})!.mass).toBe(1); // KDE path still serves these + }); + + it('degrades broken parameters to no curve, not a wrong one', () => { + const { sys } = build(['X1 ~ Uniform(0, s)', 'X2 ~ Uniform(0, 1)', 'S = X1 + X2']); + expect(sys.curve('S', { s: -1 })).toBeNull(); + expect(sys.exactProbability('S', undefined, parseExpr('1'), { s: -1 })).toBeNaN(); + expect(sys.curve('S', { s: 1 })).not.toBeNull(); + }); +}); + +describe('point masses (atoms)', () => { + it('renders a purely discrete variable as stems, not KDE bumps', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y = {X > 0: 1, 2}']); + const c = sys.curve('Y', {})!; + expect(c.pts).toEqual([]); + expect(c.atoms).toEqual([{ x: 1, p: 0.5 }, { x: 2, p: 0.5 }]); + expect(c.mean).toBeCloseTo(1.5, 12); + expect(c.sd).toBeCloseTo(0.5, 12); + }); + + it('finds the masses of a discretized uniform exactly', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'F = floor(4X)']); + expect(sys.curve('F', {})!.atoms).toEqual([ + { x: 0, p: 0.25 }, { x: 1, p: 0.25 }, { x: 2, p: 0.25 }, { x: 3, p: 0.25 }, + ]); + }); + + it('splits a mixed distribution into its atom and continuous part', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y = {X > 0: X^2, 1}']); + const c = sys.curve('Y', {})!; + expect(c.atoms).toEqual([{ x: 1, p: 0.5 }]); + let area = 0; + for (let i = 0; i + 3 < c.pts.length; i += 2) { + area += ((c.pts[i + 1] + c.pts[i + 3]) / 2) * (c.pts[i + 2] - c.pts[i]); + } + expect(area).toBeCloseTo(0.5, 1); // the continuous part carries the other half + }); + + it('leaves continuous estimates atom-free', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y ~ Normal(0, 1)', 'S = X + Y']); + expect(sys.curve('S', {})!.atoms).toBeUndefined(); + }); +}); + +describe('support edges', () => { + const curveArea = (pts: number[]) => { + let a = 0; + for (let i = 0; i + 3 < pts.length; i += 2) a += ((pts[i + 1] + pts[i + 3]) / 2) * (pts[i + 2] - pts[i]); + return a; + }; + + it('cuts a truncated variable off straight at its edge', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Y = {X > 1: X, 0}']); + const c = sys.curve('Y', {})!; + // Nothing below the edge, and the curve steps up from the axis at it. + expect(c.pts[0]).toBeGreaterThanOrEqual(1); + expect(c.pts[1]).toBe(0); + for (let i = 0; i < c.pts.length; i += 2) expect(c.pts[i]).toBeGreaterThanOrEqual(1); + // On (1, ∞) the density is φ itself: full height at the jump, no ramp. + expect(c.pts[3]).toBeCloseTo(normalpdf(1, 0, 1), 2); + let worst = 0; + for (let i = 0; i < c.pts.length; i += 2) { + if (c.pts[i] > 1.001 && c.pts[i] < 3) { + worst = Math.max(worst, Math.abs(c.pts[i + 1] - normalpdf(c.pts[i], 0, 1))); + } + } + expect(worst).toBeLessThan(0.005); + // The atom carries the other branch: P(X ≤ 1) = Φ(1). + expect(c.atoms).toHaveLength(1); + expect(c.atoms![0].x).toBe(0); + expect(c.atoms![0].p).toBeCloseTo(normalcdf(1, 0, 1), 4); + }); + + it('keeps full height at a half-normal edge', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'H = abs(X)']); + const c = sys.curve('H', {})!; + expect(c.pts[0]).toBeGreaterThanOrEqual(0); + expect(c.pts[3]).toBeCloseTo(2 * normalpdf(0, 0, 1), 2); + }); + + it('leaves a trimmed tail alone (no false edge)', () => { + const { sys } = build(['X ~ Normal(0, 1)', 'Z = X + 0']); + const c = sys.curve('Z', {})!; + // The drawn range stops in the tails, where the density is ~0 and must + // not be lifted by an edge correction. + expect(Math.abs(c.pts[1])).toBeLessThan(0.02); + let worst = 0; + for (let i = 0; i < c.pts.length; i += 2) { + if (Math.abs(c.pts[i]) <= 2.5) worst = Math.max(worst, Math.abs(c.pts[i + 1] - normalpdf(c.pts[i], 0, 1))); + } + expect(worst).toBeLessThan(0.01); + }); + + it('draws the probability of the range it covers, even at a singularity', () => { + // X² near 0 has an integrable singularity (density → ∞), where no local + // fit is meaningful; the area must still come out right. + const { sys } = build(['X ~ Normal(0, 1)', 'Y = {X > 0: X^2, 1}']); + expect(curveArea(sys.curve('Y', {})!.pts)).toBeCloseTo(0.5, 2); + const { sys: sys2 } = build(['X ~ Normal(0, 1)', 'Y = {X > 1: X, 0}']); + expect(curveArea(sys2.curve('Y', {})!.pts)).toBeCloseTo(1 - normalcdf(1, 0, 1), 2); + }); +}); + +describe('density estimation', () => { + it('recovers the standard normal density closely', () => { + const { sys } = build(['X ~ Normal(0, 1)']); + const { pts } = sys.curve('X', {})!; + let worst = 0; + for (let i = 0; i + 1 < pts.length; i += 2) { + if (Math.abs(pts[i]) > 2.5) continue; + worst = Math.max(worst, Math.abs(pts[i + 1] - normalpdf(pts[i], 0, 1))); + } + expect(worst).toBeGreaterThan(0); // the sweep saw the curve at all + expect(worst).toBeLessThan(0.02); + }); + + it('integrates to the sample mass', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y ~ Normal(0, 1)', 'S = X + Y']); + const { pts } = sys.curve('S', {})!; + let area = 0; + for (let i = 0; i + 3 < pts.length; i += 2) { + area += ((pts[i + 1] + pts[i + 3]) / 2) * (pts[i + 2] - pts[i]); + } + expect(area).toBeCloseTo(1, 1); + }); + + it('clips a shade polygon to the bounds and closes it to the axis', () => { + const { sys } = build(['X ~ Uniform(0, 1)', 'Y ~ Uniform(0, 1)', 'S = X + Y']); + const curve = sys.curve('S', {})!; + const poly = shadePolygon(curve, 0.5, 1.5)!; + expect(poly[0]).toBeCloseTo(0.5, 9); + expect(poly[1]).toBe(0); + expect(poly[poly.length - 2]).toBeCloseTo(1.5, 9); + expect(poly[poly.length - 1]).toBe(0); + for (let i = 0; i < poly.length; i += 2) { + expect(poly[i]).toBeGreaterThanOrEqual(0.5 - 1e-9); + expect(poly[i]).toBeLessThanOrEqual(1.5 + 1e-9); + } + expect(shadePolygon(curve, 5, 6)).not.toBeNull(); // empty clip yields a flat sliver + }); +}); diff --git a/lib/dist.ts b/lib/dist.ts index d57b68f..fa11244 100644 --- a/lib/dist.ts +++ b/lib/dist.ts @@ -1,25 +1,64 @@ /** * Probability distribution rows. * - * - `X ~ Normal(mean, sd)` declares a random variable; the row plots its - * density y = normalpdf(x, mean, sd). Parameters may reference constants - * (sliders) and t, so `X ~ Normal(0, a)` responds to the slider live. - * - `P(X < b)`, `P(X > b)`, `P(a < X < b)` shade the area under X's density - * over the given range (reusing the inequality-region pipeline) and report - * the numeric probability via the normal CDF. + * - `X ~ Normal(mean, sd)` (also Uniform(a, b), Exponential(rate)) declares a + * random variable; the row plots its exact density. Parameters may reference + * constants (sliders) and t, so `X ~ Normal(0, a)` responds to the slider. + * - `Y = g(X, …)` where the right side references random variables declares a + * *derived* random variable — arithmetic on distributions. `S = X1 + X2` is + * the convolution of independent summands, `X Y` the product distribution, + * and piecewise conditionals work too: `Y = {X > 0: X^2, 1}`. Derived rows + * (and bare expressions like `X + Y`) plot a density estimated from samples. + * - `P(…)` takes any inequality over the declared variables: `P(X < b)`, + * `P(a < X < b)`, `P(Y > 0.5)`, even `P(Y > X)`. Single-variable bounds on a + * base distribution stay exact (closed-form CDF + shaded region); everything + * else is estimated from the same joint samples. + * - `E(…)` takes any expression over the declared variables: `E(X)`, + * `E(X^2 + Y)`. The mean is exact when the law is (closed-form pdfs and + * uniform sums), the finite-sample mean otherwise, and the row draws a + * vertical marker at x = E under the expression's density. + * + * Sampling model: every base variable owns a deterministic stratified stream + * of standard uniforms (equal-mass quantile midpoints, shuffled by a hash of + * its name — a Latin-hypercube pairing across variables). Samples are the + * quantile transform of that stream, so distinct names are independent while + * a derived variable, evaluated per-sample over its dependencies, preserves + * the joint distribution exactly: `X + X` is 2X, `P(Y > X)` sees the + * dependence of Y on X. Streams are fixed, so results are reproducible and + * respond continuously to slider drags (common random numbers). */ -import { type Expr, erf, evaluate, freeVars, parseExpr } from './expr.ts'; +import { + EVAL_FNS, + type Expr, + FUNCTIONS, + builtinFn, + evaluate, + freeVars, + ineqComparisons, + normalcdf, + normalpdf, + parseExpr, + substVars, +} from './expr.ts'; +import { usesComplex } from './complex.ts'; +import { type GetFn, RESERVED, type ResolveOpts, resolveExpr } from './defs.ts'; +import { quadrature } from './integrate.ts'; + +// --- base distributions --- -export interface DistDef { - /** Random-variable name (the row `name ~ Normal(…)`). */ - name: string; - mean: Expr; - sd: Expr; +export type BaseKind = 'normal' | 'uniform' | 'exponential'; + +/** Argument meaning by kind — normal: [mean, sd]; uniform: [lo, hi]; exponential: [rate]. */ +export interface BaseDist { + kind: BaseKind; + args: Expr[]; } const TILDE_RE = /^\s*([A-Za-z_]\w*)\s*~\s*([\s\S]+)$/; -const DIST_RE = /^\s*([A-Za-z_]\w*)\s*\(([\s\S]*)\)\s*$/; +const DIST_RE = /^\s*([A-Za-z_]\w*)\s*(?:\(([\s\S]*)\))?\s*$/; const PROB_RE = /^\s*P\s*\(([\s\S]+)\)\s*$/; +const EXPECT_RE = /^\s*E\s*\(([\s\S]+)\)\s*$/; +const CONST_ROW_RE = /^\s*([A-Za-z_]\w*)\s*=(?!=)([\s\S]+)$/; /** Detect a `name ~ rhs` row before parsing ('~' is not an expression token). */ export function scanDistribution(text: string): { name: string; rhs: string } | null { @@ -27,35 +66,88 @@ export function scanDistribution(text: string): { name: string; rhs: string } | return m ? { name: m[1], rhs: m[2] } : null; } +interface DistSpec { + kind: BaseKind; + arity: number; + usage: string; + defaults: number[]; +} + +const DIST_SPECS = new Map([ + ...['normal', 'n'].map((a): [string, DistSpec] => + [a, { kind: 'normal', arity: 2, usage: 'Normal(mean, sd)', defaults: [0, 1] }]), + ...['uniform', 'u'].map((a): [string, DistSpec] => + [a, { kind: 'uniform', arity: 2, usage: 'Uniform(lo, hi)', defaults: [0, 1] }]), + ...['exponential', 'exp'].map((a): [string, DistSpec] => + [a, { kind: 'exponential', arity: 1, usage: 'Exponential(rate)', defaults: [1] }]), +]); + /** Parse the right side of `name ~ …`. Throws with a row-friendly message. */ -export function parseDistribution(name: string, rhs: string, fnNames: ReadonlySet): DistDef { +export function parseDistribution(rhs: string, fnNames: ReadonlySet): BaseDist { const m = DIST_RE.exec(rhs); - if (!m) throw new Error('Expected a distribution like Normal(0, 1).'); - const dist = m[1].toLowerCase(); - if (dist !== 'normal' && dist !== 'n') { - throw new Error(`Unknown distribution: ${m[1]}. Try Normal(mean, sd).`); + const spec = m && DIST_SPECS.get(m[1].toLowerCase()); + if (!m || !spec) { + throw new Error(m && !DIST_SPECS.has(m[1].toLowerCase()) + ? `Unknown distribution: ${m[1]}. Try Normal(mean, sd), Uniform(lo, hi), or Exponential(rate).` + : 'Expected a distribution like Normal(0, 1).'); + } + // A bare name takes the standard parameters: `X ~ N` is Normal(0, 1). + if (m[2] === undefined) { + return { kind: spec.kind, args: spec.defaults.map(value => ({ kind: 'num', value })) }; } let args: Expr; try { args = parseExpr(`(${m[2]})`, fnNames); } catch (e) { if (e instanceof Error && /vector components/.test(e.message)) { - throw new Error('Normal takes 2 arguments: Normal(mean, sd).'); + throw new Error(`${spec.usage} takes ${spec.arity} arguments.`); } throw e; } - if (args.kind !== 'vec' || args.items.length !== 2) { - throw new Error('Normal takes 2 arguments: Normal(mean, sd).'); + const items = args.kind === 'vec' ? args.items : [args]; + if (items.length !== spec.arity) { + throw new Error(`${spec.usage} takes ${spec.arity} argument${spec.arity > 1 ? 's' : ''}.`); } - return { name, mean: args.items[0], sd: args.items[1] }; + return { kind: spec.kind, args: items }; } const v = (name: string): Expr => ({ kind: 'var', name }); -const pdfCall = (d: DistDef): Expr => ({ kind: 'call', name: 'normalpdf', args: [v('x'), d.mean, d.sd] }); +const num = (value: number): Expr => ({ kind: 'num', value }); +const bin = (op: '+' | '-' | '*' | '/' | '^', a: Expr, b: Expr): Expr => ({ kind: 'bin', op, a, b }); +const chain = (lo: Expr, mid: Expr, hi: Expr): Expr => + ({ kind: 'ineq', op: '<', l: { kind: 'ineq', op: '<', l: lo, r: mid }, r: hi }); -/** The density curve for a random variable: y = normalpdf(x, mean, sd). */ -export function densityExpr(d: DistDef): Expr { - return { kind: 'eq', l: v('y'), r: pdfCall(d) }; +/** The exact pdf of a base distribution at `x` (piecewise where the support ends). */ +export function pdfExpr(d: BaseDist, x: Expr): Expr { + switch (d.kind) { + case 'normal': + return { kind: 'call', name: 'normalpdf', args: [x, d.args[0], d.args[1]] }; + case 'uniform': + // The condition is empty while hi <= lo (mid slider drag), so the pdf + // degrades to 0 everywhere instead of going negative. + return { + kind: 'piecewise', + cases: [{ cond: chain(d.args[0], x, d.args[1]), value: bin('/', num(1), bin('-', d.args[1], d.args[0])) }], + otherwise: num(0), + }; + case 'exponential': { + // max(rate, 0): a non-positive rate flattens to 0 rather than blowing up. + const r: Expr = { kind: 'call', name: 'max', args: [d.args[0], num(0)] }; + return { + kind: 'piecewise', + cases: [{ + cond: { kind: 'ineq', op: '>=', l: x, r: num(0) }, + value: bin('*', r, { kind: 'call', name: 'exp', args: [{ kind: 'neg', a: bin('*', r, x) }] }), + }], + otherwise: num(0), + }; + } + } +} + +/** The density curve for a base random variable: y = pdf(x). */ +export function densityExpr(d: BaseDist): Expr { + return { kind: 'eq', l: v('y'), r: pdfExpr(d, v('x')) }; } /** The inner text of a `P(…)` row, or null if the row has another shape. */ @@ -64,22 +156,44 @@ export function matchProbability(text: string): string | null { return m ? m[1] : null; } -export interface Probability { - dist: DistDef; - lo?: Expr; - hi?: Expr; +/** The inner text of an `E(…)` row, or null if the row has another shape. */ +export function matchExpectation(text: string): string | null { + const m = EXPECT_RE.exec(text); + return m ? m[1] : null; +} + +// --- probability specs --- + +export interface ProbSpec { + /** The inequality (chain) to estimate, resolved. */ + body: Expr & { kind: 'ineq' }; + /** Random variables the body references. */ + rvs: string[]; + /** + * Present when the body is constant bounds around one bare variable + * (`P(a < X < b)`): the shadeable — and for closed-form laws, exact — case. + */ + single?: { rv: string; lo?: Expr; hi?: Expr }; + /** + * Bounds around one variable-bearing *expression* (`P(0.5 < X + Y < 1.5)`): + * the same case once the caller registers the expression as an anonymous + * derived variable. + */ + inline?: { e: Expr; lo?: Expr; hi?: Expr }; } /** Interpret a parsed P(…) body against the declared random variables. */ -export function toProbability(e: Expr, dists: ReadonlyMap): Probability { - const chain: Array = []; - let node = e; - while (node.kind === 'ineq') { - chain.unshift(node); - node = node.l; - } - if (!chain.length) throw new Error('P(…) expects an inequality like P(X < 2).'); - const comps = chain.map((c, k) => ({ op: c.op, l: k === 0 ? c.l : chain[k - 1].r, r: c.r })); +export function toProbability(e: Expr, rvNames: ReadonlySet): ProbSpec { + if (e.kind !== 'ineq') throw new Error('P(…) expects an inequality like P(X < 2).'); + const frees = freeVars(e); + const rvs = [...frees].filter(n => rvNames.has(n)); + if (!rvs.length) { + throw new Error('P(…) must reference a random variable, e.g. X ~ Normal(0, 1) then P(X < 2).'); + } + for (const n of frees) { + if (/^[xyzuvw]$/.test(n)) throw new Error(`P(…) cannot use the plot coordinate ${n}.`); + } + const comps = ineqComparisons(e); if (new Set(comps.map(c => c.op[0])).size > 1) { throw new Error('Chained inequalities must point the same way.'); } @@ -87,49 +201,1388 @@ export function toProbability(e: Expr, dists: ReadonlyMap): Pro const asc = comps.map(c => (c.op[0] === '<' ? { l: c.l, r: c.r } : { l: c.r, r: c.l })); if (comps[0].op[0] === '>') asc.reverse(); const terms = [asc[0].l, ...asc.map(c => c.r)]; + const spec: ProbSpec = { body: e, rvs }; - const idx = terms.findIndex(t => t.kind === 'var' && dists.has(t.name)); - if (idx < 0) { - throw new Error('P(…) must reference a random variable, e.g. X ~ Normal(0, 1) then P(X < 2).'); + // `lo < … < hi` with exactly one variable-bearing term and variable-free + // bounds on its immediate sides shades (and computes exactly when a law is + // derivable): a bare name yields `single`, an expression `inline`. Anything + // else — P(Y > X), extra constraints beyond the bounds — samples. + const idx = terms.findIndex(t => [...freeVars(t)].some(n => rvNames.has(n))); + const others = terms.filter((_, k) => k !== idx); + if (idx >= 0 && terms.length <= 3 && idx <= 1 && idx >= terms.length - 2 + && others.every(t => [...freeVars(t)].every(n => !rvNames.has(n)))) { + const t = terms[idx]; + const lo = idx > 0 ? terms[idx - 1] : undefined; + const hi = idx < terms.length - 1 ? terms[idx + 1] : undefined; + if (t.kind === 'var' && rvNames.has(t.name)) spec.single = { rv: t.name, lo, hi }; + else spec.inline = { e: t, lo, hi }; } - terms.forEach((t, k) => { - if (k === idx) return; - if (Math.abs(k - idx) > 1) throw new Error('P(…) takes at most two bounds around the variable.'); - for (const name of freeVars(t)) { - if (dists.has(name)) throw new Error('Only one random variable may appear in P(…).'); - if (/^[xyzuvw]$/.test(name)) throw new Error(`P(…) bounds cannot use ${name}.`); - } - }); - const dist = dists.get((terms[idx] as Expr & { kind: 'var' }).name)!; - return { - dist, - lo: idx > 0 ? terms[idx - 1] : undefined, - hi: idx < terms.length - 1 ? terms[idx + 1] : undefined, - }; + return spec; +} + +export interface ExpectSpec { + /** The scalar expression to average, resolved. */ + body: Expr; + /** Random variables the body references. */ + rvs: string[]; +} + +/** Interpret a parsed E(…) body against the declared random variables. */ +export function toExpectation(e: Expr, rvNames: ReadonlySet): ExpectSpec { + if (e.kind === 'ineq') { + throw new Error('E(…) expects a value like E(X + Y); the chance of an event is P(…).'); + } + if (e.kind === 'eq' || e.kind === 'vec' || e.kind === 'list') { + throw new Error('E(…) expects a single value, like E(X + Y).'); + } + const frees = freeVars(e); + const rvs = [...frees].filter(n => rvNames.has(n)); + if (!rvs.length) { + throw new Error('E(…) must reference a random variable, e.g. X ~ Normal(0, 1) then E(X).'); + } + for (const n of frees) { + if (/^[xyzuvw]$/.test(n)) throw new Error(`E(…) cannot use the plot coordinate ${n}.`); + } + return { body: e, rvs }; } /** - * The shaded region for a probability: the area between the x-axis and the - * density, clipped to the bounds. Each part is normalized to F < 0 and + * The shaded region for an exact probability: the area between the x-axis and + * the density, clipped to the bounds. Each part is normalized to F < 0 and * combined with max() (intersection), the same shape classify() produces for * inequality chains; '<=' gives the region a drawn outline. */ -export function regionExpr(p: Probability): Expr { +export function regionExpr(d: BaseDist, lo?: Expr, hi?: Expr): Expr { const x = v('x'); const y = v('y'); - let f: Expr = { kind: 'bin', op: '-', a: y, b: pdfCall(p.dist) }; // y < pdf(x) + let f: Expr = bin('-', y, pdfExpr(d, x)); // y < pdf(x) const parts: Expr[] = [{ kind: 'neg', a: y }]; // 0 < y - if (p.lo) parts.push({ kind: 'bin', op: '-', a: p.lo, b: x }); // lo < x - if (p.hi) parts.push({ kind: 'bin', op: '-', a: x, b: p.hi }); // x < hi + if (lo) parts.push(bin('-', lo, x)); // lo < x + if (hi) parts.push(bin('-', x, hi)); // x < hi for (const part of parts) f = { kind: 'call', name: 'max', args: [f, part] }; - return { kind: 'ineq', op: '<=', l: f, r: { kind: 'num', value: 0 } }; + return { kind: 'ineq', op: '<=', l: f, r: num(0) }; +} + +/** Exact CDF of a base distribution; NaN while the parameters are invalid. */ +function cdf(d: BaseDist, x: number, env: Record): number { + const a = d.args.map(e => evaluate(e, env)); + switch (d.kind) { + case 'normal': + return a[1] > 0 ? normalcdf(x, a[0], a[1]) : NaN; + case 'uniform': + return a[1] > a[0] ? Math.min(1, Math.max(0, (x - a[0]) / (a[1] - a[0]))) : NaN; + case 'exponential': + return a[0] > 0 ? (x <= 0 ? 0 : 1 - Math.exp(-a[0] * x)) : NaN; + } +} + +/** Exact value of P(lo < X < hi) under the given constant environment. */ +export function probabilityValue( + d: BaseDist, + lo: Expr | undefined, + hi: Expr | undefined, + env: Record, +): number { + return (hi ? cdf(d, evaluate(hi, env), env) : 1) - (lo ? cdf(d, evaluate(lo, env), env) : 0); +} + +// --- row scanning --- + +/** + * Decide which rows declare random variables, before definitions are built. + * `base` rows are `name ~ …`; `derived` rows are `name = rhs` where the rhs + * mentions a random variable (transitively — `Z = Y + 1` follows `Y = X^2` + * into the set). Rows the caller has already claimed (comments, sequences) + * arrive as null. Matching is textual by design — it must run before parsing, + * because these rows must *not* become constant definitions — but it follows + * the tokenizer's identifier rule (a maximal run starting with a letter), so + * `2X` mentions X while `aX`, `X_1`, and `X2` are their own names. A \b-style + * word boundary would get `2X` wrong: 2 and X are both word characters. + */ +export function scanRandomRows(texts: readonly (string | null)[]): { + base: Map; + derived: Map; +} { + const base = new Map(); + const derived = new Map(); + const names = new Set(); + const candidates = new Map(); + texts.forEach((text, i) => { + if (!text) return; + const scan = scanDistribution(text); + if (scan) { + base.set(i, scan); + names.add(scan.name); + return; + } + const m = CONST_ROW_RE.exec(text); + // The name must be claimable as a definition (`e = X` stays an equation). + if (m && !FUNCTIONS.has(m[1]) && !RESERVED.has(m[1]) && !m[1].startsWith('u_')) { + candidates.set(i, { name: m[1], rhs: m[2] }); + } + }); + let changed = names.size > 0; + while (changed) { + changed = false; + for (const [i, c] of candidates) { + if (!(c.rhs.match(/[A-Za-z_]\w*/g) ?? []).some(t => names.has(t))) continue; + candidates.delete(i); + derived.set(i, c); + names.add(c.name); + changed = true; + } + } + return { base, derived }; +} + +/** + * Validate a resolved right-hand side as a derived random variable: a real + * scalar in random variables, constants, and t. + */ +export function checkDerived(e: Expr, rvNames: ReadonlySet, constNames: ReadonlySet): void { + if (e.kind === 'eq' || e.kind === 'ineq' || e.kind === 'vec' || e.kind === 'list') { + throw new Error('A random variable must be a single value.'); + } + if (usesComplex(e)) throw new Error('Random variables are real-valued.'); + for (const n of freeVars(e)) { + if (rvNames.has(n) || constNames.has(n) || n === 't') continue; + if (/^[xyzuv]$/.test(n)) { + throw new Error(`A random variable cannot depend on the plot coordinate ${n}.`); + } + throw new Error(`${n} is not defined.`); + } +} + +// --- the sampled system --- + +/** Joint sample count. Stratified streams keep marginals exact at any size; + * this is set by when the *derived-density* estimate looks smooth (KDE noise + * ~ 1/√(N·h), visible as low-frequency wobble on zoomed-in curves) while a + * slider drag can still resample every affected variable within a frame. */ +export const SAMPLE_COUNT = 1 << 17; + +export type RV = + | { name: string; kind: 'base'; dist: BaseDist } + | { name: string; kind: 'derived'; expr: Expr }; + +export interface DensityCurve { + /** Flat [x0, y0, x1, y1, …] polyline of the continuous part's density + * (empty when the distribution is purely discrete). */ + pts: number[]; + /** Point masses (a piecewise branch, floor, a constant): drawn as stems of + * height = probability, never smeared into the density. */ + atoms?: Array<{ x: number; p: number }>; + mean: number; + sd: number; + /** Fraction of samples that are finite (< 1 for partial support like sqrt(X)). */ + mass: number; +} + +const fnv1a = (s: string): number => { + let h = 0x811c9dc5; + for (let i = 0; i < s.length; i++) { + h ^= s.charCodeAt(i); + h = Math.imul(h, 0x01000193); + } + return h >>> 0; +}; + +const mulberry32 = (seed: number) => (): number => { + seed = (seed + 0x6d2b79f5) | 0; + let z = seed; + z = Math.imul(z ^ (z >>> 15), z | 1); + z ^= z + Math.imul(z ^ (z >>> 7), z | 61); + return ((z ^ (z >>> 14)) >>> 0) / 4294967296; +}; + +/** + * The stratified standard-uniform stream for a variable name: quantile + * midpoints (i + ½)/N shuffled by a name-seeded permutation. Memoized + * forever — the stream is a pure function of the name. + */ +const streams = new Map(); +function uniformStream(name: string): Float64Array { + let s = streams.get(name); + if (s) return s; + s = new Float64Array(SAMPLE_COUNT); + for (let i = 0; i < SAMPLE_COUNT; i++) s[i] = (i + 0.5) / SAMPLE_COUNT; + const rand = mulberry32(fnv1a(name)); + for (let i = SAMPLE_COUNT - 1; i > 0; i--) { + const j = Math.floor(rand() * (i + 1)); + const t = s[i]; + s[i] = s[j]; + s[j] = t; + } + streams.set(name, s); + return s; +} + +/** Acklam's rational approximation to the standard normal quantile (~1e-9). */ +function normalQuantile(p: number): number { + const a = [-3.969683028665376e1, 2.209460984245205e2, -2.759285104469687e2, + 1.383577518672690e2, -3.066479806614716e1, 2.506628277459239]; + const b = [-5.447609879822406e1, 1.615858368580409e2, -1.556989798598866e2, + 6.680131188771972e1, -1.328068155288572e1]; + const c = [-7.784894002430293e-3, -3.223964580411365e-1, -2.400758277161838, + -2.549732539343734, 4.374664141464968, 2.938163982698783]; + const d = [7.784695709041462e-3, 3.224671290700398e-1, 2.445134137142996, 3.754408661907416]; + const plow = 0.02425; + if (p <= 0 || p >= 1) return NaN; + if (p < plow || p > 1 - plow) { + const q = Math.sqrt(-2 * Math.log(p < plow ? p : 1 - p)); + const x = (((((c[0] * q + c[1]) * q + c[2]) * q + c[3]) * q + c[4]) * q + c[5]) + / ((((d[0] * q + d[1]) * q + d[2]) * q + d[3]) * q + 1); + return p < plow ? x : -x; + } + const q = p - 0.5; + const r = q * q; + return (((((a[0] * r + a[1]) * r + a[2]) * r + a[3]) * r + a[4]) * r + a[5]) * q + / (((((b[0] * r + b[1]) * r + b[2]) * r + b[3]) * r + b[4]) * r + 1); +} + +/** Evaluate an expression column-wise over the sample vectors. Inequalities + * yield 1/0 masks (NaN where an operand is NaN), matching evaluate(). */ +function evalCols( + e: Expr, + cols: ReadonlyMap, + env: Record, + n: number, +): Float64Array { + const alloc = () => new Float64Array(n); + switch (e.kind) { + case 'num': { + const out = alloc(); + out.fill(e.value); + return out; + } + case 'var': { + const col = cols.get(e.name); + if (col) return col; + if (!(e.name in env)) throw new Error(`Unbound variable: ${e.name}`); + const out = alloc(); + out.fill(env[e.name]); + return out; + } + case 'neg': { + const a = evalCols(e.a, cols, env, n); + const out = alloc(); + for (let i = 0; i < n; i++) out[i] = -a[i]; + return out; + } + case 'bin': { + const a = evalCols(e.a, cols, env, n); + const b = evalCols(e.b, cols, env, n); + const out = alloc(); + switch (e.op) { + case '+': for (let i = 0; i < n; i++) out[i] = a[i] + b[i]; break; + case '-': for (let i = 0; i < n; i++) out[i] = a[i] - b[i]; break; + case '*': for (let i = 0; i < n; i++) out[i] = a[i] * b[i]; break; + case '/': for (let i = 0; i < n; i++) out[i] = a[i] / b[i]; break; + case '^': for (let i = 0; i < n; i++) out[i] = Math.pow(a[i], b[i]); break; + } + return out; + } + case 'call': { + const fn = EVAL_FNS[e.name]; + if (!fn) throw new Error(`Unknown function: ${e.name}`); + const args = e.args.map(a => evalCols(a, cols, env, n)); + const out = alloc(); + if (args.length === 1) { + const a = args[0]; + for (let i = 0; i < n; i++) out[i] = fn(a[i]); + } else if (args.length === 2) { + const [a, b] = args; + for (let i = 0; i < n; i++) out[i] = fn(a[i], b[i]); + } else { + for (let i = 0; i < n; i++) out[i] = fn(...args.map(a => a[i])); + } + return out; + } + case 'eq': { + const l = evalCols(e.l, cols, env, n); + const r = evalCols(e.r, cols, env, n); + const out = alloc(); + for (let i = 0; i < n; i++) out[i] = l[i] - r[i]; + return out; + } + case 'ineq': { + const out = alloc(); + out.fill(1); + for (const { op, l, r } of ineqComparisons(e)) { + const a = evalCols(l, cols, env, n); + const b = evalCols(r, cols, env, n); + for (let i = 0; i < n; i++) { + if (Number.isNaN(out[i])) continue; + if (Number.isNaN(a[i]) || Number.isNaN(b[i])) out[i] = NaN; + else if (!(op === '<' ? a[i] < b[i] : op === '<=' ? a[i] <= b[i] + : op === '>' ? a[i] > b[i] : a[i] >= b[i])) out[i] = 0; + } + } + return out; + } + case 'piecewise': { + const out = alloc(); + out.fill(NaN); + const taken = new Uint8Array(n); + for (const c of e.cases) { + const mask = evalCols(c.cond, cols, env, n); + const val = evalCols(c.value, cols, env, n); + for (let i = 0; i < n; i++) { + if (!taken[i] && mask[i] === 1) { + out[i] = val[i]; + taken[i] = 1; + } + } + } + if (e.otherwise) { + const val = evalCols(e.otherwise, cols, env, n); + for (let i = 0; i < n; i++) if (!taken[i]) out[i] = val[i]; + } + return out; + } + case 'vec': throw new Error('Vector in scalar context.'); + case 'list': throw new Error('List in scalar context.'); + } +} + +/** Estimate a density curve from samples: point masses split off as atoms, + * the continuous remainder as a binned kernel density estimate (Silverman + * bandwidth over a robust spread) whose area equals its share of the + * finite-sample mass. Null when nothing is finite. */ +function estimateCurve(col: Float64Array): DensityCurve | null { + let finite: number[] = []; + let sum = 0; + for (let i = 0; i < col.length; i++) { + const x = col[i]; + if (isFinite(x)) { + finite.push(x); + sum += x; + } + } + const n = finite.length; + if (n < 16) return null; + const mean = sum / n; + let ss = 0; + for (const x of finite) ss += (x - mean) * (x - mean); + const sd = Math.sqrt(ss / n); + const mass = n / col.length; + + // Atoms: exactly repeated values are point masses — a piecewise branch, a + // floor, a constant — and smearing them into KDE bumps would read as + // continuous spread. Stratified streams make continuous values distinct, so + // a duplicate probe over a prefix keeps that common case on the fast path. + let atoms: Array<{ x: number; p: number }> | undefined; + const probe = new Set(); + for (let i = 0; i < Math.min(n, 4096); i++) probe.add(finite[i]); + if (probe.size < Math.min(n, 4096)) { + const counts = new Map(); + for (const x of finite) counts.set(x, (counts.get(x) ?? 0) + 1); + const minAtom = Math.max(8, col.length * 0.002); + const atomValues = new Set(); + for (const [x, count] of counts) { + if (count >= minAtom) { + atomValues.add(x); + (atoms ??= []).push({ x, p: count / col.length }); + } + } + if (atoms) { + atoms.sort((a, b) => a.x - b.x); + finite = finite.filter(x => !atomValues.has(x)); + } + } + if (finite.length < 16) return { pts: [], atoms, mean, sd, mass }; // purely discrete + // The continuous part's own count and spread size the estimate below. + const cn = finite.length; + // Quantiles from a decimated sort: plenty for a range and bandwidth. + const sub = Float64Array.from(finite.filter((_, i) => i % Math.ceil(cn / 4096) === 0)).sort(); + const q = (p: number) => sub[Math.min(sub.length - 1, Math.floor(p * sub.length))]; + const spread = Math.min(sd, (q(0.75) - q(0.25)) / 1.349); + if (!(spread > 0)) return { pts: [], atoms, mean, sd, mass }; // no continuous spread to draw + // 1.4× Silverman's rule. His 0.9 factor is MISE-optimal for i.i.d. draws; + // measured on these stratified columns, ~1.4× lowers BOTH the sup-error and + // the curve's residual wobble (second-difference energy ÷2.4) — smoothness + // is what the plotted line is judged by. + const h = 1.26 * spread * Math.pow(cn, -0.2); + // Drawn range: trim the extreme tails, but never past the observed support. + // An end the trim does not reach is the support edge itself — beyond it the + // density is truly zero, so a truncated variable like {X > 1: X, 0} has to + // cut off straight at 1 rather than ramp up to a rounded peak past it. + let x0 = Infinity; + let x1 = -Infinity; + for (const x of finite) { + if (x < x0) x0 = x; + if (x > x1) x1 = x; + } + let lo = q(0.005) - 3 * h; + let hi = q(0.995) + 3 * h; + const hardLo = lo <= x0; + const hardHi = hi >= x1; + if (hardLo) lo = x0; + if (hardHi) hi = x1; + const B = 512; + const dx = (hi - lo) / B; + const hist = new Float64Array(B + 1); + const w = 1 / (col.length * dx); + let inWindow = 0; + for (const x of finite) { + // Linear binning: split each sample between its two neighboring grid + // points, so the histogram carries no half-bin jitter into the curve. + const k = (x - lo) / dx; + const k0 = Math.floor(k); + if (k0 < 0 || k0 >= B) continue; + const f = k - k0; + hist[k0] += w * (1 - f); + hist[k0 + 1] += w * f; + inWindow++; + } + inWindow /= col.length; + // A grid point holds the density there, but the ones on a support edge + // collect from one side only — half a cell — so they read half the density. + if (hardLo) hist[0] *= 2; + if (hardHi) hist[B] *= 2; + // Gaussian smoothing of the histogram — the binned KDE. + const r = Math.min(256, Math.ceil((3 * h) / dx)); + const kernel = new Float64Array(2 * r + 1); + let ksum = 0; + for (let k = -r; k <= r; k++) ksum += kernel[k + r] = Math.exp(-0.5 * ((k * dx) / h) ** 2); + for (let k = 0; k <= 2 * r; k++) kernel[k] /= ksum; + // Boundary correction, needed only where the support ends. Straight + // truncation would halve the estimate at the edge (half the kernel hangs + // outside) and renormalizing alone still sags wherever the density is + // sloped, so fit a local *line* instead of a local mean: with kernel + // moments a₀,a₁,a₂ over the part of the window inside the support, + // f̂ = (a₂S₀ − a₁S₁)/(a₀a₂ − a₁²), which reproduces any linear density + // exactly. In the interior a₁ = 0 and a₀ = 1, so this is the plain KDE. + const hard = hardLo || hardHi; + const P0 = new Float64Array(2 * r + 2); + const P1 = new Float64Array(2 * r + 2); + const P2 = new Float64Array(2 * r + 2); + if (hard) { + for (let m = 0; m <= 2 * r; m++) { + const d = (m - r) * dx; + P0[m + 1] = P0[m] + kernel[m]; + P1[m + 1] = P1[m] + kernel[m] * d; + P2[m + 1] = P2[m] + kernel[m] * d * d; + } + } + const pts: number[] = []; + if (hardLo) pts.push(lo, 0); // the jump itself: a vertical at the edge + for (let j = 0; j <= B; j++) { + let s0 = 0; + let s1 = 0; + const k0 = Math.max(0, j - r); + const k1 = Math.min(B, j + r); + for (let k = k0; k <= k1; k++) { + const wk = hist[k] * kernel[j - k + r]; + s0 += wk; + if (hard) s1 += wk * (j - k) * dx; + } + let y = s0; + if (hard) { + // Clip the moment window only at the ends that are real support edges; + // a merely trimmed tail keeps the full window (its data continues). + const mlo = hardHi ? Math.max(0, j + r - B) : 0; + const mhi = hardLo ? Math.min(2 * r, j + r) : 2 * r; + const a0 = P0[mhi + 1] - P0[mlo]; + const a1 = P1[mhi + 1] - P1[mlo]; + const a2 = P2[mhi + 1] - P2[mlo]; + const den = a0 * a2 - a1 * a1; + // Local linear can undershoot below zero where samples are sparse; + // there, fall back to the renormalized mean, which cannot. + const ll = den > 0 ? (a2 * s0 - a1 * s1) / den : -1; + y = ll > 0 ? ll : (a0 > 0.05 ? s0 / a0 : s0); + } + pts.push(lo + j * dx, y); + } + if (hardHi) pts.push(hi, 0); + // The curve's area is the probability of the drawn range — the promise a + // density plot makes. Smoothing and the edge corrections each perturb it a + // little (and at an integrable singularity, where no local polynomial fit + // is meaningful, by more), so restore it exactly. + let area = 0; + for (let i = 0; i + 3 < pts.length; i += 2) { + area += ((pts[i + 1] + pts[i + 3]) / 2) * (pts[i + 2] - pts[i]); + } + if (area > 0) { + const k = inWindow / area; + for (let i = 1; i < pts.length; i += 2) pts[i] *= k; + } + return { pts, atoms, mean, sd, mass }; +} + +/** Linear interpolation of a density polyline at x (0 outside its range). */ +export function densityAt(curve: DensityCurve, x: number): number { + const p = curve.pts; + for (let i = 0; i + 3 < p.length; i += 2) { + if (x >= p[i] && x <= p[i + 2]) { + const f = (x - p[i]) / (p[i + 2] - p[i] || 1); + return p[i + 1] + f * (p[i + 3] - p[i + 1]); + } + } + return 0; +} + +/** Clip a density curve to [lo, hi] and close it down to the x-axis: the + * polygon a Monte Carlo `P(…)` row fills. Null when the clip is empty. */ +export function shadePolygon(curve: DensityCurve, lo?: number, hi?: number): number[] | null { + const p = curve.pts; + if (p.length < 4) return null; + const xlo = lo ?? p[0]; + const xhi = hi ?? p[p.length - 2]; + if (!(xhi > xlo)) return null; + const yAt = (x: number): number => densityAt(curve, x); + const out: number[] = [xlo, 0, xlo, yAt(xlo)]; + for (let i = 0; i < p.length; i += 2) { + if (p[i] > xlo && p[i] < xhi) out.push(p[i], p[i + 1]); + } + out.push(xhi, yAt(xhi), xhi, 0); + return out; +} + +interface CacheEntry { + /** Serialized definition + parameter values the fields were computed under. */ + sig: string; + /** Joint sample column (present once columns() ran for this sig). */ + col?: Float64Array; + curve?: DensityCurve | null; + /** Quadrature moments (present once quadMoments ran for this sig). */ + qm?: { mean: number; sd: number; mass: number } | null; } -/** Numeric value of the probability under the given constant environment. */ -export function probabilityValue(p: Probability, env: Record): number { - const mean = evaluate(p.dist.mean, env); - const sd = evaluate(p.dist.sd, env); - if (!(sd > 0)) return NaN; - const cdf = (e: Expr) => 0.5 * (1 + erf((evaluate(e, env) - mean) / (sd * Math.SQRT2))); - return (p.hi ? cdf(p.hi) : 1) - (p.lo ? cdf(p.lo) : 0); +/** The pdf and support of a base distribution at these parameter values, or + * null while the parameters are invalid. */ +function pdfClosure( + d: BaseDist, + env: Record, +): { pdf: (x: number) => number; lo: number; hi: number } | null { + const a = d.args.map(e => evaluate(e, env)); + if (!a.every(isFinite)) return null; + switch (d.kind) { + case 'normal': + return a[1] > 0 ? { pdf: x => normalpdf(x, a[0], a[1]), lo: -Infinity, hi: Infinity } : null; + case 'uniform': + return a[1] > a[0] ? { pdf: () => 1 / (a[1] - a[0]), lo: a[0], hi: a[1] } : null; + case 'exponential': + return a[0] > 0 ? { pdf: x => a[0] * Math.exp(-a[0] * x), lo: 0, hi: Infinity } : null; + } +} + +/** An expression decomposed as Σ terms[name]·name + c, coefficients free of + * random variables. Affine forms over closed families have exact laws. */ +interface Affine { + terms: Map; + c: Expr; +} + +/** An exact law: a closed-form pdf the shader draws, or a uniform-sum + * convolution evaluated as an exact piecewise polynomial per parameters. */ +export type Law = + | { kind: 'dist'; dist: BaseDist } + | { kind: 'usum'; terms: Array<{ c: Expr; lo: Expr; hi: Expr }>; d: Expr }; + +/** The numeric value of a constant-folded expression, or null (frees, NaN). */ +function numOf(e: Expr): number | null { + try { + const v = evaluate(e, {}); + return isFinite(v) ? v : null; + } catch { + return null; + } +} + +// --- exact piecewise-polynomial densities (uniform convolutions) --- +// +// Sums of independent uniforms stay piecewise polynomial forever: uniform is +// degree 0 and each convolution raises the degree by one and merges +// breakpoints (X + Y is the triangle, four terms the Irwin–Hall cubic). The +// integrands are polynomials, so "symbolic integration" here is the power +// rule; all the real work is the support bookkeeping below. Coefficients are +// plain numbers evaluated per parameter values — breakpoint *ordering* +// depends on slider values, so a symbolic form would need case analysis the +// numeric one sidesteps. + +/** Density polynomial (ascending coefficients) per interval between breaks. */ +interface PPoly { + breaks: number[]; + pieces: number[][]; +} + +const padd = (a: number[], b: number[]): number[] => { + const out = new Array(Math.max(a.length, b.length)).fill(0); + a.forEach((v, i) => { out[i] += v; }); + b.forEach((v, i) => { out[i] += v; }); + return out; +}; + +const pscale = (a: number[], k: number): number[] => a.map(v => v * k); + +const pmul = (a: number[], b: number[]): number[] => { + const out = new Array(Math.max(1, a.length + b.length - 1)).fill(0); + for (let i = 0; i < a.length; i++) { + for (let j = 0; j < b.length; j++) out[i + j] += a[i] * b[j]; + } + return out; +}; + +const peval = (a: number[], x: number): number => { + let v = 0; + for (let i = a.length - 1; i >= 0; i--) v = v * x + a[i]; + return v; +}; + +/** Antiderivative with constant 0. */ +const pint = (a: number[]): number[] => [0, ...a.map((v, i) => v / (i + 1))]; + +/** Compose: p(k·y + t) as a polynomial in y. */ +const plin = (p: number[], k: number, t: number): number[] => { + let out: number[] = [0]; + for (let j = p.length - 1; j >= 0; j--) out = padd(pmul(out, [t, k]), [p[j]]); + return out; +}; + +const uniformPP = (lo: number, hi: number): PPoly => + ({ breaks: [lo, hi], pieces: [[1 / (hi - lo)]] }); + +/** The density of c·X + t for X with density p (c ≠ 0). */ +function scalePP(p: PPoly, c: number, t: number): PPoly { + const breaks = p.breaks.map(b => c * b + t); + // g(y) = f((y − t)/c)/|c|; composing with the inverse map keeps polynomials. + const pieces = p.pieces.map(q => pscale(plin(q, 1 / c, -t / c), 1 / Math.abs(c))); + if (c < 0) { + breaks.reverse(); + pieces.reverse(); + } + return { breaks, pieces }; +} + +const binom = (n: number, k: number): number => { + let v = 1; + for (let i = 0; i < k; i++) v = (v * (n - i)) / (i + 1); + return v; +}; + +/** + * Exact convolution of two piecewise-polynomial densities. For a piece pair + * u on [a, b] and v on [c, e], h(z) = ∫ u(x)v(z−x) dx over + * x ∈ [max(a, z−e), min(b, z−c)]: the antiderivative W(x; z) is computed + * once per pair (a polynomial in x whose coefficients are polynomials in z), + * and on each output interval — the breakpoints are the pairwise support + * sums — each limit is either a constant or z + shift, so h is a polynomial. + */ +function convPP(p: PPoly, q: PPoly): PPoly { + interface Pair { a: number; b: number; c: number; e: number; Wx: number[][] } + const pairs: Pair[] = []; + const zb: number[] = []; + for (let i = 0; i + 1 < p.breaks.length; i++) { + for (let k = 0; k + 1 < q.breaks.length; k++) { + const a = p.breaks[i], b = p.breaks[i + 1]; + const c = q.breaks[k], e = q.breaks[k + 1]; + const u = p.pieces[i], v = q.pieces[k]; + // v(z−x) gathered by powers of x: Vx[r] is a polynomial in z. + const Vx: number[][] = []; + for (let kk = 0; kk < v.length; kk++) { + for (let r = 0; r <= kk; r++) { + Vx[r] ??= []; + Vx[r][kk - r] = (Vx[r][kk - r] ?? 0) + v[kk] * binom(kk, r) * (r % 2 ? -1 : 1); + } + } + // u(x)·v(z−x) by powers of x, then the antiderivative in x. + const Px: number[][] = []; + for (let i2 = 0; i2 < u.length; i2++) { + for (let r = 0; r < Vx.length; r++) { + if (!Vx[r]) continue; + Px[i2 + r] = padd(Px[i2 + r] ?? [], pscale(Vx[r], u[i2])); + } + } + const Wx: number[][] = [[]]; + for (let j = 0; j < Px.length; j++) Wx[j + 1] = Px[j] ? pscale(Px[j], 1 / (j + 1)) : []; + pairs.push({ a, b, c, e, Wx }); + zb.push(a + c, a + e, b + c, b + e); + } + } + zb.sort((x, y) => x - y); + const eps = Math.max(1e-300, (zb[zb.length - 1] - zb[0]) * 1e-12); + const zs = zb.filter((z, i) => i === 0 || z - zb[i - 1] > eps); + const breaks: number[] = [zs[0]]; + const pieces: number[][] = []; + for (let s = 0; s + 1 < zs.length; s++) { + const mid = (zs[s] + zs[s + 1]) / 2; + let acc: number[] = [0]; + for (const pr of pairs) { + if (mid <= pr.a + pr.c || mid >= pr.b + pr.e) continue; + // W at a limit that is either constant or z + shift, as a poly in z. + const wAt = (constX: number | null, shift: number): number[] => { + let out: number[] = []; + let xp = 1; + let pw: number[] = [1]; + for (let j = 0; j < pr.Wx.length; j++) { + if (pr.Wx[j].length) { + out = padd(out, constX !== null ? pscale(pr.Wx[j], xp) : pmul(pr.Wx[j], pw)); + } + if (constX !== null) xp *= constX; + else pw = pmul(pw, [shift, 1]); + } + return out; + }; + const upper = pr.b <= mid - pr.c ? wAt(pr.b, 0) : wAt(null, -pr.c); + const lower = pr.a >= mid - pr.e ? wAt(pr.a, 0) : wAt(null, -pr.e); + acc = padd(acc, padd(upper, pscale(lower, -1))); + } + breaks.push(zs[s + 1]); + pieces.push(acc); + } + return { breaks, pieces }; +} + +/** Exact CDF at x. */ +function cdfPP(p: PPoly, x: number): number { + let acc = 0; + for (let i = 0; i + 1 < p.breaks.length; i++) { + if (x <= p.breaks[i]) break; + const hi = Math.min(x, p.breaks[i + 1]); + const F = pint(p.pieces[i]); + acc += peval(F, hi) - peval(F, p.breaks[i]); + } + return acc; +} + +/** The exact curve as a polyline: pieces sampled densely, with every true + * breakpoint emitted so kinks stay corners instead of KDE shoulders. */ +function curvePP(p: PPoly): number[] { + const span = p.breaks[p.breaks.length - 1] - p.breaks[0]; + const pts: number[] = []; + for (let i = 0; i + 1 < p.breaks.length; i++) { + const x0 = p.breaks[i], x1 = p.breaks[i + 1]; + const n = Math.max(2, Math.ceil(((x1 - x0) / span) * 256)); + for (let k = 0; k <= n; k++) { + const x = x0 + ((x1 - x0) * k) / n; + pts.push(x, peval(p.pieces[i], x)); + } + } + return pts; +} + +/** + * The declared random variables of a document plus their sample columns. + * The instance persists across recompiles (reset() clears declarations, not + * caches); cache entries carry the serialized definition and the values of + * the constants the variable (transitively) references, so a slider drag + * recomputes only the variables it touches, a static scene never resamples, + * and an edited definition can never serve stale samples. + */ +export class RVSystem { + private rvs = new Map(); + private cache = new Map(); + private paramsMemo = new Map>(); + private defSigMemo = new Map(); + private affineMemo = new Map(); + private lawMemo = new Map(); + private groundedMemo = new Map(); + + /** Start a recompile: drop declarations, keep sample caches. */ + reset(): void { + this.rvs.clear(); + this.paramsMemo.clear(); + this.defSigMemo.clear(); + this.affineMemo.clear(); + this.lawMemo.clear(); + this.groundedMemo.clear(); + } + + add(rv: RV): void { + this.rvs.set(rv.name, rv); + } + + delete(name: string): void { + this.rvs.delete(name); + } + + has(name: string): boolean { + return this.rvs.has(name); + } + + get(name: string): RV | undefined { + return this.rvs.get(name); + } + + size(): number { + return this.rvs.size; + } + + /** End a recompile: drop cached samples of variables no longer declared. */ + prune(): void { + for (const k of [...this.cache.keys()]) { + if (!this.rvs.has(k)) this.cache.delete(k); + } + } + + /** Detect definition cycles; returns per-variable errors. */ + validate(): Map { + const broken = new Map(); + const state = new Map(); + const visit = (name: string, path: string[]): void => { + const rv = this.rvs.get(name); + if (!rv || state.get(name) === 'done') return; + if (state.get(name) === 'visiting') { + const cycle = path.slice(path.indexOf(name)).concat(name); + for (const cn of cycle) broken.set(cn, `${cycle.join(' → ')} is circular.`); + return; + } + state.set(name, 'visiting'); + if (rv.kind === 'derived') { + for (const dep of freeVars(rv.expr)) { + if (this.rvs.has(dep)) visit(dep, [...path, name]); + } + } + state.set(name, 'done'); + }; + for (const name of this.rvs.keys()) visit(name, []); + return broken; + } + + /** Non-random free names the variable depends on, transitively (may include 't'). */ + paramsOf(name: string): ReadonlySet { + const memo = this.paramsMemo.get(name); + if (memo) return memo; + const out = new Set(); + const seen = new Set(); + const walk = (n: string): void => { + if (seen.has(n)) return; + seen.add(n); + const rv = this.rvs.get(n); + if (!rv) return; + const frees = rv.kind === 'base' + ? rv.dist.args.reduce((s, a) => freeVars(a, s), new Set()) + : freeVars(rv.expr); + for (const f of frees) { + if (this.rvs.has(f)) walk(f); + else out.add(f); + } + }; + walk(name); + this.paramsMemo.set(name, out); + return out; + } + + /** Decompose an expression as an affine form over *base normal* names, or + * null where that fails (nonlinear use, or a non-normal base involved). */ + private affine(e: Expr): Affine | null { + const rvFree = (x: Expr): boolean => ![...freeVars(x)].some(n => this.rvs.has(n)); + if (rvFree(e)) return { terms: new Map(), c: e }; + const scale = (af: Affine | null, k: Expr): Affine | null => af && { + terms: new Map([...af.terms].map(([n, coef]) => [n, bin('*', k, coef)])), + c: bin('*', k, af.c), + }; + switch (e.kind) { + case 'var': { + const rv = this.rvs.get(e.name)!; + if (rv.kind === 'base') return { terms: new Map([[e.name, num(1)]]), c: num(0) }; + return this.affineOf(e.name); + } + case 'neg': + return scale(this.affine(e.a), num(-1)); + case 'bin': { + if (e.op === '+' || e.op === '-') { + const a = this.affine(e.a); + const b = e.op === '-' ? scale(this.affine(e.b), num(-1)) : this.affine(e.b); + if (!a || !b) return null; + const terms = new Map(a.terms); + for (const [n, coef] of b.terms) { + const prev = terms.get(n); + terms.set(n, prev ? bin('+', prev, coef) : coef); + } + return { terms, c: bin('+', a.c, b.c) }; + } + if (e.op === '*') { + if (rvFree(e.a)) return scale(this.affine(e.b), e.a); + if (rvFree(e.b)) return scale(this.affine(e.a), e.b); + return null; // X·Y: a product distribution, not affine + } + if (e.op === '/' && rvFree(e.b)) { + return scale(this.affine(e.a), bin('/', num(1), e.b)); + } + return null; // X^k and friends + } + default: + return null; // calls, piecewise, … over random variables: sampled path + } + } + + private affineOf(name: string): Affine | null { + if (this.affineMemo.has(name)) return this.affineMemo.get(name)!; + const rv = this.rvs.get(name); + const af = rv?.kind === 'derived' ? this.affine(rv.expr) : null; + this.affineMemo.set(name, af); + return af; + } + + /** + * The exact law of a variable, when one is derivable. Base declarations + * pass through. A derived variable affine in independent bases (shared + * names accumulate into one coefficient first — the covariance accounting: + * var(aX + bX) = (a+b)²σ²) reduces by family: + * + * - all normal → normal, mean Σcᵢμᵢ + d, sd √(Σ(cᵢσᵢ)²); + * - one term → the base transformed: c·U(lo,hi)+d is Uniform again (min/max + * endpoints keep a negative or slider-driven c honest), c·Exp(λ) with a + * positive literal c is Exponential(λ/c); + * - several uniform terms → 'usum', an exact piecewise-polynomial + * convolution (the triangle, Irwin–Hall, trapezoids) evaluated per + * parameter values. + * + * Null means "estimate from samples" (nonlinear transforms, products, + * mixed families). + */ + exactLaw(name: string): Law | null { + const memo = this.lawMemo.get(name); + if (memo !== undefined) return memo; + const law = this.deriveLaw(name); + this.lawMemo.set(name, law); + return law; + } + + private deriveLaw(name: string): Law | null { + const rv = this.rvs.get(name); + if (!rv) return null; + if (rv.kind === 'base') return { kind: 'dist', dist: rv.dist }; + const af = this.affineOf(name); + if (!af || !af.terms.size) return null; + const bases = [...af.terms].map(([n, coef]) => ({ + coef, + dist: (this.rvs.get(n) as RV & { kind: 'base' }).dist, + })); + if (bases.every(b => b.dist.kind === 'normal')) { + let mean = af.c; + let variance: Expr | null = null; + for (const { coef, dist } of bases) { + mean = bin('+', mean, bin('*', coef, dist.args[0])); + const term = bin('^', bin('*', coef, dist.args[1]), num(2)); + variance = variance ? bin('+', variance, term) : term; + } + const sd: Expr = { kind: 'call', name: 'sqrt', args: [variance!] }; + return { kind: 'dist', dist: { kind: 'normal', args: [mean, sd] } }; + } + if (bases.length === 1) { + const { coef, dist } = bases[0]; + if (dist.kind === 'uniform') { + const e1 = bin('+', bin('*', coef, dist.args[0]), af.c); + const e2 = bin('+', bin('*', coef, dist.args[1]), af.c); + const args: Expr[] = [ + { kind: 'call', name: 'min', args: [e1, e2] }, + { kind: 'call', name: 'max', args: [e1, e2] }, + ]; + return { kind: 'dist', dist: { kind: 'uniform', args } }; + } + if (dist.kind === 'exponential') { + // Only c·X with a positive literal c stays exponential (rate λ/c); + // a shift or flip leaves the family, and a slider c could flip live. + const c = numOf(coef); + if (c !== null && c > 0 && numOf(af.c) === 0) { + const rate = c === 1 ? dist.args[0] : bin('/', dist.args[0], num(c)); + return { kind: 'dist', dist: { kind: 'exponential', args: [rate] } }; + } + return null; + } + } + if (bases.every(b => b.dist.kind === 'uniform')) { + return { + kind: 'usum', + terms: bases.map(b => ({ c: b.coef, lo: b.dist.args[0], hi: b.dist.args[1] })), + d: af.c, + }; + } + return null; + } + + /** The exact law when it is a closed-form pdf the shader can draw. */ + exactDist(name: string): BaseDist | null { + const law = this.exactLaw(name); + return law?.kind === 'dist' ? law.dist : null; + } + + /** Non-random free names of a P(…) body, through the variables it references. */ + bodyParams(e: Expr): Set { + const out = new Set(); + for (const f of freeVars(e)) { + if (this.rvs.has(f)) for (const p of this.paramsOf(f)) out.add(p); + else out.add(f); + } + return out; + } + + /** Serialized definition of a variable *and its dependencies*, so editing + * `X ~ …` invalidates the cached samples of `Y = X + 1` too. */ + private defSig(name: string, seen = new Set()): string { + const memo = this.defSigMemo.get(name); + if (memo !== undefined) return memo; + if (seen.has(name)) return '@cycle'; // validate() reports it; keep sigs total + seen.add(name); + const rv = this.rvs.get(name); + let s: string; + if (!rv) s = '@missing'; + else if (rv.kind === 'base') s = rv.dist.kind + JSON.stringify(rv.dist.args); + else { + s = JSON.stringify(rv.expr); + for (const dep of [...freeVars(rv.expr)].sort()) { + if (this.rvs.has(dep)) s += `|${dep}:${this.defSig(dep, seen)}`; + } + } + this.defSigMemo.set(name, s); + return s; + } + + private sig(rv: RV, env: Record): string { + let s = this.defSig(rv.name); + for (const p of [...this.paramsOf(rv.name)].sort()) { + if (!(p in env)) throw new Error(`Unbound variable: ${p}`); + s += `;${p}=${env[p]}`; + } + return s; + } + + /** The cache slot for this variable at these parameter values; a stale + * signature drops every derived field at once. */ + private entry(name: string, sig: string): CacheEntry { + const hit = this.cache.get(name); + if (hit && hit.sig === sig) return hit; + const fresh: CacheEntry = { sig }; + this.cache.set(name, fresh); + return fresh; + } + + /** The sample column for a variable under the given constants. */ + columns(name: string, env: Record): Float64Array { + const rv = this.rvs.get(name); + if (!rv) throw new Error(`${name} has an error in its definition.`); + const slot = this.entry(name, this.sig(rv, env)); + if (slot.col) return slot.col; + let col: Float64Array; + if (rv.kind === 'base') { + const u = uniformStream(name); + const a = rv.dist.args.map(e => evaluate(e, env)); + col = new Float64Array(SAMPLE_COUNT); + switch (rv.dist.kind) { + case 'normal': + if (a[1] > 0) for (let i = 0; i < SAMPLE_COUNT; i++) col[i] = a[0] + a[1] * normalQuantile(u[i]); + else col.fill(NaN); + break; + case 'uniform': + if (a[1] > a[0]) for (let i = 0; i < SAMPLE_COUNT; i++) col[i] = a[0] + (a[1] - a[0]) * u[i]; + else col.fill(NaN); + break; + case 'exponential': + if (a[0] > 0) for (let i = 0; i < SAMPLE_COUNT; i++) col[i] = -Math.log(1 - u[i]) / a[0]; + else col.fill(NaN); + break; + } + } else { + const cols = new Map(); + for (const dep of freeVars(rv.expr)) { + if (this.rvs.has(dep)) cols.set(dep, this.columns(dep, env)); + } + col = evalCols(rv.expr, cols, env, SAMPLE_COUNT); + } + slot.col = col; + return col; + } + + /** Numeric piecewise polynomial of a usum law at these parameter values. */ + private usumPP(law: Law & { kind: 'usum' }, env: Record): PPoly | null { + let acc: PPoly | null = null; + const shift = evaluate(law.d, env); + if (!isFinite(shift)) return null; + for (const t of law.terms) { + const c = evaluate(t.c, env); + const lo = evaluate(t.lo, env); + const hi = evaluate(t.hi, env); + if (!isFinite(c) || !(hi > lo)) return null; + if (c === 0) continue; // a slider zeroed this term: it contributes nothing + const box = scalePP(uniformPP(lo, hi), c, 0); + acc = acc ? convPP(acc, box) : box; + } + if (!acc) return null; // every coefficient zero: a constant, nothing to draw + return shift !== 0 ? scalePP(acc, 1, shift) : acc; + } + + /** + * The density curve for a variable: the *exact* piecewise polynomial when + * its law is a uniform convolution, the sample estimate otherwise. (Rows + * whose law is a closed-form pdf never come here — they draw through the + * shader.) + */ + curve(name: string, env: Record): DensityCurve | null { + const law = this.exactLaw(name); + if (law?.kind === 'usum') { + const rv = this.rvs.get(name)!; + const slot = this.entry(name, this.sig(rv, env)); + if (slot.curve === undefined) { + const pp = this.usumPP(law, env); + const m = this.exactMoments(name, env); + slot.curve = pp && m ? { pts: curvePP(pp), mean: m.mean, sd: m.sd, mass: 1 } : null; + } + return slot.curve; + } + const col = this.columns(name, env); + const entry = this.cache.get(name)!; + if (entry.curve === undefined) entry.curve = estimateCurve(col); + return entry.curve; + } + + /** Exact mean and sd under the variable's law, or null when sampled. */ + exactMoments(name: string, env: Record): { mean: number; sd: number } | null { + const law = this.exactLaw(name); + if (!law) return null; + if (law.kind === 'dist') { + const a = law.dist.args.map(e => evaluate(e, env)); + switch (law.dist.kind) { + case 'normal': + return a[1] > 0 ? { mean: a[0], sd: a[1] } : null; + case 'uniform': + return a[1] > a[0] ? { mean: (a[0] + a[1]) / 2, sd: (a[1] - a[0]) / Math.sqrt(12) } : null; + case 'exponential': + return a[0] > 0 ? { mean: 1 / a[0], sd: 1 / a[0] } : null; + } + } + let mean = evaluate(law.d, env); + let variance = 0; + for (const t of law.terms) { + const c = evaluate(t.c, env); + const lo = evaluate(t.lo, env); + const hi = evaluate(t.hi, env); + if (!isFinite(c) || !(hi > lo)) return null; + mean += (c * (lo + hi)) / 2; + variance += (c * (hi - lo)) ** 2 / 12; + } + return isFinite(mean) ? { mean, sd: Math.sqrt(variance) } : null; + } + + /** The variable's expression with every derived dependency inlined, so + * only base variables and parameters remain — or null when the expansion + * blows up (shared subtrees duplicate under substitution). */ + private grounded(name: string): Expr | null { + const memo = this.groundedMemo.get(name); + if (memo !== undefined) return memo; + const rv = this.rvs.get(name); + let e: Expr | null = rv?.kind === 'derived' ? rv.expr : null; + // Dependencies are acyclic (validate() dropped cycles); the guard bounds + // pathological chains and substitution blowup all the same. + for (let guard = 0; e && guard < 32; guard++) { + const sub: Record = {}; + for (const n of freeVars(e)) { + const dep = this.rvs.get(n); + if (dep?.kind === 'derived') sub[n] = dep.expr; + } + if (!Object.keys(sub).length) break; + e = substVars(e, sub); + if (JSON.stringify(e).length > 200_000) e = null; + } + this.groundedMemo.set(name, e); + return e; + } + + /** + * Moments by numeric integration against the base law: for Y = g(X) with a + * single base dependency, E[g(X)] = ∫ g(x)·pdf(x) dx by adaptive + * quadrature (integrate.ts) — ~9 significant digits where the sample mean + * gives ~3. Undefined regions of g drop out of both the numerator and the + * mass, matching the sampler's "average where defined" convention. Null + * for joint dependence (E[X·Y] still samples), broken parameters, or + * integrals that fail to settle (heavy tails). + */ + quadMoments(name: string, env: Record): { mean: number; sd: number; mass: number } | null { + const rv = this.rvs.get(name); + if (rv?.kind !== 'derived') return null; // base laws: exactMoments has them + const slot = this.entry(name, this.sig(rv, env)); + if (slot.qm !== undefined) return slot.qm; + const compute = (): { mean: number; sd: number; mass: number } | null => { + const g = this.grounded(name); + if (!g) return null; + const bases = [...freeVars(g)].filter(n => this.rvs.has(n)); + const base = bases.length === 1 ? this.rvs.get(bases[0])! : null; + if (base?.kind !== 'base') return null; + let pc: ReturnType; + try { + pc = pdfClosure(base.dist, env); + } catch { + return null; + } + if (!pc) return null; + const gAt = (x: number): number => { + try { + return evaluate(g, { ...env, [base.name]: x }); + } catch { + return NaN; + } + }; + const moment = (k: 0 | 1 | 2) => quadrature(x => { + const v = gAt(x); + if (!isFinite(v)) return 0; + return (k === 0 ? 1 : k === 1 ? v : v * v) * pc!.pdf(x); + }, pc!.lo, pc!.hi); + const mass = moment(0); + if (!(mass > 1e-9)) return null; + const m1 = moment(1); + const m2 = moment(2); + if (!isFinite(m1) || !isFinite(m2)) return null; + const mean = m1 / mass; + return { mean, sd: Math.sqrt(Math.max(m2 / mass - mean * mean, 0)), mass }; + }; + slot.qm = compute(); + return slot.qm; + } + + /** The mean of a variable: exact under its law when one is derivable, + * quadrature against the base pdf for one-variable transforms, otherwise + * the finite-sample mean (NaN when nothing is finite). */ + mean(name: string, env: Record): number { + const m = this.exactMoments(name, env) ?? this.quadMoments(name, env); + if (m) return m.mean; + const col = this.columns(name, env); + let sum = 0; + let n = 0; + for (let i = 0; i < col.length; i++) { + if (isFinite(col[i])) { + sum += col[i]; + n++; + } + } + return n ? sum / n : NaN; + } + + /** Exact P(lo < name < hi) under the variable's law, or null when sampled. */ + exactProbability( + name: string, + lo: Expr | undefined, + hi: Expr | undefined, + env: Record, + ): number | null { + const law = this.exactLaw(name); + if (!law) return null; + if (law.kind === 'dist') return probabilityValue(law.dist, lo, hi, env); + const pp = this.usumPP(law, env); + if (!pp) return NaN; + const total = cdfPP(pp, pp.breaks[pp.breaks.length - 1]); + return (hi ? cdfPP(pp, evaluate(hi, env)) : total) - (lo ? cdfPP(pp, evaluate(lo, env)) : 0); + } + + /** + * Monte Carlo estimate of P(body): the fraction of joint samples where the + * inequality holds. Samples where it is undefined count as "not the event"; + * NaN when it is undefined everywhere (broken parameters). + */ + probability(body: Expr, env: Record): number { + const cols = new Map(); + for (const f of freeVars(body)) { + if (this.rvs.has(f)) cols.set(f, this.columns(f, env)); + } + const mask = evalCols(body, cols, env, SAMPLE_COUNT); + let count = 0; + let defined = 0; + for (let i = 0; i < SAMPLE_COUNT; i++) { + if (!Number.isNaN(mask[i])) { + defined++; + if (mask[i] === 1) count++; + } + } + return defined ? count / SAMPLE_COUNT : NaN; + } +} + +// --- building the system from scanned rows --- + +export interface BuildRVOpts { + fnNames: ReadonlySet; + getFn: GetFn; + ropts?: ResolveOpts; + constNames: ReadonlySet; + /** Name already claimed by a definition row (constant, function, field, …). */ + taken: (name: string) => boolean; +} + +export interface BuiltRVs { + /** Every declared name, healthy or not — the set P(…) and bare rows resolve against. */ + names: ReadonlySet; + /** Row index → the variable it declares. */ + rowRV: Map; + /** Row index → error message, for rows whose declaration failed. */ + errors: Map; +} + +/** + * Rebuild `sys` from the scanned rows: parse and resolve each declaration, + * reject name collisions, then drop definition cycles and everything that + * depends on a failed variable, reporting per-row errors. Shared by the app + * and the worker so both accept exactly the same documents. + */ +export function buildRVSystem(sys: RVSystem, scan: ReturnType, opts: BuildRVOpts): BuiltRVs { + sys.reset(); + const names = new Set([...scan.base.values(), ...scan.derived.values()].map(d => d.name)); + const rowRV = new Map(); + const errors = new Map(); + const rowOf = new Map(); + + const claim = (i: number, name: string): boolean => { + rowRV.set(i, name); + if (RESERVED.has(name) || builtinFn(name)) { + errors.set(i, `Cannot use ${name} as a random variable name.`); + } else if (sys.has(name) || opts.taken(name)) { + errors.set(i, `${name} is already defined.`); + } else { + rowOf.set(name, i); + return true; + } + return false; + }; + + for (const [i, { name, rhs }] of scan.base) { + if (!claim(i, name)) continue; + try { + const d = parseDistribution(rhs, opts.fnNames); + d.args = d.args.map(a => resolveExpr(a, opts.getFn, opts.ropts)); + for (const a of d.args) { + for (const f of freeVars(a)) { + if (names.has(f)) throw new Error('Distribution parameters cannot depend on a random variable.'); + } + checkDerived(a, new Set(), opts.constNames); + } + sys.add({ name, kind: 'base', dist: d }); + } catch (e) { + errors.set(i, e instanceof Error ? e.message : String(e)); + } + } + for (const [i, { name, rhs }] of scan.derived) { + if (!claim(i, name)) continue; + try { + const expr = resolveExpr(parseExpr(rhs, opts.fnNames), opts.getFn, opts.ropts); + checkDerived(expr, names, opts.constNames); + sys.add({ name, kind: 'derived', expr }); + } catch (e) { + errors.set(i, e instanceof Error ? e.message : String(e)); + } + } + + // Cycles, then the ripple: a variable whose dependency failed fails too. + const failed = sys.validate(); + for (const name of failed.keys()) sys.delete(name); + let changed = true; + while (changed) { + changed = false; + for (const name of rowOf.keys()) { + const rv = sys.get(name); + if (rv?.kind !== 'derived') continue; + for (const dep of freeVars(rv.expr)) { + if (names.has(dep) && !sys.has(dep)) { + failed.set(name, `${dep} has an error in its definition.`); + sys.delete(name); + changed = true; + break; + } + } + } + } + for (const [name, message] of failed) { + const row = rowOf.get(name); + if (row !== undefined && !errors.has(row)) errors.set(row, message); + } + return { names, rowRV, errors }; } diff --git a/lib/expr.ts b/lib/expr.ts index 91a6c81..8049ecc 100644 --- a/lib/expr.ts +++ b/lib/expr.ts @@ -45,8 +45,8 @@ export const FUNCTIONS = new Set([ // Small-matrix helpers (det, trace, matvec, linear solve), also lowered // symbolically — Cramer's rule for 2×2 and 3×3 (see mat.ts). 'det', 'trace', 'solve', - // Not real functions: Σ/Π binders, expanded symbolically by resolveExpr. - 'sum', 'prod', + // Not real functions: Σ/Π/∫ binders, expanded symbolically by resolveExpr. + 'sum', 'prod', 'int', // Whole-expression plot modes (see classify): domain coloring, conformal // grids, escape-time iteration, and swept tubes. 'domain', 'conformal', 'iter', 'tube', @@ -234,6 +234,7 @@ const ops = operators({ if (a?.kind !== 'var' || !isFnName(a.name)) throw new Error('Expected a function name.'); const name = canonicalFn(a.name); if (name === 'sum' || name === 'prod') return sumCall(name, b); + if (name === 'int') return intCall(b); // Tuple literals inside a call flatten into the argument list, so // tube((a, b, c)) === tube(a, b, c) and |(3, 4)| reaches abs as (3, 4); // geometry statements re-pair adjacent scalars into points (lib/geom.ts). @@ -245,6 +246,27 @@ const ops = operators({ const isRange = (e: Expr): e is Expr & { kind: 'call' } => e.kind === 'call' && e.name === '[range]'; +/** + * Shape an ∫ into a call node: args are [lo, hi] for the header form + * `int[a..b] …` (body bound from its product chain, like Σ), [body] for the + * indefinite `int(f dx)`, and [lo, hi, body] for `int(a..b, f dx)`. + * resolveExpr integrates all of them symbolically (or expands a quadrature). + */ +function intCall(b: PNode | null | undefined): Expr { + const usage = () => new Error('Expected int(f(x) dx) or int[a..b] f(x) dx.'); + if (!b || b.kind === 'popen') throw usage(); + const items = b.kind === 'series' ? b.items.map(asExpr) : [asExpr(b)]; + const ranges = items.filter(isRange); + const bodies = items.filter(x => !isRange(x)); + if (!items.length || ranges.length > 1 || bodies.length > 1) throw usage(); + const bounds = ranges.length ? [ranges[0].args[0], ranges[0].args[1]] : []; + if (!bodies.length) { + if (!bounds.length) throw usage(); + return { kind: 'call', name: 'int', args: bounds }; // header awaiting its body + } + return { kind: 'call', name: 'int', args: [...bounds, bodies[0]] }; +} + /** * Shape a Σ/Π header into a call node: args are [index, lo, hi] for the * header-only form `sum[n=1..N] …` and [index, lo, hi, body] for @@ -283,7 +305,7 @@ const syntax: PatternDict = { number: /^\d+\.?\d*$/, bar: /^\|$/, whitespace: /\s$/, - symbol: /^[A-Za-z_Σ∑Π∏][A-Za-z_0-9]*'*$/, + symbol: /^[A-Za-z_Σ∑Π∏∫∞][A-Za-z_0-9]*'*$/, operator: x => !!ops[x] || MULTI_CHAR_OPS.some(m => m.startsWith(x)), invalid(x) { throw new Error(`Invalid character: ${JSON.stringify(x)}.`); }, }; @@ -294,7 +316,8 @@ function op(str: string): Token { return { type: 'operator', str, line: -1, loc: [-1, -1] }; } -const SYMBOL_ALIASES: Record = { 'Σ': 'sum', '∑': 'sum', 'Π': 'prod', '∏': 'prod' }; +const SYMBOL_ALIASES: Record = + { 'Σ': 'sum', '∑': 'sum', 'Π': 'prod', '∏': 'prod', '∫': 'int', '∞': 'inf' }; /** * Map Σ/Π glyphs to sum/prod, and repair `1..N`: the greedy number match @@ -505,7 +528,7 @@ export function realPow(a: number, b: number): number { return NaN; // no small-denominator rational found: irrational-looking exponent } -const EVAL_FNS: Record number> = { +export const EVAL_FNS: Record number> = { sin: Math.sin, cos: Math.cos, tan: Math.tan, asin: Math.asin, acos: Math.acos, atan: Math.atan, atan2: Math.atan2, sinh: Math.sinh, cosh: Math.cosh, tanh: Math.tanh, diff --git a/lib/integrate.test.ts b/lib/integrate.test.ts new file mode 100644 index 0000000..58bf265 --- /dev/null +++ b/lib/integrate.test.ts @@ -0,0 +1,282 @@ +import { describe, expect, it } from 'vitest'; +import { antiderivative, quadrature, quadratureSum, verifyDefinite } from './integrate.ts'; +import { type Expr, evaluate, freeVars, parseExpr } from './expr.ts'; +import { diff } from './diff.ts'; +import { resolveExpr, usesIntegral } from './defs.ts'; + +/** Antiderivative of the parsed source, or null. */ +const F = (src: string, v = 'x'): Expr | null => antiderivative(parseExpr(src), v); + +/** Assert F′ = f numerically at a spread of points (the engine's own verifier + * already ran; this re-checks through an independent path). */ +function checkDeriv(src: string, v = 'x', points = [-2.3, -0.7, 0.41, 1.13, 2.9], env: Record = {}) { + const f = parseExpr(src); + const Fe = antiderivative(f, v); + expect(Fe, `no antiderivative found for ${src}`).not.toBeNull(); + const dF = diff(Fe!, v); + let compared = 0; + for (const x of points) { + let want: number; + try { + want = evaluate(f, { ...env, [v]: x }); + } catch { + continue; + } + if (!isFinite(want)) continue; + const got = evaluate(dF, { ...env, [v]: x }); + expect(got, `${src} at ${v} = ${x}`).toBeCloseTo(want, 5); + compared++; + } + expect(compared).toBeGreaterThan(2); +} + +describe('antiderivative: exact families', () => { + it('polynomials', () => { + checkDeriv('x^2'); + checkDeriv('3x^5 - 2x + 7'); + checkDeriv('(x^2 + 1)(x - 3)'); + }); + + it('polynomials with slider coefficients stay symbolic', () => { + const Fe = antiderivative(parseExpr('a x^2 + b'), 'x')!; + expect(Fe).not.toBeNull(); + expect(freeVars(Fe)).toEqual(new Set(['a', 'b', 'x'])); + // d/dx (a x³/3 + b x) = a x² + b at a = 2, b = −1. + expect(evaluate(diff(Fe, 'x'), { a: 2, b: -1, x: 1.7 })).toBeCloseTo(2 * 1.7 ** 2 - 1, 9); + }); + + it('rational functions: 1/x and friends', () => { + checkDeriv('1/x'); + checkDeriv('1/x^2'); + checkDeriv('(3x + 2)/(x^2 + 1)'); + checkDeriv('1/(x^2 - 1)'); + checkDeriv('1/(x^2 + x + 1)'); + checkDeriv('x^3/(x^2 - 4)', 'x', [-1.3, 0.4, 1.1, 3.2]); + }); + + it('rational functions with repeated factors (Hermite/Ostrogradsky part)', () => { + checkDeriv('1/(x - 2)^2', 'x', [-1, 0.5, 3.1, 4.7]); + checkDeriv('1/(x^2 - 1)^2', 'x', [-0.5, 0.3, 2.2, 3.1]); + checkDeriv('(x + 1)/(x^2 + 1)^2'); + }); + + it('table forms with linear arguments', () => { + checkDeriv('sin(x)'); + checkDeriv('cos(3x - 1)'); + checkDeriv('exp(2x)'); + checkDeriv('e^x'); + checkDeriv('tan(x)', 'x', [-0.6, 0.3, 0.9, 1.1]); + checkDeriv('sqrt(2x + 5)', 'x', [0.2, 1.3, 2.9]); + checkDeriv('ln(x)', 'x', [0.3, 1.2, 4.5]); + checkDeriv('atan(x)'); + checkDeriv('asin(x)', 'x', [-0.8, -0.2, 0.4, 0.7]); + checkDeriv('sinh(x) + cosh(2x)'); + checkDeriv('abs(x)'); + checkDeriv('erf(x)'); + checkDeriv('2^x'); + checkDeriv('normalpdf(x, 1, 2)'); + }); + + it('linear-argument coefficients may be sliders', () => { + const Fe = antiderivative(parseExpr('sin(a x + b)'), 'x')!; + expect(Fe).not.toBeNull(); + const dF = diff(Fe, 'x'); + for (const x of [-1.1, 0.4, 2.3]) { + expect(evaluate(dF, { a: 1.5, b: -0.7, x })) + .toBeCloseTo(Math.sin(1.5 * x - 0.7), 6); + } + }); + + it('gaussians via erf', () => { + checkDeriv('exp(-x^2)'); + checkDeriv('exp(-(x - 1)^2/2)'); + checkDeriv('exp(-3x^2 + 2x + 1)'); + }); + + it('u-substitution', () => { + checkDeriv('x exp(x^2)'); + checkDeriv('sin(x) cos(x)'); + checkDeriv('ln(x)/x', 'x', [0.4, 1.3, 3.7]); + checkDeriv('x/(x^2 + 1)^3'); + checkDeriv('cos(x) exp(sin(x))'); + checkDeriv('x sqrt(x^2 + 1)'); + }); + + it('integration by parts', () => { + checkDeriv('x sin(x)'); + checkDeriv('x^2 exp(x)'); + checkDeriv('x ln(x)', 'x', [0.5, 1.4, 3.3]); + checkDeriv('x exp(-x)'); + checkDeriv('x^3 cos(2x)'); + checkDeriv('x atan(x)'); + }); + + it('exp × trig', () => { + checkDeriv('exp(x) sin(x)'); + checkDeriv('exp(-x) cos(3x)'); + }); + + it('trig powers and products', () => { + checkDeriv('sin(x)^2'); + checkDeriv('cos(x)^2'); + checkDeriv('sin(x)^3'); + checkDeriv('sin(x)^2 cos(x)^2'); + checkDeriv('sin(x)^5 cos(x)^2'); + checkDeriv('sin(3x) cos(5x)'); + checkDeriv('sin(x) sin(2x)'); + }); + + it('returns null (not nonsense) where no rule applies', () => { + expect(F('exp(x^3)')).toBeNull(); + expect(F('sin(x)/x')).toBeNull(); + expect(F('exp(x)/x')).toBeNull(); + expect(F('x^x')).toBeNull(); + }); +}); + +describe('quadrature', () => { + const q = (src: string, lo: number, hi: number) => { + const e = parseExpr(src); + return quadrature(x => evaluate(e, { x }), lo, hi); + }; + + it('matches known integrals', () => { + expect(q('x^2', 0, 1)).toBeCloseTo(1 / 3, 9); + expect(q('sin(x)', 0, Math.PI)).toBeCloseTo(2, 9); + expect(q('exp(-x^2)', -6, 6)).toBeCloseTo(Math.sqrt(Math.PI), 8); + expect(q('1/x', 1, Math.E)).toBeCloseTo(1, 9); + expect(q('sin(x)/x', 0, 10)).toBeCloseTo(1.658347594218874, 7); // Si(10) + }); + + it('is sign-aware for reversed bounds', () => { + expect(q('x^2', 1, 0)).toBeCloseTo(-1 / 3, 9); + }); + + it('handles integrable endpoint singularities', () => { + expect(q('1/sqrt(x)', 0, 1)).toBeCloseTo(2, 4); + }); + + it('reports NaN for divergent integrals instead of a confident number', () => { + expect(q('1/x^2', -1, 1)).toBeNaN(); + }); + + it('transforms infinite ranges', () => { + expect(q('exp(-x^2)', -Infinity, Infinity)).toBeCloseTo(Math.sqrt(Math.PI), 7); + expect(q('exp(-x)', 0, Infinity)).toBeCloseTo(1, 8); + expect(q('1/x^2', 1, Infinity)).toBeCloseTo(1, 7); + expect(q('exp(x)', -Infinity, 0)).toBeCloseTo(1, 8); + expect(q('1/x^2', Infinity, 1)).toBeCloseTo(-1, 7); // reversed + }); +}); + +describe('verifyDefinite', () => { + it('accepts a true closed form and rejects FTC across a pole', () => { + const body = parseExpr('x^2'); + const good = parseExpr('1/3'); + expect(verifyDefinite(good, body, 'x', parseExpr('0'), parseExpr('1'))).toBe(true); + // ∫₋₁¹ x⁻² is divergent; naive FTC gives −2. + const pole = parseExpr('1/x^2'); + const naive = parseExpr('0 - 2'); + expect(verifyDefinite(naive, pole, 'x', parseExpr('0 - 1'), parseExpr('1'))).toBe(false); + }); +}); + +describe('∫ syntax through resolveExpr', () => { + const r = (s: string, fns: ReadonlySet = new Set()) => + resolveExpr(parseExpr(s, fns), () => undefined); + const val = (s: string, env: Record = {}) => evaluate(r(s), env); + + it('integrates definite forms exactly', () => { + expect(val('int[0..2] x^2 dx')).toBeCloseTo(8 / 3, 12); + expect(val('int[1..e] 1/x dx')).toBeCloseTo(1, 9); + expect(val('∫[0..pi] sin(x) dx')).toBeCloseTo(2, 12); + expect(val('int[2..0] x^2 dx')).toBeCloseTo(-8 / 3, 12); // reversed bounds + expect(val('int[0..1] dx')).toBe(1); + }); + + it('keeps slider parameters symbolic', () => { + const e = r('int[0..a] x^2 dx'); + expect(freeVars(e)).toEqual(new Set(['a'])); + expect(evaluate(e, { a: 3 })).toBeCloseTo(9, 9); + }); + + it('stays symbolic when the bound variable is also the ambient plot variable', () => { + // int[a..x] cos(x) dx: the dx binds the body's x, the bound's x is the + // plot coordinate. The closed form sin(x) − sin(a) must survive + // verification — at x = 150 a fixed-order quadrature fallback would be + // wildly wrong (this is the visible regression: a bounded sine wave, not + // hundred-unit spikes). + const e = r('int[a..x] cos(x) dx'); + expect(freeVars(e)).toEqual(new Set(['a', 'x'])); + for (const x of [1.3, 42, 150, -180]) { + expect(evaluate(e, { a: 0.5, x })).toBeCloseTo(Math.sin(x) - Math.sin(0.5), 8); + } + }); + + it('produces functions from indefinite and variable-bound forms', () => { + expect(val('int(x^2 dx)', { x: 3 })).toBeCloseTo(9, 12); + expect(val('∫ cos(x) dx', { x: 1 })).toBeCloseTo(Math.sin(1), 12); // bare chain form + expect(val('int[0..x] exp(-t^2) dt', { x: 1 })).toBeCloseTo(0.7468241328, 5); + }); + + it('falls back to a quadrature sum when no closed form is found', () => { + // Si(x) is not elementary; the expansion is an ordinary expression in x. + const si = r('int[0..x] sin(t)/t dt'); + expect(freeVars(si)).toEqual(new Set(['x'])); + expect(evaluate(si, { x: 10 })).toBeCloseTo(1.658347594218874, 5); + }); + + it('supports the classic notations', () => { + expect(val('int[0..1] dx/(1+x^2)')).toBeCloseTo(Math.PI / 4, 9); // measure first + expect(val('int[0..1] int[0..y] x dx dy')).toBeCloseTo(1 / 6, 9); // iterated, inside-out + expect(val('int[0..1] int[0..y] x dx y dy')).toBeCloseTo(1 / 8, 9); + expect(val('2 int[0..1] x dx + 1')).toBeCloseTo(2, 12); // coefficient stays outside + expect(val('int[0..1] sum[n=1..3] x^n dx')).toBeCloseTo(1 / 2 + 1 / 3 + 1 / 4, 9); + expect(val('int[0..2] (d/dt (t^2)) dt')).toBeCloseTo(4, 12); // d/dt consumes its own dt + }); + + it('supports ±inf bounds (symbolic limits, transformed quadrature fallback)', () => { + // The limit of the antiderivative replaces the infinite end… + expect(val('int[-inf..x] exp(t) dt', { x: 1.3 })).toBeCloseTo(Math.exp(1.3), 6); + expect(val('int[0..inf] exp(-x) dx')).toBeCloseTo(1, 8); + expect(val('int[1..inf] 1/x^2 dx')).toBeCloseTo(1, 6); + expect(val('int[-inf..inf] exp(-x^2) dx')).toBeCloseTo(Math.sqrt(Math.PI), 6); + expect(val('int[inf..1] 1/x^2 dx')).toBeCloseTo(-1, 6); // downhill range + // …the ∞ glyph works, and the normal cdf falls out of its pdf… + expect(val('∫[-∞..x] normalpdf(t, 0, 1) dt', { x: 0.7 })).toBeCloseTo(0.7580363, 5); + // …and a non-elementary tail still evaluates via the transformed sum: + // ∫₀^∞ e^(−x)/(1+x) dx = e·E₁(1). + expect(val('int[0..inf] exp(-x)/(1+x) dx')).toBeCloseTo(0.5963474, 4); + }); + + it('rejects malformed integrals with usable messages', () => { + expect(() => r('int[0..2]')).toThrow(/needs a body/); + expect(() => r('int(x^2, x)')).toThrow(/int\(f\(x\) dx\)/); + expect(() => r('int[0..1] x^2 dq')).not.toThrow(); // dq: integrand constant in q + expect(() => r('int(x^2)')).toThrow(/dx factor/); + }); + + it('refuses FTC across a non-integrable pole (falls back to quadrature)', () => { + // ∫₋₁¹ x⁻² diverges; the naive closed form −1/x would report −2. The + // fallback sum reports a large positive number — wrong-magnitude beats + // confidently negative for a positive integrand. + expect(val('int[-1..1] 1/x^2 dx')).toBeGreaterThan(0); + }); + + it('flags ∫ usage on the parse for row readouts', () => { + expect(usesIntegral(parseExpr('int[0..1] x^2 dx'))).toBe(true); + expect(usesIntegral(parseExpr('∫ x dx'))).toBe(true); + expect(usesIntegral(parseExpr('y = sin(x)'))).toBe(false); + }); +}); + +describe('quadratureSum', () => { + it('expands to a plain expression whose value matches the integral', () => { + const sum = quadratureSum(parseExpr('exp(-x^2)'), 'x', parseExpr('0'), parseExpr('2')); + expect(evaluate(sum, {})).toBeCloseTo(0.8820813907624215, 8); // √π/2·erf(2) + // Variable upper bound: a function of x. + const si = quadratureSum(parseExpr('sin(q)/q'), 'q', parseExpr('0'), parseExpr('x')); + expect(freeVars(si)).toEqual(new Set(['x'])); + expect(evaluate(si, { x: 10 })).toBeCloseTo(1.658347594218874, 5); + }); +}); diff --git a/lib/integrate.ts b/lib/integrate.ts new file mode 100644 index 0000000..86e5c79 --- /dev/null +++ b/lib/integrate.ts @@ -0,0 +1,1154 @@ +/** + * Symbolic integration with a numeric safety net. + * + * `antiderivative(e, v)` tries a small tower of exact methods — linearity, + * power rules, a table of elementary forms with linear arguments, the + * complete rational-function algorithm (Horowitz–Ostrogradsky reduction in + * exact rational arithmetic, log/atan terms over the certified real roots + * from poly.ts), Gaussian integrals via erf, integration by parts, trig + * product/power rewrites, and u-substitution. It is NOT a full Risch + * decision procedure: a null result means "no form found", not "no + * elementary form exists", and callers fall back to numeric quadrature. + * + * Every symbolic result is *verified before it is returned*: the candidate F + * is differentiated (symbolically when diff() can, finite differences when it + * cannot) and compared against the integrand at a spread of sample points. + * A rule with a sign slip or a domain surprise is rejected rather than + * plotted — heuristics are safe because nothing unverified escapes. + * + * For definite integrals the fundamental theorem is itself the thing that can + * lie (F(b) − F(a) is wrong across a non-integrable singularity), so + * `verifyDefinite` checks the closed form against adaptive Gauss–Kronrod + * quadrature at probe parameter values. `quadratureSum` builds the fallback: + * a fixed composite Gauss–Legendre rule expanded into an ordinary Expr, the + * same way Σ expands — so shaders, the worker VM and every other consumer + * evaluate it with no new machinery. + */ +import { type Expr, evaluate, freeVars } from './expr.ts'; +import { add, diff, div, mul, neg, pow, sub } from './diff.ts'; +import { + type FPoly, + type Frac, + exprToPoly, + fpadd, + fpderiv, + fpdivExact, + fpdivmod, + fpgcd, + fpmul, + fpscale, + frac, + primitive, + realRootsSquareFree, +} from './poly.ts'; + +const num = (value: number): Expr => ({ kind: 'num', value }); +const vr = (name: string): Expr => ({ kind: 'var', name }); +const call = (name: string, ...args: Expr[]): Expr => ({ kind: 'call', name, args }); +const ln = (x: Expr): Expr => call('ln', x); +const lnAbs = (x: Expr): Expr => call('ln', call('abs', x)); + +const isNum = (e: Expr): e is Expr & { kind: 'num' } => e.kind === 'num'; + +const key = (e: Expr): string => JSON.stringify(e); + +const isConstIn = (e: Expr, v: string): boolean => !freeVars(e).has(v); + +/** The numeric value of a v-free expression with no other frees, or null. */ +function constVal(e: Expr): number | null { + if (freeVars(e).size) return null; + try { + const x = evaluate(e, {}); + return isFinite(x) ? x : null; + } catch { + return null; + } +} + +// --- polynomials with symbolic (v-free Expr) coefficients --- + +const SYM_DEG_MAX = 8; + +/** + * Coefficients of e as a polynomial in v whose coefficients are v-free + * expressions (sliders welcome), ascending, or null. The symbolic sibling of + * poly.ts exprToPoly, for the rules that stay exact under slider parameters. + */ +export function exprPolyCoeffs(e: Expr, v: string): Expr[] | null { + if (isConstIn(e, v)) return [e]; + switch (e.kind) { + case 'var': + return e.name === v ? [num(0), num(1)] : null; + case 'neg': { + const a = exprPolyCoeffs(e.a, v); + return a && a.map(neg); + } + case 'bin': { + if (e.op === '^') { + if (!isNum(e.b) || !Number.isInteger(e.b.value) || e.b.value < 0) return null; + const a = exprPolyCoeffs(e.a, v); + if (!a || (a.length - 1) * e.b.value > SYM_DEG_MAX) return null; + let out: Expr[] = [num(1)]; + for (let i = 0; i < e.b.value; i++) out = convolve(out, a); + return out; + } + if (e.op === '/') { + if (!isConstIn(e.b, v)) return null; + const a = exprPolyCoeffs(e.a, v); + return a && a.map(c => div(c, e.b)); + } + const a = exprPolyCoeffs(e.a, v); + const b = exprPolyCoeffs(e.b, v); + if (!a || !b) return null; + if (e.op === '+' || e.op === '-') { + const out: Expr[] = []; + for (let i = 0; i < Math.max(a.length, b.length); i++) { + const x = a[i] ?? num(0); + const y = b[i] ?? num(0); + out.push(e.op === '+' ? add(x, y) : sub(x, y)); + } + return out; + } + if ((a.length - 1) + (b.length - 1) > SYM_DEG_MAX) return null; + return convolve(a, b); + } + default: + return null; + } +} + +function convolve(a: Expr[], b: Expr[]): Expr[] { + const out: Expr[] = Array.from({ length: a.length + b.length - 1 }, () => num(0)); + for (let i = 0; i < a.length; i++) { + for (let j = 0; j < b.length; j++) out[i + j] = add(out[i + j], mul(a[i], b[j])); + } + return out; +} + +/** a·v + b with a ≠ 0 (v-free a, b), or null. The "linear argument" matcher. */ +function linearIn(e: Expr, v: string): { a: Expr; b: Expr } | null { + const c = exprPolyCoeffs(e, v); + if (!c || c.length !== 2) return null; + if (isNum(c[1]) && c[1].value === 0) return null; + return { a: c[1], b: c[0] }; +} + +// --- multiplicative structure --- + +interface Factor { base: Expr; exp: number } + +/** Flatten products/quotients/negations/integer powers into factors and a sign. */ +function factorsOf(e: Expr, sign: 1 | -1, out: Factor[], flip: { neg: boolean }): void { + switch (e.kind) { + case 'neg': + flip.neg = !flip.neg; + factorsOf(e.a, sign, out, flip); + return; + case 'bin': + if (e.op === '*') { + factorsOf(e.a, sign, out, flip); + factorsOf(e.b, sign, out, flip); + return; + } + if (e.op === '/') { + factorsOf(e.a, sign, out, flip); + factorsOf(e.b, -sign as 1 | -1, out, flip); + return; + } + if (e.op === '^' && isNum(e.b) && Number.isInteger(e.b.value) + && Math.abs(e.b.value) <= 16 && e.b.value !== 0) { + out.push({ base: e.a, exp: sign * e.b.value }); + return; + } + break; + } + out.push({ base: e, exp: sign }); +} + +function flatten(e: Expr): { neg: boolean; factors: Factor[] } { + const flip = { neg: false }; + const raw: Factor[] = []; + factorsOf(e, 1, raw, flip); + // Merge repeated bases so sin(x)·sin(x) reads as sin(x)^2. + const merged = new Map(); + const order: string[] = []; + for (const f of raw) { + if (isNum(f.base) && f.base.value === 1) continue; // 1/x carries a unit factor + const k = key(f.base); + const hit = merged.get(k); + if (hit) hit.exp += f.exp; + else { + merged.set(k, { ...f }); + order.push(k); + } + } + return { neg: flip.neg, factors: order.map(k => merged.get(k)!).filter(f => f.exp !== 0) }; +} + +function rebuild(factors: Factor[], negate: boolean): Expr { + let numr: Expr | null = null; + let den: Expr | null = null; + for (const f of factors) { + const p = Math.abs(f.exp) === 1 ? f.base : pow(f.base, num(Math.abs(f.exp))); + if (f.exp > 0) numr = numr === null ? p : mul(numr, p); + else den = den === null ? p : mul(den, p); + } + let out = numr ?? num(1); + if (den) out = div(out, den); + return negate ? neg(out) : out; +} + +// --- exact rational arithmetic helpers --- + +const F1: Frac = { n: 1n, d: 1n }; + +const fracNum = (f: Frac): number => Number(f.n) / Number(f.d); + +const fpEval = (p: FPoly, x: number): number => { + let acc = 0; + for (let i = p.length - 1; i >= 0; i--) acc = acc * x + fracNum(p[i]); + return acc; +}; + +/** ∑ cᵢ vⁱ as an Expr with double coefficients. */ +function fpToExpr(p: FPoly, v: string): Expr { + let acc: Expr = num(0); + for (let i = p.length - 1; i >= 0; i--) acc = add(mul(acc, vr(v)), num(fracNum(p[i]))); + return acc; +} + +/** Exact Gaussian elimination; null when the system is singular. */ +function solveLinear(A: Frac[][], b: Frac[]): Frac[] | null { + const n = b.length; + const M = A.map((row, i) => [...row, b[i]]); + for (let col = 0; col < n; col++) { + let piv = -1; + for (let r = col; r < n; r++) { + if (M[r][col].n !== 0n) { piv = r; break; } + } + if (piv < 0) return null; + [M[col], M[piv]] = [M[piv], M[col]]; + const inv = { n: M[col][col].d, d: M[col][col].n }; + for (let c = col; c <= n; c++) M[col][c] = fmulF(M[col][c], inv); + for (let r = 0; r < n; r++) { + if (r === col || M[r][col].n === 0n) continue; + const f = M[r][col]; + for (let c = col; c <= n; c++) M[r][c] = fsubF(M[r][c], fmulF(f, M[col][c])); + } + } + return M.map(row => row[n]); +} + +const fmulF = (a: Frac, b: Frac): Frac => frac(a.n * b.n, a.d * b.d); +const fsubF = (a: Frac, b: Frac): Frac => frac(a.n * b.d - b.n * a.d, a.d * b.d); + +// --- the rational-function algorithm --- + +/** + * ∫ P/Q dv for numeric-coefficient polynomials: the polynomial part by the + * power rule; Horowitz–Ostrogradsky reduction (an exact linear solve) splits + * off the rational part P₁/Q₁; the remaining log part integrates over the + * certified real roots (residue · ln|v − r|) with one conjugate quadratic + * pair handled in closed form (ln + atan). Null when the leftover complex + * factor has degree ≥ 4 — the verifier-gated numeric fallback covers it. + */ +function integrateRational(P: FPoly, Q: FPoly, v: string): Expr | null { + if (!Q.length) return null; + const terms: Expr[] = []; + const { q: polyPart, r } = fpdivmod(P, Q); + if (polyPart.length) { + const F: FPoly = [{ n: 0n, d: 1n }]; + polyPart.forEach((c, i) => { F[i + 1] = fmulF(c, frac(1n, BigInt(i + 1))); }); + terms.push(fpToExpr(F, v)); + } + if (r.length) { + let A = r; + let D = Q; + const D1 = fpgcd(D, fpderiv(D)); + if (D1.length > 1) { + const D2 = fpdivExact(D, D1); + let H: FPoly; + try { + H = fpdivExact(fpmul(fpderiv(D1), D2), D1); + } catch { + return null; + } + // A = P₁′·D₂ − P₁·H + P₂·D₁ with deg P₁ < deg D₁, deg P₂ < deg D₂: + // one column per unknown coefficient, equated degree by degree. + const d1 = D1.length - 1; + const d2 = D2.length - 1; + const n = d1 + d2; + const rows: Frac[][] = Array.from({ length: n }, () => Array.from({ length: n }, () => ({ n: 0n, d: 1n }))); + const rhs: Frac[] = Array.from({ length: n }, (_, i) => A[i] ?? { n: 0n, d: 1n }); + if ((A.length - 1) >= n) return null; // not proper: cannot happen, be safe + const put = (colIdx: number, p: FPoly) => { + p.forEach((c, deg) => { + if (deg < n) rows[deg][colIdx] = fpadd([rows[deg][colIdx]], [c])[0] ?? { n: 0n, d: 1n }; + }); + }; + for (let j = 0; j < d1; j++) { + // P₁ = v^j: P₁′·D₂ − v^j·H + const dTerm = j > 0 ? fpmul([...Array.from({ length: j - 1 }, () => ({ n: 0n, d: 1n }) as Frac), frac(BigInt(j), 1n)], D2) : []; + const hTerm = fpscale(fpmul([...Array.from({ length: j }, () => ({ n: 0n, d: 1n }) as Frac), F1], H), { n: -1n, d: 1n }); + put(j, fpadd(dTerm, hTerm)); + } + for (let j = 0; j < d2; j++) { + put(d1 + j, fpmul([...Array.from({ length: j }, () => ({ n: 0n, d: 1n }) as Frac), F1], D1)); + } + const sol = solveLinear(rows, rhs); + if (!sol) return null; + const P1 = sol.slice(0, d1); + const P2 = sol.slice(d1); + while (P1.length && P1[P1.length - 1].n === 0n) P1.pop(); + while (P2.length && P2[P2.length - 1].n === 0n) P2.pop(); + if (P1.length) terms.push(div(fpToExpr(P1, v), fpToExpr(D1, v))); + A = P2; + D = D2; + } + if (A.length) { + const logPart = integrateLogPart(A, D, v); + if (!logPart) return null; + terms.push(logPart); + } + } + return terms.reduce((acc, t) => add(acc, t), num(0)); +} + +/** ∫ A/D with D square-free: residues at the certified real roots, plus one + * conjugate pair in closed form. */ +function integrateLogPart(A: FPoly, D: FPoly, v: string): Expr | null { + const zp = primitive(D).map(c => c.n); + const roots = realRootsSquareFree(zp); + const degLeft = (D.length - 1) - roots.length; + if (degLeft !== 0 && degLeft !== 2) return null; + const dD = fpderiv(D); + const terms: Expr[] = []; + for (const r0 of roots) { + const res = fpEval(A, r0) / fpEval(dD, r0); + if (!isFinite(res)) return null; + terms.push(mul(num(res), lnAbs(sub(vr(v), num(r0))))); + } + if (degLeft === 2) { + // Deflate the real roots numerically; the quotient is the conjugate pair. + let cur = D.map(fracNum); + for (const r0 of roots) { + const next: number[] = Array.from({ length: cur.length - 1 }, () => 0); + let carry = 0; + for (let i = cur.length - 1; i >= 1; i--) { + carry = cur[i] + carry * r0; + next[i - 1] = carry; + } + cur = next; + } + const [c, b, a] = cur; + const disc = b * b - 4 * a * c; + if (!(disc < 0) || !isFinite(disc)) return null; + const alpha = -b / (2 * a); + const beta = Math.sqrt(-disc) / (2 * a); + // Complex residue at α + βi: c = A(z)/D′(z), by complex Horner. + const horner = (p: FPoly): [number, number] => { + let re = 0; + let im = 0; + for (let i = p.length - 1; i >= 0; i--) { + const nre = re * alpha - im * beta + fracNum(p[i]); + im = re * beta + im * alpha; + re = nre; + } + return [re, im]; + }; + const [nRe, nIm] = horner(A); + const [dRe, dIm] = horner(dD); + const dd = dRe * dRe + dIm * dIm; + const resRe = (nRe * dRe + nIm * dIm) / dd; + const resIm = (nIm * dRe - nRe * dIm) / dd; + if (!isFinite(resRe) || !isFinite(resIm)) return null; + const shifted = sub(vr(v), num(alpha)); + const quadExpr = add(pow(shifted, num(2)), num(beta * beta)); + terms.push(mul(num(resRe), ln(quadExpr))); + terms.push(mul(num(-2 * resIm), call('atan', div(shifted, num(beta))))); + } + return terms.reduce((acc, t) => add(acc, t), num(0)); +} + +// --- table of elementary antiderivatives, argument u = a·v + b --- + +/** ∫ f(u) du for one call family; the caller divides by a. Null = no entry. */ +function tableCall(name: string, u: Expr): Expr | null { + switch (name) { + case 'sin': return neg(call('cos', u)); + case 'cos': return call('sin', u); + case 'tan': return neg(lnAbs(call('cos', u))); + case 'exp': return call('exp', u); + case 'sinh': return call('cosh', u); + case 'cosh': return call('sinh', u); + case 'tanh': return ln(call('cosh', u)); + case 'sech': return call('atan', call('sinh', u)); + case 'sqrt': return mul(num(2 / 3), pow(u, num(1.5))); + case 'ln': return sub(mul(u, ln(u)), u); + case 'log': return div(sub(mul(u, ln(u)), u), num(Math.LN10)); + case 'abs': return div(mul(u, call('abs', u)), num(2)); + case 'sign': return call('abs', u); + case 'asin': return add(mul(u, call('asin', u)), call('sqrt', sub(num(1), pow(u, num(2))))); + case 'acos': return sub(mul(u, call('acos', u)), call('sqrt', sub(num(1), pow(u, num(2))))); + case 'atan': return sub(mul(u, call('atan', u)), div(ln(add(num(1), pow(u, num(2)))), num(2))); + case 'asinh': return sub(mul(u, call('asinh', u)), call('sqrt', add(pow(u, num(2)), num(1)))); + case 'acosh': return sub(mul(u, call('acosh', u)), call('sqrt', sub(pow(u, num(2)), num(1)))); + case 'atanh': return add(mul(u, call('atanh', u)), div(ln(sub(num(1), pow(u, num(2)))), num(2))); + case 'erf': return add(mul(u, call('erf', u)), div(call('exp', neg(pow(u, num(2)))), num(Math.sqrt(Math.PI)))); + default: return null; + } +} + +// --- verification --- + +const V_PROBES = [-3.7, -2.1, -1.05, -0.51, 0.33, 0.77, 1.29, 2.42, 3.91]; +const PARAM_PROBES = [0.73, 1.91, -1.37, 0.42]; + +/** Probe environments for the free parameters (two spreads, cycled values). */ +function probeEnvs(frees: Iterable): Array> { + const names = [...frees]; + const mk = (offset: number): Record => { + const env: Record = {}; + names.forEach((n, i) => { env[n] = PARAM_PROBES[(i + offset) % PARAM_PROBES.length]; }); + return env; + }; + return names.length ? [mk(0), mk(1)] : [{}]; +} + +/** + * Does F′ = f? Symbolic differentiation (finite differences when diff cannot) + * compared at sample points; accepted only when enough samples are defined + * and every defined one matches. This gate is what lets the rules above be + * fearless: a wrong candidate is dropped, never returned. + */ +export function verifyAntiderivative(F: Expr, f: Expr, v: string): boolean { + let dF: Expr | null = null; + try { + dF = diff(F, v); + } catch { + dF = null; + } + const frees = new Set([...freeVars(f), ...freeVars(F)]); + frees.delete(v); + let matched = 0; + for (const env of probeEnvs(frees)) { + for (const x of V_PROBES) { + let want: number; + try { + want = evaluate(f, { ...env, [v]: x }); + } catch { + return false; // unbound name: nothing meaningful to verify against + } + if (!isFinite(want)) continue; + let got: number; + let tol = 1e-6 * (1 + Math.abs(want)); + try { + if (dF) { + got = evaluate(dF, { ...env, [v]: x }); + } else { + const h = 1e-5 * (1 + Math.abs(x)); + const fp = evaluate(F, { ...env, [v]: x + h }); + const fm = evaluate(F, { ...env, [v]: x - h }); + got = (fp - fm) / (2 * h); + tol = 1e-3 * (1 + Math.abs(want)); + } + } catch { + return false; + } + if (!isFinite(got)) continue; // F undefined there (domain edge): unjudged + if (Math.abs(got - want) > tol) return false; + matched++; + } + } + // Limited domains (asin, acosh, √) leave few in-domain probes; three + // agreeing points across the spread is already a strong certificate. + return matched >= 3; +} + +// --- adaptive quadrature (Gauss–Kronrod 7–15) --- + +const K15_X = [0, 0.2077849550078985, 0.4058451513773972, 0.5860872354676911, + 0.7415311855993945, 0.8648644233597691, 0.9491079123427585, 0.9914553711208126]; +const K15_W = [0.2094821410847278, 0.2044329400752989, 0.1903505780647854, 0.1690047266392679, + 0.1406532597155259, 0.1047900103222502, 0.0630920926299786, 0.0229353220105292]; +const G7_W = [0.4179591836734694, 0.3818300505051189, 0.2797053914892767, 0.1294849661688697]; + +/** + * Adaptive Gauss–Kronrod estimate of ∫ f over [lo, hi] (sign-aware when + * reversed; ±Infinity bounds handled by the standard rational change of + * variables). NaN when the integral fails to settle within the subdivision + * budget — divergence reads as "no answer", never a confident wrong one. + */ +export function quadrature(f: (x: number) => number, lo: number, hi: number): number { + if (Number.isNaN(lo) || Number.isNaN(hi)) return NaN; + if (lo === hi) return 0; + if (!isFinite(lo) || !isFinite(hi)) { + if (lo > hi) return -quadrature(f, hi, lo); + if (lo === -Infinity && hi === Infinity) { + // x = u/(1−u²) maps (−1, 1) onto ℝ; dx = (1+u²)/(1−u²)² du. + return quadrature(u => { + const d = 1 - u * u; + return (f(u / d) * (1 + u * u)) / (d * d); + }, -1, 1); + } + if (hi === Infinity) { + // x = lo + u/(1−u) maps (0, 1) onto (lo, ∞); dx = du/(1−u)². + return quadrature(u => { + const d = 1 - u; + return f(lo + u / d) / (d * d); + }, 0, 1); + } + // (−∞, hi): the mirror map. + return quadrature(u => { + const d = 1 - u; + return f(hi - u / d) / (d * d); + }, 0, 1); + } + const sign = lo < hi ? 1 : -1; + const a = Math.min(lo, hi); + const b = Math.max(lo, hi); + const panel = (x0: number, x1: number): { k: number; err: number } => { + const c = (x0 + x1) / 2; + const h = (x1 - x0) / 2; + let k15 = 0; + let g7 = 0; + for (let i = 0; i < 8; i++) { + const fp = f(c + h * K15_X[i]); + const fm = i === 0 ? fp : f(c - h * K15_X[i]); + const s = i === 0 ? fp : fp + fm; + k15 += K15_W[i] * s; + // The embedded G7 rule lives on the even-index K15 nodes. + if (i % 2 === 0) g7 += G7_W[i / 2] * s; + } + return { k: k15 * h, err: Math.abs((k15 - g7) * h) }; + }; + let total = 0; + let badMass = 0; + const stack: Array<{ x0: number; x1: number; depth: number }> = [{ x0: a, x1: b, depth: 0 }]; + let evals = 0; + while (stack.length) { + const { x0, x1, depth } = stack.pop()!; + const { k, err } = panel(x0, x1); + evals += 15; + const tol = 1e-10 * Math.max(1, Math.abs(total)) * ((x1 - x0) / (b - a)); + // Depth 45 lets an endpoint singularity refine down to ~1e-12 widths + // (only the singular panel keeps splitting, so the cost stays linear). + if (err <= tol || depth >= 45 || evals > 20000) { + if (isFinite(k)) { + total += k; + if (depth >= 45 || evals > 20000) badMass += err; + } + // Non-finite panels (a sampled singularity) contribute nothing: the + // neighboring subdivisions already carry the finite mass. + continue; + } + const mid = (x0 + x1) / 2; + stack.push({ x0, x1: mid, depth: depth + 1 }, { x0: mid, x1, depth: depth + 1 }); + } + if (!isFinite(total) || badMass > 1e-4 * Math.max(1, Math.abs(total))) return NaN; + return sign * total; +} + +/** + * Check a closed-form definite value against quadrature at probe parameter + * values — the guard against FTC across a singularity (∫₋₁¹ dv/v² is NOT + * v⁻¹ evaluated at the ends). True only when at least one probe is judgeable + * and every judgeable probe agrees. + */ +export function verifyDefinite(value: Expr, body: Expr, v: string, lo: Expr, hi: Expr): boolean { + // Only the BODY's occurrences of v are bound. In `int[a..x] f(x) dx` the + // upper bound's x is the ambient plot variable sharing the name — it (and + // its appearances in the substituted value) must be probed like any other + // parameter, or evaluation throws and a perfectly good closed form gets + // discarded for the far-worse fixed-order fallback. + const frees = new Set([...freeVars(lo), ...freeVars(hi), ...freeVars(value)]); + for (const f of freeVars(body)) if (f !== v) frees.add(f); + let judged = 0; + for (const env of probeEnvs(frees)) { + let a: number; + let b: number; + let want: number; + try { + a = evaluate(lo, env); + b = evaluate(hi, env); + want = evaluate(value, env); + } catch { + return false; + } + // Infinite bounds are fine (quadrature transforms them); huge FINITE + // spans are skipped — they take forever and prove little extra. + if (Number.isNaN(a) || Number.isNaN(b)) continue; + if ((isFinite(a) && Math.abs(a) > 1e6) || (isFinite(b) && Math.abs(b) > 1e6)) continue; + if (!isFinite(want)) continue; + const got = quadrature(x => { + try { + return evaluate(body, { ...env, [v]: x }); + } catch { + return NaN; + } + }, a, b); + if (!isFinite(got)) continue; + if (Math.abs(got - want) > 1e-6 * (1 + Math.abs(got))) return false; + judged++; + } + return judged > 0; +} + +// --- the numeric fallback as an ordinary expression --- + +const GL8_X = [0.1834346424956498, 0.5255324099163290, 0.7966664774136267, 0.9602898564975363]; +const GL8_W = [0.3626837833783620, 0.3137066458778873, 0.2223810344533745, 0.1012285362903763]; +const QUAD_PANELS = 5; + +/** Terms one quadratureSum expands to (for resolve-time term budgets). */ +export const QUAD_TERMS = QUAD_PANELS * 8; + +/** + * ∫ body dv over [lo, hi] as a composite Gauss–Legendre sum, expanded into a + * plain Expr (5 panels × 8 points). Exactly how Σ rows expand — so the GPU + * shaders, the worker's stack VM, sampled densities and symbolic + * differentiation all keep working with no new node kind. Fixed-order, not + * adaptive: smooth integrands are ~1e-12, kinks land near panel edges. lo + * and hi may be any expressions, including the ambient plot coordinate + * (∫₀ˣ makes a function of x). + */ +export function quadratureSum(body: Expr, v: string, lo: Expr, hi: Expr): Expr { + const width = sub(hi, lo); + let acc: Expr = num(0); + for (let p = 0; p < QUAD_PANELS; p++) { + for (let i = 0; i < 4; i++) { + for (const s of [1, -1]) { + const c = (p + 0.5 + (s * GL8_X[i]) / 2) / QUAD_PANELS; // node in (0, 1) + const xi = add(lo, mul(num(c), width)); + const term = substAll(body, v, xi); + acc = add(acc, mul(num(GL8_W[i] / (2 * QUAD_PANELS)), term)); + } + } + } + return mul(width, acc); +} + +/** + * quadratureSum over a half- or fully-infinite range: the same rational + * change of variables quadrature() uses, applied symbolically, so the result + * is still an ordinary finite-interval Gauss–Legendre expansion. A null + * bound means infinite on that side (callers normalize direction first). + */ +export function improperSum(body: Expr, v: string, lo: Expr | null, hi: Expr | null): Expr { + if (lo && hi) return quadratureSum(body, v, lo, hi); + const u = '@w'; + if (!lo && !hi) { + // x = u/(1−u²) over (−1, 1); dx = (1+u²)/(1−u²)² du. + const uu = mul(vr(u), vr(u)); + const d = sub(num(1), uu); + const x = div(vr(u), d); + const jac = div(add(num(1), uu), mul(d, d)); + return quadratureSum(mul(substAll(body, v, x), jac), u, num(-1), num(1)); + } + // Half-infinite: x = end ± u/(1−u) over (0, 1); dx = du/(1−u)². + const d = sub(num(1), vr(u)); + const step = div(vr(u), d); + const x = lo ? add(lo, step) : sub(hi!, step); + const jac = div(num(1), mul(d, d)); + return quadratureSum(mul(substAll(body, v, x), jac), u, num(0), num(1)); +} + +/** Substitute every free occurrence of variable v (no binders can capture: + * the bound variable is eliminated in the same pass). */ +function substAll(e: Expr, v: string, val: Expr): Expr { + switch (e.kind) { + case 'num': return e; + case 'var': return e.name === v ? val : e; + case 'neg': return neg(substAll(e.a, v, val)); + case 'bin': return { kind: 'bin', op: e.op, a: substAll(e.a, v, val), b: substAll(e.b, v, val) }; + case 'call': return { kind: 'call', name: e.name, args: e.args.map(a => substAll(a, v, val)) }; + case 'eq': return { kind: 'eq', l: substAll(e.l, v, val), r: substAll(e.r, v, val) }; + case 'ineq': return { kind: 'ineq', op: e.op, l: substAll(e.l, v, val), r: substAll(e.r, v, val) }; + case 'vec': return { kind: 'vec', items: e.items.map(a => substAll(a, v, val)) }; + case 'list': return { kind: 'list', items: e.items.map(a => substAll(a, v, val)) }; + case 'piecewise': return { + kind: 'piecewise', + cases: e.cases.map(c => ({ cond: substAll(c.cond, v, val), value: substAll(c.value, v, val) })), + otherwise: e.otherwise && substAll(e.otherwise, v, val), + }; + } +} + +// --- the rule tower --- + +const MAX_DEPTH = 10; + +/** + * A verified antiderivative of e with respect to v, or null. The public + * entry: candidates come from `anti`, and only survivors of + * `verifyAntiderivative` escape. + */ +export function antiderivative(e: Expr, v: string): Expr | null { + let F: Expr | null; + try { + F = anti(e, v, 0); + } catch { + return null; + } + if (!F) return null; + return verifyAntiderivative(F, e, v) ? F : null; +} + +function anti(e: Expr, v: string, depth: number): Expr | null { + if (depth > MAX_DEPTH) return null; + if (isConstIn(e, v)) return mul(e, vr(v)); + // Structural linearity keeps the current depth: splitting a sum or pulling + // a constant multiple strictly shrinks the tree, so it cannot cycle, and + // charging it depth would starve the real rules on longer expressions. + switch (e.kind) { + case 'var': + return e.name === v ? div(pow(vr(v), num(2)), num(2)) : null; + case 'neg': { + const F = anti(e.a, v, depth); + return F && neg(F); + } + case 'bin': + if (e.op === '+' || e.op === '-') { + const A = anti(e.a, v, depth); + if (!A) break; + const B = anti(e.b, v, depth); + if (!B) break; + return e.op === '+' ? add(A, B) : sub(A, B); + } + break; + default: + break; + } + + // Constant multiples pull out of the flattened product. + const { neg: negated, factors } = flatten(e); + const constFs = factors.filter(f => isConstIn(f.base, v)); + const varFs = factors.filter(f => !isConstIn(f.base, v)); + if (constFs.length) { + const coef = rebuild(constFs, negated); + const F = anti(rebuild(varFs, false), v, depth); + return F && mul(coef, F); + } + if (negated) { + const F = anti(rebuild(varFs, false), v, depth); + return F && neg(F); + } + + // A sum inside a product distributes: (A ± B)·rest splits into two + // integrals (each strictly smaller in sum-factors, so this cannot cycle). + const sumF = varFs.find(f => f.exp === 1 && f.base.kind === 'bin' + && (f.base.op === '+' || f.base.op === '-')); + if (sumF && varFs.length > 1) { + const rest = rebuild(varFs.filter(f => f !== sumF), false); + const b = sumF.base as Expr & { kind: 'bin' }; + const A = anti(mul(b.a, rest), v, depth); + if (A) { + const B = anti(mul(b.b, rest), v, depth); + if (B) return b.op === '+' ? add(A, B) : sub(A, B); + } + } + + // Exact numeric-coefficient polynomials and rational functions. + const numExpr = rebuild(varFs.filter(f => f.exp > 0), false); + const denExpr = rebuild(varFs.filter(f => f.exp < 0).map(f => ({ base: f.base, exp: -f.exp })), false); + const P = exprToPoly(numExpr, v); + const Q = exprToPoly(denExpr, v); + if (P && Q) { + const F = integrateRational(P, Q, v); + if (F) return F; + } + + // Symbolic-coefficient polynomials (slider coefficients stay exact). + const symPoly = exprPolyCoeffs(e, v); + if (symPoly) { + let F: Expr = num(0); + symPoly.forEach((c, i) => { + F = add(F, mul(div(c, num(i + 1)), pow(vr(v), num(i + 1)))); + }); + return F; + } + + const single = varFs.length === 1 ? varFs[0] : null; + + // Table forms with a linear argument: f(a·v + b) → F(u)/a. + if (single && single.exp === 1 && single.base.kind === 'call') { + const c = single.base; + if (c.args.length === 1) { + const lin = linearIn(c.args[0], v); + if (lin) { + const T = tableCall(c.name, c.args[0]); + if (T) return div(T, lin.a); + } + } + // normalpdf/normalcdf in their first argument (mean and sd v-free). + if ((c.name === 'normalpdf' || c.name === 'normalcdf') && c.args.length === 3 + && isConstIn(c.args[1], v) && isConstIn(c.args[2], v)) { + const lin = linearIn(c.args[0], v); + if (lin) { + const [x, m, s] = c.args; + const T = c.name === 'normalpdf' + ? call('normalcdf', x, m, s) + : add(mul(sub(x, m), call('normalcdf', x, m, s)), + mul(pow(s, num(2)), call('normalpdf', x, m, s))); + return div(T, lin.a); + } + } + } + + // Powers of a linear argument: (a·v + b)^p for any constant p (p = −1 → ln). + if (single && single.base.kind === 'bin' && single.base.op === '^' && single.exp === 1) { + if (isConstIn(single.base.b, v)) { + const lin = linearIn(single.base.a, v); + const p = single.base.b; + if (lin) { + if (isNum(p) && p.value === -1) return div(lnAbs(single.base.a), lin.a); + const p1 = add(p, num(1)); + return div(pow(single.base.a, p1), mul(lin.a, p1)); + } + } + // c^(a·v + b) for a v-free base: c^u / (a ln c). + const linExp = linearIn(single.base.b, v); + if (isConstIn(single.base.a, v) && linExp) { + return div(pow(single.base.a, single.base.b), mul(linExp.a, ln(single.base.a))); + } + } + if (single && single.exp === -1) { + const lin = linearIn(single.base, v); + if (lin) return div(lnAbs(single.base), lin.a); + } + // Negative powers of a linear argument: (a·v + b)^(−k). + if (single && single.exp < -1) { + const lin = linearIn(single.base, v); + if (lin) { + const p1 = single.exp + 1; + return div(pow(single.base, num(p1)), mul(lin.a, num(p1))); + } + } + + // Gaussian integrals: exp(quadratic with negative leading coefficient) → erf. + if (single && single.exp === 1 && single.base.kind === 'call' + && single.base.name === 'exp' && single.base.args.length === 1) { + const q = exprPolyCoeffs(single.base.args[0], v); + if (q && q.length === 3) { + const c2 = constVal(q[2]); + if (c2 !== null && c2 < 0) { + const k = Math.sqrt(-c2); + // exp(c₂v² + c₁v + c₀) = exp(c₀ − c₁²/4c₂) · exp(−(kv − c₁/2k)²) + const shiftK = div(q[1], num(2 * k)); // c₁/(2k) + const inner = sub(mul(num(k), vr(v)), shiftK); + const amp = call('exp', add(q[0], pow(shiftK, num(2)))); + return mul(amp, mul(num(Math.sqrt(Math.PI) / (2 * k)), call('erf', inner))); + } + } + } + + // Trig powers and products. + const trig = trigRule(varFs, v, depth); + if (trig) return trig; + + // exp(a v + b) · sin/cos(c v + d): the classic closed form. + const expTrig = expTrigRule(varFs, v); + if (expTrig) return expTrig; + + // Integration by parts: polynomial × table family. + const parts = partsRule(varFs, v, depth); + if (parts) return parts; + + // u-substitution: find u with e = g(u)·u′ structurally. + return uSubstitution(e, varFs, v, depth); +} + +// --- trig products --- + +/** sin^m(u)·cos^n(u) with a shared linear argument (odd powers via the + * complementary substitution, even ones by power reduction / product-to-sum). */ +function trigRule(factors: Factor[], v: string, depth: number): Expr | null { + if (!factors.length) return null; + if (!factors.every(f => + f.exp >= 1 && f.base.kind === 'call' && (f.base.name === 'sin' || f.base.name === 'cos') + && f.base.args.length === 1 && linearIn(f.base.args[0], v) !== null)) return null; + const sins = factors.filter(f => (f.base as Expr & { kind: 'call' }).name === 'sin'); + const coss = factors.filter(f => (f.base as Expr & { kind: 'call' }).name === 'cos'); + const argOf = (f: Factor): Expr => (f.base as Expr & { kind: 'call' }).args[0]; + + // Distinct arguments: peel one pair with a product-to-sum identity. + if (factors.length >= 2) { + const argKeys = new Set(factors.map(f => key(argOf(f)))); + if (argKeys.size > 1) { + const [f1, f2] = [factors[0], factors[1]]; + const A = argOf(f1); + const B = argOf(f2); + const n1 = (f1.base as Expr & { kind: 'call' }).name; + const n2 = (f2.base as Expr & { kind: 'call' }).name; + const rest = [ + ...(f1.exp > 1 ? [{ base: f1.base, exp: f1.exp - 1 }] : []), + ...(f2.exp > 1 ? [{ base: f2.base, exp: f2.exp - 1 }] : []), + ...factors.slice(2), + ]; + // Combine the arguments through their linear coefficients so repeated + // peels fold: cos(2x − x − x) must become the CONSTANT cos(0), not an + // expression that still looks like it depends on v. + const lA = linearIn(A, v)!; + const lB = linearIn(B, v)!; + const argSub = add(mul(sub(lA.a, lB.a), vr(v)), sub(lA.b, lB.b)); + const argAdd = add(mul(add(lA.a, lB.a), vr(v)), add(lA.b, lB.b)); + const half = (x: Expr): Expr => mul(num(0.5), x); + let sum: Expr; + if (n1 === 'sin' && n2 === 'sin') { + sum = sub(half(call('cos', argSub)), half(call('cos', argAdd))); + } else if (n1 === 'cos' && n2 === 'cos') { + sum = add(half(call('cos', argSub)), half(call('cos', argAdd))); + } else { + // sin A · cos B, arranged so the sine carries A. + const [aS, aC] = n1 === 'sin' ? [lA, lB] : [lB, lA]; + const sPlus = add(mul(add(aS.a, aC.a), vr(v)), add(aS.b, aC.b)); + const sMinus = add(mul(sub(aS.a, aC.a), vr(v)), sub(aS.b, aC.b)); + sum = add(half(call('sin', sPlus)), half(call('sin', sMinus))); + } + const restE = rebuild(rest, false); + return anti(mul(sum, restE), v, depth + 1); + } + } + + const u = argOf(factors[0]); + const lin = linearIn(u, v)!; + const m = sins.reduce((s, f) => s + f.exp, 0); + const n = coss.reduce((s, f) => s + f.exp, 0); + if (m + n > 12) return null; + + // Odd sine power: t = cos u, sin² = 1 − t². + if (m % 2 === 1) { + const F = oddTrigPoly(m, n, 'cos', u); + return div(F, lin.a); + } + // Odd cosine power: t = sin u. + if (n % 2 === 1) { + const F = oddTrigPoly(n, m, 'sin', u); + return F && div(F, lin.a); + } + // Both even: reduce one sin² or cos² by its half-angle identity and recurse. + const twoU = mul(num(2), u); + if (m >= 2) { + const rest = mul( + m - 2 >= 1 ? pow(call('sin', u), num(m - 2)) : num(1), + n >= 1 ? pow(call('cos', u), num(n)) : num(1), + ); + const rewritten = mul(sub(num(0.5), mul(num(0.5), call('cos', twoU))), rest); + return anti(rewritten, v, depth + 1); + } + if (n >= 2) { + const rest = n - 2 >= 1 ? pow(call('cos', u), num(n - 2)) : num(1); + const rewritten = mul(add(num(0.5), mul(num(0.5), call('cos', twoU))), rest); + return anti(rewritten, v, depth + 1); + } + return null; +} + +/** + * ∫ sin^m cos^n du with m odd (or the mirror): substitute t for the OTHER + * function, expand (1 − t²)^((m−1)/2)·t^n, and integrate the polynomial. + * Signs: dt = −sin u du for t = cos u, dt = +cos u du for t = sin u. + */ +function oddTrigPoly(oddPow: number, evenPow: number, tFn: 'sin' | 'cos', u: Expr): Expr { + // Polynomial in t: (1 − t²)^((odd−1)/2) · t^even. + let poly: number[] = [1]; + const half = (oddPow - 1) / 2; + for (let i = 0; i < half; i++) { + const next: number[] = Array.from({ length: poly.length + 2 }, () => 0); + poly.forEach((c, k) => { + next[k] += c; + next[k + 2] -= c; + }); + poly = next; + } + const shifted: number[] = [...Array.from({ length: evenPow }, () => 0), ...poly]; + // Integrate termwise, then substitute t = cos u (extra −1) or t = sin u. + const sign = tFn === 'cos' ? -1 : 1; + let F: Expr = num(0); + shifted.forEach((c, k) => { + if (c === 0) return; + F = add(F, mul(num((sign * c) / (k + 1)), pow(call(tFn, u), num(k + 1)))); + }); + return F; +} + +/** e^{av+b}·sin(cv+d) or ·cos(cv+d): e^u(a·trig − ...)/(a² + c²). */ +function expTrigRule(factors: Factor[], v: string): Expr | null { + if (factors.length !== 2) return null; + const ex = factors.find(f => f.exp === 1 && f.base.kind === 'call' && f.base.name === 'exp'); + const tr = factors.find(f => f.exp === 1 && f.base.kind === 'call' + && (f.base.name === 'sin' || f.base.name === 'cos')); + if (!ex || !tr) return null; + const eArg = (ex.base as Expr & { kind: 'call' }).args[0]; + const tArg = (tr.base as Expr & { kind: 'call' }).args[0]; + const la = linearIn(eArg, v); + const lc = linearIn(tArg, v); + if (!la || !lc) return null; + const a = la.a; + const c = lc.a; + const den = add(pow(a, num(2)), pow(c, num(2))); + const sinT = call('sin', tArg); + const cosT = call('cos', tArg); + const numTerm = (tr.base as Expr & { kind: 'call' }).name === 'sin' + ? sub(mul(a, sinT), mul(c, cosT)) + : add(mul(a, cosT), mul(c, sinT)); + return div(mul(ex.base, numTerm), den); +} + +/** Call families whose DERIVATIVE is algebraic: by parts these are the u to + * differentiate (∫x·ln x picks u = ln x, dv = x dx). */ +const PARTS_DIFF = new Set(['ln', 'log', 'atan', 'asin', 'acos', 'atanh', 'asinh', 'acosh', 'erf']); + +/** Integration by parts for polynomial × one table family. The polynomial is + * differentiated away for exp/trig partners; the logarithmic/inverse family + * is differentiated instead (its derivative is algebraic), the polynomial + * integrated. */ +function partsRule(factors: Factor[], v: string, depth: number): Expr | null { + if (factors.length !== 2) return null; + const polyF = factors.find(f => f.exp > 0 && f.exp <= 6 && exprPolyCoeffs(pow(f.base, num(f.exp)), v) !== null); + const other = factors.find(f => f !== polyF); + if (!polyF || !other || other.exp !== 1) return null; + const p = rebuild([polyF], false); + const g = other.base; + if (g.kind === 'call' && PARTS_DIFF.has(g.name)) { + // ∫p·g = P·g − ∫P·g′ with P = ∫p: g′ is algebraic, so the rest recurses + // into the rational/root rules instead of looping back here. + const P = anti(p, v, depth + 1); + if (!P) return null; + let dg: Expr; + try { + dg = diff(g, v); + } catch { + return null; + } + const restF = anti(mul(P, dg), v, depth + 1); + if (!restF) return null; + return sub(mul(P, g), restF); + } + const G = anti(g, v, depth + 1); + if (!G) return null; + let dp: Expr; + try { + dp = diff(p, v); + } catch { + return null; + } + const restF = anti(mul(dp, G), v, depth + 1); + if (!restF) return null; + return sub(mul(p, G), restF); +} + +// --- u-substitution --- + +/** Candidate inner functions: call arguments, non-trivial power bases and + * exponents, denominators. */ +function subCandidates(e: Expr, v: string, out: Map): void { + const consider = (u: Expr) => { + if (u.kind === 'var' || isConstIn(u, v)) return; + out.set(key(u), u); + }; + switch (e.kind) { + case 'num': + case 'var': + return; + case 'neg': + subCandidates(e.a, v, out); + return; + case 'bin': + if (e.op === '^') { + consider(e.a); + consider(e.b); + } + if (e.op === '/') consider(e.b); + subCandidates(e.a, v, out); + subCandidates(e.b, v, out); + return; + case 'call': + consider(e); // u may be the call itself: ln(x) in ln(x)/x + for (const a of e.args) { + consider(a); + subCandidates(a, v, out); + } + return; + default: + return; + } +} + +/** Structural replacement of a subexpression (by shape) with a variable. */ +function replaceExpr(e: Expr, target: string, s: Expr): Expr { + if (key(e) === target) return s; + switch (e.kind) { + case 'num': + case 'var': + return e; + case 'neg': return neg(replaceExpr(e.a, target, s)); + case 'bin': return { kind: 'bin', op: e.op, a: replaceExpr(e.a, target, s), b: replaceExpr(e.b, target, s) }; + case 'call': return { kind: 'call', name: e.name, args: e.args.map(a => replaceExpr(a, target, s)) }; + default: + return e; + } +} + +/** + * Try e = g(u)·u′: for each candidate u, divide u′ out of the factor list + * (structural cancellation, then exact polynomial division for what remains) + * and check the quotient depends on v only through u. + */ +function uSubstitution(e: Expr, factors: Factor[], v: string, depth: number): Expr | null { + const candidates = new Map(); + subCandidates(e, v, candidates); + for (const u of candidates.values()) { + let du: Expr; + try { + du = diff(u, v); + } catch { + continue; + } + const ratio = divideOut(factors, du, v); + if (!ratio) continue; + const s = `@u${depth}`; + const inS = replaceExpr(ratio, key(u), vr(s)); + if (freeVars(inS).has(v)) continue; + const G = anti(inS, s, depth + 1); + if (!G) continue; + return substAll(G, s, u); + } + return null; +} + +/** e / du as an expression, by cancelling shared factors and finishing with + * exact polynomial division; null when the quotient will not come clean. */ +function divideOut(eFactors: Factor[], du: Expr, v: string): Expr | null { + const { neg: duNeg, factors: duFs } = flatten(du); + const left = eFactors.map(f => ({ ...f })); + const still: Factor[] = []; + let coefNum = 1; + for (const f of duFs) { + if (isConstIn(f.base, v)) { + const c = constVal(f.base); + if (c === null || c === 0) return null; + coefNum *= Math.pow(c, f.exp); + continue; + } + const hit = left.find(l => key(l.base) === key(f.base)); + if (hit && hit.exp >= f.exp) hit.exp -= f.exp; + else still.push(f); + } + const remaining = left.filter(f => f.exp !== 0); + if (!still.length) { + return div(rebuild(remaining, duNeg), num(coefNum)); + } + // Whatever did not cancel structurally must divide the polynomial part. + const duRest = rebuild(still, false); + const dP = exprToPoly(duRest, v); + if (!dP || dP.length < 1) return null; + const polyFs = remaining.filter(f => f.exp > 0 && exprToPoly(pow(f.base, num(f.exp)), v) !== null); + const nonPoly = remaining.filter(f => !polyFs.includes(f)); + const eP = exprToPoly(rebuild(polyFs, false), v); + if (!eP) return null; + const { q, r } = fpdivmod(eP, dP); + if (r.length) return null; + const out = mul(fpToExpr(q, v), rebuild(nonPoly, false)); + return div(duNeg ? neg(out) : out, num(coefNum)); +} diff --git a/lib/plot.ts b/lib/plot.ts index c546cfa..cfa6c93 100644 --- a/lib/plot.ts +++ b/lib/plot.ts @@ -78,7 +78,17 @@ export type Plot = */ | { type: 'cobweb'; f: Expr; recVar: string; curveField: string; a0Name?: string } /** a_{n+1} = f(a_n, x): orbit attractor per pixel column, x as the parameter. */ - | { type: 'bifurcation'; field: string; a0Name?: string }; + | { type: 'bifurcation'; field: string; a0Name?: string } + /** A derived random variable (`S = X + Y`, or a bare expression in random + * variables): the sampled density estimate of the named variable. */ + | { type: 'density'; rv: string } + /** A `P(…)` row estimated from samples. With `shade`, the area under rv's + * density between the bounds fills in (single-variable bodies only). */ + | { type: 'prob'; body: Expr; shade?: { rv: string; lo?: Expr; hi?: Expr } } + /** An `E(…)` row: the mean lives in the row's readout; the plot is a + * vertical marker at x = E under the density of rv (the body itself, + * registered as an anonymous derived variable when not a bare name). */ + | { type: 'expect'; rv: string }; /** Symbolically differentiate each component; undefined if any is non-smooth. */ function tryGrad(exprs: Expr[], v: string): [string, string, string] | undefined { diff --git a/lib/poly.ts b/lib/poly.ts index ec746ce..5e8b16e 100644 --- a/lib/poly.ts +++ b/lib/poly.ts @@ -19,7 +19,7 @@ import { type Expr, evaluate, freeVars } from './expr.ts'; // --- exact rationals (bigint numerator / positive bigint denominator) --- -interface Frac { n: bigint; d: bigint } +export interface Frac { n: bigint; d: bigint } function bgcd(a: bigint, b: bigint): bigint { a = a < 0n ? -a : a; @@ -28,7 +28,7 @@ function bgcd(a: bigint, b: bigint): bigint { return a; } -function frac(n: bigint, d: bigint): Frac { +export function frac(n: bigint, d: bigint): Frac { if (d < 0n) { n = -n; d = -d; } const g = bgcd(n, d) || 1n; return { n: n / g, d: d / g }; @@ -54,7 +54,7 @@ function fromNumber(v: number): Frac | null { // --- polynomials as dense coefficient arrays, ascending degree --- -type FPoly = Frac[]; +export type FPoly = Frac[]; function ftrim(p: FPoly): FPoly { let n = p.length; @@ -62,17 +62,17 @@ function ftrim(p: FPoly): FPoly { return p.slice(0, n); } -function fpadd(a: FPoly, b: FPoly): FPoly { +export function fpadd(a: FPoly, b: FPoly): FPoly { const out: FPoly = []; for (let i = 0; i < Math.max(a.length, b.length); i++) out.push(fadd(a[i] ?? F0, b[i] ?? F0)); return ftrim(out); } -function fpscale(a: FPoly, s: Frac): FPoly { +export function fpscale(a: FPoly, s: Frac): FPoly { return ftrim(a.map(c => fmul(c, s))); } -function fpmul(a: FPoly, b: FPoly): FPoly { +export function fpmul(a: FPoly, b: FPoly): FPoly { if (!a.length || !b.length) return []; const out: FPoly = Array.from({ length: a.length + b.length - 1 }, () => F0); for (let i = 0; i < a.length; i++) { @@ -81,12 +81,12 @@ function fpmul(a: FPoly, b: FPoly): FPoly { return ftrim(out); } -function fpderiv(a: FPoly): FPoly { +export function fpderiv(a: FPoly): FPoly { return ftrim(a.slice(1).map((c, i) => fmul(c, { n: BigInt(i + 1), d: 1n }))); } /** Long division; returns quotient and remainder. b must be non-zero. */ -function fpdivmod(a: FPoly, b: FPoly): { q: FPoly; r: FPoly } { +export function fpdivmod(a: FPoly, b: FPoly): { q: FPoly; r: FPoly } { const q: FPoly = []; let r = a.slice(); const db = b.length - 1; @@ -104,14 +104,14 @@ function fpdivmod(a: FPoly, b: FPoly): { q: FPoly; r: FPoly } { } /** Division known to be exact (used inside Yun's algorithm). */ -function fpdivExact(a: FPoly, b: FPoly): FPoly { +export function fpdivExact(a: FPoly, b: FPoly): FPoly { const { q, r } = fpdivmod(a, b); if (r.length) throw new Error('poly: inexact division'); return q; } /** Clear denominators and divide by integer content → primitive ℤ[x]. */ -function primitive(p: FPoly): FPoly { +export function primitive(p: FPoly): FPoly { if (!p.length) return []; let l = 1n; for (const c of p) l = (l / bgcd(l, c.d)) * c.d; @@ -123,7 +123,7 @@ function primitive(p: FPoly): FPoly { } /** Polynomial gcd via Euclid with primitive-part reduction (controls blowup). */ -function fpgcd(a: FPoly, b: FPoly): FPoly { +export function fpgcd(a: FPoly, b: FPoly): FPoly { let x = primitive(a); let y = primitive(b); if (x.length < y.length) [x, y] = [y, x]; @@ -135,6 +135,28 @@ function fpgcd(a: FPoly, b: FPoly): FPoly { return x; } +/** + * Extended Euclid: g = gcd(a, b) with s·a + t·b = g (g as fpgcd computes it, + * up to a rational scale absorbed into s and t). Exact Frac arithmetic, no + * primitive-part shortcuts — the cofactors are the point (Hermite reduction + * in integrate.ts solves A = s·D + t·D′ with them). + */ +export function fpextgcd(a: FPoly, b: FPoly): { g: FPoly; s: FPoly; t: FPoly } { + let r0 = a.slice(); + let r1 = b.slice(); + let s0: FPoly = [{ n: 1n, d: 1n }]; + let s1: FPoly = []; + let t0: FPoly = []; + let t1: FPoly = [{ n: 1n, d: 1n }]; + while (r1.length) { + const { q, r } = fpdivmod(r0, r1); + [r0, r1] = [r1, r]; + [s0, s1] = [s1, fpadd(s0, fpscale(fpmul(q, s1), { n: -1n, d: 1n }))]; + [t0, t1] = [t1, fpadd(t0, fpscale(fpmul(q, t1), { n: -1n, d: 1n }))]; + } + return { g: r0, s: s0, t: t0 }; +} + // --- expression → polynomial coefficients --- const MAX_DEGREE = 128; @@ -249,7 +271,7 @@ function polyIn(e: Expr, v: string): FPoly | null { // --- Yun's square-free decomposition --- /** f = ∏ out[j].p ^ out[j].mult with each p square-free and pairwise coprime. */ -function squareFree(f: FPoly): Array<{ p: FPoly; mult: number }> { +export function squareFree(f: FPoly): Array<{ p: FPoly; mult: number }> { const df = fpderiv(f); const g = fpgcd(f, df); if (g.length <= 1) return [{ p: f, mult: 1 }]; @@ -434,7 +456,7 @@ function bitLength(x: bigint): number { } /** All real roots of a square-free primitive integer polynomial. */ -function realRootsSquareFree(p: ZPoly): number[] { +export function realRootsSquareFree(p: ZPoly): number[] { const roots: number[] = []; p = p.slice(); while (p.length && p[0] === 0n) { diff --git a/web/main.ts b/web/main.ts index 20bb30d..491d741 100644 --- a/web/main.ts +++ b/web/main.ts @@ -6,25 +6,33 @@ import { defKey, emptyDefs, evalConstEnv, - RESERVED, resolveExpr, scanDefinition, + usesIntegral, type Definition, type Defs, } from '../lib/defs.ts'; import { buildComb, buildTube, combScale, curveExtent, curveFrames } from '../lib/curve3d.ts'; import { - type DistDef, + type BaseDist, + type DensityCurve, + RVSystem, + buildRVSystem, + checkDerived, + densityAt, densityExpr, + matchExpectation, matchProbability, - parseDistribution, + pdfExpr, probabilityValue, regionExpr, - scanDistribution, + scanRandomRows, + shadePolygon, + toExpectation, toProbability, } from '../lib/dist.ts'; import { SLIDER_NUM_RE as NUM_RE, dragAxes } from '../lib/drag.ts'; -import { type Expr, builtinFn, evaluate, freeVars, parseExpr, substVars } from '../lib/expr.ts'; +import { type Expr, evaluate, freeVars, parseExpr, substVars } from '../lib/expr.ts'; import { lowerGeom, pointComps } from '../lib/geom.ts'; import { decodePayload, encodePayload } from '../lib/link.ts'; import { type GridField, angularSpacing, buildGridField, sampleGradMag } from '../lib/grid.ts'; @@ -163,6 +171,11 @@ let stateVals: Record = {}; let stateTime = 0; /** Compiled coordinate fields; non-empty replaces the Cartesian grid. */ let gridFields: GridField[] = []; +/** Declared random variables and their sample caches (persists across + * recompiles; definition-aware caching makes stale samples impossible). */ +const rvSys = new RVSystem(); +/** Every declared random-variable name, healthy or not. */ +let rvNames: ReadonlySet = new Set(); /** Click-dropped seeds for integral curves through vector fields / ODEs. */ const drops: Array<{ x: number; y: number }> = []; /** What the pointer can grab, in math coords; rebuilt by every 2D frame. */ @@ -516,6 +529,9 @@ function render() { case 'sequence': case 'cobweb': case 'bifurcation': + case 'density': + case 'prob': + case 'expect': break; // 2D-only plots (densities, flows, sequences, planar figures); skipped in 3D scenes case 'plist': { const env = { ...constEnv, t: time }; @@ -750,6 +766,52 @@ function render() { case 'bifurcation': layers.bifs.push({ field: plot.field, color, params, uniforms: { uSeed: seedOf(plot.a0Name) } }); break; + case 'density': { + let c: DensityCurve | null = null; + try { + c = rvSys.curve(plot.rv, env); + } catch { break; /* a parameter is missing this frame */ } + if (!c) break; + if (c.pts.length >= 4) extras.polylines.push({ pts: c.pts, color: css, width: 2 }); + // Point masses draw as probability stems (height = mass, not density). + for (const a of c.atoms ?? []) { + extras.polylines.push({ pts: [a.x, 0, a.x, a.p], color: css, width: 2 }); + extras.points.push({ x: a.x, y: a.p, color: css, r: 4 }); + } + break; + } + case 'prob': { + // The estimate lives in the row's readout; the plot is the shaded + // area under the variable's density, when the body has that shape. + if (!plot.shade) break; + try { + const c = rvSys.curve(plot.shade.rv, env); + if (!c) break; + const lo = plot.shade.lo ? evaluate(plot.shade.lo, env) : undefined; + const hi = plot.shade.hi ? evaluate(plot.shade.hi, env) : undefined; + const poly = shadePolygon(c, lo, hi); + if (poly) { + extras.polylines.push({ pts: poly, color: css, closed: true, fill: cssColorA(color, 0.16) }); + } + } catch { /* not evaluable this frame */ } + break; + } + case 'expect': { + // The value lives in the row's readout; the plot is a vertical + // marker at x = E under the variable's density. + try { + const m = rvSys.mean(plot.rv, env); + if (!isFinite(m)) break; + const exact = rvSys.exactDist(plot.rv); + let h = exact + ? evaluate(pdfExpr(exact, { kind: 'num', value: m }), env) + : (c => (c ? densityAt(c, m) : 0))(rvSys.curve(plot.rv, env)); + if (!isFinite(h) || h < 0) h = 0; + if (h > 0) extras.polylines.push({ pts: [m, 0, m, h], color: css, width: 2 }); + extras.points.push({ x: m, y: h, color: css, r: 4 }); + } catch { /* not evaluable this frame */ } + break; + } case 'system': // A 3-unknown system forces the 3D view, so only 2D lands here. if (plot.dim === 2) { @@ -829,10 +891,20 @@ const stateResetBtn = document.getElementById('state-reset') as HTMLButtonElemen * on every keystroke. */ function recompileAll() { + // Random-variable rows resolve outside the definition system: `X ~ …` is + // never a definition, and `Y = X^2` with X random declares a *derived* + // random variable, not a constant. The scan is transitive (`Z = Y + 1` + // follows Y into the set), so it must see the whole document first. + const rvScan = scanRandomRows(equations.map(eq => { + const text = eq.text.trim(); + return !text || text.startsWith('#') || scanSeqRec(text) ? null : text; + })); + const rvRowIdx = new Set([...rvScan.base.keys(), ...rvScan.derived.keys()]); + const raw: Definition[] = []; const defRows = new Map(); const dupRows: Equation[] = []; - for (const eq of equations) { + for (const [i, eq] of equations.entries()) { eq.cls = undefined; eq.parsed = undefined; eq.error = undefined; @@ -844,6 +916,7 @@ function recompileAll() { eq.comment = text.startsWith('#'); if (!eq.comment) eq.collapsed = undefined; if (!text || eq.comment) continue; + if (rvRowIdx.has(i)) continue; // Sequence/recurrence rows (a_n = …, a_{n+1} = …) are plots, not definitions. if (scanSeqRec(text)) continue; const d = scanDefinition(text); @@ -917,31 +990,77 @@ function recompileAll() { for (const name of defs.states.keys()) delete constVals[name]; const ropts = { consts: constVals, boundConsts: sumBoundNames }; - // Random-variable rows (`X ~ Normal(0, a)`) resolve first so P(…) rows can - // reference them regardless of row order. - const dists = new Map(); + // Random-variable rows resolve before plot rows so P(…) and bare + // expressions can reference them regardless of row order. + const builtRVs = buildRVSystem(rvSys, rvScan, { + fnNames, + getFn, + ropts, + constNames, + taken: n => defs.consts.has(n) || defs.fns.has(n) || defs.fields.has(n) + || defs.states.has(n) || defs.points.has(n) || defs.mats.has(n), + }); + rvNames = builtRVs.names; const distRows = new Set(); - for (const eq of equations) { - if (eq.def || eq.comment) continue; - const text = eq.text.trim(); - if (!text) continue; - const scan = scanDistribution(text); - if (!scan) continue; - distRows.add(eq); + // Readout environment: constants at t = 0. Animated or state-fed variables + // simply skip their readout (the sampler throws on the missing name). + let envT0: Record | null = null; + try { + envT0 = evalConstEnv(defs, 0); + } catch { /* a broken constant: rows using it already carry errors */ } + const rvInfo = (eq: Equation, name: string) => { + if (!envT0) return; try { - if (RESERVED.has(scan.name) || builtinFn(scan.name)) { - throw new Error(`Cannot use ${scan.name} as a random variable name.`); - } - if (dists.has(scan.name) || defs.consts.has(scan.name) || defs.fns.has(scan.name) || defs.fields.has(scan.name)) { - throw new Error(`${scan.name} is already defined.`); + // Exact moments where the law has a closed form; the sample estimate + // (with its ≈) everywhere else. + const m = rvSys.exactMoments(name, envT0); + if (m && isFinite(m.mean) && isFinite(m.sd)) { + eq.info = `μ = ${fmtNum(m.mean)}, σ = ${fmtNum(m.sd)}`; + return; } - const d = parseDistribution(scan.name, scan.rhs, fnNames); - d.mean = resolveExpr(d.mean, getFn, ropts); - d.sd = resolveExpr(d.sd, getFn, ropts); - dists.set(scan.name, d); - eq.cls = classify(densityExpr(d), constNames); - } catch (e) { - eq.error = e instanceof Error ? e.message : String(e); + // One-variable transforms integrate against the base pdf (quadrature): + // still ≈, but good to display precision rather than sampling noise. + const s = rvSys.quadMoments(name, envT0) ?? rvSys.curve(name, envT0); + if (!s) return; + eq.info = `μ ≈ ${s.mean.toFixed(3)}, σ ≈ ${s.sd.toFixed(3)}` + + (s.mass < 0.9995 ? `, P(defined) ≈ ${s.mass.toFixed(3)}` : ''); + } catch { /* not numerically computable right now (e.g. animated) */ } + }; + // A derived variable whose law is a closed-form pdf (affine in normals, or + // a single scaled uniform/exponential) plots exactly through the shader; a + // uniform-sum law plots its exact piecewise polynomial via curve(); only + // the rest estimate from samples. + const classifyDerived = (eq: Equation, name: string) => { + const exact = rvSys.exactDist(name); + if (exact && rvSys.get(name)!.kind === 'derived') { + eq.cls = classify(densityExpr(exact), constNames); + } else { + eq.cls = densityCls(name); + } + rvInfo(eq, name); + }; + const densityCls = (name: string): Classified => { + const ps = rvSys.paramsOf(name); + return { + plot: { type: 'density', rv: name }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(p => p !== 't'), + }; + }; + for (const [i, name] of builtRVs.rowRV) { + const eq = equations[i]; + distRows.add(eq); + const message = builtRVs.errors.get(i); + if (message) { + eq.error = message; + continue; + } + const rv = rvSys.get(name)!; + if (rv.kind === 'base') { + eq.cls = classify(densityExpr(rv.dist), constNames); + } else { + classifyDerived(eq, name); } } @@ -960,14 +1079,92 @@ function recompileAll() { } const probBody = defs.consts.has('P') || defs.fns.has('P') ? null : matchProbability(text); if (probBody !== null) { - if (!dists.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); - const p = toProbability(resolveExpr(parseExpr(probBody, fnNames), getFn, ropts), dists); - eq.cls = classify(regionExpr(p), constNames); - try { - const value = probabilityValue(p, evalConstEnv(defs, 0)); - if (isFinite(value)) eq.info = `≈ ${value.toFixed(4)}`; - } catch { - // Not numerically computable right now (e.g. animated); no readout. + if (!rvNames.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); + const p = toProbability(resolveExpr(parseExpr(probBody, fnNames), getFn, ropts), rvNames); + for (const name of p.rvs) { + if (!rvSys.has(name)) throw new Error(`${name} has an error in its definition.`); + } + // Bounds around an inline expression (`P(0.5 < X + Y < 1.5)`) become + // bounds on an anonymous derived variable, so exactness and shading + // work exactly as for a named one. + let single = p.single; + if (!single && p.inline) { + checkDerived(p.inline.e, rvNames, constNames); + const anon = `@P${eq.id}`; + rvSys.add({ name: anon, kind: 'derived', expr: p.inline.e }); + single = { rv: anon, lo: p.inline.lo, hi: p.inline.hi }; + } + // Constant bounds on one variable whose law is a closed-form pdf get + // the exact CDF and the shader-drawn region. + const exact = single ? rvSys.exactDist(single.rv) : null; + if (single && exact) { + eq.cls = classify(regionExpr(exact, single.lo, single.hi), constNames); + try { + const value = probabilityValue(exact, single.lo, single.hi, evalConstEnv(defs, 0)); + if (isFinite(value)) eq.info = `≈ ${value.toFixed(4)}`; + } catch { + // Not numerically computable right now (e.g. animated); no readout. + } + } else { + // Everything else draws/estimates through the sampled channel — + // but a uniform-sum law still gets its exact value (and its shade + // fills under the exact piecewise-polynomial curve). + const ps = rvSys.bodyParams(p.body); + eq.cls = { + plot: { type: 'prob', body: p.body, shade: single }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(v => v !== 't'), + }; + if (envT0) { + try { + const value = single + ? rvSys.exactProbability(single.rv, single.lo, single.hi, envT0) + : null; + if (value !== null) { + if (isFinite(value)) eq.info = `≈ ${value.toFixed(4)}`; + } else { + const mc = rvSys.probability(p.body, envT0); + if (isFinite(mc)) eq.info = `≈ ${mc.toFixed(3)}`; + } + } catch { /* animated or broken: no readout */ } + } + } + continue; + } + const expectBody = defs.consts.has('E') || defs.fns.has('E') ? null : matchExpectation(text); + if (expectBody !== null) { + if (!rvNames.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); + const ex = toExpectation(resolveExpr(parseExpr(expectBody, fnNames), getFn, ropts), rvNames); + for (const name of ex.rvs) { + if (!rvSys.has(name)) throw new Error(`${name} has an error in its definition.`); + } + // The body becomes the variable whose density carries the marker: a + // bare name is itself, anything else an anonymous derived variable — + // so exact laws (affine in normals, uniform sums) apply unchanged. + let name: string; + if (ex.body.kind === 'var' && rvSys.has(ex.body.name)) { + name = ex.body.name; + } else { + checkDerived(ex.body, rvNames, constNames); + name = `@E${eq.id}`; + rvSys.add({ name, kind: 'derived', expr: ex.body }); + } + const ps = rvSys.bodyParams(ex.body); + eq.cls = { + plot: { type: 'expect', rv: name }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(p => p !== 't'), + }; + if (envT0) { + try { + // Closed form and quadrature both earn full display precision; + // only the Monte Carlo fallback rounds to its noise floor. + const m = rvSys.exactMoments(name, envT0) ?? rvSys.quadMoments(name, envT0); + const value = m ? m.mean : rvSys.mean(name, envT0); + if (isFinite(value)) eq.info = `≈ ${value.toFixed(m ? 4 : 3)}`; + } catch { /* animated or broken: no readout */ } } continue; } @@ -976,7 +1173,24 @@ function recompileAll() { eq.cls = classifySeqRec(seq, fnNames, getFn, constNames, ropts); continue; } - let parsed = resolveExpr(parseExpr(text, fnNames), getFn, ropts); + const rawParsed = parseExpr(text, fnNames); + let parsed = resolveExpr(rawParsed, getFn, ropts); + // A bare expression in random variables (`X + Y`, `X^2`) plots the + // density of that derived variable — distribution arithmetic in place. + const rvRefs = [...freeVars(parsed)].filter(n => rvNames.has(n)); + if (rvRefs.length) { + for (const n of rvRefs) { + if (!rvSys.has(n)) throw new Error(`${n} has an error in its definition.`); + } + if (parsed.kind === 'ineq') { + throw new Error(`An inequality in random variables is a probability: try P(${text}).`); + } + checkDerived(parsed, rvNames, constNames); + const name = `@${eq.id}`; + rvSys.add({ name, kind: 'derived', expr: parsed }); + classifyDerived(eq, name); + continue; + } // Expand point arithmetic and geometry statements (segment, polygon, …) // into scalar expressions; a point name A becomes (A_x, A_y). parsed = lowerGeom(parsed, n => compsOf(defs, n), n => defs.mats.get(n) ?? null); @@ -985,10 +1199,19 @@ function recompileAll() { if (defs.fields.size) parsed = substVars(parsed, fieldEnv); eq.cls = classify(parsed, constNames); eq.parsed = parsed; + // A row that wrote an ∫ and resolved to a constant gets its value as a + // readout (the plot is the horizontal line y = that value). + if (envT0 && usesIntegral(rawParsed)) { + try { + const value = evaluate(parsed, envT0); + if (isFinite(value)) eq.info = `≈ ${Number(value.toPrecision(6))}`; + } catch { /* depends on plot coordinates: the curve is the answer */ } + } } catch (e) { eq.error = e instanceof Error ? e.message : String(e); } } + rvSys.prune(); // sample caches of variables that no longer exist spGen++; // queued hover recomputes predate this compile: drop them spQueue.clear(); setHover(null); @@ -1894,6 +2117,13 @@ const EXAMPLES: Array<[string, Array<[string, string]>]> = [ ['normal density', 'X ~ Normal(0, 1)'], ['P(X < b)', 'a = 1; b = 0.5; X ~ Normal(0, a); P(X < b)'], ['between two bounds', 'X ~ Normal(0, 1); P(-1 < X < 2)'], + ['uniform + exponential', 'X ~ Uniform(0, 2); Y ~ Exponential(1); P(0.5 < X < 1.5)'], + ['sum = convolution', 'X ~ Uniform(0, 1); Y ~ Uniform(0, 1); X + Y'], + ['central limit theorem', 'view(x = -0.5..4.5, y = -0.15..1.35); ' + + 'X1 ~ Uniform(0, 1); X2 ~ Uniform(0, 1); X3 ~ Uniform(0, 1); X4 ~ Uniform(0, 1); ' + + 'S = X1 + X2 + X3 + X4; Z ~ Normal(2, sqrt(1/3)); P(S > 3)'], + ['conditional variable', 'X ~ Normal(0, 1); Y = {X > 0: X^2, 1}; P(Y > 0.5); P(Y > X)'], + ['expectation', 'X ~ Uniform(0, 1); Y = X^2; E(Y); E(X + Y)'], ]], ['regions', [ ['open half-plane', 'y < x/2 + 1'], @@ -1920,6 +2150,11 @@ const EXAMPLES: Array<[string, Array<[string, string]>]> = [ ['function', 'f(x) = x^3 - 3x; y = f(x)'], ['derivative', 'y = d/dx (x^3 - 3x)'], ['tangent line', 'f(x) = x^3 - 2x; g(x) = d/dx f(x); a = 1; y = f(x); y = f(a) + g(a)(x - a)'], + ['running integral', 'view(x = -7..7, y = -1.5..4); f(x) = sin(x)^2; y = f(x); y = int[0..x] f(t) dt'], + ['antiderivative', 'f(x) = x^2 - 1; y = f(x); y = int(f(x) dx)'], + ['gaussian error fn', 'view(x = -4..4, y = -1.2..1.2); y = int[0..x] exp(-t^2) dt'], + ['normal cdf', 'view(x = -4..4, y = -0.6..1.2); y = normalpdf(x, 0, 1); y = int[-inf..x] normalpdf(t, 0, 1) dt'], + ['sine integral Si(x)', 'view(x = -20..20, y = -2.2..2.2); y = int[0..x] sin(t)/t dt'], ['orbiting charge', 'r = 2 + sin(t); ln(w - r) - ln(w + r)'], ]], ['series', [ diff --git a/web/public/llms.txt b/web/public/llms.txt index b62557f..b793159 100644 --- a/web/public/llms.txt +++ b/web/public/llms.txt @@ -4,8 +4,8 @@ > regions, scalar fields, vector fields and ODEs, complex functions (including > domain coloring, conformal maps and escape-time fractals), 3D surfaces, > parametric curves/surfaces, plane geometry (named draggable points, -> segments, polygons, squares, circles), and random variables with shaded -> probabilities. No account, no server round-trips: the graph state +> segments, polygons, squares, circles), symbolic derivatives and integrals, +> and random variables with shaded probabilities. No account, no server round-trips: the graph state > lives in the URL fragment, so you can construct a link that opens the app > with any set of equations already rendered. @@ -54,6 +54,8 @@ Examples (paste-ready): - [Double pendulum](https://equation.io/#g%20%3D%209.8;L1%20%3D%201;L2%20%3D%201;m1%20%3D%201;m2%20%3D%201;M%20%3D%20%5B((m1%2Bm2)%20L1%2C%20m2%20L2%20cos(th_1%20-%20th_2))%2C%20(L1%20cos(th_1%20-%20th_2)%2C%20L2)%5D;f%20%3D%20(-m2%20L2%20om_2%5E2%20sin(th_1%20-%20th_2)%20-%20(m1%2Bm2)%20g%20sin(th_1)%2C%20L1%20om_1%5E2%20sin(th_1%20-%20th_2)%20-%20g%20sin(th_2));th'%20%3D%20om;om'%20%3D%20solve(M%2C%20f);th(0)%20%3D%20(2.5%2C%202.4);b1%20%3D%20(L1%20sin(th_1)%2C%20-L1%20cos(th_1));b2%20%3D%20b1%20%2B%20(L2%20sin(th_2)%2C%20-L2%20cos(th_2));segment((0%2C%200)%2C%20b1);segment(b1%2C%20b2);b1;b2): `M = [((m1+m2) L1, m2 L2 cos(th_1 - th_2)), (L1 cos(th_1 - th_2), L2)]; f = (…); th' = om; om' = solve(M, f); th(0) = (2.5, 2.4); …` — the Lagrangian form M(θ)ω′ = f: th and om are 2-vector states with components th_1, th_2; solve() is Cramer's rule; segment()s between named bob points draw the linkage - [Orbit (vector gravity)](https://equation.io/#r'%20%3D%20vel;vel'%20%3D%20-9%20r%2F%7Cr%7C%5E3;r(0)%20%3D%20(2%2C%200);vel(0)%20%3D%20(0%2C%201.5);segment((0%2C%200)%2C%20r);r;(0%2C%200)): `r' = vel; vel' = -9 r/|r|^3; r(0) = (2, 0); vel(0) = (0, 1.5); segment((0, 0), r); r; (0, 0)` — vector states with point arithmetic; the bare state name draws as a moving point - [Matrix phase portrait](https://equation.io/#a%20%3D%20-1;b%20%3D%20-1%2F4;A%20%3D%20%5B(0%2C%201)%2C%20(a%2C%20b)%5D;(x'%2C%20y')%20%3D%20A%20(x%2C%20y)): `a = -1; b = -1/4; A = [(0, 1), (a, b)]; (x', y') = A (x, y)` — a linear system as its literal matrix, sliders in the entries +- [Central limit theorem](https://equation.io/#X1%20~%20Uniform%280%2C%201%29;X2%20~%20Uniform%280%2C%201%29;X3%20~%20Uniform%280%2C%201%29;X4%20~%20Uniform%280%2C%201%29;S%20%3D%20X1%20%2B%20X2%20%2B%20X3%20%2B%20X4;Z%20~%20Normal%282%2C%20sqrt%281%2F3%29%29;P%28S%20%3E%203%29): `X1 ~ Uniform(0, 1); X2 ~ Uniform(0, 1); X3 ~ Uniform(0, 1); X4 ~ Uniform(0, 1); S = X1 + X2 + X3 + X4; Z ~ Normal(2, sqrt(1/3)); P(S > 3)` — the sum of four uniforms hugs the matching normal +- [Conditional random variable](https://equation.io/#X%20~%20Normal%280%2C%201%29;Y%20%3D%20%7BX%20%3E%200%3A%20X%5E2%2C%201%7D;P%28Y%20%3E%200.5%29;P%28Y%20%3E%20X%29): `X ~ Normal(0, 1); Y = {X > 0: X^2, 1}; P(Y > 0.5); P(Y > X)` — a piecewise transform of X, with probabilities of derived and joint events ## Share links with preview images (/g/) @@ -103,6 +105,21 @@ demand instead of living in a tool description. No authentication required. outside). Bounds must be numbers or constants defined in another row (not `t`), because the sum is expanded symbolically before compiling — which makes `N` a slider that adds terms as you drag it. +- Integrals: `int[a..b] f(x) dx` (definite) and `int(f(x) dx)` + (antiderivative), also written `∫`. The trailing `dx` names the variable; + any letter works, so `y = int[0..x] exp(-t^2) dt` plots a function of x, + and bounds may use sliders, t, or x. Integration is symbolic where a + verified antiderivative exists — polynomials and rational functions + (exactly), the elementary families, u-substitution, integration by parts, + Gaussians via erf — and falls back to Gauss–Legendre quadrature expanded + in place otherwise, so non-elementary rows like `y = int[0..x] sin(t)/t dt` + (Si) still plot. A row whose integral resolves to a constant shows the + numeric value as a readout. Classic notations work: `int[0..1] dx/(1+x^2)`, + and iterated integrals pair inside-out: `int[0..1] int[0..y] x dx dy`. + Bounds may be `±inf` (or `∞`): `int[-inf..x] exp(t) dt` is exp(x), + `int[-inf..inf] exp(-x^2) dx` reads √π, and `int[-inf..x] normalpdf(t, 0, 1) dt` + is the normal cdf. Divergent limits are refused (the closed form is checked + against real quadrature), not misreported. - Comments: a row starting with `#` (e.g. `# tangent lines`) plots nothing; in the app it renders as a heading whose group (the rows until the next `#` row) can be collapsed. Use them to label sections of longer documents. @@ -162,11 +179,43 @@ demand instead of living in a tool description. No authentication required. - Probability: `P(X < 2)`, `P(X > 0)` or `P(0 < X < 2)` shades that area under X's density and shows the numeric probability (normal CDF) as a readout under the row. +- Expectation: `E(X)`, `E(X^2 + Y)` — the mean of any expression in random + variables, shown as a readout under the row and marked on the density. - Parametric curve: components in u — `(2cos(2pi u), sin(4pi u))`; three components make a 3D curve — `(2cos(6pi u), 2sin(6pi u), 4u - 2)`. - Parametric surface: three components in u and v — `(cos(2pi u)(2+cos(2pi v)), sin(2pi u)(2+cos(2pi v)), sin(2pi v))`. - 3D surface: any equation in x, y, z. +- Random variable: `X ~ Normal(mean, sd)` — also `Uniform(lo, hi)` and + `Exponential(rate)`, with short aliases `N`, `U`, `Exp`; a bare name means + the standard parameters (`X ~ N` is Normal(0, 1)). The row plots the exact + density curve; parameters may use sliders and t. +- Derived random variable: `Y = X^2`, `S = X1 + X2` — arithmetic over + declared random variables declares a new one, including piecewise + conditionals: `Y = {X > 0: X^2, 1}`. Distinct names are independent, so + `S = X1 + X2` is the convolution; the same name stays dependent, so + `X + X` is exactly 2X. A bare expression in random variables (`X + Y` as + its own row) plots that density without naming it. Affine forms stay + exact: combinations of normals (`Z = (X + Y)/2`) are the exact normal, a + scaled/shifted uniform is a uniform, and sums of uniforms (`S = X1 + X2`, + the central-limit demo) draw their exact piecewise-polynomial convolution — + the triangle's apex is a true corner. These rows report exact μ and σ. + Everything else (products, nonlinear or piecewise transforms, mixed + families) is estimated from 131072 joint samples and reports μ, σ as `≈`. + Point masses are detected and drawn as probability stems (height = mass): + `Y = {X > 0: 1, 2}` is two stems, `floor(4X)` four, and a mixed result + like `{X > 0: X^2, 1}` draws its stem plus the continuous density. +- Probability: `P(X < 2)`, `P(-1 < X < 2)`, and bounds around one expression + like `P(0.5 < X + Y < 1.5)` — exact (closed-form CDF or polynomial + integration) with the area under the density shaded, whenever the bounded + variable or expression has an exact law. Any other inequality — + `P(Y > 0.5)` for a nonlinear Y, `P(Y > X)` — is estimated from the joint + samples; the row shows the value as `≈`. +- Expectation: `E(X)`, `E(2X + 3)`, `E(X^2 + Y)` — the mean of any + expression in the declared variables, exact whenever the expression has an + exact law (closed-form pdfs, affine combinations of normals, uniform sums) + and the finite-sample mean otherwise. The row shows the value as a readout + and draws a vertical marker at x = E under the expression's density. ## Definitions (extra rows that set up the others) diff --git a/worker/graph.ts b/worker/graph.ts index a326fac..48c1216 100644 --- a/worker/graph.ts +++ b/worker/graph.ts @@ -6,27 +6,30 @@ * by using the exact same lib functions. */ import { - RESERVED, buildDefs, compsOf, defKey, evalConstEnv, resolveExpr, scanDefinition, + usesIntegral, type Definition, type Defs, } from '../lib/defs.ts'; import { - type DistDef, + RVSystem, + buildRVSystem, + checkDerived, densityExpr, + matchExpectation, matchProbability, - parseDistribution, probabilityValue, regionExpr, - scanDistribution, + scanRandomRows, + toExpectation, toProbability, } from '../lib/dist.ts'; -import { type Expr, builtinFn, parseExpr, substVars } from '../lib/expr.ts'; +import { type Expr, evaluate, freeVars, parseExpr, substVars } from '../lib/expr.ts'; import { lowerGeom } from '../lib/geom.ts'; import { type Classified, classify } from '../lib/plot.ts'; import { classifySeqRec, scanSeqRec } from '../lib/seq.ts'; @@ -45,9 +48,10 @@ export interface RowInfo { expr?: Expr; /** Set for `# label` comment rows (group headings in the app; not plotted). */ comment?: boolean; - /** Set for probability rows: `X ~ …` plots a density, `P(…)` a shaded area. */ - dist?: 'density' | 'probability'; - /** Readout shown under the row in the app (the numeric value of a P(…) row). */ + /** Set for probability rows: `X ~ …` plots a density, `P(…)` a shaded + * area, `E(…)` a mean marker. */ + dist?: 'density' | 'probability' | 'expectation'; + /** Readout shown under the row in the app (the numeric value of a P(…) or E(…) row). */ info?: string; error?: string; } @@ -57,20 +61,29 @@ export interface Analysis { defs: Defs; /** Constant values at t = 0 (for static rendering). */ constEnv: Record; + /** Declared random variables (density/prob rows sample through this). */ + rvs: RVSystem; } export function analyze(texts: string[]): Analysis { const rows: RowInfo[] = texts.map(text => ({ text: text.trim() })); + // Random-variable rows (`X ~ …`, and `Y = X^2` referencing one) resolve + // outside the definition system — mirror of web/main.ts recompileAll. + const rvScan = scanRandomRows(rows.map(r => + !r.text || r.text.startsWith('#') || scanSeqRec(r.text) ? null : r.text)); + const rvRowIdx = new Set([...rvScan.base.keys(), ...rvScan.derived.keys()]); + // Pass 1: definitions. A duplicate coordinate-field row (r = 1 + cos(theta) // after r = sqrt(x^2+y^2)) is a plot in that coordinate system, not an error. const raw: Definition[] = []; const defNames = new Set(); const dupRows: RowInfo[] = []; - for (const row of rows) { + for (const [i, row] of rows.entries()) { if (!row.text) continue; // `# label` rows are comments (collapsible group headings in the app). if (row.text.startsWith('#')) { row.comment = true; continue; } + if (rvRowIdx.has(i)) continue; // Sequence/recurrence rows (a_n = …, a_{n+1} = …) are plots, not definitions. if (scanSeqRec(row.text)) continue; const d = scanDefinition(row.text); @@ -114,37 +127,51 @@ export function analyze(texts: string[]): Analysis { if (!fn && fnNames.has(name)) throw new Error(`${name} has an error in its definition.`); return fn; }; - // Random-variable rows (`X ~ Normal(0, a)`) resolve before the plot pass so - // P(…) rows can reference them regardless of row order (as in web/main.ts). - const dists = new Map(); - const distRows = new Set(); - for (const row of rows) { - if (row.def || row.comment || row.error || !row.text) continue; - const scan = scanDistribution(row.text); - if (!scan) continue; - distRows.add(row); - try { - if (RESERVED.has(scan.name) || builtinFn(scan.name)) { - throw new Error(`Cannot use ${scan.name} as a random variable name.`); - } - if (dists.has(scan.name) || defs.consts.has(scan.name) || defs.fns.has(scan.name) || defs.fields.has(scan.name)) { - throw new Error(`${scan.name} is already defined.`); - } - const d = parseDistribution(scan.name, scan.rhs, fnNames); - d.mean = resolveExpr(d.mean, getFn); - d.sd = resolveExpr(d.sd, getFn); - dists.set(scan.name, d); - row.dist = 'density'; - row.expr = densityExpr(d); - row.cls = classify(row.expr, constNames); - } catch (e) { - row.error = e instanceof Error ? e.message : String(e); + + // Random variables next, so P(…) and bare-expression rows can reference + // them regardless of row order. + const rvs = new RVSystem(); + const builtRVs = buildRVSystem(rvs, rvScan, { + fnNames, + getFn, + constNames, + taken: n => defs.consts.has(n) || defs.fns.has(n) || defs.fields.has(n) + || defs.states.has(n) || defs.points.has(n) || defs.mats.has(n), + }); + const rvNames = builtRVs.names; + const densityCls = (name: string): Classified => { + const ps = rvs.paramsOf(name); + return { + plot: { type: 'density', rv: name }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(p => p !== 't'), + }; + }; + for (const [i, name] of builtRVs.rowRV) { + const row = rows[i]; + const message = builtRVs.errors.get(i); + if (message) { + row.error = message; + continue; + } + row.dist = 'density'; + const rv = rvs.get(name)!; + // Base declarations and derived variables with a closed form (affine in + // normal bases) draw the exact pdf; the rest estimate from samples. + const exact = rv.kind === 'base' ? rv.dist : rvs.exactDist(name); + if (exact) { + const density = densityExpr(exact); + row.cls = classify(density, constNames); + row.expr = density; + } else { + row.cls = densityCls(name); } } const seenViewKinds = new Set(); - for (const row of rows) { - if (row.def || row.comment || row.error || !row.text || distRows.has(row)) continue; + for (const [ri, row] of rows.entries()) { + if (row.def || row.comment || row.error || row.cls || !row.text) continue; try { const view = parseViewRow(row.text, constEnv); if (view) { @@ -157,35 +184,142 @@ export function analyze(texts: string[]): Analysis { // defined P themselves, in which case the row is theirs. const probBody = defs.consts.has('P') || defs.fns.has('P') ? null : matchProbability(row.text); if (probBody !== null) { - if (!dists.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); - const p = toProbability(resolveExpr(parseExpr(probBody, fnNames), getFn), dists); + if (!rvNames.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); + const p = toProbability(resolveExpr(parseExpr(probBody, fnNames), getFn), rvNames); + for (const name of p.rvs) { + if (!rvs.has(name)) throw new Error(`${name} has an error in its definition.`); + } row.dist = 'probability'; - row.expr = regionExpr(p); - row.cls = classify(row.expr, constNames); - try { - const value = probabilityValue(p, constEnv); - if (isFinite(value)) row.info = `≈ ${value.toFixed(4)}`; - } catch { - // Not numerically computable at t = 0 (e.g. animated); no readout. + // Inline bounded expressions become anonymous derived variables, so + // shading and exact laws apply — mirror of web/main.ts. + let single = p.single; + if (!single && p.inline) { + checkDerived(p.inline.e, rvNames, constNames); + const anon = `@P${ri}`; + rvs.add({ name: anon, kind: 'derived', expr: p.inline.e }); + single = { rv: anon, lo: p.inline.lo, hi: p.inline.hi }; + } + // Constant bounds on one variable with a closed form get the exact + // CDF and the shader-drawn region; the rest estimate over samples. + const exact = single ? rvs.exactDist(single.rv) : null; + if (single && exact) { + const region = regionExpr(exact, single.lo, single.hi); + row.cls = classify(region, constNames); + row.expr = region; + try { + const value = probabilityValue(exact, single.lo, single.hi, constEnv); + if (isFinite(value)) row.info = `≈ ${value.toFixed(4)}`; + } catch { + // Not numerically computable at t = 0 (e.g. animated); no readout. + } + } else { + const ps = rvs.bodyParams(p.body); + row.cls = { + plot: { type: 'prob', body: p.body, shade: single }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(v => v !== 't'), + }; + try { + // A uniform-sum law still gets its exact value (mirror of the + // app's readout); everything else estimates over joint samples. + const value = single + ? rvs.exactProbability(single.rv, single.lo, single.hi, constEnv) + : null; + if (value !== null) { + if (isFinite(value)) row.info = `≈ ${value.toFixed(4)}`; + } else { + const mc = rvs.probability(p.body, constEnv); + if (isFinite(mc)) row.info = `≈ ${mc.toFixed(3)}`; + } + } catch { /* animated or broken: no readout */ } } continue; } + // `E(…)` is the mean of an expression in random variables — unless the + // user has defined E themselves. Mirror of web/main.ts. + const expectBody = defs.consts.has('E') || defs.fns.has('E') ? null : matchExpectation(row.text); + if (expectBody !== null) { + if (!rvNames.size) throw new Error('Define a random variable first, e.g. X ~ Normal(0, 1).'); + const ex = toExpectation(resolveExpr(parseExpr(expectBody, fnNames), getFn), rvNames); + for (const name of ex.rvs) { + if (!rvs.has(name)) throw new Error(`${name} has an error in its definition.`); + } + row.dist = 'expectation'; + // A bare name is the variable itself; anything else registers as an + // anonymous derived variable, so exact laws apply unchanged. + let name: string; + if (ex.body.kind === 'var' && rvs.has(ex.body.name)) { + name = ex.body.name; + } else { + checkDerived(ex.body, rvNames, constNames); + name = `@E${ri}`; + rvs.add({ name, kind: 'derived', expr: ex.body }); + } + const ps = rvs.bodyParams(ex.body); + row.cls = { + plot: { type: 'expect', rv: name }, + animated: ps.has('t'), + needs3D: false, + params: [...ps].filter(p => p !== 't'), + }; + try { + // Closed form and quadrature both earn full display precision; + // only the Monte Carlo fallback rounds to its noise floor. + const m = rvs.exactMoments(name, constEnv) ?? rvs.quadMoments(name, constEnv); + const value = m ? m.mean : rvs.mean(name, constEnv); + if (isFinite(value)) row.info = `≈ ${value.toFixed(m ? 4 : 3)}`; + } catch { /* animated or broken: no readout */ } + continue; + } const seq = scanSeqRec(row.text); if (seq) { row.cls = classifySeqRec(seq, fnNames, getFn, constNames); continue; } - let parsed = resolveExpr(parseExpr(row.text, fnNames), getFn); + const rawParsed = parseExpr(row.text, fnNames); + let parsed = resolveExpr(rawParsed, getFn); + // A bare expression in random variables plots that derived density. + const rvRefs = [...freeVars(parsed)].filter(n => rvNames.has(n)); + if (rvRefs.length) { + for (const n of rvRefs) { + if (!rvs.has(n)) throw new Error(`${n} has an error in its definition.`); + } + if (parsed.kind === 'ineq') { + throw new Error(`An inequality in random variables is a probability: try P(${row.text}).`); + } + checkDerived(parsed, rvNames, constNames); + row.dist = 'density'; + const name = `@${ri}`; + rvs.add({ name, kind: 'derived', expr: parsed }); + const exactAnon = rvs.exactDist(name); + if (exactAnon) { + const density = densityExpr(exactAnon); + row.cls = classify(density, constNames); + row.expr = density; + } else { + row.cls = densityCls(name); + } + continue; + } // Expand point arithmetic and geometry statements (segment, polygon, …) // into scalar expressions; a point name A becomes (A_x, A_y). parsed = lowerGeom(parsed, n => compsOf(defs, n), n => defs.mats.get(n) ?? null); if (defs.fields.size) parsed = substVars(parsed, fieldEnv); row.cls = classify(parsed, constNames); row.expr = parsed; + // A row that wrote an ∫ and resolved to a constant gets its value as a + // readout (mirror of web/main.ts). + if (usesIntegral(rawParsed)) { + try { + const value = evaluate(parsed, constEnv); + if (isFinite(value)) row.info = `≈ ${Number(value.toPrecision(6))}`; + } catch { /* depends on plot coordinates: the curve is the answer */ } + } } catch (e) { row.error = e instanceof Error ? e.message : String(e); } } - return { rows, defs, constEnv }; + return { rows, defs, constEnv, rvs }; } diff --git a/worker/mcp.test.ts b/worker/mcp.test.ts index 599c8a0..4797394 100644 --- a/worker/mcp.test.ts +++ b/worker/mcp.test.ts @@ -128,6 +128,95 @@ describe('mcp endpoint', () => { expect(out.rows[0].error).toContain('X ~ Normal(0, 1)'); }); + it('validates random-variable rows: base, derived, and P(…) forms', async () => { + const { body } = await rpc('tools/call', { + name: 'create_graph', + arguments: { + equations: ['X ~ Normal(0, 1)', 'Y = {X > 0: X^2, 1}', 'P(-1 < X < 1)', 'P(Y > X)', 'X + X', + 'P(X + X < 1)'], + }, + }); + const out = body.result.structuredContent; + expect(out.valid).toBe(true); + // Every member of the family reports the same human-readable kinds the + // base rows do, whether the density is exact (X, X + X) or sampled (Y) — + // and the inline-bounded P(X + X < 1) is a probability the same way. + expect(out.rows.map((r: { kind?: string }) => r.kind)).toEqual([ + 'random variable (density curve)', + 'random variable (density curve)', + 'probability (shaded area)', + 'probability (shaded area)', + 'random variable (density curve)', + 'probability (shaded area)', + ]); + // The exact P rows read their CDF values; the Monte Carlo one estimates. + expect(out.rows[2].value).toBe('≈ 0.6827'); + expect(out.rows[3].value).toMatch(/^≈ 0\.\d{3}$/); + // X + X ~ Normal(0, 2), so P(X + X < 1) = Φ(1/2) exactly. + expect(out.rows[5].value).toBe('≈ 0.6915'); + expect(out.preview).toBe('attached'); + }); + + it('validates E(…) rows: exact and sampled means', async () => { + const { body } = await rpc('tools/call', { + name: 'create_graph', + arguments: { equations: ['X ~ Normal(2, 1)', 'Y = X^2', 'E(X)', 'E(2X + 1)', 'E(Y)'] }, + }); + const out = body.result.structuredContent; + expect(out.valid).toBe(true); + expect(out.rows.map((r: { kind?: string }) => r.kind)).toEqual([ + 'random variable (density curve)', + 'random variable (density curve)', + 'expectation (mean readout)', + 'expectation (mean readout)', + 'expectation (mean readout)', + ]); + // Exact under the law: E[X] = 2 and E[2X + 1] = 5 (affine in a normal base). + expect(out.rows[2].value).toBe('≈ 2.0000'); + expect(out.rows[3].value).toBe('≈ 5.0000'); + // E[X²] = μ² + σ² = 5, by quadrature against the base pdf. + expect(out.rows[4].value).toBe('≈ 5.0000'); + expect(out.preview).toBe('attached'); + }); + + it('validates ∫ rows: exact readouts and non-elementary curves', async () => { + const { body } = await rpc('tools/call', { + name: 'create_graph', + arguments: { + equations: ['int[0..1] exp(-x^2) dx', 'y = int[0..x] sin(t)/t dt', 'a = 2', 'y = int[0..a] t^2 dt + x'], + }, + }); + const out = body.result.structuredContent; + expect(out.valid).toBe(true); + // ∫₀¹e^(−x²) = (√π/2)erf(1), symbolically; the row reads its value. + expect(out.rows[0].value).toBe('≈ 0.746824'); + // Si(x) has no elementary form — the quadrature expansion still plots. + expect(out.rows[1].kind).toBe('implicit2d'); + expect(out.rows[3].kind).toBe('implicit2d'); // slider bound stays symbolic + expect(out.preview).toBe('attached'); + }); + + it('rejects malformed ∫ rows with a usable message', async () => { + const { body } = await rpc('tools/call', { + name: 'create_graph', + arguments: { equations: ['int(x^2)'] }, + }); + const out = body.result.structuredContent; + expect(out.valid).toBe(false); + expect(out.rows[0].error).toContain('dx'); + }); + + it('leaves E(…) rows alone when the user defines E', async () => { + const { body } = await rpc('tools/call', { + name: 'create_graph', + arguments: { equations: ['E = 3', 'X ~ Normal(0, 1)', 'E(X)'] }, + }); + const out = body.result.structuredContent; + // E is the user's constant, so E(X) is the product E·X — a derived + // density row, not an expectation readout. + expect(out.rows[2].kind).toBe('random variable (density curve)'); + }); + it('reports per-row errors without failing the call', async () => { const { body } = await rpc('tools/call', { name: 'create_graph', diff --git a/worker/mcp.ts b/worker/mcp.ts index 7a4687a..82d8fde 100644 --- a/worker/mcp.ts +++ b/worker/mcp.ts @@ -29,7 +29,7 @@ const TOOLS = [ title: 'Create a graph link', description: `Build a link that opens the equation.io grapher with the given equations already rendered, validating every row through the app's own parser. Pass the COMPLETE graph in "equations": a flat array of strings, one equation or definition per string, in display order — when editing an existing graph (see read_graph), include the unchanged rows too. -Rows can be: equations and inequalities in x,y (curves, regions; z makes it 3D), bare expressions (scalar fields; complex plots via w), points (rows report "draggable"), parametric tuples in u,v — and definitions: "a = 2" (a draggable slider), "f(x) = x^3 - a x", coordinate fields like "r = sqrt(x^2+y^2)" for polar. t animates. Also derivatives d/dx, sums sum[n=1..N], domain()/conformal()/iter() for complex analysis and fractals, y' = ... for ODE slope fields, random variables "X ~ Normal(m, s)"/"P(0 < X < 2)", and "view(x = -5..5, y = -2..2)" / "camera(theta, phi)" framing rows. That is a menu, not the syntax: before your first non-trivial graph, read the "syntax" MCP resource (also at https://equation.io/llms.txt). +Rows can be: equations and inequalities in x,y (curves, regions; z makes it 3D), bare expressions (scalar fields; complex plots via w), points (rows report "draggable"), parametric tuples in u,v — and definitions: "a = 2" (a draggable slider), "f(x) = x^3 - a x", coordinate fields like "r = sqrt(x^2+y^2)" for polar. t animates. Also derivatives d/dx, integrals int[a..b] f dx, sums sum[n=1..N], domain()/conformal()/iter() for complex plots, y' = … slope fields, random variables "X ~ Normal(m, s)"/"P(0) { ? 'random variable (density curve)' : row.dist === 'probability' ? 'probability (shaded area)' - : row.cls!.plot.type, + : row.dist === 'expectation' + ? 'expectation (mean readout)' + : row.cls!.plot.type, ...(row.cls?.animated ? { animated: true } : {}), ...(row.info ? { value: row.info } : {}), ...(drag === undefined ? {} : { draggable: drag }), diff --git a/worker/og.coverage.test.ts b/worker/og.coverage.test.ts index 6837fa4..0f0337c 100644 --- a/worker/og.coverage.test.ts +++ b/worker/og.coverage.test.ts @@ -42,6 +42,11 @@ describe('canRenderOg', () => { expect(canRenderOg(['A = (0, 0)', 'B = (4, 0)', 'C = (0, 4)', 'polygon(A, B, C)'])).toBe(true); // Random-variable rows classify to implicit2d/ineq2d, which both draw. expect(canRenderOg(['X ~ Normal(0, 1)', 'P(X < 1)'])).toBe(true); + // E(…) rows draw as a mean marker. + expect(canRenderOg(['X ~ Normal(0, 1)', 'E(X)'])).toBe(true); + // ∫ rows resolve to ordinary expressions (closed form or quadrature sum). + expect(canRenderOg(['y = int[0..x] exp(-t^2) dt'])).toBe(true); + expect(canRenderOg(['y = int[0..x] sin(t)/t dt'])).toBe(true); }); it('rejects graphs whose preview would be misleading', () => { diff --git a/worker/og.test.ts b/worker/og.test.ts index 579742e..fda0546 100644 --- a/worker/og.test.ts +++ b/worker/og.test.ts @@ -87,6 +87,80 @@ describe('og raster renderer', () => { expect(Math.max(...pixel(r, 90, 10))).toBeLessThan(245); }); + it('draws base distribution rows exactly: pdf curve + shaded P(…) region', () => { + const rows = ['view(x = -4..4, y = -0.05..0.45)', 'X ~ Normal(0, 1)', 'P(X > 1)']; + const r = renderRaster(rows, 100, 100); + // The density passes (0.6, 0.333): screen (57.5, ~48) with the view + // centered at (0, 0.2) and 0.08 units/px — off the axis gridlines, so + // only the curve itself can ink it. + const near: number[] = []; + for (const x of [57, 58]) for (const y of [47, 48, 49]) near.push(Math.min(...pixel(r, x, y))); + expect(Math.min(...near)).toBeLessThan(200); + // Inside the shaded tail (world (1.5, 0.05)) the fill tints the pixel. + expect(Math.min(...pixel(r, 69, 52))).toBeLessThan(245); + // The untouched upper-left corner stays background. + expect(Math.min(...pixel(r, 10, 10))).toBeGreaterThan(240); + }); + + it('draws the piecewise uniform pdf and its region through the VM', () => { + const rows = ['view(x = -1..3, y = -0.1..1.2)', 'X ~ Uniform(0, 2)', 'P(0.5 < X < 1.5)']; + const r = renderRaster(rows, 100, 100); + // Inside the shaded band under pdf = 0.5 (world (0.8, 0.25)) vs the same + // height outside the support (world (2.5, 0.25)). + const inside = Math.min(...pixel(r, 45, 57)); + const outside = Math.min(...pixel(r, 87, 57)); + expect(inside).toBeLessThan(outside - 5); + // The box top, pdf = 0.5 on (0, 2): world (1.4, 0.5) → screen (60, ~51). + const top = [50, 51, 52].map(y => Math.min(...pixel(r, 60, y))); + expect(Math.min(...top)).toBeLessThan(200); + }); + + it('draws sampled densities: sum of uniforms + shaded Monte Carlo P(…)', () => { + const view = 'view(x = -1..3, y = -0.1..1.2)'; + const grid = inkFraction(renderRaster([view], 100, 100)); + const rows = [view, 'X ~ Uniform(0, 1)', 'Y ~ Uniform(0, 1)', 'S = X + Y', 'P(0.5 < S < 1.5)']; + const r = renderRaster(rows, 100, 100); + expect(inkFraction(r)).toBeGreaterThan(grid + 0.02); + // Under the triangle density's peak, inside the shaded band. + expect(Math.min(...pixel(r, 50, 51))).toBeLessThan(245); + // Beyond the sum's support stays background. + expect(Math.min(...pixel(r, 90, 20))).toBeGreaterThan(240); + }); + + it('draws E(…) rows as a mean marker under the density', () => { + const rows = ['view(x = -1..3, y = -0.1..1.2)', 'X ~ Uniform(0, 1.4)', 'E(X)']; + const r = renderRaster(rows, 100, 100); + // E[X] = 0.7 → screen x ≈ 42.5; the stem runs from the axis (py ≈ 64) + // up to pdf = 1/1.4 ≈ 0.714 (py ≈ 46). + const stem = [42, 43].map(x => Math.min(...pixel(r, x, 55))); + expect(Math.min(...stem)).toBeLessThan(210); + // The same height away from the mean is inside the box but unmarked. + expect(Math.min(...pixel(r, 30, 55))).toBeGreaterThan(230); + }); + + it('draws point masses as probability stems', () => { + const rows = ['view(x = 0..3, y = -0.2..1.2)', 'X ~ Normal(0, 1)', 'Y = {X > 0: 1.3, 2.6}']; + const r = renderRaster(rows, 100, 100); + // The stem at x = 1.3 (screen ~43) runs from the axis up to p = 0.5. + expect(Math.min(...pixel(r, 43, 58))).toBeLessThan(200); + // Above the stem top there is nothing — no KDE bump smearing the atom. + expect(Math.min(...pixel(r, 43, 40))).toBeGreaterThan(230); + }); + + it('renders a non-elementary integral curve (Si) through the VM', () => { + // The quadrature-sum expansion is an ordinary expression, so the stack VM + // samples it per pixel like any other implicit curve. + const rows = ['view(x = 0..10, y = 0..2)', 'y = int[0..x] sin(t)/t dt']; + const r = renderRaster(rows, 100, 100); + // Uniform scale: the x span wins, 0.1 units/px with cy = 1. Si peaks at + // x = π with Si(π) ≈ 1.852 → screen (~31, ~41.5). + const near: number[] = []; + for (const x of [30, 31, 32]) for (const y of [40, 41, 42, 43]) near.push(Math.min(...pixel(r, x, y))); + expect(Math.min(...near)).toBeLessThan(200); + // Far from the curve stays background. + expect(Math.min(...pixel(r, 80, 80))).toBeGreaterThan(230); + }); + it('renders definitions + slider constants (tangent-line graph)', () => { const rows = ['f(x) = x^2 - 2x', 'g(x) = d/dx f(x)', 'a = 3', 'y = f(x)', 'y = f(a) + g(a)(x - a)']; expect(inkFraction(renderRaster(rows, 120, 120))).toBeGreaterThan(0.02); diff --git a/worker/og.ts b/worker/og.ts index 309aa49..8eaa703 100644 --- a/worker/og.ts +++ b/worker/og.ts @@ -8,7 +8,8 @@ * (parametric surfaces/curves, z = f(x,y) heightmaps). Output is a PNG built * with CompressionStream — no image library. */ -import { type Expr, substVars } from '../lib/expr.ts'; +import { densityAt, pdfExpr, shadePolygon } from '../lib/dist.ts'; +import { type Expr, evaluate, substVars } from '../lib/expr.ts'; import type { Plot } from '../lib/plot.ts'; import { clampPhi, fitView2D } from '../lib/view.ts'; import { type Analysis, type RowInfo, analyze } from './graph.ts'; @@ -299,10 +300,68 @@ function heightmapExpr(expr: Expr): Expr | null { return zVar(expr.l) ? expr.r : zVar(expr.r) ? expr.l : null; } -function renderRow2D(r: Raster, v: View2D, row: RowInfo, env: EvalEnv, color: [number, number, number]) { +function renderRow2D( + r: Raster, v: View2D, row: RowInfo, env: EvalEnv, color: [number, number, number], analysis: Analysis, +) { const { cls, expr } = row; if (!cls) return; const compile = (e: Expr) => compileFor(env, e); + if (cls.plot.type === 'density' || cls.plot.type === 'prob') { + // Sampled-density rows: the same estimator the app uses (lib/dist.ts), + // drawn as a polyline (density) or a filled area under it (P(…)). + const shade = cls.plot.type === 'prob' ? cls.plot.shade : undefined; + if (cls.plot.type === 'prob' && !shade) return; // readout-only row + const name = cls.plot.type === 'density' ? cls.plot.rv : shade!.rv; + const curve = analysis.rvs.curve(name, analysis.constEnv); + if (!curve) return; + if (cls.plot.type === 'density') { + // Point masses draw as probability stems (height = mass, not density). + for (const a of curve.atoms ?? []) { + const ax = toScreenX(r, v, a.x); + drawLine(r, ax, toScreenY(r, v, 0), ax, toScreenY(r, v, a.p), color); + drawDisc(r, ax, toScreenY(r, v, a.p), 3.5, color); + } + } + if (curve.pts.length < 4) return; + const pts = shade + ? shadePolygon( + curve, + shade.lo ? evaluate(shade.lo, analysis.constEnv) : undefined, + shade.hi ? evaluate(shade.hi, analysis.constEnv) : undefined, + ) + : curve.pts; + if (!pts) return; + const sx: number[] = [], sy: number[] = []; + for (let i = 0; i + 1 < pts.length; i += 2) { + sx.push(toScreenX(r, v, pts[i])); + sy.push(toScreenY(r, v, pts[i + 1])); + } + if (shade) fillPolygon(r, sx, sy, color, 0.16); + for (let i = 0; i + 1 < sx.length; i++) drawLine(r, sx[i], sy[i], sx[i + 1], sy[i + 1], color); + if (shade) drawLine(r, sx[sx.length - 1], sy[sy.length - 1], sx[0], sy[0], color); + return; + } + if (cls.plot.type === 'expect') { + // The mean marker the app draws: a stem from the axis to the density at + // x = E, capped with a dot (web/main.ts case 'expect'). + const name = cls.plot.rv; + const m = analysis.rvs.mean(name, analysis.constEnv); + if (!Number.isFinite(m)) return; + const exact = analysis.rvs.exactDist(name); + let h: number; + try { + h = exact + ? evaluate(pdfExpr(exact, { kind: 'num', value: m }), analysis.constEnv) + : (c => (c ? densityAt(c, m) : 0))(analysis.rvs.curve(name, analysis.constEnv)); + } catch { + return; + } + if (!Number.isFinite(h) || h < 0) h = 0; + const mx = toScreenX(r, v, m); + if (h > 0) drawLine(r, mx, toScreenY(r, v, 0), mx, toScreenY(r, v, h), color); + drawDisc(r, mx, toScreenY(r, v, h), 3.5, color); + return; + } if (cls.plot.type === 'cobweb') { // A recurrence a_{n+1} = f(a_n) draws the same three pieces as the app // (web/main.ts): the map's curve y = f(x), the y = x diagonal the orbit @@ -510,6 +569,13 @@ export const OG_COVERAGE: Record = { // Solution marks require running the numeric solver, which is not wired // into this backend yet. system: 'fallback', + // Sampled densities run the same lib/dist.ts estimator on the CPU: the + // curve is a polyline, a shaded P(…) a filled polygon, an E(…) row a + // vertical mean marker. A readout-only P(…) row (no shade) draws nothing, + // matching the app's canvas. + density: 'draws', + prob: 'draws', + expect: 'draws', }; /** @@ -617,7 +683,7 @@ export function renderRaster(texts: string[], w = OG_WIDTH, h = OG_HEIGHT): Rast const view: View2D = box?.kind === 'view' ? fitView2D(box, w, h) : { cx: 0, cy: 0, upp: 12 / h }; drawGrid2D(raster, view); for (const row of plotRows) { - try { renderRow2D(raster, view, row, env, colorOf(row)); } catch { /* skip row */ } + try { renderRow2D(raster, view, row, env, colorOf(row), analysis); } catch { /* skip row */ } } } return raster; diff --git a/worker/vm.ts b/worker/vm.ts index d3682b1..168780a 100644 --- a/worker/vm.ts +++ b/worker/vm.ts @@ -5,9 +5,9 @@ * sample is too slow and Workers forbid dynamic codegen (`new Function`), so * expressions compile once to opcode arrays run by a small stack machine. */ -import { type Expr, erf, normalcdf, normalpdf } from '../lib/expr.ts'; +import { type Expr, erf, ineqComparisons, normalcdf, normalpdf } from '../lib/expr.ts'; -const enum Op { Const, Var, Add, Sub, Mul, Div, Pow, Neg, Fn1, Fn2, Fn3 } +const enum Op { Const, Var, Add, Sub, Mul, Div, Pow, Neg, Fn1, Fn2, Fn3, Lt, Le, Gt, Ge, Sel } const FN1: Record number> = { sin: Math.sin, cos: Math.cos, tan: Math.tan, @@ -95,6 +95,41 @@ export function compileProg(e: Expr, slots: ReadonlyMap): Prog { } return; } + case 'piecewise': { + // No jumps: every case value evaluates eagerly and Sel keeps the first + // whose condition holds. A NaN in a discarded branch costs nothing. + const emitCond = (cond: Expr): void => { + if (cond.kind !== 'ineq') throw new Error('Piecewise conditions must be inequalities.'); + ineqComparisons(cond).forEach(({ op, l, r }, k) => { + emit(l); + emit(r); + code.push(op === '<' ? Op.Lt : op === '<=' ? Op.Le : op === '>' ? Op.Gt : Op.Ge, 0); + push(-1); + if (k > 0) { + code.push(Op.Mul, 0); // AND of 0/1 masks + push(-1); + } + }); + }; + const emitCases = (k: number): void => { + if (k === node.cases.length) { + if (node.otherwise) emit(node.otherwise); + else { + code.push(Op.Const, consts.length); + consts.push(NaN); + push(1); + } + return; + } + emitCond(node.cases[k].cond); + emit(node.cases[k].value); + emitCases(k + 1); + code.push(Op.Sel, 0); + push(-2); + }; + emitCases(0); + return; + } default: throw new Error(`Cannot evaluate a ${node.kind} node numerically.`); } @@ -124,6 +159,13 @@ export function run(p: Prog, vars: ArrayLike, stack: Float64Array): numb case Op.Fn1: stack[sp - 1] = FN1_TABLE[arg](stack[sp - 1]); break; case Op.Fn2: sp--; stack[sp - 1] = FN2_TABLE[arg](stack[sp - 1], stack[sp]); break; case Op.Fn3: sp -= 2; stack[sp - 1] = FN3_TABLE[arg](stack[sp - 1], stack[sp], stack[sp + 1]); break; + // Comparisons yield 1/0 masks (0 for NaN operands, like a false branch). + case Op.Lt: sp--; stack[sp - 1] = stack[sp - 1] < stack[sp] ? 1 : 0; break; + case Op.Le: sp--; stack[sp - 1] = stack[sp - 1] <= stack[sp] ? 1 : 0; break; + case Op.Gt: sp--; stack[sp - 1] = stack[sp - 1] > stack[sp] ? 1 : 0; break; + case Op.Ge: sp--; stack[sp - 1] = stack[sp - 1] >= stack[sp] ? 1 : 0; break; + // [cond, then, else] → the first matching case wins. + case Op.Sel: sp -= 2; stack[sp - 1] = stack[sp - 1] === 1 ? stack[sp] : stack[sp + 1]; break; } } return stack[sp - 1];