Multilayer coil WIP

6 files changed, 166 insertions, 25 deletions

diff --git a/gerbonara/cad/kicad/base_types.py b/gerbonara/cad/kicad/base_types.py
index 001138c..11eeb6d 100644
--- a/gerbonara/cad/kicad/base_types.py
+++ b/gerbonara/cad/kicad/base_types.py
@@ -208,6 +208,9 @@ class XYCoord:
         else:
             self.x, self.y = x, y
 
+    def within_distance(self, x, y, dist):
+        return math.dist((x, y), (self.x, self.y)) < dist
+
     def isclose(self, other, tol=1e-3):
         return math.isclose(self.x, other.x, tol) and math.isclose(self.y, other.y, tol)
 
diff --git a/gerbonara/cad/kicad/footprints.py b/gerbonara/cad/kicad/footprints.py
index 47d474a..5c27855 100644
--- a/gerbonara/cad/kicad/footprints.py
+++ b/gerbonara/cad/kicad/footprints.py
@@ -409,12 +409,47 @@ class Pad:
         x, y = rotate_point(self.at.x, self.at.y, math.radians(pr))
         return x+px, y+py, self.at.rotation, False
 
+    @property
+    def layer_mask(self):
+        return layer_mask(self.layers)
+
     def offset(self, x=0, y=0):
         self.at = self.at.with_offset(x, y)
 
-    def find_connected(self, **filters):
+    def find_connected_footprints(self, **filters):
         """ Find footprints connected to the same net as this pad """
         return self.footprint.board.find_footprints(net=self.net.name, **filters)
+
+    def find_same_net(self, include_vias=True):
+        """ Find traces and vias of the same net as this pad. """ 
+        return self.footprint.board.find_traces(self.net.name, include_vias=include_vias)
+
+    def find_connected_traces(self, consider_candidates=5):
+        board = self.footprint.board
+
+        found = set()
+        search_frontier = [(self.at, 0, self.layer_mask)]
+        while search_frontier:
+            coord, size, layers = search_frontier.pop()
+            x, y = coord.x, coord.y
+
+            for cand, attr, cand_size in self.footprint.board.query_trace_index((x, x, y, y), layers,
+                                                               n=consider_candidates):
+                if cand in found:
+                    continue
+
+                cand_coord = getattr(cand, attr)
+                cand_x, cand_y = cand_coord.x, cand_coord.y
+                if math.dist((x, y), (cand_x, cand_y)) <= size/2 + cand_size/2:
+                    found.add(cand)
+                    yield cand
+
+                    if hasattr(cand, 'at'): # via or pad
+                        search_frontier.append((cand.at, getattr(cand, 'size', 0), cand.layer_mask))
+                    else:
+                        mask = cand.layer_mask
+                        search_frontier.append((cand.start, cand.width, mask))
+                        search_frontier.append((cand.end, cand.width, mask))
     
     def render(self, variables=None, margin=None, cache=None):
         #if self.type in (Atom.connect, Atom.np_thru_hole):
diff --git a/gerbonara/cad/kicad/graphical_primitives.py b/gerbonara/cad/kicad/graphical_primitives.py
index f1d13e6..568d1d2 100644
--- a/gerbonara/cad/kicad/graphical_primitives.py
+++ b/gerbonara/cad/kicad/graphical_primitives.py
@@ -78,6 +78,10 @@ class Line:
     stroke: Stroke = field(default_factory=Stroke)
     tstamp: Timestamp = None
 
+    def rotate(self, angle, cx=None, cy=None):
+        self.start.x, self.start.y = rotate_point(self.start.x, self.start.y, angle, cx, cy)
+        self.end.x, self.end.y = rotate_point(self.end.x, self.end.y, angle, cx, cy)
+
     def render(self, variables=None):
         if self.angle:
             raise NotImplementedError('Angles on lines are not implemented. Please raise an issue and provide an example file.')
@@ -176,6 +180,19 @@ class Arc:
     width: Named(float) = None
     stroke: Stroke = field(default_factory=Stroke)
     tstamp: Timestamp = None
+    _: SEXP_END = None
+    center: XYCoord = None
+
+    def __post_init__(self):
+        self.start = XYCoord(self.start)
+        self.end = XYCoord(self.end)
+        self.mid = XYCoord(self.mid) if self.mid else center_arc_to_kicad_mid(XYCoord(self.center), self.start, self.end)
+        self.center = None
+
+    def rotate(self, angle, cx=None, cy=None):
+        self.start.x, self.start.y = rotate_point(self.start.x, self.start.y, angle, cx, cy)
+        self.mid.x, self.mid.y = rotate_point(self.mid.x, self.mid.y, angle, cx, cy)
+        self.end.x, self.end.y = rotate_point(self.end.x, self.end.y, angle, cx, cy)
 
     def render(self, variables=None):
         # FIXME stroke support
diff --git a/gerbonara/cad/kicad/pcb.py b/gerbonara/cad/kicad/pcb.py
index 2ad126e..ccda43f 100644
--- a/gerbonara/cad/kicad/pcb.py
+++ b/gerbonara/cad/kicad/pcb.py
@@ -4,7 +4,7 @@ Library for handling KiCad's PCB files (`*.kicad_mod`).
 
 import math
 from pathlib import Path
