AI Flow

#!/usr/bin/env python3
"""Independent reference values for crates/aiflow-router/src/impedance.rs.

A 2-D finite-volume Laplace solve of a trace cross-section on a graded rectilinear grid
(numpy + scipy). Conductors are fixed-potential node sets, the ground planes are rows at 0 V,
the open sides and top of a microstrip are a 0 V box about 25 dielectric heights away. The
capacitance per metre comes from the field energy, once with the dielectric and once in air:

    Z0 = 1 / (c * sqrt(C * C_air)),  eeff = C / C_air

Pairs are solved in odd mode (+1 V / -1 V) and even mode (+1 V / +1 V); Zdiff = 2 * Zodd.
No solder mask, no loss, no dispersion: the same quasi-static problem the closed forms solve.
With `cpw` a same-layer ground (0 V, the strip's thickness) starts that far beside the outer
copper edges and runs to the box side (or `pour` mm wide): grounded coplanar waveguide, the
pour stitched to the plane under it.

    python3 tools/impedance_reference.py            # prints the table used by the Rust tests
    python3 tools/impedance_reference.py --coplanar-sweep > tests/fixtures/coplanar-reference.jsonl

Each case is solved on two grids (finest cell f and f/2) and extrapolated to zero cell size
(Richardson, first order: the strip corners dominate the grid error); the two grids differ by
under 0.5 %.
"""
import math
import sys

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla

EPS0 = 8.8541878128e-12
C0 = 299792458.0


def graded(points, fine, coarse, growth=1.15):
    """Node coordinates through every point in `points` (sorted), at most `fine` apart near
    them and growing by `growth` up to `coarse` away from them."""
    pts = sorted(set(points))
    out = [pts[0]]
    for a, b in zip(pts, pts[1:]):
        span = b - a
        if span <= fine * 1.0001:
            out.append(b)
            continue
        # grow from both ends toward the middle
        left, right = [a], [b]
        step = fine
        while True:
            if left[-1] + step >= right[-1] - step:
                break
            left.append(left[-1] + step)
            right.append(right[-1] - step)
            step = min(step * growth, coarse)
        mid = left + right[::-1]
        out.extend(mid[1:])
    return np.array(sorted(set(round(x, 9) for x in out)))


def solve(xs, ys, eps_cell, fixed):
    """Node potentials and energy. eps_cell[j, i] is the relative permittivity of the cell
    between nodes (i, j) and (i+1, j+1); fixed maps node index -> potential."""
    nx, ny = len(xs), len(ys)
    dx, dy = np.diff(xs), np.diff(ys)
    idx = lambda i, j: j * nx + i
    rows, cols, vals = [], [], []
    edges = []  # (a, b, coefficient)
    for j in range(ny):
        for i in range(nx):
            # east edge
            if i + 1 < nx:
                c = 0.0
                if j - 1 >= 0:
                    c += eps_cell[j - 1, i] * dy[j - 1] / 2
                if j < ny - 1:
                    c += eps_cell[j, i] * dy[j] / 2
                edges.append((idx(i, j), idx(i + 1, j), c / dx[i]))
            if j + 1 < ny:
                c = 0.0
                if i - 1 >= 0:
                    c += eps_cell[j, i - 1] * dx[i - 1] / 2
                if i < nx - 1:
                    c += eps_cell[j, i] * dx[i] / 2
                edges.append((idx(i, j), idx(i, j + 1), c / dy[j]))
    n = nx * ny
    e = np.array(edges)
    a, b, c = e[:, 0].astype(int), e[:, 1].astype(int), e[:, 2]
    diag = np.zeros(n)
    np.add.at(diag, a, c)
    np.add.at(diag, b, c)
    K = sp.coo_matrix((np.concatenate([-c, -c, diag]), (np.concatenate([a, b, np.arange(n)]), np.concatenate([b, a, np.arange(n)]))), shape=(n, n)).tocsr()
    phi = np.zeros(n)
    is_fixed = np.zeros(n, bool)
    for k, v in fixed.items():
        phi[k] = v
        is_fixed[k] = True
    free = ~is_fixed
    rhs = -K[free][:, is_fixed] @ phi[is_fixed]
    phi[free] = spla.spsolve(K[free][:, free].tocsc(), rhs)
    energy = 0.5 * EPS0 * np.sum(c * (phi[a] - phi[b]) ** 2)
    return energy


