1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 7
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def retention_sweep():
 10    # f(s)=s, reset=0, constant gate g => m_n=g^n*m_0 exactly.
 11    n = 20; rows = []
 12    for g in [0.0, 0.25, 0.5, 0.8, 0.9, 0.99]:
 13        observed = float(g ** n)
 14        if 0 < g < 1:
 15            half_pred = math.log(0.5) / math.log(g)
 16            half_obs = next(k for k in range(101) if g ** k <= 0.5)
 17        elif g == 0:
 18            half_pred, half_obs = 1.0, 1
 19        else:
 20            half_pred, half_obs = float('inf'), None
 21        rows.append({'g':g, 'steps':n, 'observed':observed,
 22                     'predicted':float(g**n), 'half_life_observed':half_obs,
 23                     'half_life_predicted':half_pred})
 24    return rows
 25
 26
 27def stability_sweep():
 28    # f(s)=lambda*s then retention s+=g*s gives multiplier q=lambda*g.
 29    rows=[]
 30    for lam in [0.8, 1.0, 1.05, 1.2]:
 31        for g in [0.5, 0.8, 0.99, 1.0, 1.1]:
 32            q=lam*g; s=1.0
 33            for _ in range(40): s*=q
 34            rows.append({'lambda':lam,'g':g,'multiplier':q,'final_abs':abs(s),
 35                         'observed_stable':abs(s)<=1.0,
 36                         'predicted_stable':q<=1.0})
 37    return rows
 38
 39
 40def overwrite_sweep():
 41    # s+=g*s-+(1-g)*z, old=1,z=0 => overwrite amount 1-g.
 42    rows=[]
 43    for g in [0.0,.25,.5,.75,1.0]:
 44        out=g
 45        rows.append({'g':g,'observed_overwrite':1-out,
 46                     'predicted_overwrite':1-g,'resulting_state':out})
 47    return rows
 48
 49
 50def delayed_bit_train():
 51    # Stream: bit, distractors, query. State is partitioned into a two-value
 52    # write port and persistent workspace. Only the query output is scored.
 53    try:
 54        import torch
 55        import torch.nn as nn
 56        torch.manual_seed(SEED)
 57        try:
 58            device='cuda' if torch.cuda.is_available() else 'cpu'
 59            # Small tensors only; fall back if CUDA is unavailable/unhealthy.
 60            torch.zeros(1, device=device)
 61        except Exception:
 62            device='cpu'
 63
 64        class Workspace(nn.Module):
 65            def __init__(self, hidden=8, gated=True):
 66                super().__init__(); self.hidden=hidden; self.gated=gated
 67                self.enc=nn.Linear(2,hidden)
 68                self.trans=nn.Sequential(nn.Linear(hidden,hidden),nn.Tanh())
 69                self.read=nn.Linear(hidden,1)
 70                self.gate=nn.Linear(2*hidden,hidden) if gated else None
 71                self.reset=nn.Parameter(torch.zeros(hidden))
 72            def forward(self,x):
 73                b,time,_=x.shape
 74                # state starts with write-port zeros and workspace zeros
 75                s=torch.zeros(b,self.hidden,device=x.device)
 76                outs=[]
 77                for t in range(time):
 78                    z=self.enc(x[:,t])
 79                    # Functional overwrite: write coordinates are replaced;
 80                    # coordinates 2: are retained from the prior state.
 81                    s=torch.cat([z[:,:2],s[:,2:]],dim=-1)
 82                    s=self.trans(s)
 83                    outs.append(self.read(s).squeeze(-1))
 84                    if self.gated:
 85                        g=torch.sigmoid(self.gate(torch.cat([s,z],dim=-1)))
 86                        s=g*s+(1-g)*self.reset
 87                return torch.stack(outs,dim=1)
 88
 89        def make_batch(batch,gap):
 90            time=gap+2; x=torch.zeros(batch,time,2,device=device)
 91            bit=torch.randint(0,2,(batch,),device=device).float()
 92            x[:,0,0]=bit; x[:,-1,1]=1.
 93            return x,bit
 94
 95        def run(gated,gap):
 96            torch.manual_seed(SEED+gap+int(gated))
 97            model=Workspace(gated=gated).to(device)
 98            opt=torch.optim.Adam(model.parameters(),lr=.01)
 99            for _ in range(600):
100                x,y=make_batch(64,gap); pred=model(x)[:,-1]
101                loss=nn.functional.binary_cross_entropy_with_logits(pred,y)
102                opt.zero_grad(); loss.backward()
103                torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
104            with torch.no_grad():
105                x,y=make_batch(1024,gap)
106                acc=((model(x)[:,-1]>0).float()==y).float().mean().item()
107            return acc
108
109        return {'device':device,'results':{
110            str(g):{'baseline_no_retention':run(False,g),
111                    'gated_workspace':run(True,g)} for g in [2,8,16]}}
112    except Exception as exc:
113        return {'device':'cpu','error':repr(exc)}
114
115
116def main():
117    out={'retention_prediction':retention_sweep(),
118         'stability_prediction':stability_sweep(),
119         'overwrite_prediction':overwrite_sweep(),
120         'delayed_bit':delayed_bit_train()}
121    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
122
123if __name__=='__main__': main()