OCuLink Riser for the Supermicro AIOM Slot

#!/usr/bin/env python3
"""After the maze: GND fan-out vias for SMD ground pads, GND stitching and return vias, pours and keepouts,
board text. Writes design/routes.json (everything build_board.py final needs)."""
import json, math, os, sys, time
import numpy as np
from scipy import ndimage
from shapely.geometry import Polygon, box, Point
sys.path.insert(0, os.path.dirname(__file__))
import route_all as R
import maze, hs
from maze import G, LI, NX, NY

ROOT = R.ROOT
ls = json.load(open(os.path.join(ROOT, "design", "routes_ls.json")))
B = R.B
T, V = R.T, R.V
_ts = set(map(tuple, T)); _vs = set(map(tuple, V))
for t in map(tuple, ls["tracks"]):
    if t not in _ts:
        T.append(t); B.track(*t); _ts.add(t)
for v in map(tuple, ls["vias"]):
    if v not in _vs:
        V.append(v); B.via(v[0], v[1], v[2], v[3]); _vs.add(v)
gnd = B.nid("GND")


def holes_mask(gap=0.62):
    """cells too close to any existing drilled hole (via or plated pad) for hole-to-hole spacing"""
    m = np.zeros((NY, NX), bool)
    pts = [(v[0], v[1], v[4] if len(v) > 4 else 0.2) for v in V]
    pts += [(p["x"], p["y"], max(p["drill"])) for p in R.pads if p["drill"][0] > 0]
    for (x, y, h) in pts:
        r = gap + h / 2
        j0, j1 = max(0, int((y - r) / G)), min(NY - 1, int((y + r) / G) + 1)
        i0, i1 = max(0, int((x - r) / G)), min(NX - 1, int((x + r) / G) + 1)
        yy_, xx_ = np.mgrid[j0:j1 + 1, i0:i1 + 1]
        m[j0:j1 + 1, i0:i1 + 1] |= (xx_ * G - x) ** 2 + (yy_ * G - y) ** 2 < r * r
    return m


def via_ok_mask(clr=0.25, clr_hs=0.35, d=maze.VIA_D):
    ok = ~holes_mask()
    for li in range(4):
        lab = B.lab[li]
        other = (lab >= 0) & (lab != gnd)
        ok &= ndimage.distance_transform_edt(~other) * G >= clr + d / 2
        oh = other & B.hs[li]
        ok &= ndimage.distance_transform_edt(~oh) * G >= clr_hs + d / 2
        ok &= ~B.block_via[li]
    # keep top/bottom copper out of the 2.75 mm long-edge strips (OCP: component and trace free)
    ok &= (B.X > 3.0) & (B.X < 73.0)
    return ok


# ---------------------------------------------------------------- probe-pad stubs (design/tp_stubs.json)
_st = os.path.join(ROOT, "design", "tp_stubs.json")
if os.path.exists(_st):
    for t in json.load(open(_st))["tracks"]:
        t = tuple(t); T.append(t); B.track(*t)

# ---------------------------------------------------------------- GND fan-out for SMD ground pads
fan = 0
FAN_T, FAN_V = [], []
FIXED_NETS = {t[1] for t in ls.get("fixed_tracks", [])} | {v[2] for v in ls.get("fixed_vias", [])}


