Auto-thin each chain to its measured autocorrelation time - concept-collection/timeseries-entropy

concept-collection / timeseries-entropy

Auto-thin each chain to its measured autocorrelation time

With thin=1 the slow-mixing narrowband large-sigma cells violate the r=1.5 variance condition: E[Delta_m^2] stays flat instead of decaying faster than 2^(-r m), so the estimator is heavy-tailed and a rare deep truncation level returns hundreds of bits. A level-8 draw on bp300-6000_s32 (run_index 1, past 2, rep 5) produced 305 bits, inflating the cached cell to 3.63 +/- 0.81 against a predicted 2.99. thin='auto' (now the default) probes each past's integrated autocorrelation time (Sokal windowing), thins by ceil(tau) capped at 64, and treats reps as a per-past budget so per-past cost stays roughly flat. Sweep records now store the per-past values, resolved thin, and tau, so any future outlier can be traced to its seed and replayed. The estimates cache is reset; draws made before this change are discarded.

Jeremy Magland <jmagland@flatironinstitute.org> committed 18 days ago commit efed18a14c03 parent 4a9d41f Browse files

9 changed files+200−43

.github/workflows/estimates.ymlmodified+5−2View file

@@ -22,7 +22,10 @@ concurrency:

2222 jobs:

2323 sweep:

2424 runs-on: ubuntu-latest

25- timeout-minutes: 60

25+ # Auto-thinning makes the slow-mixing cells legitimately expensive, and a

26+ # deep truncation-level draw can add tens of minutes on top; a run that

27+ # exceeds the timeout pushes nothing, which is safe but wasted.

28+ timeout-minutes: 120

2629 steps:

2730 - uses: actions/checkout@v4

2831

@@ -44,7 +47,7 @@ jobs:

4447 "https://x-access-token:${GH_TOKEN}@github.com/${GITHUB_REPOSITORY}.git"

4548 if git fetch -q --depth 1 origin estimates; then

4649 git checkout -q -B estimates FETCH_HEAD

47- echo "continuing $(wc -l < runs.jsonl) cached records"

50+ echo "continuing $(wc -l < runs.jsonl 2>/dev/null || echo 0) cached records"

4851 else

4952 git checkout -q -b estimates

5053 echo "estimates branch does not exist yet — starting it"

README.mdmodified+20−6View file

@@ -22,7 +22,10 @@ it there too.

2222 stationarity — no burn-in bias. Each Gibbs conditional is a box-truncated

2323 normal; after each sweep the free tail latent is drawn fresh, emitting one

2424 exact sample of z_{M+1}. The samples form a stationary, autocorrelated

25- discrete chain.

25+ discrete chain. By default each chain is auto-thinned: a short probe

26+ measures the draws' integrated autocorrelation time tau and the chain

27+ then takes ceil(tau) sweeps per emitted sample (capped), so downstream

28+ levels see approximately independent draws.

2629 2. **Unbiased entropy of the chain's marginal.** Plug-in entropies of blocks

2730 whose sizes double per level are combined by Rhee–Glynn randomized

2831 telescoping with antithetic half-block corrections

@@ -71,9 +74,17 @@ Filters match the web app: `none`, `moving-average`, `lowpass`, `bandpass`,

7174 expected work needs r > 1; finite variance needs E[Delta_m^2] to decay

7275 faster than 2^(-r m). Run `--pilot 6` to see the RMS Delta_m decay before

7376 trusting a value.

74-- `n0` (default 128) is the base block size; `thin` inserts extra Gibbs

75- sweeps per emitted sample to cut autocorrelation for slowly mixing

76- (narrowband, large-sigma) settings.

77+- `thin` (default `'auto'`) inserts extra Gibbs sweeps per emitted sample to

78+ cut autocorrelation. This is not optional for slowly mixing (narrowband,

79+ large-sigma) settings: without thinning their Delta_m decay too slowly for

80+ r = 1.5 and the estimator's variance is infinite — still unbiased, but a

