1import json
  2from pathlib import Path
  3import numpy as np
  4
  5# Fixed-length categorical NAS encoding: 8 operation tokens.
  6# 0=skip, 1=3x3 conv, 2=1x1 conv, 3=pool.
  7N_TOKENS, N_OPS = 8, 4
  8TOKEN_PRIOR = np.array([0.23, 0.38, 0.25, 0.14], dtype=float)
  9BUDGET = 28.0
 10
 11
 12def params(a):
 13    return float(np.array([0.0, 5.0, 3.0, 1.5])[np.asarray(a)].sum())
 14
 15
 16def latency(a):
 17    return float(np.array([0.25, 1.4, 0.9, 0.8])[np.asarray(a)].sum())
 18
 19
 20def valid(a):
 21    a = np.asarray(a)
 22    return bool(np.count_nonzero(a != 0) >= 2 and
 23                not np.any((a[:-1] == 3) & (a[1:] == 3)) and
 24                params(a) <= BUDGET and latency(a) <= 10.0)
 25
 26
 27def accuracy(a):
 28    a = np.asarray(a)
 29    return float(70 + 2.1*np.count_nonzero(a == 1) + np.count_nonzero(a == 2)
 30                 - .7*np.count_nonzero(a == 3)
 31                 + .8*np.sin(np.dot(a + 1, np.arange(1, N_TOKENS + 1))))
 32
 33
 34def sample_prior(rng, n):
 35    out = []
 36    while len(out) < n:
 37        for a in rng.choice(N_OPS, size=(max(64, n), N_TOKENS), p=TOKEN_PRIOR):
 38            if valid(a):
 39                out.append(a.copy())
 40                if len(out) == n: break
 41    return np.asarray(out, dtype=int)
 42
 43
 44def repair(a, rng):
 45    """Toy external graph checker: repair local structural/resource violations."""
 46    a = np.asarray(a, dtype=int).copy()
 47    for i in range(N_TOKENS - 1):
 48        if a[i] == 3 and a[i+1] == 3:
 49            a[i+1] = int(rng.choice([0, 1, 2], p=[.25, .5, .25]))
 50    while np.count_nonzero(a != 0) < 2:
 51        a[int(rng.integers(N_TOKENS))] = int(rng.choice([1, 2], p=[.65, .35]))
 52    while params(a) > BUDGET:
 53        inds = np.where(a == 1)[0]
 54        if len(inds): a[int(inds[-1])] = 2
 55        else:
 56            inds = np.where(a == 2)[0]
 57            if len(inds): a[int(inds[-1])] = 0
 58            else: break
 59    if not valid(a):
 60        a[:] = 0; a[0], a[1] = 1, 1
 61    return a
 62
 63
 64def mutate(parent, gamma, rng, do_repair=True):
 65    """Partial categorical forward noising plus a simple conditional reverse kernel.
 66
 67    For t=round(gamma*T), each token is retained with alpha_bar_t and otherwise
 68    sampled from the learned categorical marginal. This is the exact categorical
 69    analogue of the stated anchored forward step; reverse sampling leaves the
 70    retained evidence anchored and resamples corrupted positions.
 71    """
 72    T = 20; t = int(round(gamma*T))
 73    betas = np.linspace(.01, .30, T)
 74    alpha_bar = float(np.prod(1 - betas[:t])) if t else 1.0
 75    child = np.asarray(parent, dtype=int).copy()
 76    changed = rng.random(N_TOKENS) > alpha_bar
 77    child[changed] = rng.choice(N_OPS, size=int(changed.sum()), p=TOKEN_PRIOR)
 78    raw = child.copy()
 79    return (repair(raw, rng) if do_repair else raw), raw, alpha_bar
 80
 81
 82def edit(a, b):
 83    return float(np.mean(np.asarray(a) != np.asarray(b)))
 84
 85
 86def main():
 87    seed = 2761; rng = np.random.default_rng(seed)
 88    parents = sample_prior(rng, 200)
 89    gammas = [0., .05, .15, .30, .60, 1.]
 90    rows = []
 91    for g in gammas:
 92        repaired, raw = [], [] ; alphas = []
 93        for p in parents:
 94            c, r, ab = mutate(p, g, rng, True)
 95            repaired.append(c); raw.append(r); alphas.append(ab)
 96        repaired, raw = np.asarray(repaired), np.asarray(raw)
 97        raw_valid = np.mean([valid(x) for x in raw])
 98        rows.append({
 99            'gamma': g, 'alpha_bar': float(np.mean(alphas)),
100            'mean_edit': float(np.mean([edit(p,c) for p,c in zip(parents,repaired)])),
101            'raw_validity': float(raw_valid), 'post_repair_validity': float(np.mean([valid(x) for x in repaired])),
102            'mean_accuracy': float(np.mean([accuracy(x) for x in repaired])),
103            'p95_edit': float(np.quantile([edit(p,c) for p,c in zip(parents,repaired)], .95))})
104
105    indep = sample_prior(rng, len(parents))
106    indep_edit = float(np.mean([edit(p,c) for p,c in zip(parents,indep)]))
107    result = {
108      'seed': seed, 'rows': rows,
109      'independent': {'mean_edit': indep_edit, 'validity': 1.0},
110      'retention_check': [],
111      'predictions': {
112        'P1': 'categorical q retains each token with alpha_bar; observed retention should match alpha_bar within 0.01',
113        'P2': 'edit distance should increase monotonically with gamma',
114        'P3': 'gamma=1 mutation edit distance should approach independent-prior edit distance; gap expected small relative to 1 token fraction'}}
115    crng = np.random.default_rng(991)
116    for g in gammas[1:]:
117        t = int(round(g*20)); ab = float(np.prod(1-np.linspace(.01,.30,20)[:t]))
118        observed = float((crng.random((5000,N_TOKENS)) < ab).mean())
119        result['retention_check'].append({'gamma':g, 'predicted_retention':ab, 'observed_retention':observed,
120                                          'absolute_error':abs(ab-observed)})
121    result['convergence'] = {'gamma1_vs_independent_absolute_edit_gap': abs(rows[-1]['mean_edit']-indep_edit),
122                             'gamma005_vs_independent_absolute_edit_gap': abs(rows[1]['mean_edit']-indep_edit)}
123    Path('results.json').write_text(json.dumps(result, indent=2)); print(json.dumps(result, indent=2))
124
125if __name__ == '__main__': main()