consequence-ledger 0.3.0__py3-none-any.whl

ledger/estimate.py

@@ -0,0 +1,412 @@
+"""Off-policy estimates of each arm's bad-outcome rate.
+
+For every arm a we estimate E[Y(a)]: the 30-day bad-outcome rate if *every*
+decision had gone to a. Estimators:
+
+  naive  -- plain mean of Y among decisions that got a (biased when routing
+            depends on the task, e.g. hard tasks go to the "strong" agent)
+  ipw    -- inverse-propensity weighting with the logged probabilities
+  snipw  -- self-normalised IPW (lower variance, small bias)
+  dm     -- direct method: per-stratum outcome model, averaged over the population
+  dr     -- doubly robust: dm plus an IPW correction of its residuals
+
+Confidence intervals come from a percentile bootstrap over decisions.
+"""
+import hashlib
+import json
+import math
+import random
+from collections import defaultdict
+
+ESTIMATORS = ("naive", "ipw", "snipw", "dm", "dr")
+# The outcome model is fitted on other folds than the rows it corrects (cross-fitting). In a simulation with
+# 100 strata and 400 decisions the doubly-robust interval covered the truth 87.7% of the time fitted in-sample
+# and 95.3% cross-fitted. Below MIN_CROSSFIT rows a fold is too small to fit anything, so it is switched off.
+DEFAULT_FOLDS = 5
+MIN_CROSSFIT = 50
+
+
+def load_rows(con, stratum_key="stratum", cluster_key=None, window="full"):
+    """Attributed decisions whose outcome window has closed.
+
+    stratum_key: one context key, or several joined by commas ("difficulty,language"); the outcome
+    model then works per combination. cluster_key names a context field (e.g. "developer") whose
+    values share outcomes; intervals then resample whole clusters instead of single decisions.
+    window: "full" uses the main horizon; "early" the early-warning window; "both" keeps decisions
+    with both windows closed and adds the early label as "y_early".
+    """
+    keys = [k.strip() for k in (stratum_key or "").split(",") if k.strip()]
+    where = {"full": "o.mature = 1", "early": "o.early_mature = 1",
+             "both": "o.mature = 1 AND o.early_mature = 1"}[window]
+    rows = []
+    query = f"""SELECT d.id, d.chosen, d.propensities, d.context, o.bad, o.early_bad
+                FROM decisions d JOIN outcomes o ON o.decision_id = d.id
+                WHERE {where} AND o.kind != 'no_commit'
+                ORDER BY d.ts, d.id"""
+    for r in con.execute(query):
+        context = json.loads(r["context"])
+        row = {
+            "id": r["id"],
+            "arm": r["chosen"],
+            "y": r["early_bad"] if window == "early" else r["bad"],
+            "p": json.loads(r["propensities"]),
+            "s": "|".join(str(context.get(k, "all")) for k in keys) if keys else "all",
+            "c": str(context.get(cluster_key, r["id"])) if cluster_key else r["id"],
+        }
+        if window == "both":
+            row["y_early"] = r["early_bad"]
+        rows.append(row)
+    return rows
+
+
+def early_readout(con, baseline, stratum_key="stratum", cluster_key=None, planned_n=2000):
+    """First look from the early window, plus how well early differences have tracked final ones.
+
+    Early labels see only the bad outcomes that surface fast, so early differences are typically
+    smaller than final ones. The calibration compares both on decisions old enough to have both.
+    """
+    rows = load_rows(con, stratum_key, cluster_key, "early")
+    if not rows:
+        return None
+    arms = sorted({r["arm"] for r in rows} | {a for r in rows for a in r["p"]})
+    if baseline not in arms:
+        baseline = arms[0]
+    early = analytic(rows, arms, baseline, t_star=planned_n)
+    result = {"n": len(rows), "baseline": baseline,
+              "arms": {a: {k: early[a]["dr"][k] for k in ("value", "diff", "diff_lo", "diff_hi", "cs_lo", "cs_hi")}
+                       for a in arms}}
+    both = load_rows(con, stratum_key, cluster_key, "both")
+    total_bad = sum(r["y"] for r in both)
+    if both and total_bad:
+        early_rows = [dict(r, y=r["y_early"]) for r in both]
+        e = point_estimates(early_rows, arms)
+        f = point_estimates(both, arms)
+        pairs = [(e[a]["dr"] - e[baseline]["dr"], f[a]["dr"] - f[baseline]["dr"]) for a in arms if a != baseline]
+        result["calibration"] = {
+            "n": len(both),
+            # Share of the final bad outcomes that had already surfaced in the early window.
+            "surfaced_early": sum(r["y_early"] for r in both) / total_bad,
+            "pairs": {a: {"early_diff": pe, "final_diff": pf} for a, (pe, pf) in zip([a for a in arms if a != baseline], pairs)},
+            "same_direction": sum((pe > 0) == (pf > 0) for pe, pf in pairs if pe and pf),
+            "compared": sum(1 for pe, pf in pairs if pe and pf),
+        }
+    return result
+
+
+def merge_balance(con):
+    """Per-arm share of decisions that never merged, and a chi-square test that it is the same for all arms.
+
+    Unmerged decisions are excluded from estimation. If one arm's work is abandoned more often,
+    the merged sets are no longer comparable and the estimates are biased.
+    """
+    counts = {}
+    for r in con.execute("""SELECT d.chosen arm, SUM(o.kind = 'no_commit') unmerged, COUNT(*) n
+                            FROM decisions d JOIN outcomes o ON o.decision_id = d.id GROUP BY d.chosen"""):
+        counts[r["arm"]] = (r["unmerged"], r["n"])
+    total_u = sum(u for u, _ in counts.values())
+    total_n = sum(n for _, n in counts.values())
+    chi2 = 0.0
+    if 0 < total_u < total_n:
+        rate = total_u / total_n
+        for u, n in counts.values():
+            for observed, expected in ((u, n * rate), (n - u, n * (1 - rate))):
+                if expected > 0:
+                    chi2 += (observed - expected) ** 2 / expected
+    df = len(counts) - 1
+    return {"arms": {a: {"unmerged": u, "n": n, "rate": u / n if n else float("nan")} for a, (u, n) in counts.items()},
+            "chi2": chi2, "df": df, "p": chi2_sf(chi2, df) if df > 0 else None}
+
+
+def chi2_sf(x, df):
+    """Upper tail of the chi-square distribution (Wilson-Hilferty approximation; fine for a warning)."""
+    if x <= 0:
+        return 1.0
+    z = ((x / df) ** (1 / 3) - (1 - 2 / (9 * df))) / math.sqrt(2 / (9 * df))
+    return 0.5 * math.erfc(z / math.sqrt(2))
+
+
+def outcome_model(rows):
+    """Mean Y per (stratum, arm), falling back to the arm mean for empty cells."""
+    cell = defaultdict(lambda: [0, 0])
+    arm_tot = defaultdict(lambda: [0, 0])
+    for r in rows:
+        cell[(r["s"], r["arm"])][0] += r["y"]
+        cell[(r["s"], r["arm"])][1] += 1
+        arm_tot[r["arm"]][0] += r["y"]
+        arm_tot[r["arm"]][1] += 1
+
+    def model(s, a):
+        y, k = cell.get((s, a), (0, 0))
+        if k:
+            return y / k
+        y, k = arm_tot.get(a, (0, 0))
+        return y / k if k else float("nan")
+
+    return model
+
+
+def fold_of(row, folds, position=0):
+    """Stable fold for a row. Rows of one cluster share a fold, so correlated outcomes never leak across folds."""
+    key = str(row.get("c", row.get("id", position)))
+    return int(hashlib.sha1(key.encode("utf-8")).hexdigest()[:8], 16) % folds
+
+
+def _tally(rows):
+    cell, arm_tot = defaultdict(lambda: [0, 0]), defaultdict(lambda: [0, 0])
+    for r in rows:
+        for table, key in ((cell, (r["s"], r["arm"])), (arm_tot, r["arm"])):
+            table[key][0] += r["y"]
+            table[key][1] += 1
+    return cell, arm_tot
+
+
+def _model_from(cell, arm_tot):
+    def model(s, a):
+        y, k = cell.get((s, a), (0, 0))
+        if k:
+            return y / k
+        y, k = arm_tot.get(a, (0, 0))
+        return y / k if k else float("nan")
+
+    return model
+
+
+def row_models(rows, folds=None):
+    """The outcome model to use for each row: fitted on the other folds when folds >= 2, else on all rows.
+
+    Counts are tallied once per fold and subtracted from the totals, so cross-fitting costs one pass.
+    """
+    full = outcome_model(rows)
+    if not folds or folds < 2 or len(rows) < folds:
+        return [full] * len(rows)
+    assign = []
+    for i, r in enumerate(rows):
+        if r.get("_folds") != folds:       # cached on the row: bootstrap resamples reuse the same row objects
+            r["_fold"], r["_folds"] = fold_of(r, folds, i), folds
+        assign.append(r["_fold"])
+    per_fold = [(defaultdict(lambda: [0, 0]), defaultdict(lambda: [0, 0])) for _ in range(folds)]
+    total_cell, total_arm = defaultdict(lambda: [0, 0]), defaultdict(lambda: [0, 0])
+    for r, f in zip(rows, assign):
+        key = (r["s"], r["arm"])
+        for table, k in ((per_fold[f][0], key), (per_fold[f][1], r["arm"]), (total_cell, key), (total_arm, r["arm"])):
+            table[k][0] += r["y"]
+            table[k][1] += 1
+    models = []
+    for k in range(folds):
+        cell_k, arm_k = per_fold[k]
+        cell = {key: (v[0] - cell_k[key][0], v[1] - cell_k[key][1]) for key, v in total_cell.items()} if cell_k else total_cell
+        arm = {key: (v[0] - arm_k[key][0], v[1] - arm_k[key][1]) for key, v in total_arm.items()} if arm_k else total_arm
+        fitted = _model_from(cell, arm)
+
+        def model(s, a, fitted=fitted):
+            value = fitted(s, a)
+            return full(s, a) if value != value else value   # an arm unseen in the training folds
+
+        models.append(model)
+    return [models[f] for f in assign]
+
+
+def contributions(rows, arms, clip=None, folds=None):
+    """Per-decision terms whose means are the IPW and DR estimates (their influence functions)."""
+    models = row_models(rows, folds)
+    out = {}
+    for a in arms:
+        ipw, dr, dm, w_all = [], [], [], []
+        for r, model in zip(rows, models):
+            m = model(r["s"], a)
+            w = 0.0
+            if r["arm"] == a:
+                w = 1.0 / r["p"][a]
+                if clip:
+                    w = min(w, clip)
+            ipw.append(w * r["y"])
+            dr.append(m + w * (r["y"] - m))
+            dm.append(m)
+            w_all.append(w)
+        out[a] = {"ipw": ipw, "dr": dr, "dm": dm, "w": w_all}
+    return out
+
+
+def point_estimates(rows, arms, clip=None, folds=None):
+    n = len(rows)
+    nan = float("nan")
+    out = {}
+    for a, c in contributions(rows, arms, clip, folds).items():
+        taken = [r["y"] for r in rows if r["arm"] == a]
+        w_sum = sum(c["w"])
+        out[a] = {
+            "naive": sum(taken) / len(taken) if taken else nan,
+            "ipw": sum(c["ipw"]) / n if n else nan,
+            "snipw": sum(c["ipw"]) / w_sum if w_sum else nan,
+            "dm": sum(c["dm"]) / n if n else nan,
+            "dr": sum(c["dr"]) / n if n else nan,
+        }
+    return out
+
+
+def _mean_se(values, clusters=None):
+    """Mean and standard error; with clusters, the cluster-robust (sandwich) standard error."""
+    n = len(values)
+    mean = sum(values) / n
+    if clusters is None:
+        var = sum((v - mean) ** 2 for v in values) / (n - 1) if n > 1 else float("nan")
+        return mean, math.sqrt(var / n)
+    sums = defaultdict(float)
+    for v, g in zip(values, clusters):
+        sums[g] += v - mean
+    k = len(sums)
+    if k < 2:
+        return mean, float("nan")
+    return mean, math.sqrt(k / (k - 1) * sum(t * t for t in sums.values())) / n
+
+
+def cs_radius(sigma, t, alpha=0.05, t_star=2000):
+    """Half-width of the Gaussian-mixture asymptotic confidence sequence.
+
+    Waudby-Smith, Arbour, Sinha, Kennedy & Ramdas, "Time-uniform central limit theory and
+    asymptotic confidence sequences", Ann. Statist. 52(6), 2024, Eq. 13. Unlike a fixed-n CI it
+    stays valid when the report is re-run every week. rho is fixed in advance so the sequence is
+    tightest near t_star decisions; it must not be tuned on the data.
+    """
+    if t < 2 or math.isnan(sigma):
+        return float("nan")
+    rho2 = (-2 * math.log(alpha) + math.log(-2 * math.log(alpha) + 1)) / t_star
+    return sigma * math.sqrt(2 * (t * rho2 + 1) / (t * t * rho2) * math.log(math.sqrt(t * rho2 + 1) / alpha))
+
+
+def floor_se(counts_a, counts_b=None):
+    """Smallest standard error believable for a rate (or a difference of two) given how few bad outcomes it saw.
+
+    Adds one pseudo-bad and one pseudo-clean outcome per arm (Agresti and Caffo 2000). In simulation, intervals
+    for a 3% bad rate from 300 decisions covered 92.5% of the time without this floor and 96.5% with it, and
+    it changes nothing visible once there are hundreds of bad outcomes. counts are (bad, n) pairs.
+    """
+    var = 0.0
+    for counts in (counts_a, counts_b):
+        if counts and counts[1] > 0:
+            p = (counts[0] + 1) / (counts[1] + 2)
+            var += p * (1 - p) / counts[1]
+    return math.sqrt(var)
+
+
+def _counts(rows, arms):
+    out = {a: [0, 0] for a in arms}
+    for r in rows:
+        if r["arm"] in out:
+            out[r["arm"]][0] += r["y"]
+            out[r["arm"]][1] += 1
+    return out
+
+
+def analytic(rows, arms, baseline, clip=None, z=1.959964, t_star=2000, folds=None):
+    """Normal-approximation CIs from influence functions: fast, no resampling.
+
+    Treats the outcome model as fixed, which is accurate when strata are few and well populated.
+    """
+    c = contributions(rows, arms, clip, folds)
+    clusters = [r.get("c") for r in rows] if is_clustered(rows) else None
+    counts = _counts(rows, arms)
+    result = {}
+    for a in arms:
+        result[a] = {}
+        for e in ("ipw", "dr"):
+            m, se = _mean_se(c[a][e], clusters)
+            dm, dse = _mean_se([x - y for x, y in zip(c[a][e], c[baseline][e])], clusters)
+            se = max(se, floor_se(counts[a]))
+            dse = max(dse, floor_se(counts[a], counts[baseline]) if a != baseline else 0.0)
+            # se is sigma/sqrt(n) (cluster-robust if clustered), so sigma = se * sqrt(n).
+            n = len(rows)
+            r = cs_radius(dse * math.sqrt(n), n, 1 - math.erf(z / math.sqrt(2)), t_star)
+            result[a][e] = {"value": m, "se": se, "lo": m - z * se, "hi": m + z * se,
+                            "diff": dm, "diff_se": dse, "diff_lo": dm - z * dse, "diff_hi": dm + z * dse,
+                            "cs_lo": dm - r, "cs_hi": dm + r}
+    return result
+
+
+def is_clustered(rows):
+    return len({r.get("c", i) for i, r in enumerate(rows)}) < len(rows)
+
+
+def _percentile(values, q):
+    values = sorted(v for v in values if not math.isnan(v))
+    if not values:
+        return float("nan")
+    i = (len(values) - 1) * q
+    lo, hi = math.floor(i), math.ceil(i)
+    return values[lo] + (values[hi] - values[lo]) * (i - lo)
+
+
+def diagnostics(rows, arms):
+    diag = {}
+    for a in arms:
+        w = [1.0 / r["p"][a] for r in rows if r["arm"] == a]
+        diag[a] = {
+            "n": len(w),
+            "bad": sum(r["y"] for r in rows if r["arm"] == a),
+            "min_propensity": min((r["p"][a] for r in rows), default=float("nan")),
+            "max_weight": max(w, default=float("nan")),
+            # Kish effective sample size of the IPW weights.
+            "ess": (sum(w) ** 2 / sum(x * x for x in w)) if w else 0.0,
+        }
+    return diag
+
+
+def _widen(value, bound, pick):
+    return bound if value != value else pick(value, bound)
+
+
+def estimate(rows, arms=None, baseline=None, boot=500, clip=None, seed=0, alpha=0.05, planned_n=2000,
+             folds=DEFAULT_FOLDS):
+    if not rows:
+        raise ValueError("no mature, attributed decisions to estimate from")
+    arms = arms or sorted({r["arm"] for r in rows} | {a for r in rows for a in r["p"]})
+    baseline = baseline or arms[0]
+    if baseline not in arms:
+        raise ValueError(f"baseline {baseline!r} is not one of {arms}")
+    folds = folds if folds and len(rows) >= MIN_CROSSFIT else 0
+    point = point_estimates(rows, arms, clip, folds)
+
+    rng = random.Random(seed)
+    samples = defaultdict(list)  # (arm, est) and ("diff", arm, est)
+    n = len(rows)
+    groups = defaultdict(list)
+    for i, r in enumerate(rows):
+        groups[r.get("c", i)].append(r)
+    groups = list(groups.values())
+    for _ in range(boot):
+        # Resample whole clusters so correlated outcomes (same developer, same service) widen the CI.
+        resample = [r for _ in range(len(groups)) for r in groups[rng.randrange(len(groups))]]
+        b = point_estimates(resample, arms, clip, folds)
+        for a in arms:
+            for e in ESTIMATORS:
+                samples[(a, e)].append(b[a][e])
+                samples[("diff", a, e)].append(b[a][e] - b[baseline][e])
+
+    lo_q, hi_q = alpha / 2, 1 - alpha / 2
+    result = {"baseline": baseline, "arms": {}, "n": n, "clusters": len(groups), "boot": boot, "clip": clip,
+              "diagnostics": diagnostics(rows, arms),
+              "analytic": analytic(rows, arms, baseline, clip, t_star=planned_n, folds=folds),
+              "planned_n": planned_n, "folds": folds}
+    counts = _counts(rows, arms)
+    z = 1.959964
+    for a in arms:
+        entry = {"estimates": {}, "vs_baseline": {}}
+        half = z * floor_se(counts[a])
+        half_diff = z * floor_se(counts[a], counts[baseline]) if a != baseline else 0.0
+        for e in ESTIMATORS:
+            entry["estimates"][e] = {
+                "value": point[a][e],
+                "lo": _widen(_percentile(samples[(a, e)], lo_q), point[a][e] - half, min),
+                "hi": _widen(_percentile(samples[(a, e)], hi_q), point[a][e] + half, max),
+            }
+            diffs = samples[("diff", a, e)]
+            entry["vs_baseline"][e] = {
+                "value": point[a][e] - point[baseline][e],
+                "lo": _widen(_percentile(diffs, lo_q), point[a][e] - point[baseline][e] - half_diff, min),
+                "hi": _widen(_percentile(diffs, hi_q), point[a][e] - point[baseline][e] + half_diff, max),
+                # Two-sided bootstrap p-value for "no difference".
+                "p": (min(1.0, 2 * min(sum(d <= 0 for d in diffs), sum(d >= 0 for d in diffs)) / len(diffs))
+                      if a != baseline and diffs else None),
+            }
+        result["arms"][a] = entry
+    return result

