1"""Run the cached-estimate grid and fold the results into the data branch.
3 python scripts/run_sweep.py --data-dir data
5Appends one line per cell to <data-dir>/runs.jsonl and rebuilds
6<data-dir>/estimates.json from the whole log. Each dispatch draws fresh,
7independent pasts (the seed advances with the number of runs already recorded
8for that cell), so the pooled mean tightens run after run.
10Driven by .github/workflows/estimates.yml, but works standalone: point it at
11any directory and run it as often as you like.
12"""
14import argparse
15import json
16import math
17import os
18import time
19import zlib
20from datetime import datetime, timezone
21from pathlib import Path
23import numpy as np
25import grid # first: also puts src/ on sys.path for a bare checkout
26from timeseries_entropy import estimate_conditional_entropy
27from timeseries_entropy.theory import predict_entropy_rate
29RUNS = 'runs.jsonl'
30SUMMARY = 'estimates.json'
33def _thin_arg(s):
34 return s if s == 'auto' else int(s)
37def main():
38 ap = argparse.ArgumentParser(description=__doc__)
39 ap.add_argument('--data-dir', required=True,
40 help='checkout of the estimates branch')
41 ap.add_argument('--pasts', type=int, default=8,
42 help='independent pasts added per cell per run')
43 ap.add_argument('--reps', type=int, default=8,
44 help='per-past realization budget at thin=1')
45 ap.add_argument('--n0', type=int, default=128)
46 ap.add_argument('--r', type=float, default=1.5)
47 ap.add_argument('--thin', type=_thin_arg, default='auto',
48 help="Gibbs sweeps per draw, or 'auto' to match each "
49 "chain's measured autocorrelation time (default)")
50 ap.add_argument('--workers', type=int)
51 ap.add_argument('--only', help='substring of cell id, for local testing')
52 args = ap.parse_args()
54 data_dir = Path(args.data_dir)
55 data_dir.mkdir(parents=True, exist_ok=True)
56 records = read_records(data_dir / RUNS)
57 runs_so_far = {}
58 for rec in records:
59 runs_so_far[rec['cell']] = runs_so_far.get(rec['cell'], 0) + 1
61 cells = [c for c in grid.grid()
62 if args.only is None or args.only in grid.cell_id(c)]
63 print(f'{len(cells)} cells x {args.pasts} pasts x {args.reps} reps '
64 f'({len(records)} records already cached)')
66 new = []
67 t_all = time.time()
68 for i, params in enumerate(cells):
69 cid = grid.cell_id(params)
70 run_index = runs_so_far.get(cid, 0)
71 kernel = grid.kernel_from_params(params)
72 M = max(512, 4 * kernel.size)
73 seed = [zlib.crc32(cid.encode()), run_index]
74 t0 = time.time()
75 est = estimate_conditional_entropy(
76 kernel, params['sigma'], past=M, pasts=args.pasts, reps=args.reps,
77 n0=args.n0, r=args.r, thin=args.thin, seed=seed,
78 workers=args.workers)
79 v = np.asarray(est.per_past, dtype=float)
80 rec = {
81 'cell': cid,
82 'params': params,
83 'M': M,
84 'L': int(kernel.size),
85 'pasts': int(v.size),
86 'reps': args.reps,
87 'n0': args.n0,
88 'r': args.r,
89 'thin': args.thin,
90 'sum': float(v.sum()),
91 'sumsq': float((v ** 2).sum()),
92 # Per-past forensics: the values themselves (so outliers are
93 # visible directly), the resolved thinning and realization
94 # counts, and each probe's autocorrelation-time estimate.
95 'values': [float(x) for x in v],
96 'thin_resolved': [int(t) for t in est.thin],
97 'reps_resolved': [int(k) for k in est.reps],
98 'tau': [None if math.isnan(t) else round(float(t), 1)
99 for t in est.tau],
100 'seed': seed,
101 'run_index': run_index,
102 'utc': utc_now(),
103 'code_sha': os.environ.get('GITHUB_SHA'),
104 'workflow_run': os.environ.get('GITHUB_RUN_ID'),
105 }
106 new.append(rec)
107 thins = est.thin
108 print(f' [{i + 1:2d}/{len(cells)}] {cid:<16} '
109 f'H = {est.mean:.4f} +/- {est.se:.4f} '
110 f'thin {int(thins.min())}-{int(thins.max())} '
111 f'({time.time() - t0:.0f}s)', flush=True)
113 append_records(data_dir / RUNS, new)
114 summary = summarize(records + new, args)
115 write_json(data_dir / SUMMARY, summary)
116 write_readme(data_dir / 'README.md', summary)
117 print(f'\n{len(new)} records appended in {time.time() - t_all:.0f}s; '
118 f'{len(summary["cells"])} cells summarized')
119 step_summary(summary)
122def read_records(path):
123 if not path.exists():
124 return []
125 with path.open() as f:
126 return [json.loads(line) for line in f if line.strip()]
129def append_records(path, records):
130 with path.open('a') as f:
131 for rec in records:
132 f.write(json.dumps(rec, sort_keys=True) + '\n')
135def summarize(records, args):
136 """Pool every recorded per-past value, per cell, into a mean and an SE."""
137 cells = {}
138 for rec in records:
139 c = cells.setdefault(rec['cell'], {
140 'id': rec['cell'], 'n_pasts': 0, 'n_runs': 0,
141 '_sum': 0.0, '_sumsq': 0.0})
142 c['params'] = rec['params']
143 c['M'] = rec['M']
144 c['L'] = rec['L']
145 c['n_pasts'] += rec['pasts']
146 c['n_runs'] += 1
147 c['_sum'] += rec['sum']
148 c['_sumsq'] += rec['sumsq']
149 c['last_utc'] = max(rec['utc'], c.get('last_utc', ''))
151 out = []
152 for c in sorted(cells.values(), key=lambda c: c['id']):
153 n, s, ss = c['n_pasts'], c.pop('_sum'), c.pop('_sumsq')
154 mean = s / n
155 # Sample variance of the pooled per-past values; the pasts are
156 # independent across runs as well as within one, so the SE of their
157 # mean is the honest error bar on the whole cache.
158 var = max((ss - s * s / n) / (n - 1), 0.0) if n > 1 else float('nan')
159 pred = predict_entropy_rate(grid.kernel_from_params(c['params']),
160 c['params']['sigma'])
161 c.update({
162 'mean': mean,
163 'se': math.sqrt(var / n) if n > 1 else None,
164 'sd': math.sqrt(var) if n > 1 else None,
165 'predicted': pred['corrected'],
166 'highres_szego': pred['highres'] if math.isfinite(pred['highres'])
167 else None,
168 })
169 out.append(c)
171 return {
172 'generated': utc_now(),
173 'estimand': 'H(z_{M+1} | z_1..z_M) in bits/sample, for '
174 'x ~ iid N(0, sigma^2) -> y = h * x -> z = round(y)',
175 'total_runs': max((c['n_runs'] for c in out), default=0),
176 'settings': {'pasts_per_run': args.pasts, 'reps': args.reps,
177 'n0': args.n0, 'r': args.r, 'thin': args.thin},
178 'cells': out,
179 }
182def write_json(path, obj):
183 path.write_text(json.dumps(obj, indent=2, sort_keys=True) + '\n')
186def utc_now():
187 return datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
190def rows(summary):
191 for c in summary['cells']:
192 se = '' if c['se'] is None else f"{c['se']:.4f}"
193 yield (c['id'], f"{c['params']['sigma']:g}", str(c['M']),
194 f"{c['mean']:.4f}", se, str(c['n_pasts']),
195 f"{c['predicted']:.4f}")
198HEAD = ('cell', 'sigma', 'M', 'mean', 'se', 'pasts', 'predicted')
201def table(summary):
202 body = [HEAD, tuple('-' * 3 for _ in HEAD)] + list(rows(summary))
203 return '\n'.join('| ' + ' | '.join(r) + ' |' for r in body)
206def write_readme(path, summary):
207 path.write_text(f"""\
208# Cached entropy estimates
210Machine-generated by [`.github/workflows/estimates.yml`][wf] on `main` — do not
211edit by hand. Every dispatch runs the whole grid again with fresh, independent
212pasts, appends one record per cell to `runs.jsonl`, and rebuilds
213`estimates.json` from the full log, so the means tighten run after run.
215- **`runs.jsonl`** — append-only, one JSON object per (cell, dispatch): the
216 cell's parameters, the per-past values with their sum and sum of squares,
217 the per-past resolved Gibbs thinning and probe autocorrelation times
218 (`thin='auto'` matches thinning to each chain's measured mixing), the
219 seed, and the code commit that produced them.
220- **`estimates.json`** — pooled `mean` and `se` per cell over every past ever
221 drawn for it, alongside the analytic `predicted` rate for comparison.
223The estimand is H(z_(M+1) | z_1..z_M) in bits/sample — an upper bound on the
224entropy rate of z that tightens as the conditioning window M grows.
226Fetch it from a browser or a script (raw.githubusercontent.com sends
227`access-control-allow-origin: *`):
229 https://raw.githubusercontent.com/concept-collection/timeseries-entropy/estimates/estimates.json
231## Current state ({summary['total_runs']} runs, generated {summary['generated']})
233{table(summary)}
235[wf]: https://github.com/concept-collection/timeseries-entropy/blob/main/.github/workflows/estimates.yml
236""")
239def step_summary(summary):
240 path = os.environ.get('GITHUB_STEP_SUMMARY')
241 if not path:
242 return
243 with open(path, 'a') as f:
244 f.write(f"## Cached estimates after {summary['total_runs']} runs\n\n"
245 f'{table(summary)}\n')
248if __name__ == '__main__':
249 main()