pcb/scripts/stitch_plane_pads.py

#!/usr/bin/env python3
"""Connect SMT power/ground pads locally to their intended solid planes.

This deliberately does not replace power-stage routing.  It only installs
short pad-to-via connections where a filled, same-net plane is physically
present below the candidate via.  High-current converter loops and plane
bridges are handled separately.

It runs *after* route_power_stage.py: a candidate via is only checked against
copper that already exists, so any stage that lays down wide trunks has to go
first or its copper is invisible here.
"""

from __future__ import annotations

import argparse
import math
import shutil
from pathlib import Path

import pcbnew


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INPUT = ROOT / "hardware" / "board_power_routed.kicad_pcb"
DEFAULT_OUTPUT = ROOT / "hardware" / "board_plane_stitched.kicad_pcb"
PLANE_NETS = {
    "GND", "SYS5", "VCC0V95", "VCC1V8", "VCC3V3", "VCC3V3_AUX",
    "VCC3V0_CLK", "VCC3V0_DIG", "VCC3V0_RF", "VCC3V3_ANT",
}
VIA_DIAMETER = 0.45
VIA_DRILL = 0.225
CLEARANCE = 0.10


def mm(value: float) -> int:
    return pcbnew.FromMM(value)


def distance_mm(a: pcbnew.VECTOR2I, b: pcbnew.VECTOR2I) -> float:
    return math.hypot(pcbnew.ToMM(a.x - b.x), pcbnew.ToMM(a.y - b.y))


def add_track(board: pcbnew.BOARD, pad: pcbnew.PAD,
              end: pcbnew.VECTOR2I, width: float) -> pcbnew.PCB_TRACK:
    track = pcbnew.PCB_TRACK(board)
    track.SetStart(pad.GetPosition())
    track.SetEnd(end)
    track.SetWidth(mm(width))
    track.SetLayer(pad.GetLayer())
    track.SetNet(pad.GetNet())
    track.SetLocked(True)
    board.Add(track)
    return track


def already_connected(pad: pcbnew.PAD, tracks: list[pcbnew.PCB_TRACK],
                      vias: list[pcbnew.PCB_VIA]) -> bool:
    """Report whether an earlier stage already attached copper to this pad."""
    net = pad.GetNetname()
    box = pad.GetBoundingBox()

    def inside(position: pcbnew.VECTOR2I) -> bool:
        return (box.GetX() <= position.x <= box.GetRight() and
                box.GetY() <= position.y <= box.GetBottom())

    for track in tracks:
        if track.GetNetname() == net and (inside(track.GetStart()) or inside(track.GetEnd())):
            return True
    return any(via.GetNetname() == net and inside(via.GetPosition()) for via in vias)


def add_via(board: pcbnew.BOARD, pad: pcbnew.PAD,
            position: pcbnew.VECTOR2I) -> pcbnew.PCB_VIA:
    via = pcbnew.PCB_VIA(board)
    via.SetPosition(position)
    via.SetWidth(mm(VIA_DIAMETER))
    via.SetDrill(mm(VIA_DRILL))
    via.SetLayerPair(pcbnew.F_Cu, pcbnew.B_Cu)
    via.SetNet(pad.GetNet())
    via.SetLocked(True)
    board.Add(via)
    return via


def plane_at(zones_by_net: dict[str, list[pcbnew.ZONE]], net: str,
             position: pcbnew.VECTOR2I) -> bool:
    return any(zone.HitTestFilledArea(zone.GetLayer(), position)
               for zone in zones_by_net.get(net, ()))


def inside_temporary_keepout(position: pcbnew.VECTOR2I,
                             keepouts: list[pcbnew.ZONE]) -> bool:
    """Reject a through-via whose copper can touch an RF/USB rule area."""
    radius = VIA_DIAMETER / 2 + CLEARANCE
    probes = [(0.0, 0.0)]
    for index in range(8):
        angle = index * math.pi / 4
        probes.append((radius * math.cos(angle), radius * math.sin(angle)))
    for dx, dy in probes:
        probe = pcbnew.VECTOR2I(position.x + mm(dx), position.y + mm(dy))
        if any(zone.Outline().Contains(probe) for zone in keepouts):
            return True
    return False


def segment_distance_mm(point: pcbnew.VECTOR2I, start: pcbnew.VECTOR2I,
                        end: pcbnew.VECTOR2I) -> float:
    px, py = pcbnew.ToMM(point.x), pcbnew.ToMM(point.y)
    ax, ay = pcbnew.ToMM(start.x), pcbnew.ToMM(start.y)
    bx, by = pcbnew.ToMM(end.x), pcbnew.ToMM(end.y)
    dx, dy = bx - ax, by - ay
    length = dx * dx + dy * dy
    if length == 0.0:
        return math.hypot(px - ax, py - ay)
    t = max(0.0, min(1.0, ((px - ax) * dx + (py - ay) * dy) / length))
    return math.hypot(px - (ax + t * dx), py - (ay + t * dy))


