1"""Rhee-Glynn (randomized telescoping) unbiased entropy estimation.
3Input: a stationary, ergodic sequence of discrete draws, possibly
4autocorrelated, each with the target marginal p. Plug-in entropies of blocks
5whose sizes double from level to level form a telescoping sum via the
6antithetic correction
8 Delta_m = h(B_m) - [h(B_m^1) + h(B_m^2)] / 2,
10where B_m^1, B_m^2 are the two halves of B_m. Truncating the sum at a random
11level N with P(N >= m) = 2^(-r m) and reweighting by the survival
12probabilities gives an estimator whose expectation is exactly the entropy of
13p, despite the bias of every finite-block plug-in estimate.
15All entropies are in bits.
16"""
18import numpy as np
21def plugin_entropy(samples):
22 """Plug-in entropy (bits) of a block of discrete samples."""
23 _, counts = np.unique(np.asarray(samples), return_counts=True)
24 return _entropy(counts.astype(float), counts.sum())
27def _entropy(counts, n):
28 return float(np.log2(n) - (counts * np.log2(counts)).sum() / n)
31def _dict_entropy(counts, n):
32 return _entropy(np.fromiter(counts.values(), dtype=float), n)
35def _telescope(draw, n0, levels):
36 """h(B_0) and [Delta_1, ..., Delta_levels] over one growing block.
38 Counts are merged upward (each sample is counted once), so the cost is
39 linear in the n0 * 2**levels samples drawn.
40 """
41 counts = {}
42 _count_into(counts, draw(n0))
43 size = n0
44 h0 = _dict_entropy(counts, size)
45 h_prev = h0
46 deltas = []
47 for _ in range(levels):
48 half = {}
49 _count_into(half, draw(size))
50 h2 = _dict_entropy(half, size)
51 for v, c in half.items():
52 counts[v] = counts.get(v, 0) + c
53 size *= 2
54 h_full = _dict_entropy(counts, size)
55 deltas.append(h_full - 0.5 * (h_prev + h2))
56 h_prev = h_full
57 return h0, deltas
60def _count_into(counts, seg):
61 vals, cnts = np.unique(np.asarray(seg), return_counts=True)
62 for v, c in zip(vals.tolist(), cnts.tolist()):
63 counts[v] = counts.get(v, 0) + c
66def unbiased_entropy(draw, n0=128, r=1.5, rng=None):
67 """One randomized-telescoping realization of the marginal entropy (bits).
69 draw(k) must return the next k consecutive samples of a stationary
70 discrete chain; successive calls continue the chain. The realization
71 consumes n0 * 2**N samples with P(N >= m) = 2^(-r m). Average many
72 realizations (they may continue one chain back-to-back) to reduce
73 variance; each has expectation exactly H(p).
75 r trades expected work against variance: levels must decay like
76 E[Delta_m^2] = O(2^(-r'm)) with r' > r > 1 for both to be finite.
77 r = 1.5 suits the typical second-moment decay 2^(-2m); check with
78 level_corrections when in doubt.
79 """
80 rng = np.random.default_rng() if rng is None else rng
81 rho = 2.0 ** -r
82 N = 0
83 while rng.random() < rho:
84 N += 1
85 h0, deltas = _telescope(draw, n0, N)
86 return h0 + sum(d * 2.0 ** (r * m) for m, d in enumerate(deltas, start=1))
89def level_corrections(draw, n0=128, levels=6):
90 """Deterministic pilot: [Delta_1, ..., Delta_levels] from one block.
92 Draws n0 * 2**levels samples. Repeat over fresh chains, look at the decay
93 of mean(Delta_m^2) with m, and pick r below the decay exponent.
94 """
95 _, deltas = _telescope(draw, n0, levels)
96 return np.array(deltas)
99def integrated_autocorr_time(x, c=5.0):
100 """Integrated autocorrelation time of a stationary sequence, in samples.
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)