def fan_pad(p):
    vok = via_ok_mask(0.25, 0.32)
    cx, cy = p["x"], p["y"]
    j0, i0 = int(round(cy / G)), int(round(cx / G))
    r = int(2.5 / G)
    jj, ii = np.nonzero(vok[max(0, j0 - r):j0 + r + 1, max(0, i0 - r):i0 + r + 1])
    goal = {(LI["L1"], int(j + max(0, j0 - r)), int(i + max(0, i0 - r))) for j, i in zip(jj, ii)}
    x0, y0, x1, y1 = p["geom"].bounds
    from shapely import contains_xy
    sj0, sj1, si0, si1 = int(y0 / G), int(y1 / G) + 1, int(x0 / G), int(x1 / G) + 1
    sub = contains_xy(p["geom"].buffer(-0.02), B.X[sj0:sj1 + 1, si0:si1 + 1], B.Y[sj0:sj1 + 1, si0:si1 + 1])
    sjs, sis = np.nonzero(sub)
    start = {(LI["L1"], int(j + sj0), int(i + si0)) for j, i in zip(sjs, sis)}
    if not goal or not start:
        return False
    paths, msg = maze.route_net(B, "GND", [start, goal], w=0.25, clr=0.2, clr_hs=0.32, allow=("L1",), max_expand=400000)
    if paths is None:
        # no via site nearby: a short top-layer trace to GND copper already there (another GND pad, a fan-out)
        own = p["geom"].buffer(0.25)
        gj, gi = np.nonzero(B.lab[LI["L1"], max(0, j0 - r):j0 + r + 1, max(0, i0 - r):i0 + r + 1] == gnd)
        goal2 = {(LI["L1"], int(j + max(0, j0 - r)), int(i + max(0, i0 - r))) for j, i in zip(gj, gi)}
        goal2 = {c for c in goal2 if not own.contains(Point(c[2] * G, c[1] * G))}
        if goal2:
            paths, msg = maze.route_net(B, "GND", [start, goal2], w=0.2, clr=0.15, clr_hs=0.3, allow=("L1",), max_expand=400000)
        if paths is not None:
            tr, _ = maze.path_to_copper(paths[0], "GND", 0.2)
            for t in tr:
                T.append(t); B.track(*t); FAN_T.append(t)
            print("GND trace to nearby GND copper for", p["ref"], p["num"])
            return True
    if paths is None:
        # no room beside the pad: a filled, capped via in the pad (JLC 6-layer POFV, as at the OCuLink pads)
        d_ip, h_ip = 0.3, 0.15
        vin = ~holes_mask(0.35)
        for li in range(4):
            other = (B.lab[li] >= 0) & (B.lab[li] != gnd)
            vin &= ndimage.distance_transform_edt(~other) * G >= 0.15 + d_ip / 2
            oh = other & B.hs[li]
            vin &= ndimage.distance_transform_edt(~oh) * G >= 0.28 + d_ip / 2
        cand = [(math.hypot(c[2] * G - cx, c[1] * G - cy), c) for c in start if vin[c[1], c[2]]
                and p["geom"].buffer(-0.03).contains(Point(c[2] * G, c[1] * G).buffer(d_ip / 2))]
        if not cand:
            cand = [(math.hypot(c[2] * G - cx, c[1] * G - cy), c) for c in start if vin[c[1], c[2]]]
        if not cand:
            return False
        c = min(cand)[1]
        v = (round(c[2] * G, 4), round(c[1] * G, 4), "GND", d_ip, h_ip)
        V.append(v); B.via(v[0], v[1], "GND", d_ip); FAN_V.append(v)
        print("GND via-in-pad for", p["ref"], p["num"], v[:2])
        return True
    tr, _ = maze.path_to_copper(paths[0], "GND", 0.25)
    end = paths[0][-1]
    for t in tr:
        T.append(t); B.track(*t); FAN_T.append(t)
    v = (round(end[2] * G, 4), round(end[1] * G, 4), "GND", maze.VIA_D, maze.VIA_H)
    V.append(v); B.via(v[0], v[1], "GND", v[3]); FAN_V.append(v)
    return True


def unpaint(net):
    nid = B.nid(net)
    B.lab[B.lab == nid] = -1
    for q in R.pads:
        if q["net"] == net:
            B.paint_geom(q["lis"], q["geom"], nid)


