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 main():
34 ap = argparse.ArgumentParser(description=__doc__)
35 ap.add_argument('--data-dir', required=True,
36 help='checkout of the estimates branch')
37 ap.add_argument('--pasts', type=int, default=8,
38 help='independent pasts added per cell per run')
39 ap.add_argument('--reps', type=int, default=8,
40 help='randomized realizations averaged per past')
41 ap.add_argument('--n0', type=int, default=128)
42 ap.add_argument('--r', type=float, default=1.5)
43 ap.add_argument('--thin', type=int, default=1)
44 ap.add_argument('--workers', type=int)
45 ap.add_argument('--only', help='substring of cell id, for local testing')
46 args = ap.parse_args()
48 data_dir = Path(args.data_dir)
49 data_dir.mkdir(parents=True, exist_ok=True)
50 records = read_records(data_dir / RUNS)
51 runs_so_far = {}
52 for rec in records:
53 runs_so_far[rec['cell']] = runs_so_far.get(rec['cell'], 0) + 1
55 cells = [c for c in grid.grid()
56 if args.only is None or args.only in grid.cell_id(c)]
57 print(f'{len(cells)} cells x {args.pasts} pasts x {args.reps} reps '
58 f'({len(records)} records already cached)')
60 new = []
61 t_all = time.time()
62 for i, params in enumerate(cells):
63 cid = grid.cell_id(params)
64 run_index = runs_so_far.get(cid, 0)
65 kernel = grid.kernel_from_params(params)
66 M = max(512, 4 * kernel.size)
67 seed = [zlib.crc32(cid.encode()), run_index]
68 t0 = time.time()
69 est = estimate_conditional_entropy(
70 kernel, params['sigma'], past=M, pasts=args.pasts, reps=args.reps,
71 n0=args.n0, r=args.r, thin=args.thin, seed=seed,
72 workers=args.workers)
73 v = np.asarray(est.per_past, dtype=float)
74 rec = {
75 'cell': cid,
76 'params': params,
77 'M': M,
78 'L': int(kernel.size),
79 'pasts': int(v.size),
80 'reps': args.reps,
81 'n0': args.n0,
82 'r': args.r,
83 'thin': args.thin,
84 'sum': float(v.sum()),
85 'sumsq': float((v ** 2).sum()),
86 'seed': seed,
87 'run_index': run_index,
88 'utc': utc_now(),
89 'code_sha': os.environ.get('GITHUB_SHA'),
90 'workflow_run': os.environ.get('GITHUB_RUN_ID'),
91 }
92 new.append(rec)
93 print(f' [{i + 1:2d}/{len(cells)}] {cid:<16} '
94 f'H = {est.mean:.4f} +/- {est.se:.4f} '
95 f'({time.time() - t0:.0f}s)', flush=True)
97 append_records(data_dir / RUNS, new)
98 summary = summarize(records + new, args)
99 write_json(data_dir / SUMMARY, summary)
100 write_readme(data_dir / 'README.md', summary)
101 print(f'\n{len(new)} records appended in {time.time() - t_all:.0f}s; '
102 f'{len(summary["cells"])} cells summarized')
103 step_summary(summary)
106def read_records(path):
107 if not path.exists():
108 return []
109 with path.open() as f:
110 return [json.loads(line) for line in f if line.strip()]
113def append_records(path, records):
114 with path.open('a') as f:
115 for rec in records:
116 f.write(json.dumps(rec, sort_keys=True) + '\n')
119def summarize(records, args):
120 """Pool every recorded per-past value, per cell, into a mean and an SE."""
121 cells = {}
122 for rec in records:
123 c = cells.setdefault(rec['cell'], {
124 'id': rec['cell'], 'n_pasts': 0, 'n_runs': 0,
125 '_sum': 0.0, '_sumsq': 0.0})
126 c['params'] = rec['params']
127 c['M'] = rec['M']
128 c['L'] = rec['L']
129 c['n_pasts'] += rec['pasts']
130 c['n_runs'] += 1
131 c['_sum'] += rec['sum']
132 c['_sumsq'] += rec['sumsq']
133 c['last_utc'] = max(rec['utc'], c.get('last_utc', ''))
135 out = []
136 for c in sorted(cells.values(), key=lambda c: c['id']):
137 n, s, ss = c['n_pasts'], c.pop('_sum'), c.pop('_sumsq')
138 mean = s / n
139 # Sample variance of the pooled per-past values; the pasts are
140 # independent across runs as well as within one, so the SE of their
141 # mean is the honest error bar on the whole cache.
142 var = max((ss - s * s / n) / (n - 1), 0.0) if n > 1 else float('nan')
143 pred = predict_entropy_rate(grid.kernel_from_params(c['params']),
144 c['params']['sigma'])
145 c.update({
146 'mean': mean,
147 'se': math.sqrt(var / n) if n > 1 else None,
148 'sd': math.sqrt(var) if n > 1 else None,
149 'predicted': pred['corrected'],
150 'highres_szego': pred['highres'] if math.isfinite(pred['highres'])
151 else None,
152 })
153 out.append(c)
155 return {
156 'generated': utc_now(),
157 'estimand': 'H(z_{M+1} | z_1..z_M) in bits/sample, for '
158 'x ~ iid N(0, sigma^2) -> y = h * x -> z = round(y)',
159 'total_runs': max((c['n_runs'] for c in out), default=0),
160 'settings': {'pasts_per_run': args.pasts, 'reps': args.reps,
161 'n0': args.n0, 'r': args.r, 'thin': args.thin},
162 'cells': out,
163 }
166def write_json(path, obj):
167 path.write_text(json.dumps(obj, indent=2, sort_keys=True) + '\n')
170def utc_now():
171 return datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
174def rows(summary):
175 for c in summary['cells']:
176 se = '' if c['se'] is None else f"{c['se']:.4f}"
177 yield (c['id'], f"{c['params']['sigma']:g}", str(c['M']),
178 f"{c['mean']:.4f}", se, str(c['n_pasts']),
179 f"{c['predicted']:.4f}")
182HEAD = ('cell', 'sigma', 'M', 'mean', 'se', 'pasts', 'predicted')
185def table(summary):
186 body = [HEAD, tuple('-' * 3 for _ in HEAD)] + list(rows(summary))
187 return '\n'.join('| ' + ' | '.join(r) + ' |' for r in body)
190def write_readme(path, summary):
191 path.write_text(f"""\
192# Cached entropy estimates
194Machine-generated by [`.github/workflows/estimates.yml`][wf] on `main` — do not
195edit by hand. Every dispatch runs the whole grid again with fresh, independent
196pasts, appends one record per cell to `runs.jsonl`, and rebuilds
197`estimates.json` from the full log, so the means tighten run after run.
199- **`runs.jsonl`** — append-only, one JSON object per (cell, dispatch): the
200 cell's parameters, how many pasts it contributed, their sum and sum of
201 squares, the seed, and the code commit that produced them.
202- **`estimates.json`** — pooled `mean` and `se` per cell over every past ever
203 drawn for it, alongside the analytic `predicted` rate for comparison.
205The estimand is H(z_(M+1) | z_1..z_M) in bits/sample — an upper bound on the
206entropy rate of z that tightens as the conditioning window M grows.
208Fetch it from a browser or a script (raw.githubusercontent.com sends
209`access-control-allow-origin: *`):
211 https://raw.githubusercontent.com/concept-collection/timeseries-entropy/estimates/estimates.json
213## Current state ({summary['total_runs']} runs, generated {summary['generated']})
215{table(summary)}
217[wf]: https://github.com/concept-collection/timeseries-entropy/blob/main/.github/workflows/estimates.yml
218""")
221def step_summary(summary):
222 path = os.environ.get('GITHUB_STEP_SUMMARY')
223 if not path:
224 return
225 with open(path, 'a') as f:
226 f.write(f"## Cached estimates after {summary['total_runs']} runs\n\n"
227 f'{table(summary)}\n')
230if __name__ == '__main__':
231 main()