Closes #9

.vscode/settings.json

@@ -1,5 +1,9 @@
 {
     "python.pythonPath": "/Users/tlima/envs/zeropdk/bin/python",
     "python.formatting.provider": "black",
-    "editor.formatOnSave": true
+    "editor.formatOnSave": true,
+    "python.analysis.typeCheckingMode": "basic",
+    "python.linting.mypyEnabled": true,
+    "python.linting.enabled": true,
+    "python.linting.pylintEnabled": false
 }

setup.py

@@ -46,7 +46,7 @@ def main():
             "Operating System :: OS Independent",
             "Topic :: Scientific/Engineering",
         ),
-        install_requires=["numpy", "klayout", "scipy"],
+        install_requires=["numpy>=1.20", "klayout>=0.28.3", "scipy"],
     )
 
     setup(**metadata)

zeropdk/klayout_helper/extend.py

@@ -1,7 +1,3 @@
-from functools import wraps
-from typing import Type
-import klayout.db as kdb
-
 from . import point, cell, layout, polygon  # noqa
 
 cell.patch_cell()

zeropdk/klayout_helper/layout.py

@@ -29,4 +29,4 @@ def layout_read_cell(layout: Layout, cell_name: str, filepath: str) -> Cell:
 
 
 def patch_layout():
-    Layout.read_cell = layout_read_cell
+    Layout.read_cell = layout_read_cell  # type: ignore

zeropdk/klayout_helper/point.py

@@ -105,8 +105,9 @@ def make_points_picklable():
 
 
 def patch_points():
-    make_points_picklable()
     for klass in KLayoutPoints:
+        klass.__getstate__ = pyaPoint__getstate__
+        klass.__setstate__ = pyaPoint__setstate__
         klass.__init__ = pyaPoint__init__
         klass.__rmul__ = pyaPoint__rmul__
         klass.__mul__ = pyaPoint__mul__

zeropdk/layout/algorithms/sampling.py