ledger/gitio.py

@@ -0,0 +1,102 @@
+"""Thin wrappers around the git CLI."""
+import hashlib
+import subprocess
+from pathlib import Path
+
+
+class GitError(RuntimeError):
+    pass
+
+
+def git(repo, *args, input=None):
+    """Run git in repo and return its stdout. `input` (str) is fed to stdin."""
+    data = input.encode("utf-8", "replace") if isinstance(input, str) else input
+    proc = subprocess.run(["git", "-C", str(repo), *args], capture_output=True, input=data)
+    if proc.returncode != 0:
+        raise GitError(f"git {' '.join(args)}: {proc.stderr.decode('utf-8', 'replace').strip()}")
+    return proc.stdout.decode("utf-8", "replace")
+
+
+def repo_root(path="."):
+    """Root of the main working tree for any folder inside a repository or one of its linked worktrees.
+
+    The ledger lives there, so every worktree and subfolder shares one. A folder that is not inside
+    a git repository is returned unchanged and each command reports its own error.
+    """
+    path = Path(path)
+    try:
+        common = git(path, "rev-parse", "--git-common-dir").strip()
+    except (GitError, OSError):
+        return path
+    common = Path(common) if Path(common).is_absolute() else path / common
+    common = common.resolve()
+    if common.name == ".git":
+        return common.parent
+    return Path(git(path, "rev-parse", "--show-toplevel").strip())  # bare or custom gitdir layouts
+
+
+def current_branch(repo):
+    """The checked-out branch name, or None when HEAD is detached."""
+    try:
+        return git(repo, "symbolic-ref", "--short", "-q", "HEAD").strip() or None
+    except GitError:
+        return None
+
+
+def branch_key(branch):
+    """Filename for a branch's active decision: git's own blob hash of the name, so any branch name
+    (slashes, unicode, 'a/b' vs 'a_b') maps to a distinct, filesystem-safe key. The commit hook
+    computes the same value with `git hash-object --stdin`."""
+    # A detached HEAD gets a name no real branch can have: git forbids spaces in ref names.
+    data = (branch if branch is not None else "(detached HEAD)").encode("utf-8")
+    return hashlib.sha1(b"blob " + str(len(data)).encode() + b"\0" + data).hexdigest()
+
+
+def safe_rev(rev):
+    """A revision from user input, refused if git could mistake it for an option (e.g. --output=file)."""
+    rev = str(rev)
+    if not rev or rev.startswith("-") or "\0" in rev or "\n" in rev:
+        raise ValueError(f"{rev!r} is not a valid revision (it must not start with '-')")
+    return rev
+
+
+def log_commits(repo, rev="HEAD", all_refs=False):
+    """All commits reachable from rev, oldest first.
+
+    With all_refs, every commit git still has: all branches, remote-tracking refs and reflog entries.
+    That includes work on branches that were squash-merged and deleted.
+    """
+    scope = ["--all", "--reflog"] if all_refs else [safe_rev(rev)]
+    out = git(repo, "log", "--reverse", "--format=%H%x1f%P%x1f%ct%x1f%an <%ae>%x1f%B%x1e", *scope)
+    commits = []
+    for record in out.split("\x1e"):
+        record = record.strip("\n")
+        if not record:
+            continue
+        sha, parents, ts, author, body = record.split("\x1f", 4)
+        commits.append({
+            "sha": sha,
+            "parents": parents.split(),
+            "ts": int(ts),
+            "author": author,
+            "body": body,
+            "subject": body.split("\n", 1)[0],
+        })
+    return commits
+
+
+def lines_changed(repo, rev="HEAD"):
+    """{sha: lines added + deleted} for non-merge commits (binary files count as 0)."""
+    out = git(repo, "log", "--no-merges", "--format=%x1e%H", "--numstat", safe_rev(rev))
+    sizes = {}
+    for record in out.split("\x1e"):
+        parts = record.strip().split("\n")
+        if not parts or not parts[0]:
+            continue
+        total = 0
+        for line in parts[1:]:
+            cols = line.split("\t")
+            if len(cols) >= 3:
+                total += sum(int(c) for c in cols[:2] if c.isdigit())
+        sizes[parts[0]] = total
+    return sizes