OCuLink Riser for the Supermicro AIOM Slot

#!/usr/bin/env python3
"""Multi-layer A* maze router for the riser's low-speed nets, on a 0.1 mm grid.

Copper of every net is rasterised per layer into a label grid; a net is routed through cells whose distance to
other nets' copper is at least (clearance + width/2) (scipy distance transform per net). Vias need a free disk
on every copper layer. Routed copper is rasterised back before the next net."""
import heapq, json, math, os, sys
import numpy as np
from scipy import ndimage
from shapely.geometry import Polygon, LineString, Point, box
from shapely import affinity

G = 0.1
NX, NY = 760, 1150
LAYERS = ["L1", "L3", "L5", "L6"]          # routing/obstacle layers (L2/L4 are solid GND planes)
LI = {l: i for i, l in enumerate(LAYERS)}
VIA_D, VIA_H = 0.40, 0.20


class Board:
    def __init__(self):
        self.lab = np.full((4, NY, NX), -1, np.int32)     # net id of copper in the cell, -1 none
        self.block = np.zeros((4, NY, NX), bool)          # static keepouts per layer
        self.nets = {}
        self.netname = []
        self.hs = np.zeros((4, NY, NX), bool)             # high-speed copper (gets extra clearance)
        yy, xx = np.mgrid[0:NY, 0:NX]
        self.X = xx * G
        self.Y = yy * G

    def nid(self, name):
        if name not in self.nets:
            self.nets[name] = len(self.netname)
            self.netname.append(name)
        return self.nets[name]

    def paint_geom(self, layer_idx, geom, net, hs=False):
        """Rasterise a shapely geometry (cells whose centre lies inside), dilated by 0.06 mm so the painted cells
        cover the true copper edge (a cell centre can sit up to ~0.07 mm inside it)."""
        geom = geom.buffer(0.06)
        x0, y0, x1, y1 = geom.bounds
        i0, i1 = max(0, int(x0 / G) - 1), min(NX - 1, int(x1 / G) + 2)
        j0, j1 = max(0, int(y0 / G) - 1), min(NY - 1, int(y1 / G) + 2)
        if i1 < i0 or j1 < j0:
            return
        from shapely import contains_xy
        sub = contains_xy(geom, self.X[j0:j1 + 1, i0:i1 + 1], self.Y[j0:j1 + 1, i0:i1 + 1])
        # thin geometry may contain no cell centre: paint the nearest cells of its centreline
        if not sub.any():
            c = geom.centroid
            sub = np.zeros_like(sub)
            jj, ii = int(round(c.y / G)) - j0, int(round(c.x / G)) - i0
            if 0 <= jj < sub.shape[0] and 0 <= ii < sub.shape[1]:
                sub[jj, ii] = True
        for li in ([layer_idx] if isinstance(layer_idx, int) else layer_idx):
            region = self.lab[li, j0:j1 + 1, i0:i1 + 1]
            region[sub] = net
            if hs:
                self.hs[li, j0:j1 + 1, i0:i1 + 1][sub] = True

    def track(self, layer, net, x1, y1, x2, y2, w, hs=False):
        g = LineString([(x1, y1), (x2, y2)]).buffer(w / 2, cap_style=1) if (x1, y1) != (x2, y2) else Point(x1, y1).buffer(w / 2)
        self.paint_geom(LI[layer], g, self.nid(net), hs)

    def via(self, x, y, net, d, hs=False):
        self.paint_geom([0, 1, 2, 3], Point(x, y).buffer(d / 2), self.nid(net), hs)


def free_mask(B, net, w, clr, clr_hs, li):
    """cells where a track of width w for `net` may have its centreline on layer li."""
    lab = B.lab[li]
    other = (lab >= 0) & (lab != net)
    d = ndimage.distance_transform_edt(~other) * G
    ok = d >= clr + w / 2 - 1e-9
    if clr_hs > clr:
        oh = other & B.hs[li]
        dh = ndimage.distance_transform_edt(~oh) * G
        ok &= dh >= clr_hs + w / 2 - 1e-9
    ok &= ~B.block[li]
    return ok