def clear_of_tracks(position: pcbnew.VECTOR2I, net: str,
                    tracks: list[pcbnew.PCB_TRACK], extra: float = 0.0) -> bool:
    """Keep a via barrel clear of copper installed by the routing stages."""
    for track in tracks:
        if track.GetNetname() == net or isinstance(track, pcbnew.PCB_VIA):
            continue
        limit = (VIA_DIAMETER / 2 + extra + pcbnew.ToMM(track.GetWidth()) / 2
                 + CLEARANCE + 0.03)
        if segment_distance_mm(position, track.GetStart(), track.GetEnd()) < limit:
            return False
    return True


def clear_track_of_tracks(start: pcbnew.VECTOR2I, end: pcbnew.VECTOR2I, net: str,
                          tracks: list[pcbnew.PCB_TRACK], layer: int,
                          width: float) -> bool:
    for track in tracks:
        if track.GetNetname() == net or isinstance(track, pcbnew.PCB_VIA):
            continue
        if track.GetLayer() != layer:
            continue
        limit = width / 2 + pcbnew.ToMM(track.GetWidth()) / 2 + CLEARANCE + 0.03
        for step in range(0, 13):
            sample = pcbnew.VECTOR2I(
                round(start.x + (end.x - start.x) * step / 12),
                round(start.y + (end.y - start.y) * step / 12),
            )
            if segment_distance_mm(sample, track.GetStart(), track.GetEnd()) < limit:
                return False
    return True


def clear_of_other_copper(position: pcbnew.VECTOR2I, net: str,
                          pads: list[pcbnew.PAD], vias: list[pcbnew.PCB_VIA]) -> bool:
    x, y = pcbnew.ToMM(position.x), pcbnew.ToMM(position.y)
    margin = VIA_DIAMETER / 2 + CLEARANCE + 0.03
    for other in pads:
        if other.GetNetname() == net:
            continue
        box = other.GetBoundingBox()
        x0, y0 = pcbnew.ToMM(box.GetX()) - margin, pcbnew.ToMM(box.GetY()) - margin
        x1, y1 = pcbnew.ToMM(box.GetRight()) + margin, pcbnew.ToMM(box.GetBottom()) + margin
        if x0 <= x <= x1 and y0 <= y <= y1:
            return False
    for other in vias:
        minimum = (VIA_DIAMETER + pcbnew.ToMM(other.GetWidth(pcbnew.F_Cu))) / 2 + CLEARANCE
        if distance_mm(position, other.GetPosition()) < minimum:
            return False
    return True


def clear_track_path(start: pcbnew.VECTOR2I, end: pcbnew.VECTOR2I, net: str,
                     source: pcbnew.PAD, pads: list[pcbnew.PAD], width: float) -> bool:
    """Conservatively reject short dogbones that cross another pad bbox."""
    margin = width / 2 + CLEARANCE + 0.02
    for step in range(1, 12):
        fraction = step / 12
        sample = pcbnew.VECTOR2I(
            round(start.x + (end.x - start.x) * fraction),
            round(start.y + (end.y - start.y) * fraction),
        )
        x, y = pcbnew.ToMM(sample.x), pcbnew.ToMM(sample.y)
        for other in pads:
            if other is source or other.GetNetname() == net:
                continue
            box = other.GetBoundingBox()
            if (pcbnew.ToMM(box.GetX()) - margin <= x <= pcbnew.ToMM(box.GetRight()) + margin and
                    pcbnew.ToMM(box.GetY()) - margin <= y <= pcbnew.ToMM(box.GetBottom()) + margin):
                return False
    return True


def candidates(pad: pcbnew.PAD, footprint: pcbnew.FOOTPRINT) -> list[pcbnew.VECTOR2I]:
    pad_pos = pad.GetPosition()
    fp_pos = footprint.GetPosition()
    away = math.atan2(pad_pos.y - fp_pos.y, pad_pos.x - fp_pos.x)
    if pad_pos == fp_pos:
        away = 0.0
    radius = max(pcbnew.ToMM(pad.GetSizeX()), pcbnew.ToMM(pad.GetSizeY())) / 2
    radius += VIA_DIAMETER / 2 + CLEARANCE + 0.08
    # Freerouting treats fixed non-45-degree traces as invalid geometry. Use
    # only global orthogonal/45-degree directions, ordered by how closely they
    # point away from the footprint centre.
    directions = [(math.cos(math.radians(angle)), math.sin(math.radians(angle)))
                  for angle in range(0, 360, 45)]
    directions.sort(key=lambda direction: -(
        direction[0] * math.cos(away) + direction[1] * math.sin(away)
    ))
    result = []
    for scale in (1.0, 1.35, 1.75):
        for dx, dy in directions:
            result.append(pcbnew.VECTOR2I(
                pad_pos.x + mm(radius * scale * dx),
                pad_pos.y + mm(radius * scale * dy),
            ))
    return result


