1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED=251
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device=torch.device('cpu')
 13
 14def entropy(p):
 15    if p<=0 or p>=1: return 0.0
 16    return -p*math.log2(p)-(1-p)*math.log2(1-p)
 17def insertion_capacity(delta):
 18    return (1+delta)*(1-entropy(delta/(1+delta)))
 19
 20def markov_sample(n,q,rng):
 21    z=np.empty(n,dtype=np.int64); z[0]=rng.integers(2)
 22    for i in range(1,n): z[i]=z[i-1] if rng.random()>=q else 1-z[i-1]
 23    return z
 24
 25def is_subsequence(c,y):
 26    j=0
 27    for x in y:
 28        if j<len(c) and c[j]==x: j+=1
 29    return j==len(c)
 30
 31def insert_bits(c,delta,rng):
 32    k=int(round(delta*len(c))); positions=sorted(rng.choice(len(c)+k,k,replace=False)) if k else []
 33    out=[]; j=0; posset=set(positions)
 34    for i in range(len(c)+k):
 35        if i in posset: out.append(int(rng.integers(2)))
 36        else: out.append(int(c[j])); j+=1
 37    return np.array(out,dtype=np.int64)
 38
 39def codebook(kind,n,M,q,rng):
 40    if kind=='iid': return [rng.integers(2,size=n,dtype=np.int64) for _ in range(M)]
 41    return [markov_sample(n,q,rng) for _ in range(M)]
 42
 43def list_probe(codes,delta,trials,rng):
 44    lists=[]
 45    for _ in range(trials):
 46        c=codes[int(rng.integers(len(codes)))]; y=insert_bits(c,delta,rng)
 47        lists.append(sum(is_subsequence(x,y) for x in codes))
 48    return float(np.mean(lists)), int(np.max(lists)), float(np.mean(np.array(lists)>1))
 49
 50def make_data(codes,delta,count,rng):
 51    xs=[]; ys=[]
 52    for _ in range(count):
 53        label=int(rng.integers(len(codes))); xs.append(insert_bits(codes[label],delta,rng)); ys.append(label)
 54    return torch.tensor(np.stack(xs),dtype=torch.long), torch.tensor(ys,dtype=torch.long)
 55
 56class Decoder(nn.Module):
 57    def __init__(self,M,h=48):
 58        super().__init__(); self.emb=nn.Embedding(2,16); self.rnn=nn.GRU(16,h,batch_first=True); self.fc=nn.Linear(h,M)
 59    def forward(self,x):
 60        _,h=self.rnn(self.emb(x)); return self.fc(h[-1])
 61
 62def train_eval(codes,delta,rng):
 63    # Fixed small setup, identical optimizer and number of updates.
 64    global device
 65    tr_x,tr_y=make_data(codes,delta,3200,rng)
 66    te={d:make_data(codes,d,1000,rng) for d in [0,.25,.5]}
 67    def run(dev):
 68        model=Decoder(len(codes)).to(dev)
 69        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 70        model.train(); bs=64
 71        for epoch in range(12):
 72            perm=torch.randperm(len(tr_y))
 73            for st in range(0,len(tr_y),bs):
 74                ix=perm[st:st+bs]
 75                opt.zero_grad(set_to_none=True)
 76                loss=nn.functional.cross_entropy(model(tr_x[ix].to(dev)),tr_y[ix].to(dev))
 77                loss.backward(); opt.step()
 78        out={}; model.eval()
 79        with torch.no_grad():
 80            for d,(x,y) in te.items():
 81                pred=model(x.to(dev)).argmax(1).cpu()
 82                out[str(d)]=float((pred==y).float().mean())
 83        return out
 84    try:
 85        return run(device)
 86    except Exception as exc:
 87        print('device failure, falling back to CPU:', repr(exc))
 88        if device.type == 'cuda':
 89            torch.cuda.empty_cache()
 90        device=torch.device('cpu')
 91        return run(device)
 92
 93def main():
 94    delta=.25; n=12; M=8; q=.25; rng=np.random.default_rng(SEED)
 95    cap=insertion_capacity(delta); rate=math.log2(M)/n
 96    mathcheck={'delta':delta,'capacity_formula':cap,'code_rate':rate,'safe_rate_(eps=.1)':.9*cap,'rate_satisfies_margin':rate<=.9*cap}
 97    # Several independently drawn codebooks reduce dependence on one random draw.
 98    rows=[]
 99    for kind,kq in [('iid',.5),('markov',.25),('markov',.4)]:
100        lp=[]; acc=[]; t0=time.time()
101        for rep in range(2):
102            rr=np.random.default_rng(SEED+100*rep+(0 if kind=='iid' else 1))
103            codes=codebook(kind,n,M,kq,rr)
104            lp.append(list_probe(codes,delta,800,rr))
105            acc.append(train_eval(codes,delta,rr))
106        rows.append({'kind':kind,'q':kq,'list_mean':float(np.mean([x[0] for x in lp])),'list_max_mean':float(np.mean([x[1] for x in lp])),'collision_fraction':float(np.mean([x[2] for x in lp])),'accuracy':{d:float(np.mean([a[d] for a in acc])) for d in ['0','0.25','0.5']},'seconds':time.time()-t0})
107    result={'seed':SEED,'device':str(device),'mathcheck':mathcheck,'results':rows,'notes':'iid q=.5 is the symmetric iid Bernoulli source; Markov q=.25 has positive persistence. Lists are sampled by inserting random bits and counting codewords that are subsequences.'}
108    Path('results.json').write_text(json.dumps(result,indent=2))
109    print(json.dumps(result,indent=2))
110if __name__=='__main__': main()