def section(w, t, h_below, h_above=None, er=4.3, gap=None, mode="single", fine=None, cpw=None, pour=None):
    """Capacitance per metre of a strip (or pair) of width w, thickness t (mm) at h_below over a
    ground plane; a second plane h_above over the top of the strip makes it a stripline.
    cpw: a same-layer ground (0 V, thickness t) at this gap beside the outer edge(s), out to the
    box side (or `pour` mm wide): grounded coplanar waveguide, single or edge-coupled."""
    mm = 1e-3
    fine = fine or min(t / 4, w / 20, h_below / 20, (gap / 8) if gap else 1, (cpw / 8) if cpw else 1)
    strips = [(-w / 2, w / 2)] if gap is None else [(-gap / 2 - w, -gap / 2), (gap / 2, gap / 2 + w)]
    span = (w if gap is None else 2 * w + gap)
    hb = h_below
    top = (hb + t + h_above) if h_above is not None else hb + t + 25 * hb + 5 * span
    side = span / 2 + (8 * (hb + (h_above or 0)) + 3 * span if h_above is not None else 25 * hb + 5 * span)
    grounds = []
    if cpw is not None:
        e = span / 2 + cpw
        far = side if pour is None else min(side, e + pour)
        grounds = [(-far, -e), (e, far)]
    xk = [-side, side] + [x for s in strips + grounds for x in s] + [0.0]
    yk = [0.0, hb, hb + t, top]
    coarse = max(span, hb) / 2
    xs = graded(xk, fine, coarse)
    ys = graded(yk, fine, coarse)
    nx, ny = len(xs), len(ys)
    def energy(er_d, potentials):
        eps = np.ones((ny - 1, nx - 1))
        ymid = (ys[:-1] + ys[1:]) / 2
        if h_above is None:
            eps[ymid < hb, :] = er_d
        else:
            eps[:, :] = er_d
        fixed = {}
        for i in range(nx):
            fixed[i] = 0.0  # bottom plane
            fixed[(ny - 1) * nx + i] = 0.0  # top plane or box
        for j in range(ny):
            fixed[j * nx] = 0.0
            fixed[j * nx + nx - 1] = 0.0
        for (x0, x1), v in list(zip(strips, potentials)) + [(g, 0.0) for g in grounds]:
            for j in range(ny):
                if hb - 1e-12 <= ys[j] <= hb + t + 1e-12:
                    for i in range(nx):
                        if x0 - 1e-12 <= xs[i] <= x1 + 1e-12:
                            fixed[j * nx + i] = v
        return solve(xs * mm, ys * mm, eps, fixed)
    pots = {"single": [1.0], "odd": [1.0, -1.0], "even": [1.0, 1.0]}[mode]
    e_d, e_a = energy(er, pots), energy(1.0, pots)
    # per-line capacitance: single C = 2E; odd and even (two lines at +-1 V) C_line = E
    k = 2.0 if mode == "single" else 1.0
    cd, ca = k * e_d, k * e_a
    return 1 / (C0 * math.sqrt(cd * ca)), cd / ca