-from dataclasses import field, KW_ONLY
+from dataclasses import field, KW_ONLY, fields
 from itertools import chain
 import re
 import fnmatch
@@ -14,6 +14,7 @@ from .base_types import *
 from .primitives import *
 from .footprints import Footprint
 from . import graphical_primitives as gr
+import rtree.index
 
 from .. import primitives as cad_pr
 
@@ -164,6 +165,10 @@ class TrackSegment:
         self.start = XYCoord(self.start)
         self.end = XYCoord(self.end)
 
+    @property
+    def layer_mask(self):
+        return layer_mask([self.layer])
+
     def render(self, variables=None, cache=None):
         if not self.width:
             return
@@ -193,29 +198,18 @@ class TrackArc:
     locked: Flag() = False
     net: Named(int) = 0
     tstamp: Timestamp = field(default_factory=Timestamp)
-    _: KW_ONLY
+    _: SEXP_END = None
     center: XYCoord = None
 
     def __post_init__(self):
         self.start = XYCoord(self.start)
         self.end = XYCoord(self.end)
-        if self.center is not None:
-            # Convert normal p1/p2/center notation to the insanity that is kicad's midpoint notation
-            center = XYCoord(self.center)
-            cx, cy = center.x, center.y
-            x1, y1 = self.start.x - cx, self.start.y - cy
-            x2, y2 = self.end.x - cx, self.end.y - cy
-            # Get a vector pointing towards the middle between "start" and "end"
-            dx, dy = (x1 + x2)/2, (y1 + y2)/2
-            # normalize vector, and multiply by radius to get final point
-            r = math.hypot(x1, y1)
-            l = math.hypot(dx, dy)
-            mx = cx + dx / l * r
-            my = cy + dy / l * r
-            self.mid = XYCoord(mx, my)
-            self.center = None
-        else:
-            self.mid = XYCoord(self.mid)
+        self.mid = XYCoord(self.mid) if self.mid else center_arc_to_kicad_mid(XYCoord(self.center), self.start, self.end)
+        self.center = None
+
+    @property
+    def layer_mask(self):
+        return layer_mask([self.layer])
 
     def render(self, variables=None, cache=None):
         if not self.width:
@@ -228,9 +222,6 @@ class TrackArc:
         yield go.Arc(x1, y1, x2, y2, cx-x1, cy-y1, aperture=aperture, clockwise=True, unit=MM)
 
     def rotate(self, angle, cx=None, cy=None):
-        if cx is None or cy is None:
-            cx, cy = self.mid.x, self.mid.y
-
         self.start.x, self.start.y = rotate_point(self.start.x, self.start.y, angle, cx, cy)
         self.mid.x, self.mid.y = rotate_point(self.mid.x, self.mid.y, angle, cx, cy)
         self.end.x, self.end.y = rotate_point(self.end.x, self.end.y, angle, cx, cy)
@@ -259,6 +250,14 @@ class Via:
     def abs_pos(self):
         return self.at.x, self.at.y, 0, False
 
+    @property
+    def layer_mask(self):
+        return layer_mask(self.layers)
+
+    @property
+    def width(self):
+        return self.size
+
     def __post_init__(self):
         self.at = XYCoord(self.at)
 
@@ -314,8 +313,47 @@ class Board:
     _ : SEXP_END = None
     original_filename: str = None
     _bounding_box: tuple = None
+    _trace_index: rtree.index.Index = None
+    _trace_index_map: dict = None
 
 
+    def rebuild_trace_index(self):
+        idx = self._trace_index = rtree.index.Index()
+        id_map = self._trace_index_map = {}
+        for obj in chain(self.track_segments, self.track_arcs):
+            for field in ('start', 'end'):
+                obj_id = id(obj)
+                coord = getattr(obj, field)
+                _trace_index_map[obj_id] = obj, field, obj.width, obj.layer_mask
+                idx.insert(obj_id, (coord.x, coord.x, coord.y, coord.y))
+
+        for fp in self.footprints:
+            for pad in fp.pads:
+                obj_id = id(pad)
+                _trace_index_map[obj_id] = pad, 'at', 0, pad.layer_mask
+                idx.insert(obj_id, (pad.at.x, pad.at.x, pad.at.y, pad.at.y))
+
+        for via in self.vias:
+            obj_id = id(via)
+            _trace_index_map[obj_id] = via, 'at', via.size, via.layer_mask
+            idx.insert(obj_id, (via.at.x, via.at.x, via.at.y, via.at.y))
+    
+    def query_trace_index(self, point, layers='*.Cu', n=5):
+        if self._trace_index is None:
+            self.rebuild_trace_index()
+
+        if isinstance(layers, str):
+            layers = [l.strip() for l in layers.split(',')]
+
+        if not isinstance(layers, int):
+            layers = layer_mask(layers)
+
+        x, y = point
+        for obj_id in self._trace_index.nearest((x, x, y, y), n):
+            entry = obj, attr, size, mask = _trace_index_map[obj_id]
+            if layers & mask:
+                yield entry
+
     def __after_parse__(self, parent):
         self.properties = {prop.key: prop.value for prop in self.properties}
 
