1"""Small, exact interval CEGAR abstraction for a 2-D neural-state-space toy.
  2
  3The affine map is an intentionally transparent surrogate for a saturated RNN cell:
  4    h' = A h + b.
  5For axis-aligned boxes, interval propagation is exact for each coordinate's
  6independent extrema.  Therefore every edge generated by intersecting the
  7resulting box with partition cells is conservative (up to floating point).
  8"""
  9import json, math, random
 10from dataclasses import dataclass
 11from pathlib import Path
 12import numpy as np
 13
 14A = np.array([[.85, .15], [-.15, .85]], dtype=float)
 15b = np.array([.20, 0.0])
 16DOMAIN = (-1.0, 1.0, -1.0, 1.0)
 17UNSAFE_X = .72
 18
 19@dataclass(frozen=True)
 20class Cell:
 21    x0: float; x1: float; y0: float; y1: float
 22    def bounds(self): return np.array([self.x0,self.x1,self.y0,self.y1])
 23    def center(self): return np.array([(self.x0+self.x1)/2,(self.y0+self.y1)/2])
 24    def width(self): return max(self.x1-self.x0,self.y1-self.y0)
 25    def intersects(self, box, eps=1e-12):
 26        return (self.x1 >= box[0]-eps and self.x0 <= box[1]+eps and
 27                self.y1 >= box[2]-eps and self.y0 <= box[3]+eps)
 28    def unsafe(self): return self.x1 >= UNSAFE_X
 29
 30def affine_box(cell, margin=0.0):
 31    lo, hi = cell.bounds()[::2], cell.bounds()[1::2]
 32    # extrema of each affine coordinate over a rectangle
 33    outlo=[]; outhi=[]
 34    for row, bias in zip(A,b):
 35        mn = bias + sum(c*l if c >= 0 else c*h for c,l,h in zip(row,lo,hi))
 36        mx = bias + sum(c*h if c >= 0 else c*l for c,l,h in zip(row,lo,hi))
 37        outlo.append(mn-margin); outhi.append(mx+margin)
 38    return np.array([outlo[0], outhi[0], outlo[1], outhi[1]])
 39
 40def uniform_partition(n):
 41    xs=np.linspace(DOMAIN[0],DOMAIN[1],n+1); ys=np.linspace(DOMAIN[2],DOMAIN[3],n+1)
 42    return [Cell(xs[i],xs[i+1],ys[j],ys[j+1]) for i in range(n) for j in range(n)]
 43
 44def split(c):
 45    xm=(c.x0+c.x1)/2; ym=(c.y0+c.y1)/2
 46    return [Cell(c.x0,xm,c.y0,ym),Cell(c.x0,xm,ym,c.y1),Cell(xm,c.x1,c.y0,ym),Cell(xm,c.x1,ym,c.y1)]
 47
 48def build_graph(cells, margin=0.0):
 49    edges=[]
 50    for c in cells:
 51        box=affine_box(c,margin)
 52        edges.append([j for j,d in enumerate(cells) if d.intersects(box)])
 53    return edges
 54
 55def reachable(cells, edges, initial=(.01,.01), steps=12):
 56    # Boundary ownership is immaterial because all intersecting cells are kept.
 57    cur={next(i for i,c in enumerate(cells) if c.x0 <= initial[0] <= c.x1 and c.y0 <= initial[1] <= c.y1)}
 58    levels=[set(cur)]; parents=[]
 59    for _ in range(steps):
 60        nxt=set(); par={}
 61        for i in cur:
 62            for j in edges[i]:
 63                nxt.add(j); par.setdefault(j,i)
 64        parents.append(par); levels.append(nxt); cur=nxt
 65    unsafe=sorted({i for lev in levels for i in lev if cells[i].unsafe()})
 66    return levels, parents, unsafe
 67
 68def path_to_unsafe(levels, parents, target):
 69    # recover one shortest-ish abstract counterexample
 70    t=next(k for k,lev in enumerate(levels) if target in lev)
 71    path=[target]
 72    for k in range(t-1,-1,-1): path.append(parents[k][path[-1]])
 73    return list(reversed(path))
 74
 75def concrete_rollout(h=(.01,.01), steps=30):
 76    z=np.array(h,float); out=[z.copy()]
 77    for _ in range(steps): z=A@z+b; out.append(z.copy())
 78    return np.array(out)
 79
 80def cegar(initial_n=4, rounds=8, margin=1e-10):
 81    cells=uniform_partition(initial_n); refinements=[]
 82    for _ in range(rounds):
 83        edges=build_graph(cells,margin); levels,parents,unsafe=reachable(cells,edges)
 84        if not unsafe: break
 85        path=path_to_unsafe(levels,parents,unsafe[0])
 86        # Refine the widest cell on the counterexample, prioritizing its prefix.
 87        k=max(range(len(path)), key=lambda q: cells[path[q]].width())
 88        chosen=path[k]; old=cells[chosen]; cells=cells[:chosen]+split(old)+cells[chosen+1:]
 89        refinements.append(old.width())
 90    edges=build_graph(cells,margin); levels,parents,unsafe=reachable(cells,edges)
 91    return cells, levels, unsafe, refinements
 92
 93def sampled_soundness(n=4, margin=0.0, samples=20000, seed=7):
 94    rng=np.random.default_rng(seed); cells=uniform_partition(n); edges=build_graph(cells,margin)
 95    failures=0; checked=0
 96    for i,c in enumerate(cells):
 97        pts=rng.uniform([c.x0,c.y0],[c.x1,c.y1],size=(max(1,samples//len(cells)),2))
 98        nxt=pts@A.T+b
 99        for z in nxt:
100            # The abstraction is over DOMAIN; successors outside it are
101            # intentionally excluded from this conditional soundness check.
102            if not (DOMAIN[0]-1e-12 <= z[0] <= DOMAIN[1]+1e-12 and DOMAIN[2]-1e-12 <= z[1] <= DOMAIN[3]+1e-12):
103                continue
104            js=[j for j,d in enumerate(cells) if d.x0-1e-12<=z[0]<=d.x1+1e-12 and d.y0-1e-12<=z[1]<=d.y1+1e-12]
105            checked+=1
106            if not any(j in edges[i] for j in js): failures+=1
107    return checked, failures
108
109def run():
110    np.set_printoptions(precision=4,suppress=True)
111    exact=concrete_rollout(); concrete_unsafe=bool(np.max(exact[:,0])>=UNSAFE_X)
112    sweep=[]
113    for n in [2,4,8,16]:
114        cells=uniform_partition(n); ed=build_graph(cells); lev,_,u=reachable(cells,ed)
115        area=sum((cells[i].x1-cells[i].x0)*(cells[i].y1-cells[i].y0) for i in u)
116        sweep.append({'n':n,'cells':len(cells),'abstract_unsafe_cells':len(u),
117                      'abstract_unsafe_area':area,
118                      'abstract_unsafe_levels':sum(any(cells[i].unsafe() for i in q) for q in lev),
119                      'concrete_unsafe':concrete_unsafe})
120    ce_cells,ce_lev,ce_u,ref=cegar()
121    ce_area=sum((c.x1-c.x0)*(c.y1-c.y0) for i,c in enumerate(ce_cells) if i in ce_u)
122    margins=[]
123    for m in [0,.001,.01,.05]:
124        checked,fail=sampled_soundness(8,m)
125        edge_count=sum(len(x) for x in build_graph(uniform_partition(8),m))
126        margins.append({'margin':m,'checked':checked,'failures':fail,'coverage':1-fail/checked,'edges':edge_count})
127    result={'map_A':A.tolist(),'map_b':b.tolist(),'unsafe_x':UNSAFE_X,
128            'concrete_max_x':float(np.max(exact[:,0])),'predictions':{
129              'P1_soundness': 'conditional in-domain sampled omissions are zero for every nonnegative margin',
130              'P2_uniform_resolution': 'area of reachable unsafe overapproximation decreases as cell width shrinks',
131              'P3_margin': 'adding interval margin cannot reduce edge count',
132              'P4_CEGAR_focus': 'CEGAR reaches comparable unsafe area with fewer cells than uniform n=16'},
133            'uniform_sweep':sweep,
134            'cegar':{'final_cells':len(ce_cells),'unsafe_cells':len(ce_u),'unsafe_area':ce_area,
135                     'unsafe_levels':sum(any(ce_cells[i].unsafe() for i in q) for q in ce_lev),
136                     'refined_widths':ref},'margin_sweep':margins}
137    Path('results.json').write_text(json.dumps(result,indent=2))
138    print(json.dumps(result,indent=2))
139
140if __name__=='__main__': run()