81+ rare deep truncation level then returns a value of hundreds of bits (one

82+ such draw was traced producing 305 bits on the bandpass sigma=32 grid

83+ cell). `'auto'` probes each chain's integrated autocorrelation time and

84+ thins by ceil(tau), capped at `thin_cap` (default 64); `reps` then acts as

85+ a per-past budget (realizations = max(1, reps / thin)), so per-past cost

86+ stays roughly flat and accuracy accumulates over pasts instead.

87+- `n0` (default 128) is the base block size.

7788 - `--past` sets M; increase it until the estimate stops moving to approach

7889 the rate.

7990 - `--workers` caps the process pool (default: all cores).

@@ -88,9 +99,12 @@ bandpass 300-6000 Hz at 30 kHz) crossed with sigma in {1, 2, 4, 8, 16, 32}.

8899

89100 The **estimates** workflow fills it every two hours, and can also be dispatched

90101 by hand. Each run draws 8 fresh, independent pasts per cell, appends one record

91-per cell to `runs.jsonl`, and

102+per cell to `runs.jsonl` (including the per-past values and resolved

103+thinning, so any outlier can be traced to its seed and replayed), and

92104 rebuilds `estimates.json` by pooling every past ever drawn — so the means keep

93-tightening the more often it runs. The grid lives in

105+tightening the more often it runs. The cache was reset on 2026-08-01: draws

106+made before auto-thinning landed had heavy-tailed outliers on the

107+narrowband large-sigma cells (see Tuning) and were discarded. The grid lives in

94108 [scripts/grid.py](scripts/grid.py); adding a cell does not invalidate the

95109 cache, since each record stores its own parameters. To fill the cache locally

96110 instead:

scripts/run_sweep.pymodified+22−4View file

@@ -30,6 +30,10 @@ RUNS = 'runs.jsonl'

3030 SUMMARY = 'estimates.json'

3131

3232

33+def _thin_arg(s):

34+ return s if s == 'auto' else int(s)

35+

36+

3337 def main():

3438 ap = argparse.ArgumentParser(description=__doc__)

3539 ap.add_argument('--data-dir', required=True,

@@ -37,10 +41,12 @@ def main():

3741 ap.add_argument('--pasts', type=int, default=8,

3842 help='independent pasts added per cell per run')

3943 ap.add_argument('--reps', type=int, default=8,

40- help='randomized realizations averaged per past')

44+ help='per-past realization budget at thin=1')

4145 ap.add_argument('--n0', type=int, default=128)

4246 ap.add_argument('--r', type=float, default=1.5)

43- ap.add_argument('--thin', type=int, default=1)

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)")

4450 ap.add_argument('--workers', type=int)

4551 ap.add_argument('--only', help='substring of cell id, for local testing')

4652 args = ap.parse_args()

@@ -83,6 +89,14 @@ def main():

8389 'thin': args.thin,

8490 'sum': float(v.sum()),

8591 '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],

86100 'seed': seed,

87101 'run_index': run_index,

88102 'utc': utc_now(),

@@ -90,8 +104,10 @@ def main():

90104 'workflow_run': os.environ.get('GITHUB_RUN_ID'),

91105 }

92106 new.append(rec)

107+ thins = est.thin

93108 print(f' [{i + 1:2d}/{len(cells)}] {cid:<16} '

94109 f'H = {est.mean:.4f} +/- {est.se:.4f} '

110+ f'thin {int(thins.min())}-{int(thins.max())} '

95111 f'({time.time() - t0:.0f}s)', flush=True)

96112

97113 append_records(data_dir / RUNS, new)

@@ -197,8 +213,10 @@ pasts, appends one record per cell to `runs.jsonl`, and rebuilds

197213 `estimates.json` from the full log, so the means tighten run after run.

198214

