imrishabh18/rp2040-motor-controller

scripts/find-encoder-cap-slots.py

import json,math
from pathlib import Path
from shapely.geometry import Point,LineString,Polygon,box
from shapely.affinity import rotate,translate
from shapely.strtree import STRtree
from copper_geometry import trace_segments
j=json.load(open('dist/baseline/circuit.json'));sc={e['source_component_id']:e['name'] for e in j if e['type']=='source_component'};pc={e['pcb_component_id']:e for e in j if e['type']=='pcb_component'}
parent={}
def root(x):
 parent.setdefault(x,x)
 if parent[x]!=x:parent[x]=root(parent[x])
 return parent[x]
for e in j:
 ids=e.get('connected_source_port_ids',[])+e.get('connected_source_net_ids',[]) if e['type']=='source_trace' else e.get('source_port_ids',[]) if e['type']=='source_component_internal_connection' else []
 for v in ids[1:]:parent[root(v)]=root(ids[0])
sp={e['source_port_id']:e for e in j if e['type']=='source_port'};pp={e['pcb_port_id']:e for e in j if e['type']=='pcb_port'};st={e['source_trace_id']:e for e in j if e['type']=='source_trace'};pt={e['pcb_trace_id']:e for e in j if e['type']=='pcb_trace'}
def tnet(e):
 if e.get('source_net_id'):return root(e['source_net_id'])
 if e.get('source_trace_id') in st:
  s=st[e['source_trace_id']];return root((s['connected_source_port_ids']+s['connected_source_net_ids'])[0])
 if e.get('pcb_trace_id') in pt and e['type']!='pcb_trace':return tnet(pt[e['pcb_trace_id']])
 raise ValueError(e)
portnet={id:root(e['source_port_id']) for id,e in pp.items() if e.get('source_port_id')}
ground=root(next(e['source_net_id'] for e in j if e['type']=='source_net' and e['name']=='GND'))
v33=root(next(e['source_net_id'] for e in j if e['type']=='source_net' and e['name']=='V3V3'))
courts=[]
for e in j:
 if e['type'] not in ['pcb_courtyard_rect','pcb_courtyard_outline'] or e['layer']!='top':continue
 name=sc.get(pc[e['pcb_component_id']].get('source_component_id'),'')
 if name.startswith('TP_'):continue
 if e['type']=='pcb_courtyard_outline':g=Polygon([(p['x'],p['y']) for p in e['outline']])
 else:g=translate(rotate(box(-e['width']/2,-e['height']/2,e['width']/2,e['height']/2),e.get('ccw_rotation',0)),e['center']['x'],e['center']['y'])
 courts.append(g)
ct=STRtree(courts)
shapes=[];nets=[];viaflags=[];owners=[];protected=set()
for t in pt.values():
 if max((p.get("width",0) for p in t["route"]),default=0)>=.35 or st.get(t.get("source_trace_id"),{}).get("name","").startswith(("XIN","XOUT","CXIN","CXOUT","CURRENT_A_SENSE","CURRENT_B_SENSE")):protected.add(t["pcb_trace_id"])
for e in j:
 if e['type']=='pcb_trace':
  for a,b,l,w in trace_segments(e['route']):
   if l=='top':shapes.append(LineString([a,b]).buffer(w/2));nets.append(tnet(e));viaflags.append(False);owners.append(e["pcb_trace_id"])
 elif e['type']=='pcb_via':
  shapes.append(Point(e['x'],e['y']).buffer(e['outer_diameter']/2));nets.append(None);viaflags.append(True);owners.append(e.get("pcb_trace_id","manual"))
trace_tree=STRtree(shapes);candidates=[]
for angle in [0,90,180,270]:
 base_court=rotate(box(-.93,-.47,.93,.47),angle)
 base_pads=[rotate(translate(box(-.27,-.32,.27,.32),x,0),angle,origin=(0,0)) for x in [-.51,.51]]
 for xi in range(-140,171):
  for yi in range(-200,101):
   x,y=xi*.05,yi*.05
   px=x-.51*math.cos(math.radians(angle));py=y-.51*math.sin(math.radians(angle))
   if math.hypot(px-.635,py+2.7051)>7.3:continue
   court=translate(base_court,x,y)
   if any(court.intersects(courts[i]) for i in ct.query(court)):continue
   valid=True;gap=1;hits=set()
   for pad,net in zip(base_pads,[v33,ground]):
    g=translate(pad,x,y)
    for i in trace_tree.query(g.buffer(.151)):
     if nets[i]==net:continue
     d=g.distance(shapes[i]);limit=.126 if viaflags[i] else .106
     if d<limit:
      if owners[i] in protected or owners[i]=='manual':valid=False;break
      hits.add(owners[i])
     gap=min(gap,d)
    if not valid:break
   if valid:
    score=100*len(hits)+((px-.635)**2+(py+2.7051)**2)/100
    candidates.append((score,x,y,angle,court,sorted(hits)))
candidates.sort(key=lambda p:p[0]);candidates=candidates[:600];print('legal slots',len(candidates))
pairs=[]
for a in candidates:
 for b in candidates:
  if b[0]<a[0] or a[4].distance(b[4])<.04:continue
  pairs.append((100*len(set(a[5]+b[5]))+(a[0]%100)+(b[0]%100),a[:4]+(a[5],),b[:4]+(b[5],)))
print('best pairs',sorted(pairs)[:15])
Path('reports/cap-slot-candidates.json').write_text(json.dumps({'candidates':[p[:4]+(p[5],) for p in candidates],'pairs':sorted(pairs)[:30]},indent=2))