def route_net(B, net, terminals, w=0.15, clr=0.15, clr_hs=0.3, allow=("L1", "L3", "L5"), via_cost=12.0,
              layer_cost=None, l6_region=None, max_expand=6_000_000):
    """terminals: list of sets of (li, j, i) cells (each a pad). Returns list of paths (lists of (li,j,i))."""
    nid = B.nid(net)
    layer_cost = layer_cost or {"L1": 1.0, "L3": 1.15, "L5": 1.15, "L6": 3.0}
    ok = np.zeros((4, NY, NX), bool)
    for l in allow:
        ok[LI[l]] = free_mask(B, nid, w, clr, clr_hs, LI[l])
    if "L6" in allow and l6_region is not None:
        ok[LI["L6"]] &= l6_region
    # via allowed where a via disk is clear on every layer (incl. L6, never in keepouts)
    vok = np.ones((NY, NX), bool)
    for li in range(4):
        lab = B.lab[li]
        other = (lab >= 0) & (lab != nid)
        d = ndimage.distance_transform_edt(~other) * G
        vok &= d >= clr + VIA_D / 2 - 1e-9
        oh = other & B.hs[li]
        dh = ndimage.distance_transform_edt(~oh) * G
        vok &= dh >= clr_hs + VIA_D / 2 - 1e-9
        vok &= ~B.block_via[li] if hasattr(B, "block_via") else True
    # terminals' own cells are always enterable on their layers
    for t in terminals:
        for (li, j, i) in t:
            ok[li, j, i] = True
    tree = set(terminals[0])
    paths = []
    remaining = [set(t) for t in terminals[1:]]
    allowed_li = [LI[l] for l in allow]
    lc = np.array([layer_cost.get(l, 1.0) for l in LAYERS])
    while remaining:
        goal = set().union(*remaining)
        gj = np.array([c[1] for c in goal]); gi = np.array([c[2] for c in goal])
        gcj, gci = gj.mean(), gi.mean()
        rad = max(np.hypot(gj - gcj, gi - gci).max(), 0)

        def h(j, i):
            dj, di = abs(j - gcj), abs(i - gci)
            return max(0.0, (max(dj, di) + 0.414 * min(dj, di)) - rad)

        start = list(tree)
        dist = {}
        prev = {}
        heap = []
        for s in start:
            dist[s] = 0.0
            heapq.heappush(heap, (h(s[1], s[2]), 0.0, s))
        found = None
        n = 0
        while heap:
            f, g, c = heapq.heappop(heap)
            if g > dist.get(c, 1e18) + 1e-9:
                continue
            if c in goal:
                found = c; break
            n += 1
            if n > max_expand:
                break
            li, j, i = c
            for dj, di, cost in ((0, 1, 1), (0, -1, 1), (1, 0, 1), (-1, 0, 1), (1, 1, 1.4142), (1, -1, 1.4142), (-1, 1, 1.4142), (-1, -1, 1.4142)):
                jj, ii = j + dj, i + di
                if 0 <= jj < NY and 0 <= ii < NX and ok[li, jj, ii]:
                    if dj and di and not (ok[li, j, ii] and ok[li, jj, i]):
                        continue
                    nc = (li, jj, ii)
                    ng = g + cost * lc[li]
                    if ng < dist.get(nc, 1e18):
                        dist[nc] = ng; prev[nc] = c
                        heapq.heappush(heap, (ng + h(jj, ii), ng, nc))
            if vok[j, i]:
                for l2 in allowed_li:
                    if l2 != li and ok[l2, j, i]:
                        nc = (l2, j, i)
                        ng = g + via_cost
                        if ng < dist.get(nc, 1e18):
                            dist[nc] = ng; prev[nc] = c
                            heapq.heappush(heap, (ng + h(j, i), ng, nc))
        if found is None:
            return None, f"no path ({n} expanded)"
        path = [found]
        while path[-1] in prev:
            path.append(prev[path[-1]])
        path.reverse()
        paths.append(path)
        tree |= set(path)
        for r in remaining:
            if found in r:
                tree |= r
                remaining.remove(r)
                break
    return paths, "ok"


def path_to_copper(path, net, w):
    """Cells -> merged straight tracks and vias."""
    tracks, vias = [], []
    seg = [path[0]]

    def flush(seg):
        if len(seg) < 2:
            return
        pts = [(c[2] * G, c[1] * G) for c in seg]
        # merge collinear
        out = [pts[0]]
        for k in range(1, len(pts) - 1):
            a, b, c = out[-1], pts[k], pts[k + 1]
            if abs((b[0] - a[0]) * (c[1] - b[1]) - (b[1] - a[1]) * (c[0] - b[0])) < 1e-9:
                continue
            out.append(b)
        out.append(pts[-1])
        for a, b in zip(out, out[1:]):
            tracks.append((LAYERS[seg[0][0]], net, round(a[0], 4), round(a[1], 4), round(b[0], 4), round(b[1], 4), w))

    for c in path[1:]:
        if c[0] != seg[-1][0]:
            flush(seg)
            vias.append((round(c[2] * G, 4), round(c[1] * G, 4), net, VIA_D, VIA_H))
            seg = [c]
        else:
            seg.append(c)
    flush(seg)
    return tracks, vias