round_tracks_action.py

# Original example Copyright [2017] [Miles McCoo]
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# extensively modified by Julian Loiacono, Oct 2018
# updated and repacked as action plugin by Antoine Pintout, May 2019
# updated subdivision algorithm + native arc support by mitxela, 2021-2024

import pcbnew
import math
import os
import wx
import time
from .round_tracks_utils import *
from .round_tracks_gui import RoundTracksDialog

RADIUS_DEFAULT = 2.0
PASSES_DEFAULT = 3

class RoundTracks(RoundTracksDialog):

    def __init__(self, board, action):
        super(RoundTracks, self).__init__(None)
        self.board = board
        self.basefilename = os.path.splitext(board.GetFileName())[0]
        if self.basefilename.endswith('-rounded'):
            self.basefilename = self.basefilename[:-len('-rounded')]
        self.configfilepath = self.basefilename+".round-tracks-config"
        self.action = action
        self.config = {}
        self.netClassCount = 1
        self.load_config()
        if 'checkboxes' not in self.config:
            self.config['checkboxes'] = {'new_file':False, 'native':True, 'avoid_junctions':False}
        self.do_create.SetValue( self.config['checkboxes']['new_file'])
        self.use_native.SetValue( self.config['checkboxes']['native'])
        self.avoid_junctions.SetValue( self.config['checkboxes']['avoid_junctions'])

        c = self.config['classes']
        if "Default" not in c:
            self.netclasslist.AppendItem( ["Default", True, str(RADIUS_DEFAULT), str(PASSES_DEFAULT)])
        else:
            self.netclasslist.AppendItem( ["Default", c['Default']['do_round'],  str(c['Default']['scaling']),  str(c['Default']['passes'])])
        for class_id in self.board.GetNetClasses():
            classname = str(class_id)
            self.netClassCount += 1
            if classname not in c:
                self.netclasslist.AppendItem( [classname, True, str(RADIUS_DEFAULT), str(PASSES_DEFAULT)])
            else:
                self.netclasslist.AppendItem( [classname, c[classname]['do_round'],  str(c[classname]['scaling']),  str(c[classname]['passes'])])
        self.validate_all_data()


    def run( self, event ):
        start = time.time()
        self.apply.SetLabel("Working...")
        self.validate_all_data()
        self.save_config()

        self.prog = wx.ProgressDialog("Processing", "Starting...", 100, self, wx.PD_AUTO_HIDE|wx.PD_APP_MODAL|wx.PD_ELAPSED_TIME)

        if self.do_create.IsChecked():
            new_name = self.basefilename+"-rounded.kicad_pcb"
            self.board.SetFileName(new_name)

        anySelected = False
        for t in self.board.GetTracks():
            if t.IsSelected():
                anySelected = True
                break

        avoid = self.avoid_junctions.IsChecked()
        classes = self.config['classes']
        for classname in classes:
            if classes[classname]['do_round']:
                if self.use_native.IsChecked():
                    self.addIntermediateTracks(scaling = classes[classname]['scaling'], netclass = classname, native = True, onlySelection = anySelected, avoid_junctions = avoid)
                else:
                    for i in range(classes[classname]['passes']):
                        self.addIntermediateTracks(scaling = classes[classname]['scaling'], netclass = classname, native = False, onlySelection = anySelected, avoid_junctions = avoid, msg=f", pass {i+1}")

        # if m_AutoRefillZones is set, we should skip here, but PCBNEW_SETTINGS is not exposed to swig
        # ZONE_FILLER has SetProgressReporter, but PROGRESS_REPORTER is also not available, so we can't use it
        # even zone.SetNeedRefill(False) doesn't prevent it running twice
        #self.prog.Pulse("Rebuilding zones...")
        #wx.Yield()
        #filler = pcbnew.ZONE_FILLER(self.board)
        #filler.Fill(self.board.Zones())

        if bool(self.prog):
            self.prog.Destroy()
            wx.Yield()
        dt = time.time()-start
        if dt>0.1:
            wx.MessageBox("Done, took {:.3f} seconds".format(time.time()-start), parent=self)
        self.EndModal(wx.ID_OK)

    def unrun( self, event ):

        anySelected = False
        for t in self.board.GetTracks():
            if t.IsSelected():
                anySelected = True
                break

        board = self.board
        netcodes = board.GetNetsByNetcode()
        allTracks = board.GetTracks()

        tracksStartsToExtend = []
        tracksEndsToExtend = []
        tracksToAdd = []
        tracksToRemove = []
        fail=False

        def findExtendTrack( pos, angle, intersection, tracksInNet, arc ):
            angle = normalizeAngle(angle)
            layer = arc.GetLayer()
            for t in tracksInNet:
                if t.GetLayer() != layer:
                    continue
                if similarPoints(t.GetStart(), pos):
                    if similarAngle(angle, getTrackAngle(t)):
                        #t.SetStart(intersection)
                        tracksStartsToExtend.append( (t, intersection) )
                        return True
                if similarPoints(t.GetEnd(), pos):
                    if similarAngle(angle, getTrackAngle(t)+math.pi):
                        #t.SetEnd(intersection)
                        tracksEndsToExtend.append( (t, intersection) )
                        return True
            # before creating new track, look for new co-linear tracks
            width = arc.GetWidth()
            net = arc.GetNetCode()
            for i,t in enumerate(tracksToAdd):
                (sp,ep,w,l,n) = t
                if (n==net and l==layer and w==width and similarPoints(sp, pos) and
                    similarAngle(math.atan2(ep.y-sp.y,ep.x-sp.x), math.atan2(pos.y-intersection.y,pos.x-intersection.x))):
                    tracksToAdd[i]= (intersection, ep ,w,l,n)
                    return False

            tracksToAdd.append((pos, intersection, width, layer, net))
            return False

        for netcode, net in netcodes.items():
            tracksInNet = []
            arcsInNet = []
            intersectionStack = []

            for t in allTracks:
                if t.GetNetCode() == netcode:
                    if t.GetClass() == 'PCB_ARC' and (not anySelected or t.IsSelected()):
                        arcsInNet.append(t)
                    # we may need to extend tracks outside of selection
                    elif t.GetClass() == 'PCB_TRACK':
                        tracksInNet.append(t)

            # check arc angle < pi ? oversize arcs shouldn't exist

            for a in arcsInNet:
                start  = a.GetStart()
                end    = a.GetEnd()
                center = a.GetCenter()
                startAngle = a.GetArcAngleStart().AsRadians()
                endAngle = a.GetArcAngleEnd().AsRadians()
                halfAngle = a.GetAngle().AsRadians()/2
                interdistance = a.GetRadius() / math.cos(halfAngle)
                if interdistance > 2**32:
                    fail=True
                    tracksToRemove.append(a)
                    continue
                if a.IsCCW():
                    interAngle = startAngle +halfAngle
                else:
                    interAngle = endAngle -halfAngle

                intersection = pcbnew.VECTOR2I(
                    center.x + int(interdistance*math.cos(interAngle)),
                    center.y + int(interdistance*math.sin(interAngle))
                    )
                adjust = math.pi/2 if a.IsCCW() else -math.pi/2

                intersectionStack.append( (intersection, a.GetLayer()) )

                findExtendTrack( start, startAngle-adjust, intersection, tracksInNet, a )
                findExtendTrack( end, endAngle+adjust, intersection, tracksInNet, a )
                tracksToRemove.append(a)

            # find junctions
            for i,iL in intersectionStack:
                junction = []
                for j,jL in intersectionStack:
                    if i==j or iL!=jL: continue
                    if ( ((i.x > j.x - unpickJunctionTolerance)
                     and  (i.x < j.x + unpickJunctionTolerance))
                     and ((i.y > j.y - unpickJunctionTolerance)
                     and  (i.y < j.y + unpickJunctionTolerance))):
                        junction.append(j)
                if len(junction)>0:
                    junction.append(i)
                    x = 0
                    y = 0
                    for j in junction:
                        x += j.x
                        y += j.y
                    newCenter = pcbnew.VECTOR2I(int(x/len(junction)), int(y/len(junction)))
                    extremities = []
                    # correct positions of new track ends
                    for idx,(t,s) in enumerate(tracksStartsToExtend):
                        if s in junction:
                            tracksStartsToExtend[idx] = t,newCenter
                            extremities.append(t.GetStart())
                    for idx,(t,s) in enumerate(tracksEndsToExtend):
                        if s in junction:
                            tracksEndsToExtend[idx] = t,newCenter
                            extremities.append(t.GetEnd())
                    for idx, (sp, ep, width, layer, n) in enumerate(tracksToAdd):
                        if layer==iL and net==n and sp in junction:
                            tracksToAdd[idx] = newCenter, ep, width, layer, n
                            extremities.append(ep)

                    # remove unmodified, straight tracks that pass through the junction
                    for t in tracksInNet:
                        if t.GetLayer()==iL and t.GetStart() in extremities and t.GetEnd() in extremities:
                            tracksToRemove.append(t)

            # junctions may stack the same track start/end twice
            for t, s in tracksStartsToExtend:
                t.SetStart(s)
            tracksStartsToExtend = []

            for t, s in tracksEndsToExtend:
                t.SetEnd(s)
            tracksEndsToExtend = []

        for sp, ep, width, layer, net in tracksToAdd:
            track = pcbnew.PCB_TRACK(board)
            track.SetStart(sp)
            track.SetEnd(ep)
            track.SetWidth(width)
            track.SetLayer(layer)
            board.Add(track)
            track.SetNetCode(net)
            if anySelected:
                track.SetSelected()

        for t in tracksToRemove:
            board.Remove(t)

        if fail:
            wx.MessageBox("Some arcs failed to unround.\n\nThis is usually caused by very, very small arcs which have nonsensical angles.", style=wx.ICON_ERROR)

        wx.MessageBox("Unrounded!\n\nPlease note, due to rounding errors, a round trip may not exactly match the original.", parent=self)

        self.EndModal(wx.ID_OK)


    def on_close( self, event ):
        self.EndModal(wx.ID_OK)

    def on_item_editing( self, event ):
        if bool(self.netclasslist):
            self.validate_all_data()

    def load_config(self):
        new_config = {}
        if os.path.isfile(self.configfilepath):
            with open(self.configfilepath, "r") as configfile:
                for line in configfile.readlines():
                    params = line[:-1].split("\t")
                    new_config_line = {}
                    try:
                        new_config_line['do_round'] = params[1] == "True"
                        new_config_line['scaling'] = float(params[2])
                        new_config_line['passes'] = int(params[3])
                        new_config[params[0]] = new_config_line
                    except Exception as e:
                        try:
                            new_config_line['new_file'] = params[0] == "True"
                            new_config_line['native'] = params[1] == "True"
                            new_config_line['avoid_junctions'] = params[2] == "True"
                            self.config['checkboxes'] = new_config_line
                        except Exception as e:
                            pass
        self.config['classes'] = new_config

    def save_config(self):
        classes = self.config['classes']
        try:
            with open(self.configfilepath, "w") as configfile:
                for classname in classes:
                    configfile.write('%s\t%s\t%s\t%s\n' % (classname, str(classes[classname]['do_round']), str(classes[classname]['scaling']), str(classes[classname]['passes'])))
                configfile.write('%s\t%s\t%s\n' % (str(self.config['checkboxes']['new_file']), str(self.config['checkboxes']['native']), str(self.config['checkboxes']['avoid_junctions'])))
        except PermissionError:
            pass

    def validate_all_data (self):
        new_config = {}
        for i in range(self.netClassCount):
            for j in range(5):
                if j == 2:
                    # param should be between 0 and 1
                    try:
                        tested_val = float(self.netclasslist.GetTextValue(i, j))
                        if tested_val < 0:
                            self.netclasslist.SetTextValue(str(RADIUS_DEFAULT), i, j)
                    except Exception as e:
                        self.netclasslist.SetTextValue(str(RADIUS_DEFAULT), i, j)
                if j == 3:
                    # param should be between int 1 and 5
                    try:
                        tested_val = int(self.netclasslist.GetTextValue(i, j))
                        if tested_val < 0 or tested_val > 5:
                            self.netclasslist.SetTextValue(str(PASSES_DEFAULT), i, j)
                    except Exception as e:
                        self.netclasslist.SetTextValue(str(PASSES_DEFAULT), i, j)
            new_config[self.netclasslist.GetTextValue(i, 0)] = {
                'do_round' : self.netclasslist.GetToggleValue(i, 1),
                'scaling' : float(self.netclasslist.GetTextValue(i, 2)),
                'passes' : int(self.netclasslist.GetTextValue(i, 3))
            }
        self.config['classes'] = new_config
        self.config['checkboxes'] = {'new_file':self.do_create.IsChecked(), 'native':self.use_native.IsChecked(), 'avoid_junctions':self.avoid_junctions.IsChecked()}

    def addIntermediateTracks( self, scaling = RADIUS_DEFAULT, netclass = None, native = False, onlySelection = False, avoid_junctions = False, msg=""):

        # A 90 degree bend will get a maximum radius of this amount
        RADIUS = pcbnew.FromMM(scaling /(math.sin( math.pi/4 )+1))

        board = self.board
        netcodes = board.GetNetsByNetcode()
        allTracks = board.GetTracks()
        allPads = board.GetPads()

        tracksToRemove = []

        for netcode, net in netcodes.items():

            if netclass is not None and netclass == net.GetNetClassName():

                # BOARD::TracksInNet casts everything to a PCB_TRACK, even PCB_VIA and PCB_ARC
                # tracksInNet = board.TracksInNet(net.GetNetCode())
                tracksInNet = []
                viasInNet = []
                for t in allTracks:
                    if t.GetNetCode() == netcode and (not onlySelection or t.IsSelected()):
                        if t.GetClass() == 'PCB_VIA':
                            viasInNet.append(t)
                        else:
                            tracksInNet.append(t)

                tracksPerLayer = {}
                # separate track by layer
                for t in tracksInNet:
                    layer = t.GetLayer()
                    if layer not in tracksPerLayer:
                        tracksPerLayer[layer] = []
                    tracksPerLayer[layer].append(t)

                for v in viasInNet:
                    # a buried/blind via will report only layers affected
                    # a through via will return all 32 possible layers
                    layerSet = v.GetLayerSet().CuStack()
                    for layer in tracksPerLayer:
                        if layer in layerSet:
                            tracksPerLayer[layer].append(v)

                # TH pads cover all layers
                # SMD/CONN pads only touch F.Cu and B.Cu (layers 0 and 31)
                # Due to glitch in KiCad, pad.GetLayer() always returns 0. Need to use GetLayerSet().Contains() to actually check

                padsInNet = []
                FCuPadsInNet = []
                BCuPadsInNet = []

                for p in allPads:
                    if p.GetNetCode() == netcode and (not onlySelection or t.IsSelected()):
                        attr = p.GetAttribute()
                        if attr == pcbnew.PAD_ATTRIB_NPTH or attr == pcbnew.PAD_ATTRIB_PTH:
                            padsInNet.append(p)
                        else:
                            if p.GetLayerSet().Contains(31):
                                BCuPadsInNet.append(p)
                            else:
                                FCuPadsInNet.append(p)

                for layer in tracksPerLayer:
                    tracks = tracksPerLayer[layer]

                    # add all the possible intersections to a unique set, for iterating over later
                    intersections = set();  
                    for t1 in range(len(tracks)):
                        for t2 in range(t1+1, len(tracks)):
                            # check if these two tracks share an endpoint
                            # reduce it to a 2-part tuple so there are not multiple objects of the same point in the set
                            if(tracks[t1].IsPointOnEnds(tracks[t2].GetStart())): 
                                intersections.add((tracks[t2].GetStart().x, tracks[t2].GetStart().y))
                            if(tracks[t1].IsPointOnEnds(tracks[t2].GetEnd())):
                                intersections.add((tracks[t2].GetEnd().x, tracks[t2].GetEnd().y))

                    # for each remaining intersection, shorten each track by the same amount, and place a track between.
                    tracksToAdd = []
                    arcsToAdd = []
                    trackLengths = {}
                    for ip in intersections:
                        (newX, newY) = ip;
                        intersection = pcbnew.VECTOR2I(newX, newY)
                        tracksHere = [];
                        for t1 in tracks:
                            if similarPoints(t1.GetStart(), intersection):
                                tracksHere.append(t1)
                            elif similarPoints(t1.GetEnd(), intersection):
                                # flip track such that all tracks start at the IP
                                reverseTrack(t1)
                                tracksHere.append(t1)

                        if len(tracksHere) == 0 or (avoid_junctions and len(tracksHere)>2):
                            continue

                        # if there are any arcs or vias present, skip the intersection entirely
                        skip = False
                        for t1 in tracksHere:
                            if t1.GetClass() != 'PCB_TRACK':
                                skip = True
                                break

                        # If the intersection is within a pad, but none of the tracks end within the pad, skip
                        for p in padsInNet:
                            if withinPad(p, intersection, tracksHere):
                                skip = True
                                break

                        if layer == 0:
                            for p in FCuPadsInNet:
                                if withinPad(p, intersection, tracksHere):
                                    skip = True
                                    break
                        elif layer == 31:
                            for p in BCuPadsInNet:
                                if withinPad(p, intersection, tracksHere):
                                    skip = True
                                    break

                        if skip:
                            continue

                        shortest=1e9
                        for t1 in tracksHere:
                            if id(t1) not in trackLengths:
                                trackLengths[id(t1)] = t1.GetLength()
                            if trackLengths[id(t1)]<shortest:
                                shortest = trackLengths[id(t1)]

                        #sort these tracks by angle, so new tracks can be drawn between them
                        tracksHere.sort(key = getTrackAngle)

                        if native:
                            halfTrackAngle = {} # cache this, because after shortening the length may end up zero
                            for t1 in range(len(tracksHere)):
                                halfTrackAngle[t1] = getTrackAngleDifference( tracksHere[t1], tracksHere[(t1+1)%len(tracksHere)] )/2

                            shortenAmount = shortest * 0.5
                            for t1 in range(len(tracksHere)):
                                f = math.sin( halfTrackAngle[t1] )+1
                                shortenAmount = min( RADIUS *f, shortenAmount )
                            for t1 in range(len(tracksHere)):
                                if shortenTrack(tracksHere[t1], shortenAmount):
                                    tracksToRemove.append(tracksHere[t1])

                            for t1 in range(len(tracksHere)):
                                if not (len(tracksHere) == 2 and t1 == 1):
                                    theta = math.pi/2 - halfTrackAngle[t1]
                                    f = 1/(2*math.cos(theta) +2)

                                    sp = cloneVECTOR2I(tracksHere[t1].GetStart())
                                    ep = cloneVECTOR2I(tracksHere[(t1+1)%len(tracksHere)].GetStart())
                                    if halfTrackAngle[t1]> math.pi/2 -0.001:
                                        tracksToAdd.append((sp, ep, tracksHere[t1].GetWidth(), tracksHere[t1].GetLayer(), tracksHere[t1].GetNetCode()))
                                    else:
                                        mp = pcbnew.VECTOR2I(int(newX*(1-f*2)+sp.x*f+ep.x*f), int(newY*(1-f*2)+sp.y*f+ep.y*f))
                                        arcsToAdd.append((sp, ep, mp, tracksHere[t1].GetWidth(), tracksHere[t1].GetLayer(), tracksHere[t1].GetNetCode()))

                        else:
                            shortenAmount = RADIUS
                            for t1 in range(len(tracksHere)):
                                theta = math.pi/2 - getTrackAngleDifference( tracksHere[t1], tracksHere[(t1+1)%len(tracksHere)] )/2
                                f = 1/(2*math.cos(theta) +2)
                                shortenAmount = min(shortenAmount, shortest * f)
                            for t1 in range(len(tracksHere)):
                                shortenTrack(tracksHere[t1], shortenAmount)

                            #connect the new startpoints in a circle around the old center point
                            for t1 in range(len(tracksHere)):
                                #dont add 2 new tracks in the 2 track case
                                if not (len(tracksHere) == 2 and t1 == 1):
                                    newPoint1 = cloneVECTOR2I(tracksHere[t1].GetStart())
                                    newPoint2 = cloneVECTOR2I(tracksHere[(t1+1)%len(tracksHere)].GetStart())
                                    tracksToAdd.append((newPoint1, newPoint2, tracksHere[t1].GetWidth(), tracksHere[t1].GetLayer(), tracksHere[t1].GetNetCode()))

                    #add all the new tracks in post, so as not to cause problems with set iteration
                    for trackpoints in tracksToAdd:
                        (sp, ep, width, layer, net) = trackpoints

                        track = pcbnew.PCB_TRACK(board)
                        track.SetStart(sp)
                        track.SetEnd(ep)
                        track.SetWidth(width)
                        track.SetLayer(layer)
                        board.Add(track)
                        track.SetNetCode(net)
                        if onlySelection:
                            track.SetSelected()

                    for trackpoints in arcsToAdd:
                        (sp, ep, mp, width, layer, net) = trackpoints

                        arc = pcbnew.PCB_ARC(board)
                        arc.SetStart(sp)
                        arc.SetMid(mp)
                        arc.SetEnd(ep)
                        arc.SetWidth(width)
                        arc.SetLayer(layer)
                        board.Add(arc)
                        arc.SetNetCode(net)
                        if onlySelection:
                            arc.SetSelected()

            self.prog.Pulse(f"Netclass: {netclass}, {netcode+1} of {len(netcodes)}{msg}")

        for t in tracksToRemove:
            board.Remove(t)

class ActionRoundTracks( pcbnew.ActionPlugin ):
 
    def defaults( self ):
        self.name = "Round tracks"
        self.category = "Modify PCB"
        self.description = "Subdivision-based track rounding"
        self.show_toolbar_button = True
        self.icon_file_name = os.path.join(os.path.dirname(__file__), "./round_tracks.png")

    def Run( self ):
        if pcbnew.GetBuildVersion()=='(7.0.0)':
            wx.MessageBox("Sorry, this plugin is not compatible with KiCad 7.0.0\n\nIt should work with 7.0.1, if that's not out yet you can try KiCad nightly.")
            return
        board = pcbnew.GetBoard()
        rt = RoundTracks(board, self)
        rt.ShowModal()
        rt.Destroy()

        # this shouldn't be needed, but without it, arcs are sometimes not acknowledged as tangential until they've been moved or the file re-opened
        # as of kicad 6.0.0, doing pcbnew.Refresh() here causes the arcs to glitch out and create duplicate shapes that can't be deleted
        pcbnew.UpdateUserInterface()