1import json
 2import numpy as np
 3from dissipative_softmax import generator, stationary, softmax
 4
 5
 6def entropy_production(pi, Q):
 7    s = 0.0
 8    for i in range(len(pi)):
 9        for j in range(i + 1, len(pi)):
10            f, b = pi[i] * Q[i, j], pi[j] * Q[j, i]
11            if f > 0 and b > 0:
12                s += (f - b) * np.log(f / b)
13    return float(s)
14
15
16def baseline_and_layer():
17    rng = np.random.default_rng(2750)
18    K, n = 8, 500
19    logits = rng.normal(size=(n, K))
20    labels = np.array([rng.choice(K, p=softmax(x)) for x in logits])
21    base_nll = -np.mean(np.log([softmax(x)[y] for x, y in zip(logits, labels)]))
22    out = []
23    for ratio in [1, 10, 100]:
24        vals = []
25        for x, y in zip(logits, labels):
26            Q, _ = generator(x, ratio, 1.0, 0.15)
27            pi = stationary(Q); q = pi[1:] / pi[1:].sum()
28            vals.append(-np.log(max(q[y], 1e-300)))
29        out.append((ratio, float(np.mean(vals))))
30    return base_nll, out
31
32
33def weak_affinity_check():
34    """Build a cycle which is detailed-balanced at A=0, then tilt it by A."""
35    X = np.array([1.0, .2, -.5, -1.0]); K = len(X); r = 30.; k = 1.; c = .2
36    Q0, p = generator(X, r, k, drive=0.)
37    pi0 = stationary(Q0)
38    edges = [(0, 1)] + [(i, i+1) for i in range(1, K)] + [(K, 0)]
39    results = []
40    L = len(edges)
41    for A in [0.005, .01, .02, .04, .08, .16]:
42        Q = Q0.copy()
43        for i, j in edges:
44            # At A=0, added rates obey pi0_i w_ij = pi0_j w_ji.
45            base_f = c * np.exp((np.log(pi0[j])-np.log(pi0[i]))/2)
46            base_b = c * np.exp((np.log(pi0[i])-np.log(pi0[j]))/2)
47            Q[i, j] += base_f * np.exp(A/(2*L))
48            Q[j, i] += base_b * np.exp(-A/(2*L))
49        np.fill_diagonal(Q, 0.)
50        np.fill_diagonal(Q, -Q.sum(axis=1))
51        pi = stationary(Q)
52        Js = [pi[i]*Q[i,j] - pi[j]*Q[j,i] for i,j in edges]
53        J = float(np.mean(Js)); sigma = entropy_production(pi, Q)
54        results.append((A, J, sigma))
55    slope_j = np.polyfit(np.log([x[0] for x in results]), np.log(np.abs([x[1] for x in results])), 1)[0]
56    slope_s = np.polyfit(np.log([x[0] for x in results]), np.log([x[2] for x in results]), 1)[0]
57    return results, slope_j, slope_s
58
59if __name__ == '__main__':
60    base, layer = baseline_and_layer()
61    aff, sj, ss = weak_affinity_check()
62    print(json.dumps({'baseline_nll': base, 'dissipative_nll': layer,
63                      'weak_affinity': aff, 'loglog_slope_J_vs_A': sj,
64                      'loglog_slope_sigma_vs_A': ss}, indent=2))