CASES = [
    # name, kwargs
    ("microstrip 1.6 mm FR4 er4.5 w2.95 t35", dict(w=2.95, t=0.035, h_below=1.6, er=4.5)),
    ("microstrip JLC7628 h0.2104 er4.4 w0.35 t35", dict(w=0.35, t=0.035, h_below=0.2104, er=4.4)),
    ("microstrip h0.2 er4.2 w0.3 t35", dict(w=0.30, t=0.035, h_below=0.2, er=4.2)),
    ("pair microstrip JLC7628 w0.2 s0.2 odd", dict(w=0.2, t=0.035, h_below=0.2104, er=4.4, gap=0.2, mode="odd")),
    ("pair microstrip JLC7628 w0.25 s0.15 odd", dict(w=0.25, t=0.035, h_below=0.2104, er=4.4, gap=0.15, mode="odd")),
    ("pair microstrip 1.6 mm er4.5 w1.0 s0.5 odd", dict(w=1.0, t=0.035, h_below=1.6, er=4.5, gap=0.5, mode="odd")),
    ("pair microstrip 1.6 mm er4.5 w1.0 s0.2 odd (s/h 0.125: refused)", dict(w=1.0, t=0.035, h_below=1.6, er=4.5, gap=0.2, mode="odd")),
    ("stripline sym b0.5 er4.3 w0.2 t35", dict(w=0.2, t=0.035, h_below=0.2325, h_above=0.2325, er=4.3)),
    ("stripline offset 0.2/0.4 er4.3 w0.15 t35", dict(w=0.15, t=0.035, h_below=0.2, h_above=0.4, er=4.3)),
    ("stripline offset 0.2/0.6 er4.3 w0.15 t35", dict(w=0.15, t=0.035, h_below=0.2, h_above=0.6, er=4.3)),
    ("stripline offset 0.2/1.03 er4.3 w0.15 t35 (5:1: refused)", dict(w=0.15, t=0.035, h_below=0.2, h_above=1.03, er=4.3)),
    ("pair stripline sym b0.5 er4.3 w0.15 s0.15 odd", dict(w=0.15, t=0.035, h_below=0.2325, h_above=0.2325, er=4.3, gap=0.15, mode="odd")),
    # grounded coplanar (a same-layer ground pour at gap cpw beside the copper, out to the box side)
    ("GCPW 1.6 mm FR4 er4.5 w0.802 g0.15 t35", dict(w=0.802, t=0.035, h_below=1.6, er=4.5, cpw=0.15)),
    ("pair GCPW 1.51 er4.5 w0.546 s0.2 g0.2 odd (StickHub)", dict(w=0.546, t=0.035, h_below=1.51, er=4.5, cpw=0.2, gap=0.2, mode="odd")),
    ("pair GCPW 1.51 er4.5 w0.379 s0.15 g0.15 odd", dict(w=0.379, t=0.035, h_below=1.51, er=4.5, cpw=0.15, gap=0.15, mode="odd")),
    ("pair GCPW 1.51 er4.5 w0.8485 s0.3515 g0.2 odd (grid pitch)", dict(w=0.8485, t=0.035, h_below=1.51, er=4.5, cpw=0.2, gap=0.3515, mode="odd")),
    ("GCPW JLC7628 h0.2104 er4.4 w0.3095 g0.15", dict(w=0.3095, t=0.035, h_below=0.2104, er=4.4, cpw=0.15)),
    ("pair GCPW JLC7628 w0.2688 s0.15 g0.15 odd", dict(w=0.2688, t=0.035, h_below=0.2104, er=4.4, cpw=0.15, gap=0.15, mode="odd")),
    ("GCPW h1.0 er4.4 w0.2 g1.0 (range edge, g/h 1)", dict(w=0.2, t=0.035, h_below=1.0, er=4.4, cpw=1.0)),
    ("pair GCPW h1.0 er4.4 w0.2 s0.6 g0.5 odd (range edge)", dict(w=0.2, t=0.035, h_below=1.0, er=4.4, cpw=0.5, gap=0.6, mode="odd")),
    ("GCPW h0.508 er3.66 w0.8 g0.2", dict(w=0.8, t=0.035, h_below=0.508, er=3.66, cpw=0.2)),
    ("pair GCPW h0.508 er3.66 w0.3 s0.2 g0.2 odd", dict(w=0.3, t=0.035, h_below=0.508, er=3.66, cpw=0.2, gap=0.2, mode="odd")),
    ("GCPW h0.635 er9.8 w0.3 g0.15 t18", dict(w=0.3, t=0.018, h_below=0.635, er=9.8, cpw=0.15)),
    ("GCPW h0.5 er4.4 w0.2 g1.0 (g/h 2: refused)", dict(w=0.2, t=0.035, h_below=0.5, er=4.4, cpw=1.0)),
    ("pair GCPW 1.51 er4.5 w0.546 s0.2 g0.2 odd, pour 0.3 mm wide", dict(w=0.546, t=0.035, h_below=1.51, er=4.5, cpw=0.2, gap=0.2, mode="odd", pour=0.3)),
    ("GCPW 1.6 mm FR4 er4.5 w0.802 g0.15, pour 0.3 mm wide", dict(w=0.802, t=0.035, h_below=1.6, er=4.5, cpw=0.15, pour=0.3)),
]