@@ -365,6 +403,12 @@ class Board:
             case _:
                 raise TypeError('Can only remove KiCad objects, cannot map generic gerbonara.cad objects for removal')
 
+    def remove_many(self, iterable):
+        iterable = {id(obj) for obj in iterable}
+        for field in fields(self):
+            if field.default_factory is list and field.name not in ('nets', 'properties'):
+                setattr(self, field.name, [obj for obj in getattr(self, field.name) if id(obj) not in iterable])
+
     def add(self, obj):
         match obj:
             case gr.Text():
@@ -481,6 +525,13 @@ class Board:
                 continue
             yield fp
 
+    def find_traces(self, net=None, include_vias=True):
+        net_id = self.net_id(net, create=False)
+        match = lambda obj: obj.net == net_id
+        for obj in chain(self.track_segments, self.track_arcs, self.vias):
+            if obj.net == net_id:
+                yield obj
+
     @property
     def version(self):
         return self._version
diff --git a/gerbonara/cad/kicad/primitives.py b/gerbonara/cad/kicad/primitives.py
index 58a5b2c..d5ee205 100644
--- a/gerbonara/cad/kicad/primitives.py
+++ b/gerbonara/cad/kicad/primitives.py
@@ -1,5 +1,6 @@
 
 import enum
+import math
 import re
 
 from .sexp import *
@@ -12,6 +13,7 @@ def unfuck_layers(layers):
     else:
         return layers
 
+
 def fuck_layers(layers):
     if layers and 'F.Cu' in layers and 'B.Cu' in layers and not any(re.match(r'^In[0-9]+\.Cu$', l) for l in layers):
         return ['F&B.Cu', *(l for l in layers if l not in ('F.Cu', 'B.Cu'))]
@@ -19,6 +21,38 @@ def fuck_layers(layers):
         return layers
 
 
+def layer_mask(layers):
+    mask = 0
+    for layer in layers:
+        match layer:
+            case '*.Cu':
+                return 0xff
+            case 'F.Cu':
+                mask |= 1<<0
+            case 'B.Cu':
+                mask |= 1<<31
+            case _:
+                if (m := re.match(f'In([0-9]+)\.Cu', layer)):
+                    mask |= 1<<int(m.group(1))
+    return mask
+
+
+def center_arc_to_kicad_mid(center, start, end):
+    # Convert normal p1/p2/center notation to the insanity that is kicad's midpoint notation
+    cx, cy = center.x, center.y
+    x1, y1 = start.x - cx, start.y - cy
+    x2, y2 = end.x - cx, end.y - cy
+    # Get a vector pointing from the center to the "mid" point.
+    dx, dy = x1 - x2, y1 - y2 # Get a vector pointing from "end" to "start"
+    dx, dy = -dy, dx # rotate by 90 degrees counter-clockwise
+    # normalize vector, and multiply by radius to get final point
+    r = math.hypot(x1, y1)
+    l = math.hypot(dx, dy)
+    mx = cx + dx / l * r
+    my = cy + dy / l * r
+    return XYCoord(mx, my)
+
+
 @sexp_type('hatch')
 class Hatch:
     style: AtomChoice(Atom.none, Atom.edge, Atom.full) = Atom.edge
diff --git a/gerbonara/cad/kicad/schematic.py b/gerbonara/cad/kicad/schematic.py
index 5f4e920..390716e 100644
--- a/gerbonara/cad/kicad/schematic.py
+++ b/gerbonara/cad/kicad/schematic.py
@@ -633,7 +633,7 @@ class Schematic:
     
 # From: https://jakevdp.github.io/blog/2012/10/07/xkcd-style-plots-in-matplotlib/
 #def xkcd_line(x, y, xlim=None, ylim=None, mag=1.0, f1=30, f2=0.05, f3=15):
-def xkcd_line(x, y, xlim=None, ylim=None, mag=1.0, f1=10, f2=0.10, f3=5):
+def xkcd_line(x, y, xlim=None, ylim=None, mag=1.0, f1=10, f2=0.05, f3=5):
     """
     Mimic a hand-drawn line from (x, y) data
 
@@ -745,9 +745,10 @@ def wonkify(path):
             for p0, p1, p2 in zip(pts[0::], pts[1::], pts[2::]):
                 dx1, dy1 = p1[0] - p0[0], p1[1] - p0[1]
                 dx2, dy2 = p2[0] - p1[0], p2[1] - p1[1]
+                l1, l2 = math.hypot(dx1, dy1), math.hypot(dx2, dy2)
                 a1, a2 = math.atan2(dy1, dx1), math.atan2(dy2, dx2)
                 da = (a2 - a1 + math.pi) % (2*math.pi) - math.pi
-                if abs(da) > math.pi/4:
+                if abs(da) > math.pi/4 and l1+l2 > 3:
                     seg.append(p1)
                     seg = [p1, p2]
                     segs.append(seg)