1#!/usr/bin/env python3
2"""Monte Carlo estimate of the true entropy rate of the quantized filtered
3Gaussian pipeline, alongside the closed-form approximation the UI plots.
5Implements docs/mc-true-rate.md: sequential Monte Carlo (Genz separation of
6variables with resampling) over the latent Gaussian path constrained to the
7observed integer boxes. All systematic errors are upward, so the estimate
8converges to the true R from above; double --particles and compare to check
9convergence.
11The UI shows the exact command for the current parameter set. Requires
12numpy and scipy.
14Examples:
15 python scripts/true_rate.py --sigma 5 --filter bandpass --low-hz 300 \
16 --high-hz 6000 --taps 101 --sample-rate 30000
17 python scripts/true_rate.py --sigma 0.5 --filter none
18 python scripts/true_rate.py --selftest
19"""
21import argparse
22import math
23import sys
25import numpy as np
26from scipy.special import log_ndtr, logsumexp, ndtr, ndtri
28# ---------------------------------------------------------------- kernels
30def windowed_sinc_lowpass(fc: float, taps: int) -> np.ndarray:
31 n = taps if taps % 2 == 1 else taps + 1
32 i = np.arange(n)
33 t = i - (n - 1) / 2
34 with np.errstate(invalid="ignore", divide="ignore"):
35 sinc = np.where(t == 0, 2 * fc, np.sin(2 * np.pi * fc * t) / (np.pi * t))
36 w = 0.54 - 0.46 * np.cos(2 * np.pi * i / (n - 1))
37 h = sinc * w
38 return h / h.sum()
41def design_kernel(args) -> np.ndarray:
42 f = args.filter
43 if f == "none":
44 return np.array([1.0])
45 if f == "first-difference":
46 return np.array([1.0, -1.0])
47 if f == "moving-average":
48 return np.full(args.width, 1.0 / args.width)
49 if f == "lowpass":
50 return windowed_sinc_lowpass(args.cutoff_hz / args.sample_rate, args.taps)
51 if f == "bandpass":
52 lo = windowed_sinc_lowpass(args.low_hz / args.sample_rate, args.taps)
53 hi = windowed_sinc_lowpass(args.high_hz / args.sample_rate, args.taps)
54 return hi - lo
55 raise ValueError(f)
57# ------------------------------------------------- closed-form quantities
59def h_delta(s: float) -> float:
60 """Exact entropy (bits) of round(N(0, s^2)) on the unit lattice."""
61 if s <= 0:
62 return 0.0
63 zmax = int(np.ceil(8 * s + 4))
64 z = np.arange(-zmax, zmax + 1)
65 p = ndtr((z + 0.5) / s) - ndtr((z - 0.5) / s)
66 p = p[p > 1e-300]
67 return float(-(p @ np.log2(p)))
70def magnitude_sq(h: np.ndarray, f: np.ndarray) -> np.ndarray:
71 n = np.arange(len(h))
72 e = np.exp(-2j * np.pi * np.outer(f, n))
73 H = e @ h
74 return np.abs(H) ** 2
77def r_approx(h: np.ndarray, sigma: float, dither: bool, points: int = 8192) -> float:
78 """The UI's formula: noise-floored Szego prediction error through H_delta."""
79 floor = 1 / 6 if dither else 1 / 12
80 f = (np.arange(points) + 0.5) * 0.5 / points
81 s_z = sigma**2 * magnitude_sq(h, f) + floor
82 integral = float(np.log2(s_z).mean() * 0.5)
83 return h_delta(2.0**integral)
85# ----------------------------------------------------------- MC estimator
87def autocovariance(h: np.ndarray, sigma: float, kmax: int) -> np.ndarray:
88 full = sigma**2 * np.correlate(h, h, "full")
89 r = np.zeros(kmax + 1)
90 take = min(kmax + 1, len(h))
91 r[:take] = full[len(h) - 1 : len(h) - 1 + take]
92 return r
95def levinson_all(r: np.ndarray, kmax: int):
96 """Prediction coefficients for every order 0..kmax and innovation variances.
98 A[p] holds a_1..a_p with yhat_t = sum_j a_j y_{t-j}; v[p] is the order-p
99 prediction error variance.
100 """
101 A = [np.zeros(0)]
102 v = np.zeros(kmax + 1)
103 v[0] = r[0]
104 a = np.zeros(kmax)
105 for p in range(1, kmax + 1):
106 acc = r[p] - (a[: p - 1] @ r[p - 1 : 0 : -1] if p > 1 else 0.0)
107 k = acc / v[p - 1]
108 if p > 1:
109 a[: p - 1] = a[: p - 1] - k * a[p - 2 :: -1]
110 a[p - 1] = k
111 v[p] = max(v[p - 1] * (1 - k * k), 1e-30)
112 A.append(a[:p].copy())
113 return A, v
116def log_phi_diff(alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
117 """log(Phi(beta) - Phi(alpha)) elementwise, safe in both tails."""
118 out = np.empty_like(alpha)
119 hi_tail = alpha >= 0 # reflect to the lower tail
120 a = np.where(hi_tail, -beta, alpha)
121 b = np.where(hi_tail, -alpha, beta)
122 straddle = b >= 0 # a < 0 <= b: safe in linear space
123 with np.errstate(divide="ignore"):
124 out[straddle] = np.log(ndtr(b[straddle]) - ndtr(a[straddle]))
125 lo = ~straddle # both below 0: log-space difference
126 la, lb = log_ndtr(a[lo]), log_ndtr(b[lo])
127 out[lo] = lb + np.log1p(-np.exp(np.minimum(la - lb, -1e-12)))
128 return out
131def sample_truncated(mu, s, lo, hi, rng):
132 """Sample N(mu, s^2) truncated to [lo, hi], vectorized, tail-safe by
133 reflection. Callers guarantee the interval has nonzero mass."""
134 alpha = (lo - mu) / s
135 beta = (hi - mu) / s
136 flip = alpha > 0
137 a = np.where(flip, -beta, alpha)
138 b = np.where(flip, -alpha, beta)
139 pa, pb = ndtr(a), ndtr(b)
140 u = pa + rng.random(len(mu)) * (pb - pa)
141 x = ndtri(np.clip(u, 1e-320, 1 - 1e-16))
142 x = np.where(flip, -x, x)
143 return mu + s * np.clip(x, alpha, beta)
146def generate_z(h, sigma, dither, T, rng):
147 L = len(h)
148 x = rng.standard_normal(T + L - 1) * sigma
149 y = np.convolve(x, h, "valid")
150 if dither:
151 y = y + rng.uniform(-0.5, 0.5, T)
152 return np.round(y).astype(np.int64)
155def smc_replicate(h, sigma, dither, T, burn, N, kmax, rng):
156 """One replicate: -(1/(T-burn)) sum log2 phat(z_t | z_<t) after burn-in."""
157 A, v = levinson_all(autocovariance(h, sigma, kmax), kmax)
158 s_by_order = np.sqrt(v)
159 z = generate_z(h, sigma, dither, T, rng)
161 W = np.zeros((N, kmax), dtype=np.float32) # each particle's last kmax values
162 log_p = np.zeros(T)
163 bad_steps = 0
164 for t in range(T):
165 p = min(t, kmax)
166 a = A[p]
167 s = s_by_order[p]
168 mu = (W[:, kmax - p :] @ a[::-1].astype(np.float32)).astype(np.float64) if p else np.zeros(N)
169 d = rng.uniform(-0.5, 0.5, N) if dither else 0.0
170 lo = z[t] - 0.5 - d
171 hi = z[t] + 0.5 - d
172 logw = log_phi_diff((lo - mu) / s, (hi - mu) / s)
173 lse = logsumexp(logw)
174 if not np.isfinite(lse):
175 raise RuntimeError(
176 f"particle collapse at step {t}: no particle is consistent with "
177 f"the observation — rerun with more --particles"
178 )
179 log_p[t] = lse - math.log(N)
180 # A per-step surprisal beyond ~40 bits means the particle cloud has
181 # drifted away from every path consistent with the data — the
182 # genealogical-collapse failure mode of docs/mc-true-rate.md §5.3,
183 # not a property of the data. Fail loudly rather than average it in.
184 if -log_p[t] / math.log(2) > 40:
185 bad_steps += 1
186 if bad_steps > 25:
187 raise RuntimeError(
188 "the sequential filter degenerated (near-deterministic dynamics; "
189 "see docs/mc-true-rate.md §5.3) — this kernel needs the "
190 "lookahead/twisted-proposal extension. Reduce --taps to study "
191 "the trend with a shallower stopband."
192 )
193 # Systematic resampling, then extend the chosen ancestors.
194 probs = np.exp(logw - logw.max())
195 cdf = np.cumsum(probs)
196 cdf /= cdf[-1]
197 ancestors = np.searchsorted(cdf, (rng.random() + np.arange(N)) / N)
198 mu_a = mu[ancestors]
199 lo_a = lo[ancestors] if dither else np.full(N, lo)
200 hi_a = hi[ancestors] if dither else np.full(N, hi)
201 y_new = sample_truncated(mu_a, s, lo_a, hi_a, rng)
202 W = W[ancestors]
203 W[:, :-1] = W[:, 1:]
204 W[:, -1] = y_new.astype(np.float32)
205 return float(-log_p[burn:].mean() / math.log(2))
208def estimate_true_rate(h, sigma, dither, args, seed):
209 rates = []
210 for j in range(args.replicates):
211 rng = np.random.default_rng(seed + j)
212 r = smc_replicate(h, sigma, dither, args.steps, args.burn, args.particles, args.kmax, rng)
213 rates.append(r)
214 print(f" replicate {j + 1}/{args.replicates}: {r:.4f}", flush=True)
215 rates = np.array(rates)
216 se = rates.std(ddof=1) / math.sqrt(len(rates)) if len(rates) > 1 else float("nan")
217 return rates.mean(), se
219# ------------------------------------------------------------- self-test
221def selftest() -> int:
222 failures = 0
224 def check(name, got, want, tol):
225 nonlocal failures
226 ok = abs(got - want) <= tol
227 failures += 0 if ok else 1
228 print(f" {'PASS' if ok else 'FAIL'} {name}: got {got:.5f}, want {want:.5f} (tol {tol})")
230 print("H_delta against reference values:")
231 for s, want in [(0.1, 0.00001), (0.3, 0.55042), (1, 2.10483), (5, 4.37142), (50, 7.69098)]:
232 check(f"H_delta({s})", h_delta(s), want, 2e-4)
234 print("Formula against reference values (iid):")
235 none = np.array([1.0])
236 for s, want in [(0.1, 0.57611), (1, 2.15829), (5, 4.37382), (100, 8.69096)]:
237 check(f"R_approx none sigma={s}", r_approx(none, s, False), want, 1e-3)
239 print("Formula against reference values (default bandpass, end-to-end kernel check):")
240 lo = windowed_sinc_lowpass(300 / 30000, 101)
241 hi = windowed_sinc_lowpass(6000 / 30000, 101)
242 bp = hi - lo
243 check("||h||_2", float(np.sqrt((bp**2).sum())), 0.60216, 1e-4)
244 for s, want in [(0.5, 0.8620), (5, 1.9397), (20, 2.7282), (100, 3.7166)]:
245 check(f"R_approx bandpass sigma={s}", r_approx(bp, s, False), want, 1e-3)
246 check("R_approx bandpass sigma=5 dither", r_approx(bp, 5, True), 2.2111, 1e-3)
248 print("Quick MC on the exact iid case (truth = H_delta(sigma)):")
249 rng = np.random.default_rng(1)
250 est = smc_replicate(none, 5.0, False, T=1500, burn=300, N=512, kmax=8, rng=rng)
251 check("MC none sigma=5", est, h_delta(5.0), 0.05)
252 est = smc_replicate(none, 0.5, False, T=1500, burn=300, N=512, kmax=8, rng=rng)
253 check("MC none sigma=0.5", est, h_delta(0.5), 0.05)
255 print("FAILED" if failures else "all tests passed")
256 return 1 if failures else 0
258# ------------------------------------------------------------------ main
260def main() -> int:
261 ap = argparse.ArgumentParser(
262 description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
263 )
264 ap.add_argument("--sigma", type=float, default=5.0, help="input std in quantization steps")
265 ap.add_argument(
266 "--filter",
267 choices=["none", "moving-average", "lowpass", "bandpass", "first-difference"],
268 default="none",
269 )
270 ap.add_argument("--low-hz", type=float, default=300, help="bandpass low edge")
271 ap.add_argument("--high-hz", type=float, default=6000, help="bandpass high edge")
272 ap.add_argument("--cutoff-hz", type=float, default=6000, help="lowpass cutoff")
273 ap.add_argument("--taps", type=int, default=101, help="windowed-sinc kernel length")
274 ap.add_argument("--width", type=int, default=8, help="moving-average width")
275 ap.add_argument("--sample-rate", type=float, default=30000)
276 ap.add_argument("--dither", action="store_true")
277 ap.add_argument("--particles", type=int, default=2048)
278 ap.add_argument("--steps", type=int, default=3000)
279 ap.add_argument("--burn", type=int, default=800)
280 ap.add_argument("--replicates", type=int, default=8)
281 ap.add_argument("--kmax", type=int, default=0, help="predictor memory; 0 = auto (4L, capped 512)")
282 ap.add_argument("--seed", type=int, default=0)
283 ap.add_argument("--selftest", action="store_true")
284 args = ap.parse_args()
286 if args.selftest:
287 return selftest()
289 h = design_kernel(args)
290 L = len(h)
291 if args.kmax == 0:
292 args.kmax = min(max(4 * L, 64), 512)
293 args.burn = max(args.burn, 2 * args.kmax)
294 if args.steps <= args.burn + 500:
295 args.steps = args.burn + 2000
297 norm = float(np.sqrt((h**2).sum()))
298 formula = r_approx(h, args.sigma, args.dither)
299 _, v = levinson_all(autocovariance(h, args.sigma, args.kmax), args.kmax)
300 s_inn = math.sqrt(v[args.kmax])
301 print(f"kernel: {args.filter}, L={L}, ||h||2={norm:.5f}, sigma_y={args.sigma * norm:.3f}, "
302 f"innovation std={s_inn:.4f}")
303 if s_inn < 0.12:
304 print("warning: innovation std << quantization step — near-deterministic dynamics. "
305 "The plain sequential filter (docs/mc-true-rate.md §5.3) will likely degenerate "
306 "here; reduce --taps for a shallower stopband, or implement the "
307 "lookahead/twisted-proposal extension.")
308 print(f"formula (as in the UI): R = {formula:.4f} bits/sample (ratio {16 / formula:.2f}x)"
309 if formula > 0 else f"formula (as in the UI): R = {formula:.4f} bits/sample")
310 if args.filter == "none":
311 exact = h_delta(args.sigma) if not args.dither else None
312 if exact is not None:
313 print(f"exact truth (iid closed form): R = {exact:.4f} bits/sample")
315 print(f"MC (N={args.particles} particles, T={args.steps} steps, burn={args.burn}, "
316 f"kmax={args.kmax}, {args.replicates} replicates):")
317 try:
318 mean, se = estimate_true_rate(h, args.sigma, args.dither, args, args.seed)
319 except RuntimeError as e:
320 print(f"aborted: {e}")
321 return 2
322 print(f"MC true rate: R = {mean:.4f} +/- {se:.4f} bits/sample "
323 f"(ratio {16 / mean:.2f}x)" if mean > 0 else f"MC true rate: R = {mean:.4f}")
324 print(f"difference (formula - MC): {formula - mean:+.4f} bits/sample")
325 print("note: finite window and inner MC both bias the estimate upward; "
326 "double --particles and compare to confirm convergence.")
327 return 0
330if __name__ == "__main__":
331 sys.exit(main())