def stitch(board: pcbnew.BOARD) -> tuple[dict[str, int], dict[str, int]]:
    """Connect every reachable SMD power pad to its plane on the given board.

    Callable twice: once before signal routing, so the planes already have
    real connection points, and again from finalize_routing.py once the RF and
    USB corridors are no longer reserved by temporary rule areas.  Pads that
    are already stitched are skipped because the candidate location is then
    occupied by the earlier via.
    """
    zones_by_net: dict[str, list[pcbnew.ZONE]] = {}
    for zone in board.Zones():
        if not zone.GetIsRuleArea() and zone.GetNetname() in PLANE_NETS:
            zones_by_net.setdefault(zone.GetNetname(), []).append(zone)
    temporary_keepouts = [zone for zone in board.Zones()
                          if zone.GetIsRuleArea() and zone.GetZoneName().startswith("TEMP_")]

    all_pads = [pad for fp in board.GetFootprints() for pad in fp.Pads()]
    all_vias = [item for item in board.GetTracks() if isinstance(item, pcbnew.PCB_VIA)]
    all_tracks = [item for item in board.GetTracks() if not isinstance(item, pcbnew.PCB_VIA)]
    connected: dict[str, int] = {}
    skipped: dict[str, int] = {}

    for fp in board.GetFootprints():
        # U1 already has deterministic dogbones.  J2 is an edge-launch part;
        # its ground lands are completed by finalize_routing.py's RF ground
        # zone/fence, not by ordinary tracks that would run along Edge.Cuts.
        # U13 ground/thermal vias are part of the deterministic high-current
        # converter geometry in route_power_stage.py. Generic radial dogbones
        # here can intrude into its SW and VIN paths.
        if fp.GetReference() in {"U1", "U13", "J2"}:
            continue
        for pad in fp.Pads():
            net = pad.GetNetname()
            if net not in PLANE_NETS or pad.GetAttribute() != pcbnew.PAD_ATTRIB_SMD:
                continue
            if already_connected(pad, all_tracks, all_vias):
                continue
            chosen = None
            pad_min = min(pcbnew.ToMM(pad.GetSizeX()), pcbnew.ToMM(pad.GetSizeY()))
            width = max(0.15, min(0.30, pad_min * 0.65))
            for candidate in candidates(pad, fp):
                x, y = pcbnew.ToMM(candidate.x), pcbnew.ToMM(candidate.y)
                if not (20.5 < x < 174.5 and 20.5 < y < 109.5):
                    continue
                if not plane_at(zones_by_net, net, candidate):
                    continue
                if inside_temporary_keepout(candidate, temporary_keepouts):
                    continue
                if not clear_of_other_copper(candidate, net, all_pads, all_vias):
                    continue
                if not clear_track_path(pad.GetPosition(), candidate, net, pad, all_pads, width):
                    continue
                if not clear_of_tracks(candidate, net, all_tracks):
                    continue
                if not clear_track_of_tracks(pad.GetPosition(), candidate, net,
                                             all_tracks, pad.GetLayer(), width):
                    continue
                chosen = candidate
                break
            if chosen is None:
                skipped[net] = skipped.get(net, 0) + 1
                continue
            all_tracks.append(add_track(board, pad, chosen, width))
            all_vias.append(add_via(board, pad, chosen))
            connected[net] = connected.get(net, 0) + 1

    return connected, skipped


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", nargs="?", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("output", nargs="?", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()

    board = pcbnew.LoadBoard(str(args.input))
    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Could not fill planes before stitching")
    connected, skipped = stitch(board)

    if not pcbnew.ZONE_FILLER(board).Fill(board.Zones()):
        raise RuntimeError("Could not refill planes after stitching")
    pcbnew.SaveBoard(str(args.output), board)
    for suffix in (".kicad_pro", ".kicad_dru"):
        source = args.input.with_suffix(suffix)
        if source.exists():
            shutil.copyfile(source, args.output.with_suffix(suffix))
    print("Plane connections:", ", ".join(f"{n}={k}" for n, k in sorted(connected.items())))
    print("Pads without a local same-net plane:",
          ", ".join(f"{n}={k}" for n, k in sorted(skipped.items())))
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()