@@ -165,7 +165,7 @@ def _sample_function(
 
         # compute the length of each line segment in the path
         dp = np.diff(p, axis=-1)
-        s = np.sqrt((dp**2).sum(axis=0))
+        s = np.sqrt((dp ** 2).sum(axis=0))
         s_tot = s.sum()
 
         # compute the angle between consecutive line segments

zeropdk/layout/cache.py

@@ -6,7 +6,7 @@
 import logging
 from hashlib import sha256
 from functools import partial, wraps
-from typing import Any, Type, Union, Callable, Dict
+from typing import Any, Optional, Type, Union, Callable, Dict
 
 import klayout.db as pya
 from zeropdk.klayout_helper.point import make_points_picklable
@@ -130,7 +130,7 @@ def read_layout(layout: pya.Layout, gds_filename: str, disambiguation_name: str
 
 
 def cache_cell(
-    cls: Type[PCell] = None, *, extra_hash: Any = None, cache_dir: str = CACHE_DIR
+    cls: Optional[Type[PCell]] = None, *, extra_hash: Any = None, cache_dir: str = CACHE_DIR
 ) -> Union[Type[PCell], Callable]:
     """Caches results of pcell call to save build time.
 
@@ -239,7 +239,7 @@ def wrapper_draw(self, cell):
             cell.insert(
                 pya.DCellInstArray(
                     retrieved_cell.cell_index(),
-                    pya.DTrans(pya.DTrans.R0, pya.DPoint(0, 0)),
+                    pya.DTrans(pya.DTrans.R0, pya.DVector(0, 0)),
                 )
             )
             return cell, ports

zeropdk/layout/geometry.py

@@ -124,7 +124,7 @@ def coords2d_func(t) -> NDArray[np.float_]:
             coords2d_func, [t0, t1], tol=0.0001 / scale, min_points=100
         )  # 1000 times more precise than the scale
         dp = np.diff(sampled_coords, axis=-1)
-    ds = np.sqrt((dp**2).sum(axis=0))
+    ds = np.sqrt((dp ** 2).sum(axis=0))
     return ds.sum()
 
 
@@ -225,8 +225,8 @@ def bezier_line(P0, P1, P2, P3):
     curve_func = (
         lambda t: (1 - t) ** 3 * P0
         + 3 * (1 - t) ** 2 * t * P1
-        + 3 * (1 - t) * t**2 * P2
-        + t**3 * P3
+        + 3 * (1 - t) * t ** 2 * P2
+        + t ** 3 * P3
     )
     return curve_func
 
@@ -244,7 +244,7 @@ def curvature_bezier(P0, P1, P2, P3):
     b_prime = (
         lambda t: 3 * (1 - t) ** 2 * (P1 - P0)
         + 6 * (1 - t) * t * (P2 - P1)
-        + 3 * t**2 * (P3 - P2)
+        + 3 * t ** 2 * (P3 - P2)
     )
     b_second = lambda t: 6 * (1 - t) * (P2 - 2 * P1 + P0) + 6 * t * (P3 - 2 * P2 + P1)
     dx = lambda t: b_prime(t).x
@@ -340,7 +340,7 @@ def __str__(self):
         return f"Point({self.x}, {self.y})"
 
     def norm(self):
-        return sqrt(self.x**2 + self.y**2)
+        return sqrt(self.x ** 2 + self.y ** 2)
 
 
 class _Line(_Point):

zeropdk/layout/routing.py

@@ -64,6 +64,7 @@ def ensure_layer(layout: kdb.Layout, layer: GeneralLayer) -> Optional[int]:
         return layer
     else:
         logger.error(f"{layer} not recognized")
+    return None
 
 
 def common_layout_manhattan_traces(

zeropdk/layout/waveguide_rounding.py

@@ -1,12 +1,14 @@
 """ Straight waveguide rounding algorithms"""
 from functools import lru_cache
 from math import atan2, tan, inf
-from typing import List, Sequence, Tuple, Union
+from numbers import Number
+from typing import Callable, List, Sequence, Tuple, Union
 import warnings
 import numpy as np
+import numpy.typing as npt
 import klayout.db as kdb
 from zeropdk.klayout_helper import as_point
-from zeropdk.layout import GeneralLayer
+from zeropdk.layout import GeneralLayer, insert_shape
 from zeropdk.layout.geometry import rotate, fix_angle, cross_prod
 from zeropdk.layout.algorithms.sampling import sample_function
 from zeropdk.layout.polygons import layout_path
@@ -109,7 +111,9 @@ def get_points(self):
             theta_start = (theta_start - theta_end) % (2 * pi) + theta_end
             theta_start, theta_end = theta_end, theta_start
 
-        arc_function = lambda t: np.array([r * np.cos(t), r * np.sin(t)])
+        arc_function: Callable[[npt.ArrayLike], npt.NDArray] = lambda t: np.array(
+            [r * np.cos(t), r * np.sin(t)]
+        )
 
         # in the function below, theta_start must be smaller than theta_end
         t, coords = sample_function(arc_function, [theta_start, theta_end], tol=0.002 / r)
@@ -136,7 +140,7 @@ def __repr__(self):
 
 
 class _Line:
-    def __init__(self, P1, P2):
+    def __init__(self, P1: kdb.DPoint, P2: kdb.DPoint):
         self.P1 = P1
         self.P2 = P2
 
@@ -473,8 +477,8 @@ def compute_rounded_path(points: Sequence[kdb.DPoint], radius: float) -> Rounded
 class _Path:
     """Object holding path plus width information"""
 
-    def __init__(self, points, widths):
-        self.points = points
+    def __init__(self, points: Sequence[kdb.DPoint], widths: Union[Number, Sequence[Number]]):
+        self.points: Sequence[kdb.DPoint] = points
 
         # This can be a single width or a list of widths, just like in layout_waveguide()
         self.widths = widths
@@ -504,7 +508,7 @@ def __repr__(self):
         )
 
 
-def _compute_tapered_line(line, waveguide_width, taper_width, taper_length):
+def _compute_tapered_line(line: _Line, waveguide_width, taper_width, taper_length):
     """Takes a _Line object and computes two tapers with taper_width and taper_length"""
 
     minimum_length = 30 + 2 * taper_length  # don't bother tapering waveguides beyond this length
@@ -524,11 +528,11 @@ def _compute_tapered_line(line, waveguide_width, taper_width, taper_length):
     ]
 
 
-def compute_untapered_path(path, waveguide_width):
+def compute_untapered_path(path, waveguide_width) -> List[_Path]:
     return [_Path(element.get_points(), waveguide_width) for element in path]
 
 
-def compute_tapered_path(path, waveguide_width, taper_width, taper_length):
+def compute_tapered_path(path, waveguide_width, taper_width, taper_length) -> List[_Path]:
     tapered_path = []
     for element in path:
         if isinstance(element, _Line):
@@ -557,15 +561,23 @@ def unique_points(point_list: Sequence[kdb.DPoint]) -> List[kdb.DPoint]:
 
 
 def waveguide_from_points(
-    cell, layer, points, width, radius, taper_width=None, taper_length=None
+    cell: kdb.Cell,
+    layer: kdb.LayerInfo,
+    points: Sequence[kdb.DPoint],
+    width: float,
+    radius,
+    taper_width=None,
+    taper_length=None,
 ) -> Tuple[List[kdb.DPoint], List[float]]:
     """Draws a waveguide with rounded corners given a path of manhattan-like points.
 
     Returns:
         - points: list of DPoints
-        - widths: list of widths with same length as points.
+        - width: floating point
+        - radius: floating point (must be greater than half the width)
 
     """
+    assert isinstance(width, Number), "This function takes a single width"
     assert radius > width / 2, "Please use a radius larger than the half-width"
     points = unique_points(points)
 
@@ -590,23 +602,26 @@ def waveguide_from_points(
         waveguide_path = compute_untapered_path(rounded_path, width)
 
     # creating a single path
-    _draw_points = []
-    _draw_widths = []
+    _draw_points: List[kdb.DPoint] = []
+    _draw_widths: List[float] = []
     for element in waveguide_path:
-        points, width = element.points, element.widths
+        points, _width = element.points, element.widths
         n_points = len(points)
         try:
-            if len(width) == n_points:
+            if isinstance(_width, Number):
                 _draw_points.extend(points)
-                _draw_widths.extend(width)
-            elif len(width) == 2:
-                _draw_widths.extend(np.linspace(width[0], width[1], n_points))
+                _draw_widths.extend(np.ones(n_points) * _width)
+            elif len(_width) == n_points:
+                _draw_points.extend(points)
+                _draw_widths.extend(_width)
+            elif len(_width) == 2:
+                _draw_widths.extend(np.linspace(_width[0], _width[1], n_points))
                 _draw_points.extend(points)
             else:
                 raise RuntimeError("Internal error detected. Debug please.")
         except TypeError:
             _draw_points.extend(points)
-            _draw_widths.extend(np.ones(n_points) * width)
+            _draw_widths.extend(np.ones(n_points) * _width)
 
     # deleting repeated points
     _cur_point = None
@@ -642,18 +657,17 @@ def layout_waveguide_from_points(
 
 
 def main():
-    def trace_rounded_path(cell, layer, rounded_path, width):
+    def trace_rounded_path(cell: kdb.Cell, layer: kdb.LayerInfo, rounded_path, width):
         points = []
         for item in rounded_path:
             points.extend(item.get_points())
 
         dpath = kdb.DPath(points, width, 0, 0)
+        insert_shape(cell, layer, dpath)
 
-        cell.shapes(layer).insert(dpath)
-
-    def trace_reference_path(cell, layer, points, width):
+    def trace_reference_path(cell: kdb.Cell, layer: kdb.LayerInfo, points, width):
         dpath = kdb.DPath(points, width, 0, 0)
-        cell.shapes(layer).insert(dpath)
+        insert_shape(cell, layer, dpath)
 
     layout = kdb.Layout()
     TOP = layout.create_cell("TOP")

zeropdk/layout/waveguides.py

@@ -10,7 +10,7 @@
 """
 
 from itertools import repeat
-from numbers import Real
+from numbers import Number, Real
 from typing import List, Sequence, Tuple, TypeVar, Union
 import numpy as np
 from numpy import cos, sin, pi, sqrt
@@ -264,7 +264,7 @@ def smooth_append(point_list: List[kdb.DPoint], point: kdb.DPoint) -> List[kdb.D
 
             # Only add new point if the area of the triangle built with
             # current edge and previous edge is greater than dbu^2/2
-            if abs(cross_prod(prev_edge, curr_edge)) > dbu**2 / 2:
+            if abs(cross_prod(prev_edge, curr_edge)) > dbu ** 2 / 2:
                 if smooth:
                     # avoid corners when smoothing
                     if cos_angle(curr_edge, prev_edge) > cos(130 / 180 * pi):
@@ -298,7 +298,7 @@ def layout_waveguide(
     cell: kdb.Cell,
     layer: GeneralLayer,
     points_list: Sequence[kdb.DPoint],
-    width: Union[float, Sequence[float]],
+    width: Union[Number, Sequence[Number]],
     smooth=False,
 ) -> ZeroPDKDSimplePolygon:
     """Lays out a waveguide (or trace) with a certain width along given points.