def rip_and_fan(p):
    """Last resort for a boxed-in GND pad: rip the low-speed nets crowding it on L1, fan the pad out, then
    route those nets again around the new copper (restoring them as they were if that fails)."""
    near = p["geom"].buffer(0.7)
    x0, y0, x1, y1 = near.bounds
    j0, j1, i0, i1 = int(y0 / G), int(y1 / G) + 1, int(x0 / G), int(x1 / G) + 1
    ids = set(np.unique(B.lab[:, j0:j1 + 1, i0:i1 + 1])) - {-1, gnd}
    rips = sorted(B.netname[i] for i in ids if B.netname[i] not in FIXED_NETS and not B.netname[i].startswith(("+", "PE", "REFCLK")))
    if not rips:
        return False
    saved = {n: ([t for t in T if t[1] == n], [v for v in V if v[2] == n]) for n in rips}
    for n in rips:
        T[:] = [t for t in T if t[1] != n]; V[:] = [v for v in V if v[2] != n]
        unpaint(n)
    nT, nV = len(FAN_T), len(FAN_V)
    ok = fan_pad(p)
    if ok:
        newc = {}
        for n in rips:
            w, allow = {r_[0]: (r_[1], r_[2]) for r_ in R.ROUTE}.get(n, (0.15, ("L1", "L3", "L5")))
            paths, msg = maze.route_net(B, n, R.terminals(n), w=w, clr=R.CLR, clr_hs=R.CLR_HS, allow=allow, l6_region=R.L6_REGION)
            if paths is None:
                ok = False; print("  reroute failed:", n, msg); break
            tr, vi = [], []
            for pth in paths:
                a_, b_ = maze.path_to_copper(pth, n, w); tr += a_; vi += b_
            for t in tr: B.track(*t)
            for v in vi: B.via(v[0], v[1], n, v[3])
            newc[n] = (tr, vi)
        if ok:
            for n, (tr, vi) in newc.items():
                T.extend(tr); V.extend(vi)
            print("  ripped and re-routed", rips, "for", p["ref"], p["num"])
            return True
        # undo: drop the fan-out and the partial reroutes, restore the old copper
        drop_t, drop_v = set(map(tuple, FAN_T[nT:])), set(map(tuple, FAN_V[nV:]))
        del FAN_T[nT:]; del FAN_V[nV:]
        T[:] = [t for t in T if tuple(t) not in drop_t]; V[:] = [v for v in V if tuple(v) not in drop_v]
        unpaint("GND")
        for t in T:
            if t[1] == "GND": B.track(*t)
        for v in V:
            if v[2] == "GND": B.via(v[0], v[1], "GND", v[3])
    for n in rips:
        unpaint(n)
        T[:] = [t for t in T if t[1] != n]; V[:] = [v for v in V if v[2] != n]
        tr, vi = saved[n]
        T.extend(tr); V.extend(vi)
        for t in tr: B.track(*t)
        for v in vi: B.via(v[0], v[1], n, v[3])
    return False


for p in R.pads:
    if p["net"] != "GND" or p["ref"] in ("J1", "J2", "J3", "J4", "J5", "H1", "H2") or p["drill"][0] > 0:
        continue
    if fan_pad(p) or rip_and_fan(p):
        fan += 1
    else:
        print("no GND fan-out for", p["ref"], p["num"])
print("GND fan-out vias:", fan)

# ---------------------------------------------------------------- return vias beside every high-speed via pair
vok = via_ok_mask(0.25, 0.3)
ret = 0
hsv = [v for v in V if v[2] != "GND" and (v[2].startswith(("PET", "PER", "REFCLK")))]
taken = []
for (x, y, net, d, h) in hsv:
    # look for a GND via spot within 1.1 mm, prefer the closest, at most one per signal via
    j0, i0 = int(round(y / G)), int(round(x / G))
    cands = []
    for dj in range(-11, 12):
        for di in range(-11, 12):
            j, i = j0 + dj, i0 + di
            if 0 <= j < NY and 0 <= i < NX and vok[j, i]:
                rr = math.hypot(dj, di) * G
                if 0.55 <= rr <= 1.1:
                    cands.append((rr, j, i))
    cands.sort()
    for rr, j, i in cands:
        if all(math.hypot(i * G - a, j * G - b) > 0.65 for (a, b) in taken):
            v = (i * G, j * G, "GND", maze.VIA_D, maze.VIA_H)
            V.append(v); B.via(i * G, j * G, "GND", maze.VIA_D)
            taken.append((i * G, j * G)); ret += 1
            vok = None
            break
    if vok is None:
        vok = via_ok_mask(0.25, 0.3)
print("return vias:", ret)

# ---------------------------------------------------------------- stitching grid
vok = via_ok_mask(0.25, 0.35)
st = 0
pitch = 2.0
for yy in np.arange(1.2, 112.0, pitch):
    for xx in np.arange(3.2, 73.0, pitch):
        j, i = int(round(yy / G)), int(round(xx / G))
        if vok[j, i]:
            V.append((round(xx, 3), round(yy, 3), "GND", maze.VIA_D, maze.VIA_H)); st += 1
            # keep the mask honest for the next candidates
            r = int(0.8 / G)
            vok[max(0, j - r):j + r + 1, max(0, i - r):i + r + 1] = False
print("stitching vias:", st)

# ---------------------------------------------------------------- pours
outline = Polygon(R.lay["outline"])
outer = outline.buffer(-0.3).intersection(box(2.75, 0.3, 76 - 2.75, 111.95))
inner = outline.buffer(-1.0).intersection(box(1.0, 0.3, 75.0, 113.6))


