1import numpy as np
  2from scipy.special import logsumexp
  3
  4SEED = 3114
  5L = 8
  6BETA = 0.65
  7J = 1.0
  8N = 4096
  9
 10
 11def energy(x):
 12    x = np.asarray(x)
 13    s = 2*x - 1
 14    return -J * np.sum(s * np.roll(s, -1, axis=-1), axis=-1)
 15
 16
 17def state_id(x):
 18    return np.asarray(x, dtype=np.int64) @ (1 << np.arange(L))
 19
 20
 21def exact_target():
 22    states = ((np.arange(1 << L)[:, None] >> np.arange(L)) & 1).astype(np.int64)
 23    e = energy(states)
 24    lp = -BETA * e
 25    p = np.exp(lp - logsumexp(lp))
 26    return states, e, p
 27
 28
 29def normalized_ess(logw):
 30    # Paper's ESS: (sum exp A)^2 / (M sum exp(2A)); computed stably.
 31    m = len(logw)
 32    return float(np.exp(2*logsumexp(logw) - np.log(m) - logsumexp(2*logw)))
 33
 34
 35def sample_proposal(rng, n):
 36    # Collapsed autoregressive special case: independent Bernoulli conditionals.
 37    return (rng.random((n, L)) < 0.90).astype(np.int64)
 38
 39
 40def log_q(x):
 41    x = np.asarray(x)
 42    return np.sum(x*np.log(.9) + (1-x)*np.log(.1), axis=1)
 43
 44
 45def rate(y, i):
 46    z = y.copy()
 47    z[i] = 1 - z[i]
 48    de = float(energy(z) - energy(y))
 49    return min(1.0, np.exp(-BETA * de))
 50
 51
 52def ctmc_one(y, horizon, rng):
 53    y = y.copy()
 54    logpath = 0.0
 55    t = 0.0
 56    jumps = 0
 57    while t < horizon:
 58        rates = np.array([rate(y, i) for i in range(L)])
 59        lam = rates.sum()
 60        dt = rng.exponential(1.0 / lam)
 61        if t + dt >= horizon:
 62            logpath -= lam * (horizon - t)
 63            break
 64        t += dt
 65        i = int(rng.choice(L, p=rates / lam))
 66        logpath += np.log(rates[i]) - lam * dt
 67        y[i] = 1-y[i]
 68        jumps += 1
 69    return y, logpath, jumps
 70
 71
 72def ctmc_batch(x, horizon, rng):
 73    ys, lps, js = [], [], []
 74    for y in x:
 75        z, lp, j = ctmc_one(y, horizon, rng)
 76        ys.append(z); lps.append(lp); js.append(j)
 77    return np.array(ys), np.array(lps), np.array(js)
 78
 79
 80def empirical_distribution(x):
 81    counts = np.bincount(state_id(x), minlength=1 << L).astype(float)
 82    return counts / len(x)
 83
 84
 85def kl(p, q):
 86    # q is exact target and strictly positive.
 87    mask = p > 0
 88    return float(np.sum(p[mask] * (np.log(p[mask]) - np.log(q[mask]))))
 89
 90
 91def generator_check():
 92    y = np.zeros(L, dtype=np.int64)
 93    Q = np.zeros((1 << L, 1 << L))
 94    for sid in range(1 << L):
 95        x = ((sid >> np.arange(L)) & 1).astype(np.int64)
 96        for i in range(L):
 97            z=x.copy(); z[i]=1-z[i]
 98            Q[state_id(z), sid] = rate(x, i)
 99        Q[sid, sid] = -np.sum(Q[:, sid])
100    colsum = np.max(np.abs(Q.sum(axis=0)))
101    # Small-time transition check from all-zero state, using expm-free first order.
102    sid = 0
103    eps = 1e-7
104    first_order = np.eye(1 << L)[:, sid] + eps*Q[:, sid]
105    valid_nonnegative = bool(first_order.min() >= -1e-12)
106    return float(colsum), valid_nonnegative
107
108
109def path_check(rng):
110    y=np.zeros(L,dtype=np.int64)
111    z, lp, jumps=ctmc_one(y, 0.5, rng)
112    # Every accepted jump contributes log(rate)-lambda*waiting; survival contributes too.
113    return bool(np.isfinite(lp) and jumps >= 0), float(lp), int(jumps)
114
115
116def main():
117    states, energies, target = exact_target()
118    rng = np.random.default_rng(SEED)
119    x0 = sample_proposal(rng, N)
120    a0 = -BETA*energy(x0) - log_q(x0)
121    proposal_ess = normalized_ess(a0)
122    proposal_kl = kl(empirical_distribution(x0), target)
123
124    # Uniform mutation baseline: same number of local updates as the CTMC's mean.
125    # One mutation per site-independent random event, with no energy-aware rates.
126    uniform = x0.copy()
127    for _ in range(8):
128        rows=rng.integers(N,size=N); cols=rng.integers(L,size=N)
129        uniform[rows,cols] = 1-uniform[rows,cols]
130    uniform_kl = kl(empirical_distribution(uniform), target)
131
132    # CTMC refinement, then importance weights based on the proposal (the stable option
133    # recommended by the idea for fixed augmentation). Compare raw sample distribution
134    # and weighted ESS; path likelihood is also recorded as a diagnostic.
135    refined, path_lp, jumps = ctmc_batch(x0, 1.0, rng)
136    refined_kl = kl(empirical_distribution(refined), target)
137    refined_a0 = -BETA*energy(refined) - log_q(refined)
138    refined_ess = normalized_ess(refined_a0)
139    corrected_ess = normalized_ess(refined_a0 - path_lp)
140
141    # Mode recovery: count exact states represented among samples; target modes are
142    # the two ground states of the ferromagnetic ring.
143    ground = set(np.flatnonzero(energies == energies.min()).tolist())
144    def ground_hits(x): return int(np.isin(state_id(x), list(ground)).sum())
145    print('CONFIG', {'L':L,'beta':BETA,'N':N,'seed':SEED})
146    print('MATH', {'generator_max_column_sum':generator_check()[0],
147                   'generator_first_order_nonnegative':generator_check()[1],
148                   'path_finite_and_valid':path_check(np.random.default_rng(SEED+1))})
149    print('RESULT', {
150        'proposal_ess':proposal_ess,
151        'proposal_kl':proposal_kl,
152        'uniform_mutation_kl':uniform_kl,
153        'ctmc_kl':refined_kl,
154        'ctmc_ess_A0':refined_ess,
155        'ctmc_ess_path_corrected':corrected_ess,
156        'mean_ctmc_jumps':float(jumps.mean()),
157        'ground_hits_proposal':ground_hits(x0),
158        'ground_hits_uniform':ground_hits(uniform),
159        'ground_hits_ctmc':ground_hits(refined),
160        'unique_proposal':int(np.unique(state_id(x0)).size),
161        'unique_uniform':int(np.unique(state_id(uniform)).size),
162        'unique_ctmc':int(np.unique(state_id(refined)).size),
163    })
164
165if __name__ == '__main__':
166    main()