1"""Unbiased Monte-Carlo entropy estimation for quantized filtered Gaussian
2time series: x iid N(0, sigma^2) -> y = h * x -> z = round(y)."""
4import math
5import os
6from concurrent.futures import ProcessPoolExecutor, as_completed
7from dataclasses import dataclass
9import numpy as np
11from .estimator import (plugin_entropy, unbiased_entropy, level_corrections,
12 integrated_autocorr_time)
13from .model import ConditionalChain
14from . import kernels
16__all__ = [
17 'estimate_conditional_entropy', 'Estimate', 'ConditionalChain',
18 'plugin_entropy', 'unbiased_entropy', 'level_corrections',
19 'integrated_autocorr_time', 'kernels',
20]
24class Estimate:
25 """mean +/- se (over independent pasts) of H(z_{M+1} | z_1..z_M), bits.
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 """
31 mean: float
32 se: float
33 per_past: np.ndarray
34 thin: np.ndarray
35 reps: np.ndarray
36 tau: np.ndarray
39def _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
49def _one_past(kernel, sigma, M, thin, n0, r, reps, probe, thin_cap, seed_seq):
50 rng = np.random.default_rng(seed_seq)
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
59def estimate_conditional_entropy(kernel, sigma, past=None, pasts=24, reps=8,
60 n0=128, r=1.5, thin='auto', probe=512,
61 thin_cap=64, seed=None, progress=None,
62 workers=None):
63 """Unbiased estimate of H(z_{M+1} | z_1..z_M) in bits.
65 For each of `pasts` independent pasts, a stationary Gibbs chain of
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.
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.
84 past defaults to max(512, 4 * len(kernel)); the estimand decreases toward
85 the entropy rate as it grows. progress, if given, is called as
86 progress(i, values) after each past finishes (completion order when
87 parallel).
89 Pasts run in parallel across `workers` processes (default: all cores).
90 Each past gets its own spawned RNG stream, so a given seed yields the
91 same result for any worker count.
92 """
93 kernel = np.asarray(kernel, dtype=float)
94 M = int(past) if past is not None else max(512, 4 * kernel.size)
95 seeds = np.random.SeedSequence(seed).spawn(pasts)
96 if workers is None:
97 workers = min(pasts, os.cpu_count() or 1)
98 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)
102 values = []
104 def record(i, result):
105 per_past[i], thins[i], n_reps[i], taus[i] = result
106 values.append(per_past[i])
108 if workers <= 1:
109 for i in range(pasts):
110 record(i, _one_past(kernel, sigma, M, thin, n0, r, reps, probe,
111 thin_cap, seeds[i]))
112 if progress is not None:
113 progress(i, values)
114 else:
115 with ProcessPoolExecutor(max_workers=workers) as pool:
116 futures = {
117 pool.submit(_one_past, kernel, sigma, M, thin, n0, r, reps,
118 probe, thin_cap, seeds[i]): i
119 for i in range(pasts)}
120 for done, fut in enumerate(as_completed(futures)):
121 record(futures[fut], fut.result())
122 if progress is not None:
123 progress(done, values)
124 se = (float(per_past.std(ddof=1) / math.sqrt(len(per_past)))
125 if len(per_past) > 1 else float('nan'))
126 return Estimate(float(per_past.mean()), se, per_past, thins, n_reps, taus)