def poly_pts(g):
    if g.geom_type == "MultiPolygon":
        g = max(g.geoms, key=lambda q: q.area)
    return [(round(x, 4), round(y, 4)) for (x, y) in list(g.exterior.coords)[:-1]]



def prune_dangling():
    """Drop signal stubs that end nowhere (a fixed bus start the maze tapped further along): a via touching
    fewer than two tracks and no pad, and a track whose free end touches no pad, via or other track."""
    from shapely.geometry import LineString
    padg = {}
    for p in R.pads:
        padg.setdefault(p["net"], []).append(p["geom"])
    removed = 0
    while True:
        changed = False
        for net in sorted({t[1] for t in T} - {"GND"}):
            tr = [t for t in T if t[1] == net]
            vi = [v for v in V if v[2] == net]
            pg = padg.get(net, [])

            def on_track(x, y, layer, skip):
                pt = Point(x, y)
                return any(t is not skip and t[0] == layer and
                           LineString([(t[2], t[3]), (t[4], t[5])]).distance(pt) <= (t[6] + skip[6]) / 2 - 0.005 for t in tr)

            def at_via(x, y, w):
                return any(math.hypot(v[0] - x, v[1] - y) <= (v[3] + w) / 2 - 0.005 for v in vi)

            def in_pad(x, y, r=0.0):
                return any(g.distance(Point(x, y)) <= max(r - 0.005, 1e-3) for g in pg)
            for v in vi:
                n = sum(1 for t in tr if LineString([(t[2], t[3]), (t[4], t[5])]).distance(Point(v[0], v[1])) <= (v[3] + t[6]) / 2 - 0.005)
                if n < 2 and not in_pad(v[0], v[1], v[3] / 2):
                    V.remove(v); removed += 1; changed = True
            vi = [v for v in V if v[2] == net]
            for t in tr:
                for end in (0, 1):
                    x, y = (t[4], t[5]) if end else (t[2], t[3])
                    if in_pad(x, y, t[6] / 2) or at_via(x, y, t[6]) or on_track(x, y, t[0], t):
                        continue
                    # free end: cut the segment back to the last tap along it (another track's end, a via or
                    # a pad touching it), or drop it when nothing touches it but its other end
                    sx, sy = (t[2], t[3]) if end else (t[4], t[5])
                    L = math.hypot(x - sx, y - sy)
                    seg = LineString([(sx, sy), (x, y)])
                    us = []
                    for o in tr:
                        if o is t or o[0] != t[0]:
                            continue
                        for (ox, oy) in ((o[2], o[3]), (o[4], o[5])):
                            if seg.distance(Point(ox, oy)) <= (o[6] + t[6]) / 2 - 0.005:
                                us.append(seg.project(Point(ox, oy)))
                    for v in vi:
                        if seg.distance(Point(v[0], v[1])) <= (v[3] + t[6]) / 2 - 0.005:
                            us.append(seg.project(Point(v[0], v[1])))
                    for g in pg:
                        if g.distance(seg) <= t[6] / 2 - 0.005:
                            us.append(seg.project(g.centroid))
                    us = [u for u in us if 1e-3 < u < L - 1e-3]
                    T.remove(t); removed += 1; changed = True
                    if us:
                        u = max(us)
                        nx, ny = sx + (x - sx) * u / L, sy + (y - sy) * u / L
                        nt = (t[0], t[1], round(sx, 4), round(sy, 4), round(nx, 4), round(ny, 4), t[6])
                        T.append(nt)
                    break
        if not changed:
            return removed


print("dangling stubs pruned:", prune_dangling())

zones = []
for layer, g, clr in (("L1", outer, 0.35), ("L6", outer, 0.35), ("L3", inner, 0.3), ("L5", inner, 0.3),
                      ("L2", inner, 0.25), ("L4", inner, 0.25)):
    zones.append({"layer": layer, "net": "GND", "poly": poly_pts(g), "clearance": clr, "priority": 0,
                  "thermal": False})

texts = [
    {"text": "OCP NIC 3.0 SFF  ->  4x OCuLink 4i", "x": 38.0, "y": 104.6, "layer": "F.SilkS", "size": 1.0},
]
json.dump({"tracks": T, "vias": V, "zones": zones, "texts": [], "fanout_tracks": FAN_T, "fanout_vias": FAN_V, "stats": {"fanout": fan, "return": ret, "stitch": st}},
          open(os.path.join(ROOT, "design", "routes.json"), "w"))
print("routes.json:", len(T), "tracks,", len(V), "vias,", len(zones), "zones")