def coplanar_sweep():
    """The grounded coplanar cross-sections the validated range in impedance.rs comes from
    (tests/fixtures/coplanar-reference.jsonl): single lines and odd-mode pairs, er 4.4."""
    import itertools
    out = []
    for h, w, g in itertools.product([0.1, 0.2104, 0.5, 1.0, 1.6], [0.1, 0.2, 0.4, 0.8, 1.5], [0.1, 0.15, 0.25, 0.5, 1.0]):
        out.append(dict(w=w, t=0.035, h_below=h, er=4.4, cpw=g))
    for h, w, s, g in itertools.product([0.1, 0.2104, 1.0, 1.6], [0.1, 0.2, 0.4, 0.8], [0.1, 0.15, 0.3, 0.6], [0.1, 0.2, 0.5]):
        out.append(dict(w=w, t=0.035, h_below=h, er=4.4, cpw=g, gap=s, mode="odd"))
    for h, w, g, t in itertools.product([0.2104, 1.6], [0.2, 0.5], [0.1, 0.2], [0.018, 0.07]):
        out.append(dict(w=w, t=t, h_below=h, er=4.4, cpw=g))
        out.append(dict(w=w, t=t, h_below=h, er=4.4, cpw=g, gap=g, mode="odd"))
    for h, w, g in itertools.product([0.2, 1.0], [0.1, 0.4], [0.1, 0.2, 0.4, 1.0]):
        out.append(dict(w=w, t=0.004, h_below=h, er=4.4, cpw=g))
    return out


def extrapolated(kw):
    """Two grids (finest cell f and f/2), Richardson to zero cell size: (Z, eeff)."""
    f = min(kw["t"] / 4, kw["w"] / 20, kw["h_below"] / 20, (kw["gap"] / 8) if kw.get("gap") else 1, (kw["cpw"] / 8) if kw.get("cpw") else 1)
    z1, e1 = section(**kw, fine=f)
    z2, e2 = section(**kw, fine=f / 2)
    return 2 * z2 - z1, 2 * e2 - e1


def _sweep_line(kw):
    import json
    z, _ = extrapolated(kw)
    odd = kw.get("mode") == "odd"
    return json.dumps({"h": kw["h_below"], "w": kw["w"], "g": kw["cpw"], "s": kw.get("gap"), "t": kw["t"], "er": kw["er"], "z": round(2 * z if odd else z, 3)})

if __name__ == "__main__":
    if sys.argv[1:2] == ["--coplanar-sweep"]:
        # python3 tools/impedance_reference.py --coplanar-sweep > tests/fixtures/coplanar-reference.jsonl
        from multiprocessing import Pool
        with Pool() as p:
            for line in p.map(_sweep_line, coplanar_sweep()):
                print(line, flush=True)
        sys.exit(0)
    only = sys.argv[1:]
    for name, kw in CASES:
        if only and not any(o in name for o in only):
            continue
        # first-order grid error (the strip corners): extrapolate from two grids (Richardson)
        z, ee = extrapolated(kw)
        extra = f"  Zdiff {2 * z:.2f}" if kw.get("mode") == "odd" else ""
        print(f"{name:52s} Z {z:7.2f} ohm  eeff {ee:5.3f}{extra}", flush=True)