199215 - **`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.

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.

202220 - **`estimates.json`** — pooled `mean` and `se` per cell over every past ever

203221 drawn for it, alongside the analytic `predicted` rate for comparison.

204222

src/timeseries_entropy/__init__.pymodified+63−20View file

@@ -8,41 +8,78 @@ from dataclasses import dataclass

88

99 import numpy as np

1010

11-from .estimator import plugin_entropy, unbiased_entropy, level_corrections

11+from .estimator import (plugin_entropy, unbiased_entropy, level_corrections,

12+ integrated_autocorr_time)

1213 from .model import ConditionalChain

1314 from . import kernels

1415

1516 __all__ = [

1617 'estimate_conditional_entropy', 'Estimate', 'ConditionalChain',

17- 'plugin_entropy', 'unbiased_entropy', 'level_corrections', 'kernels',

18+ 'plugin_entropy', 'unbiased_entropy', 'level_corrections',

19+ 'integrated_autocorr_time', 'kernels',

1820 ]

1921

2022

2123 @dataclass

2224 class Estimate:

23- """mean +/- se (over independent pasts) of H(z_{M+1} | z_1..z_M), bits."""

25+ """mean +/- se (over independent pasts) of H(z_{M+1} | z_1..z_M), bits.

26+

27+ per_past holds one unbiased value per past; thin, reps, and tau record,

28+ per past, the resolved Gibbs sweeps per draw, the realizations averaged,

29+ and the probe's autocorrelation-time estimate (nan when thin was fixed).

30+ """

2431 mean: float

2532 se: float

2633 per_past: np.ndarray

34+ thin: np.ndarray

35+ reps: np.ndarray

36+ tau: np.ndarray

37+

38+

39+def _resolve_thin(chain, thin, probe, thin_cap):

40+ """Set chain.thin; returns (resolved thin, probe tau or nan)."""

41+ if thin != 'auto':

42+ chain.thin = int(thin)

43+ return int(thin), float('nan')

44+ tau = integrated_autocorr_time(chain.draw(probe))

45+ chain.thin = int(min(thin_cap, math.ceil(tau)))

46+ return chain.thin, tau

2747

2848

29-def _one_past(kernel, sigma, M, thin, n0, r, reps, seed_seq):

49+def _one_past(kernel, sigma, M, thin, n0, r, reps, probe, thin_cap, seed_seq):

3050 rng = np.random.default_rng(seed_seq)

31- chain = ConditionalChain(kernel, sigma, M, rng, thin)

32- return float(np.mean(

33- [unbiased_entropy(chain.draw, n0, r, rng) for _ in range(reps)]))

51+ chain = ConditionalChain(kernel, sigma, M, rng, 1)

52+ t, tau = _resolve_thin(chain, thin, probe, thin_cap)

53+ n_reps = max(1, round(reps / t)) if thin == 'auto' else reps

54+ val = float(np.mean(

55+ [unbiased_entropy(chain.draw, n0, r, rng) for _ in range(n_reps)]))

56+ return val, t, n_reps, tau

3457

3558

3659 def estimate_conditional_entropy(kernel, sigma, past=None, pasts=24, reps=8,

37- n0=128, r=1.5, thin=1, seed=None,

38- progress=None, workers=None):

60+ n0=128, r=1.5, thin='auto', probe=512,

61+ thin_cap=64, seed=None, progress=None,

62+ workers=None):

3963 """Unbiased estimate of H(z_{M+1} | z_1..z_M) in bits.

4064

4165 For each of `pasts` independent pasts, a stationary Gibbs chain of

42- z_{M+1} draws feeds `reps` randomized-telescoping realizations (on

43- consecutive segments of the chain); their average is one unbiased value

44- per past. Returns the mean and standard error over pasts — valid because

45- pasts are independent.

66+ z_{M+1} draws feeds randomized-telescoping realizations (on consecutive

67+ segments of the chain); their average is one unbiased value per past.

68+ Returns the mean and standard error over pasts — valid because pasts are

69+ independent.

70+

71+ thin='auto' (the default) matches the chain's thinning to its measured

72+ mixing: each past first draws `probe` samples at thin=1, estimates their

73+ integrated autocorrelation time tau, and thins by ceil(tau), capped at

74+ `thin_cap`. This keeps the level corrections Delta_m decaying fast

75+ enough for the r exponent — with slowly mixing chains (narrowband

76+ kernels x large sigma) and no thinning, E[Delta_m^2] can decay slower

77+ than 2^(-r m), which makes the estimator's variance infinite: still

78+ unbiased, but with rare enormous realizations and meaningless standard

79+ errors. Under 'auto', `reps` is a per-past budget at thin=1: the

80+ realization count becomes max(1, round(reps / thin)), so per-past cost

81+ stays roughly flat and slow cells trade realizations for better draws.

82+ Pass an integer thin to control both knobs explicitly.

4683

4784 past defaults to max(512, 4 * len(kernel)); the estimand decreases toward

4885 the entropy rate as it grows. progress, if given, is called as

@@ -59,25 +96,31 @@ def estimate_conditional_entropy(kernel, sigma, past=None, pasts=24, reps=8,

5996 if workers is None:

6097 workers = min(pasts, os.cpu_count() or 1)

6198 per_past = np.empty(pasts)

99+ thins = np.empty(pasts, dtype=int)

100+ n_reps = np.empty(pasts, dtype=int)

101+ taus = np.empty(pasts)

62102 values = []

103+

104+ def record(i, result):

105+ per_past[i], thins[i], n_reps[i], taus[i] = result

106+ values.append(per_past[i])

107+

63108 if workers <= 1:

64109 for i in range(pasts):

65- per_past[i] = _one_past(kernel, sigma, M, thin, n0, r, reps,

66- seeds[i])

67- values.append(per_past[i])

110+ record(i, _one_past(kernel, sigma, M, thin, n0, r, reps, probe,

111+ thin_cap, seeds[i]))

68112 if progress is not None:

69113 progress(i, values)

70114 else:

71115 with ProcessPoolExecutor(max_workers=workers) as pool:

72116 futures = {

73117 pool.submit(_one_past, kernel, sigma, M, thin, n0, r, reps,

74- seeds[i]): i

118+ probe, thin_cap, seeds[i]): i

75119 for i in range(pasts)}

76120 for done, fut in enumerate(as_completed(futures)):

77- per_past[futures[fut]] = fut.result()

78- values.append(per_past[futures[fut]])

121+ record(futures[fut], fut.result())

79122 if progress is not None:

80123 progress(done, values)

81124 se = (float(per_past.std(ddof=1) / math.sqrt(len(per_past)))

82125 if len(per_past) > 1 else float('nan'))

83- return Estimate(float(per_past.mean()), se, per_past)

126+ return Estimate(float(per_past.mean()), se, per_past, thins, n_reps, taus)

src/timeseries_entropy/cli.pymodified+25−9View file

@@ -6,11 +6,16 @@ from concurrent.futures import ProcessPoolExecutor

66

77 import numpy as np

88

9-from . import estimate_conditional_entropy, level_corrections, kernels

9+from . import (estimate_conditional_entropy, level_corrections, kernels,

10+ _resolve_thin)

1011 from .model import ConditionalChain

1112 from .theory import predict_entropy_rate

1213

1314

15+def _thin_arg(s):

16+ return s if s == 'auto' else int(s)

17+

18+

1419 def design_kernel(args):

1520 if args.filter == 'none':

1621 return kernels.identity()

@@ -53,12 +58,13 @@ def main():

5358 ap.add_argument('--pasts', type=int, default=24,

5459 help='independent pasts to average')

5560 ap.add_argument('--reps', type=int, default=8,

56- help='randomized realizations per past')

61+ help='per-past realization budget at thin=1')

5762 ap.add_argument('--n0', type=int, default=128, help='base block size')

5863 ap.add_argument('--r', type=float, default=1.5,

5964 help='truncation exponent: P(N >= m) = 2^(-r m)')

60- ap.add_argument('--thin', type=int, default=1,

61- help='Gibbs sweeps per emitted sample')

65+ ap.add_argument('--thin', type=_thin_arg, default='auto',

66+ help="Gibbs sweeps per emitted sample, or 'auto' "

67+ '(default) to match the measured mixing per past')

6268 ap.add_argument('--workers', type=int,

6369 help='parallel processes over pasts (default: all cores)')

6470 ap.add_argument('--seed', type=int, default=0)

@@ -95,6 +101,10 @@ def main():

95101 kernel, args.sigma, past=M, pasts=args.pasts, reps=args.reps,

96102 n0=args.n0, r=args.r, thin=args.thin, seed=args.seed,

97103 progress=progress, workers=args.workers)

104+ if args.thin == 'auto':

105+ print(f'resolved thin {est.thin.min()}-{est.thin.max()} '

106+ f'(tau {est.tau.min():.1f}-{est.tau.max():.1f}), '

107+ f'{est.reps.min()}-{est.reps.max()} reps/past')

98108 ratio = f' (ratio vs int16: {16 / est.mean:.3f}x)' if est.mean > 0 else ''

99109 print(f'\nH(z_next | {M} past samples) = {est.mean:.4f} +/- {est.se:.4f} '

100110 f'bits/sample{ratio}')

@@ -104,19 +114,25 @@ def main():

104114

105115 def _one_pilot(kernel, sigma, M, thin, n0, levels, seed_seq):

106116 rng = np.random.default_rng(seed_seq)

107- return level_corrections(

108- ConditionalChain(kernel, sigma, M, rng, thin).draw, n0, levels)

117+ chain = ConditionalChain(kernel, sigma, M, rng, 1)

118+ t, _ = _resolve_thin(chain, thin, probe=512, thin_cap=64)

119+ return t, level_corrections(chain.draw, n0, levels)

109120

110121

111122 def run_pilot(kernel, args, M):

112123 seeds = np.random.SeedSequence(args.seed).spawn(args.pasts)

113124 workers = args.workers or min(args.pasts, os.cpu_count() or 1)

114- print(f'pilot: {args.pasts} pasts, levels 1..{args.pilot}, n0={args.n0}')

125+ print(f'pilot: {args.pasts} pasts, levels 1..{args.pilot}, n0={args.n0}, '

126+ f'thin={args.thin}')

115127 with ProcessPoolExecutor(max_workers=workers) as pool:

116- deltas = np.array(list(pool.map(

128+ results = list(pool.map(

117129 _one_pilot,

118130 *zip(*[(kernel, args.sigma, M, args.thin, args.n0, args.pilot, s)

119- for s in seeds]))))

131+ for s in seeds])))

132+ thins = np.array([t for t, _ in results])

133+ deltas = np.array([d for _, d in results])

134+ if args.thin == 'auto':

135+ print(f' resolved thin {thins.min()}-{thins.max()}')

120136 rms = np.sqrt((deltas ** 2).mean(axis=0))

121137 for m in range(args.pilot):

122138 note = ''

src/timeseries_entropy/estimator.pymodified+27−0View file

@@ -94,3 +94,30 @@ def level_corrections(draw, n0=128, levels=6):

9494 """

9595 _, deltas = _telescope(draw, n0, levels)

9696 return np.array(deltas)

97+

98+

99+def integrated_autocorr_time(x, c=5.0):

100+ """Integrated autocorrelation time of a stationary sequence, in samples.

101+

102+ tau = 1 + 2 sum_k rho_k with Sokal's automatic windowing: the sum stops

103+ at the smallest lag W >= c * tau(W). Resolving tau needs len(x) >> c *

104+ tau; longer times saturate near len(x) / (2 c), so cap the result when

105+ the sequence may mix slower than the probe can see. Returns >= 1.

106+ """

107+ x = np.asarray(x, dtype=float)

108+ n = x.size

109+ if n < 2:

110+ return 1.0

111+ x = x - x.mean()

112+ denom = float(x @ x)

113+ if denom == 0.0:

114+ return 1.0

115+ f = np.fft.rfft(x, 2 * n)

116+ acf = np.fft.irfft(f * f.conj())[:n] / denom

117+ csum = np.cumsum(acf[1:n // 2 + 1])

118+ tau = 1.0

119+ for w in range(1, csum.size + 1):

120+ tau = 1.0 + 2.0 * float(csum[w - 1])

121+ if w >= c * tau:

122+ break

123+ return max(tau, 1.0)

src/timeseries_entropy/model.pymodified+3−1View file

@@ -34,7 +34,9 @@ class ConditionalChain:

3434 themselves an exact draw from p(x | z), so the Gibbs chain starts in

3535 stationarity — no burn-in bias, only autocorrelation. draw(k) advances the

3636 chain k steps (thin sweeps each) and returns the k sampled z_{M+1} values,

37- each marginally distributed exactly as z_{M+1} | z_1..z_M.

37+ each marginally distributed exactly as z_{M+1} | z_1..z_M. thin may be

38+ reassigned between draws (e.g. probe at thin=1, then thin by the measured

39+ autocorrelation time); stationarity is unaffected.

3840

3941 Each Gibbs conditional x_i | rest is N(0, sigma^2) truncated to the

4042 interval read off the <= L constraint boxes x_i appears in. Coordinates a

tests/test_estimator.pymodified+23−1View file

@@ -3,7 +3,8 @@ import math

33 import numpy as np

44 import pytest

55

6-from timeseries_entropy import plugin_entropy, unbiased_entropy, level_corrections

6+from timeseries_entropy import (plugin_entropy, unbiased_entropy,

7+ level_corrections, integrated_autocorr_time)

78

89

910 def test_plugin_entropy_exact():

@@ -46,6 +47,27 @@ def test_unbiased_on_autocorrelated_chain():

4647 assert se < 0.02

4748

4849

50+def test_autocorr_time_iid():

51+ rng = np.random.default_rng(6)

52+ tau = integrated_autocorr_time(rng.integers(0, 4, size=4096))

53+ assert 0.8 < tau < 1.5

54+

55+

56+def test_autocorr_time_sticky_markov():

57+ # Two-state chain flipping w.p. eps: rho_k = (1 - 2 eps)^k, so

58+ # tau = 1 + 2 rho / (1 - rho) = 19 at eps = 0.05.

59+ rng = np.random.default_rng(7)

60+ flips = rng.random(200000) < 0.05

61+ x = np.cumsum(flips) % 2

62+ tau = integrated_autocorr_time(x)

63+ assert 12 < tau < 28

64+

65+

66+def test_autocorr_time_degenerate():

67+ assert integrated_autocorr_time([3]) == 1.0

68+ assert integrated_autocorr_time([2, 2, 2, 2]) == 1.0

69+

70+

4971 def test_level_corrections_decay():

5072 rng = np.random.default_rng(3)

5173 p = np.array([0.5, 0.3, 0.2])

tests/test_model.pymodified+12−0View file

@@ -39,6 +39,18 @@ def test_gibbs_preserves_constraints():

3939 assert np.array_equal(np.floor(y + 0.5), z0)

4040

4141

42+def test_auto_thin_deterministic_across_workers():

43+ kwargs = dict(sigma=4.0, past=32, pasts=3, reps=4, n0=32, thin='auto',

44+ probe=64, seed=9)

45+ kernel = np.full(4, 0.25)

46+ a = estimate_conditional_entropy(kernel, workers=1, **kwargs)

47+ b = estimate_conditional_entropy(kernel, workers=3, **kwargs)

48+ assert np.array_equal(a.per_past, b.per_past)

49+ assert np.array_equal(a.thin, b.thin)

50+ assert np.array_equal(a.reps, b.reps)

51+ assert np.all(a.thin >= 1) and np.all(a.reps >= 1)

52+

53+

4254 def test_no_filter_matches_exact_entropy():

4355 # kernel [1]: z is iid round(N(0, sigma^2)), so the conditional entropy

4456 # equals the exact marginal entropy for any past length.