imrishabh18/rp2040-motor-controller

scripts/audit-ground-copper.py

import json,collections,hashlib,sys
from pathlib import Path
from shapely import make_valid
from shapely.geometry import Point,Polygon,LineString,box
from shapely.affinity import rotate,translate
from shapely.strtree import STRtree
from copper_geometry import trace_segments, subtract_drills
input_path=Path(sys.argv[1] if len(sys.argv)>1 else 'dist/index/circuit.json')
layers=next((a.split('=',1)[1].split(',') for a in sys.argv if a.startswith('--layers=')),['top','inner1','inner2','bottom'])
assert set(layers) <= {'top','inner1','inner2','bottom'} and layers
output_path=Path(next((a.split('=',1)[1] for a in sys.argv if a.startswith('--output=')),'reports/ground-copper-connectivity.json'))
j=json.loads(input_path.read_text()); parent={}
def root(x):
 parent.setdefault(x,x)
 if parent[x]!=x:parent[x]=root(parent[x])
 return parent[x]
def join(a,b):parent[root(a)]=root(b)
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 x in ids[1:]:join(ids[0],x)
gnd=root(next(e['source_net_id'] for e in j if e['type']=='source_net' and e['name']=='GND'))
ports={e['pcb_port_id']:root(e['source_port_id'])==gnd for e in j if e['type']=='pcb_port' and e.get('source_port_id')}
traces={e['source_trace_id']:any(root(x)==gnd for x in e['connected_source_port_ids']+e['connected_source_net_ids']) for e in j if e['type']=='source_trace'}
ground_vias={(round(r['x'],6),round(r['y'],6)) for e in j if e['type']=='pcb_trace' and traces.get(e.get('source_trace_id')) for r in e['route'] if r['route_type']=='via'}
shapes=[]; meta=[]; barrels=[]
def add(g,l,id):
 if l not in layers:return
 if g.is_empty:return
 if g.geom_type in ("MultiPolygon", "GeometryCollection"):
  for piece in g.geoms:add(piece,l,id)
  return
 shapes.append(g);meta.append((l,id))
def pad(e,outer=False,hole=False):
 if e.get("shape")=="polygon":return make_valid(Polygon([(p["x"],p["y"]) for p in e["points"]]))
 prefix='hole_' if hole else 'outer_' if outer else ''
 if e['shape']=='circle':g=Point(0,0).buffer(e.get(prefix+'diameter',2*e.get('radius',0))/2,quad_segs=64)
 else:
  w,h=e[prefix+'width'],e[prefix+'height']; r=min(w,h)/2 if 'pill' in e['shape'] else e.get('corner_radius',0)
  if r:
   if abs(w-2*r)<1e-8 or abs(h-2*r)<1e-8:g=LineString([(-w/2+r,-h/2+r),(w/2-r,h/2-r)]).buffer(r,quad_segs=64)
   else:g=box(-w/2+r,-h/2+r,w/2-r,h/2-r).buffer(r,quad_segs=64)
  else:g=box(-w/2,-h/2,w/2,h/2)
 return translate(rotate(g,e.get('ccw_rotation',0),origin=(0,0)),e['x'],e['y'])
for e in j:
 t=e['type'];id=e.get(t+'_id')
 if t=='pcb_copper_pour' and root(e['source_net_id'])==gnd:
  b=e['brep_shape']; xy=lambda r:[(v['x'],v['y']) for v in r['vertices']]
  g=Polygon(xy(b['outer_ring']),[xy(r) for r in b['inner_rings']]);g=make_valid(g);add(g,e['layer'],id)
 elif t=='pcb_smtpad' and ports.get(e.get('pcb_port_id')):add(pad(e),e['layer'],id)
 elif t=='pcb_plated_hole' and ports.get(e.get('pcb_port_id')):
  ids=[];g=pad(e,outer=True).difference(pad(e,hole=True))
  for l in e['layers']:
   if l in layers:ids.append(len(shapes));add(g,l,id)
  barrels.append(ids)
 elif t=='pcb_via' and (traces.get(e.get('source_trace_id')) or (round(e['x'],6),round(e['y'],6)) in ground_vias):
  g=Point(e['x'],e['y']).buffer(e['outer_diameter']/2,quad_segs=64).difference(Point(e['x'],e['y']).buffer(e['hole_diameter']/2,quad_segs=64));ids=[]
  for l in e['layers']:
   if l in layers:ids.append(len(shapes));add(g,l,id)
  barrels.append(ids)
 elif t=='pcb_trace' and traces.get(e.get('source_trace_id')):
  for a,b,layer,width in trace_segments(e['route']):
   add(LineString([a,b]).buffer(width/2,quad_segs=64),layer,id)
drills=[]
for e in j:
 if e['type']=='pcb_via':drills.append((Point(e['x'],e['y']).buffer(e['hole_diameter']/2,quad_segs=64),e['layers']))
 elif e['type']=='pcb_plated_hole':drills.append((pad(e,hole=True),e['layers']))
shapes,meta,barrels=subtract_drills(shapes,meta,barrels,drills,lambda m:m[0])
parent={}
for i in range(len(shapes)):root(i)
for ids in barrels:
 for i in ids[1:]:join(ids[0],i)
for l in layers:
 indices=[i for i,m in enumerate(meta) if m[0]==l];tree=STRtree([shapes[i] for i in indices])
 for i in indices:
  for k in tree.query(shapes[i].buffer(1e-7)):
   other=indices[k]
   if shapes[i].distance(shapes[other])<=1e-7:join(i,other)
groups=collections.defaultdict(list)
for i in range(len(shapes)):groups[root(i)].append(i)
result={'circuitJsonSha256':hashlib.sha256(input_path.read_bytes()).hexdigest(),'groundCopperConnectedComponents':len(groups),'groundPourCount':sum(e['type']=='pcb_copper_pour' for e in j),'geometryPieceCount':len(shapes),'groups':[{'pieces':len(v),'areaSumMm2':sum(shapes[i].area for i in v),'elements':sorted(set(meta[i][1] for i in v))} for v in sorted(groups.values(),key=len,reverse=True)]}
Path('reports').mkdir(exist_ok=True)
result['auditedLayers']=layers
output_path.write_text(json.dumps(result,indent=2)+'\n')
print(json.dumps({**result,'groups':[dict(g,elements=g['elements'] if g['pieces']<30 else ['main ground network']) for g in result['groups']]},indent=2))

if len(groups) != 1:sys.exit("Ground copper is not one connected network")