path: root/kisim.py

"""kisim -- run a KiCad schematic in ngspice from Python and read signals by name.

Single file, no project knowledge. Needs: kicad-cli, ngspice, numpy, spicelib.

From a script in a subfolder of sim/:

    import sys; from pathlib import Path; sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    from kisim import KiSim

    sim = KiSim("softstart.kicad_sch", tran="10u 400m")      # path relative to the calling script
    r = sim.run()                                            # or sim.run(params={"X1": {"TACT": "250m"}})
    r = sim.run(subs={r"^V1 .*$": "V1 +28V GND DC 35"}, tag="v35")   # regex edits of netlist lines
    r = sim.run(sources={"V1": pwl((0, 0), (1e-3, 28))}, tag="ramp")  # replace a source's value
    KiSim(..., save=["+28V", "/VG", "i(V1)"])                # only keep these signals (KiCad names)
    r.v("+28V"), r.v("Net-(Q2-D)"), r.v("Net-(D1-A)", "+28V")  # KiCad net names; 2nd arg = reference node
    r.i("V1"), r.i("Q2", "d"), r["v(x1.p12r)"], r.find("q7")  # currents (pins of subckts) / raw names / search
    KiSim(..., power=["Q2", "R1"]).run().p("Q2")             # device power, W (must be requested)
    r.at("i(v1)", 0.1), r.window(0.015, 0.03)                 # value at a time / slice in time
    run_cases(sim, [Case("ov", sources={...}, analyses=[peak], limits={...})])  # test cases, see Case
    Wave(8).ramp(12, 1e-3).hold(2e-3).ramp(20, 0.5e-3)       # piecewise stimulus, str() -> PWL
    Case("x", **stim(piece_a, piece_b))                       # merge reusable scenario pieces
    run_cases(sim, corners(case, {"R4": tol("10.7k", .01), "X1.PACT": ["9", "11"]}))  # WCA corners

Signal names: ngspice lower-cases everything and keeps KiCad's net names
(e.g. "Net-(Q2-D)" -> "v(net-_q2-d_)"; KiCad's SPICE export swaps "(" ")" for "_").
Subcircuit internals are "v(x1.p12r)" style. r.names lists all of them.
Saving: by default ngspice keeps every node voltage and every voltage-source
current. save=[...] keeps only the listed signals: KiCad net names ("+12V",
"/TC_OFF", "Net-(Q2-D)") become v(...); entries with a "(" ("i(V1)",
"v(x1.p12r)") pass through. Use it for big vendor models -- thousands of
internal nodes make huge, slow raw files.
Currents: currents=True adds '.probe alli', giving every device terminal
current, e.g. "i(r1)", "i(xq2:d)". Off by default: it inserts a 0 V probe
source in every device terminal (bigger matrix, and it has broken convergence
of marginal models); turn it on only when you need device pin currents.
Power: ngspice-47 ignores
'.probe allp', so power is probed per device: power=["Q2"] -> "xq2:power".
KiCad prefixes subcircuit instances with X in the netlist (Q2 -> XQ2); the
helpers accept the schematic reference and resolve that.
"""
from __future__ import annotations

import inspect
import os
import re
import subprocess
from pathlib import Path

import numpy as np
from spicelib import RawRead


def pwl(*points) -> str:
    """PWL source value from (time_s, value) pairs: pwl((0, 0), (1e-3, 5))."""
    return "PWL(" + " ".join(f"{t:g} {v:g}" for t, v in points) + ")"


class KiSim:
    def __init__(self, schematic, tran: str = "10u 100m", options: str = "method=gear reltol=1e-4",
                 extra: list[str] | None = None, currents: bool = False, power: list[str] | None = None,
                 workdir=None, ngspice: str = "ngspice", compat: str | None = None,
                 save: list[str] | None = None):
        """schematic: .kicad_sch path (relative paths resolve against the calling script).
        tran: '.tran' arguments. options: '.options' line contents. extra: raw netlist lines
        to add (e.g. '.ic', '.param'). currents: also save all device terminal currents
        ('.probe alli', off by default). save: only keep these signals (KiCad net names
        or raw ngspice names like "i(V1)"); default keeps all node voltages.
        power: schematic references whose dissipated power to save.
        workdir: where run.cir/run.raw go (default: results/).
        compat: ngspice compatibility mode, e.g. "psa" for PSpice vendor models (TI)."""
        caller = Path(inspect.stack()[1].filename).resolve().parent
        self.sch = (caller / schematic).resolve()
        self.tran, self.options, self.extra, self.ngspice = tran, options, extra or [], ngspice
        self.compat = ["-D", f"ngbehavior={compat}"] if compat else []
        self.currents, self.power, self.save = currents, power or [], save or []
        self.workdir = Path(workdir) if workdir else self.sch.parent / "results"
        self.workdir.mkdir(parents=True, exist_ok=True)

    def netlist(self, params: dict | None = None, tag: str = "run", subs: dict | None = None,
                sources: dict | None = None, tran: str | None = None) -> Path:
        """Export the schematic to a SPICE netlist, apply parameter overrides, source values
        (sources={"V1": "PWL(...)"}: everything after the two nodes is replaced) and regex
        substitutions (subs={pattern: replacement}, multiline), add the analysis."""
        cir = self.workdir / f"{tag}.cir"
        subprocess.run(["kicad-cli", "sch", "export", "netlist", "--format", "spice", "-o", cir, self.sch],
                       check=True, capture_output=True)
        lines = cir.read_text().splitlines()
        def find(ref):                                   # schematic ref -> netlist line (Q2 may be XQ2)
            for i, l in enumerate(lines):
                if l.split()[:1] and l.split()[0].upper() in (ref.upper(), "X" + ref.upper()):
                    return i
            raise KeyError(f"{ref} not in netlist")
        for ref, kv in (params or {}).items():          # {"X1": {"TACT": "250m"}} -> edit/append KEY=VAL
            idx = find(ref)
            for k, val in kv.items():
                line, n = re.subn(rf"\b{k}=\S+", f"{k}={val}", lines[idx], flags=re.I)
                lines[idx] = line if n else f"{line} {k}={val}"
        for ref, val in (sources or {}).items():         # "V4 +5VSC GND2 DC 5" -> "V4 +5VSC GND2 <val>"
            idx = find(ref)
            lines[idx] = " ".join(lines[idx].split()[:3] + [str(val)])   # str(): Wave -> PWL
        if subs:
            text = "\n".join(lines)
            for pat, rep in subs.items():
                text, n = re.subn(pat, rep, text, flags=re.M)
                if not n:
                    raise ValueError(f"subs pattern {pat!r} matched nothing")
            lines = text.splitlines()
        end = max(i for i, l in enumerate(lines) if l.strip().lower() == ".end")
        probes = [".save all", ".probe alli"] if self.currents else []
        if self.save:                                    # KiCad names -> v(node); raw names pass through
            probes += [".save " + " ".join(n.lower() if "(" in n else f"v({Result.node(n)})" for n in self.save)]
        probes += [f".probe p({lines[find(ref)].split()[0]})" for ref in self.power]
        lines[end:end] = [f".options {self.options}", *probes, *self.extra, f".tran {tran or self.tran}"]
        cir.write_text("\n".join(lines) + "\n")
        return cir

    def run(self, params: dict | None = None, tag: str = "run", cache: bool = True,
            subs: dict | None = None, sources: dict | None = None, tran: str | None = None) -> "Result":
        """Simulate and return a Result. With cache=True an unchanged netlist is not re-simulated.
        params/subs/sources: see netlist(). tran: '.tran' for this run only (default: self.tran)."""
        cir, raw = self.workdir / f"{tag}.cir", self.workdir / f"{tag}.raw"
        old = cir.read_text() if cir.exists() else None
        self.netlist(params, tag, subs, sources, tran)
        if not (cache and raw.exists() and old == cir.read_text()):
            tmp = raw.with_suffix(".raw.tmp")          # only a finished run becomes the cached .raw
            raw.unlink(missing_ok=True)
            env = {**os.environ, "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS", "1")}
            # one thread per run: as fast as ngspice's OpenMP threads on these circuits, and
            # parallel case runs (run_cases) no longer oversubscribe the CPU
            p = subprocess.run([self.ngspice, *self.compat, "-b", "-r", tmp, cir], capture_output=True,
                               text=True, env=env)
            if p.returncode or not tmp.exists():
                raise RuntimeError(f"ngspice failed:\n{p.stdout[-2000:]}{p.stderr[-2000:]}")
            tmp.replace(raw)
        return Result(raw)


class Result:
    """Transient results. All waveforms are numpy arrays aligned with .t (seconds)."""

    def __init__(self, raw: Path):
        self.raw = RawRead(raw, verbose=False)
        self.names = self.raw.get_trace_names()
        self.t = self["time"]

    def __getitem__(self, name: str) -> np.ndarray:
        key = name.lower()
        if key not in self.names:
            raise KeyError(f"{name!r} not found; try .find({name.split('(')[-1].strip(')')!r})")
        return np.real(self.raw.get_trace(key).get_wave())

    @staticmethod
    def node(net: str) -> str:
        """KiCad net name -> ngspice node name."""
        return re.sub(r"[() ]", "_", net.lower())  # '/Sheet name/LABEL' keeps its slashes

    def v(self, net: str, ref: str | None = None) -> np.ndarray:
        """Voltage of a KiCad net (or ngspice node), optionally relative to another net."""
        vv = self[f"v({self.node(net)})"]
        return vv - self.v(ref) if ref else vv

    def _dev(self, ref: str, fmt: str) -> np.ndarray:
        for name in (ref.lower(), "x" + ref.lower()):     # Q2 is XQ2 in the netlist
            if fmt.format(name) in self.names:
                return self[fmt.format(name)]
        raise KeyError(f"{fmt.format(ref.lower())!r} not found; try .find({ref.lower()!r})")

    def i(self, ref: str, pin: str | None = None) -> np.ndarray:
        """Current of a device, e.g. 'V1', 'R1', or a subcircuit pin: ('Q2', 'd').
        Sign is ngspice's: positive into the terminal (for V sources: into the + terminal)."""
        return self._dev(ref, "i({}:%s)" % pin.lower() if pin else "i({})")

    def p(self, ref: str) -> np.ndarray:
        """Power dissipated in a device listed in KiSim(power=[...]), in W."""
        return self._dev(ref, "{}:power")

    def find(self, pattern: str) -> list[str]:
        """Signal names containing pattern (case-insensitive)."""
        return [n for n in self.names if pattern.lower() in n]

    def at(self, name: str, time: float) -> float:
        """Signal value at a time (linear interpolation)."""
        return float(np.interp(time, self.t, self[name]))

    def plot(self, rows, title: str | None = None, file=None, xlim=None, show: bool = False):
        """Stacked plot, one axis per entry. rows: list of entries; a sub-list puts several
        entries on one axis. An entry is any of:
          - a KiCad net name: "+12V", "/TC_OFF", "Net-(Q2-D)"
          - a raw trace name: "i(V1)", "xq1:power", "v(x1.p12r)" (case-insensitive); device
            quantities also by schematic ref: "Q1:power" -> "xq1:power", "D2:power" -> "xd2:power"
          - a (label, signal) tuple; signal is a name as above or an array, e.g.
            ("IL", (r.v("+12V") - r.v("Net-(U2-ISNS-)")) / 68e-3)
          - an array (numpy or array-like), plotted as-is against r.t: r.i("V1"), r.p("Q2")
        e.g. r.plot(["+12V", ["/OVP_COMP", "/OVP_COMP_FILT"], "i(V1)", r.p("Q1")]).
        Time in ms. Returns (fig, axes)."""
        import matplotlib.pyplot as plt
        if not isinstance(rows, list):
            rows = [rows]
        rows = [[x] if isinstance(x, str) or not isinstance(x, list) else list(x) for x in rows]
        fig, axes = plt.subplots(len(rows), 1, sharex=True, figsize=(9, 1.8 * len(rows) + 0.6), squeeze=False)
        axes = axes[:, 0]
        for ax, row in zip(axes, rows):
            labeled = False
            for spec in row:
                label, y = self._curve(spec)
                ax.plot(self.t * 1e3, y, label=label)
                labeled = labeled or label is not None
            ax.grid(alpha=0.3)
            if labeled:
                ax.legend(fontsize=8, loc="upper left")
        axes[-1].set_xlabel("time (ms)")
        if xlim:
            axes[-1].set_xlim(*xlim)
        if title:
            axes[0].set_title(title)
        fig.tight_layout()
        if file:
            fig.savefig(file, dpi=110)
        if show:
            plt.show()
        return fig, axes

    def _curve(self, spec) -> tuple[str | None, np.ndarray]:
        """(label, y) for one plot entry: net/trace name, (label, signal) tuple, or array.
        A string is looked up as a raw trace name first ("i(V1)", "xq1:power"), then as a
        KiCad net name."""
        label, sig = spec if isinstance(spec, tuple) else (spec, spec)
        if isinstance(sig, str):
            key = sig.lower()
            if key not in self.names and ":" in key and "x" + key in self.names:
                key = "x" + key                      # schematic ref: "D2:power" -> "xd2:power"
            y = self[key] if key in self.names else self.v(sig)
        else:
            y = np.asarray(sig)
        return (label if isinstance(label, str) else None), y

    def window(self, t0: float, t1: float) -> slice:
        """Index slice for t0 <= t <= t1, e.g. r.t[w], r.v('+28V')[w]."""
        i0, i1 = np.searchsorted(self.t, [t0, t1])
        return slice(i0, i1 + 1)


# --------------------------------------------------------------------------------------
# Test cases: stimulus + analyses + limits, run in parallel, summarised in one table.
# --------------------------------------------------------------------------------------
from dataclasses import dataclass, field


@dataclass
class Case:
    """One simulation case.

    name      tag of the run (results/<name>.cir/.raw/.png) and row label in the table
    sources / params / subs / tran   as for KiSim.run(), applied to this case only
    analyses  functions f(r: Result) -> dict. Scalar values are metrics (table columns);
              array values (same length as r.t) are extra curves added to the case plot.
    limits    {metric: (op, value)} with op one of "<", "<=", ">", ">=", "==", or a function
              f(value) -> bool. Checked metrics get a ✓/✗ in the table.
    plot      rows for Result.plot(); None = no plot. Array curves from analyses are added.
    xlim      plot x-range in ms
    """
    name: str
    sources: dict | None = None
    params: dict | None = None
    subs: dict | None = None
    tran: str | None = None
    analyses: list = field(default_factory=list)
    limits: dict = field(default_factory=dict)
    plot: list | None = None
    xlim: tuple | None = None


_OPS = {"<": lambda a, b: a < b, "<=": lambda a, b: a <= b, ">": lambda a, b: a > b,
        ">=": lambda a, b: a >= b, "==": lambda a, b: a == b}


def _check(value, limit) -> bool:
    if callable(limit):
        return bool(limit(value))
    op, ref = limit
    return bool(_OPS[op](value, ref))


def run_cases(sim: "KiSim", cases: list, parallel: bool = True, workers: int | None = None,
              cache: bool = True, show: bool = False) -> dict:
    """Run cases (in parallel by default: each case is its own ngspice process), apply each
    case's analyses, check limits, save a plot per case to sim.workdir/<name>.png and print a
    summary table. A case whose simulation fails is reported, the others still run.
    Returns {name: {"result": Result | None, "metrics": dict, "ok": bool | None, "error": str}}."""
    from concurrent.futures import ThreadPoolExecutor
    import matplotlib
    if not show:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    def simulate(c):
        try:
            return c, sim.run(params=c.params, tag=c.name, cache=cache, subs=c.subs,
                              sources=c.sources, tran=c.tran), ""
        except Exception as e:                       # report, keep going with the other cases
            msg = str(e)
            m = re.search(r"Timestep too small[^\n]*|Error[^\n]*", msg)
            return c, None, (m.group(0) if m else msg.strip().splitlines()[-1] if msg.strip() else "failed")

    if parallel and len(cases) > 1:
        with ThreadPoolExecutor(max_workers=workers or len(cases)) as ex:
            done = list(ex.map(simulate, cases))
    else:
        done = [simulate(c) for c in cases]

    out = {}
    for c, r, err in done:
        metrics, curves, ok = {}, [], None
        if r is not None:
            for f in c.analyses:
                try:
                    res = f(r) or {}
                except Exception as e:
                    res = {f"{f.__name__} ERROR": f"{type(e).__name__}: {e}"}
                for k, v in res.items():
                    if np.ndim(v) == 1 and len(v) == len(r.t):
                        curves.append((k, np.asarray(v)))
                    else:
                        metrics[k] = v
            checks = [_check(metrics[k], lim) for k, lim in c.limits.items() if k in metrics]
            ok = all(checks) if checks else None
            if c.plot is not None or curves:
                fig, _ = r.plot(list(c.plot or []) + curves, title=c.name, xlim=c.xlim,
                                file=sim.workdir / f"{c.name}.png")
                if not show:
                    plt.close(fig)
        out[c.name] = {"result": r, "metrics": metrics, "ok": ok, "error": err}

    # summary table
    cols = list(dict.fromkeys(k for v in out.values() for k in v["metrics"]))
    def fmt(name, k):
        v = out[name]["metrics"].get(k)
        if v is None:
            return "-"
        s = f"{v:.4g}" if isinstance(v, (int, float, np.floating)) and not isinstance(v, bool) else str(v)
        lim = next((c.limits.get(k) for c in cases if c.name == name), None)
        return s + ((" ✓" if _check(v, lim) else " ✗") if lim is not None else "")
    w0 = max(len(n) for n in out) + 2
    widths = [max(len(k), *(len(fmt(n, k)) for n in out)) + 2 for k in cols]
    print("case".ljust(w0) + "".join(k.rjust(w) for k, w in zip(cols, widths)) + "  result")
    for n, v in out.items():
        res = "FAILED: " + v["error"] if v["error"] else {True: "pass", False: "FAIL", None: ""}[v["ok"]]
        print(n.ljust(w0) + "".join(fmt(n, k).rjust(w) for k, w in zip(cols, widths)) + "  " + res)
    if show:
        plt.show()
    return out


# --------------------------------------------------------------------------------------
# Stimulus building: piecewise waveforms and reusable scenario pieces.
# --------------------------------------------------------------------------------------
class Wave:
    """Piecewise-linear waveform built step by step; str(w) / w.pwl gives the PWL string.

        Wave(8).ramp(12, 1e-3).hold(1.5e-3).ramp(20, 0.5e-3).hold(1.5e-3).ramp(12, 0.5e-3)
        Wave(12).hold(3e-3).step(0, 1e-6).hold(2e-3).step(12, 1e-6)      # power cycle

    hold(d): stay for d seconds. ramp(v, d): go linearly to v over d seconds.
    step(v, rise=1e-6): fast edge to v. at(t): jump the clock to absolute time t (holding).
    After the last point the value is held forever (PWL semantics)."""

    def __init__(self, v0: float):
        self.pts = [(0.0, float(v0))]

    @property
    def t(self) -> float:
        return self.pts[-1][0]

    @property
    def v(self) -> float:
        return self.pts[-1][1]

    def hold(self, d: float) -> "Wave":
        self.pts.append((self.t + d, self.v))
        return self

    def ramp(self, v: float, d: float) -> "Wave":
        self.pts.append((self.t + d, float(v)))
        return self

    def step(self, v: float, rise: float = 1e-6) -> "Wave":
        return self.ramp(v, rise)

    def at(self, t: float) -> "Wave":
        if t > self.t:
            self.hold(t - self.t)
        return self

    @property
    def pwl(self) -> str:
        return pwl(*self.pts)

    def __str__(self) -> str:
        return self.pwl


def stim(*parts: dict, **kw) -> dict:
    """Merge scenario pieces into Case keyword arguments. Each piece is a dict with any of
    sources / params / subs / tran; later pieces override earlier ones per key.
        Case("ov_fast", **stim(supply_ov(rate=16e3), db_active(9)), analyses=[...])"""
    out: dict = {}
    for p in (*parts, kw):
        for k, v in p.items():
            if k in ("sources", "params", "subs") and isinstance(v, dict):
                merged = dict(out.get(k) or {})
                for kk, vv in v.items():          # params: merge per instance
                    if k == "params" and isinstance(vv, dict) and isinstance(merged.get(kk), dict):
                        merged[kk] = {**merged[kk], **vv}
                    else:
                        merged[kk] = vv
                out[k] = merged
            else:
                out[k] = v
    return out


# --------------------------------------------------------------------------------------
# Corners: expand one case into tolerance / parameter variants (worst-case analysis).
# --------------------------------------------------------------------------------------
_SI = {"f": 1e-15, "p": 1e-12, "n": 1e-9, "u": 1e-6, "m": 1e-3, "k": 1e3, "g": 1e9, "t": 1e12}


def spice_num(s) -> float:
    """'10.7k' -> 10700.0, '4.7n' -> 4.7e-9, '1meg' -> 1e6, numbers pass through."""
    if isinstance(s, (int, float)):
        return float(s)
    m = re.match(r"^\s*([-+]?[\d.]+(?:e[-+]?\d+)?)\s*(meg|[fpnumkgt])?", s.lower())
    if not m:
        raise ValueError(f"not a SPICE number: {s!r}")
    return float(m.group(1)) * (1e6 if m.group(2) == "meg" else _SI.get(m.group(2) or "", 1.0))


def tol(nominal, rel: float) -> list[str]:
    """[min, nom, max] for a value with relative tolerance: tol("10.7k", 0.01)."""
    n = spice_num(nominal)
    return [f"{n * (1 - rel):.6g}", f"{n:.6g}", f"{n * (1 + rel):.6g}"]


def corners(case: "Case", vary: dict, mode: str = "extremes") -> list:
    """Expand a Case into variants. vary maps what to change to a list of values:
        "R4": ["10.59k", "10.7k", "10.81k"]          two-terminal part value (via sources)
        "X1.PACT": ["9", "11"]                         instance parameter (via params)
        "V1": [Wave(...), Wave(...)]                   source waveform
    mode: "extremes"  every combination of each list's first and last value (2^n runs)
          "factorial" every combination of all listed values
          "oat"       one at a time: each value of each parameter, the others at their
                      middle (nominal) value
    Variant names: <case>__R4=10.81k__X1.PACT=11 (unique tags, safe as file names)."""
    import itertools
    keys = list(vary)
    vals = {k: list(v) for k, v in vary.items()}
    nom = {k: v[(len(v) - 1) // 2] for k, v in vals.items()}    # middle; first of an even list
    if mode == "extremes":
        combos = list(itertools.product(*[[v[0], v[-1]] if len(v) > 1 else v for v in vals.values()]))
    elif mode == "factorial":
        combos = list(itertools.product(*vals.values()))
    elif mode == "oat":
        combos, seen = [], set()
        for k in keys:
            for v in vals[k]:
                c = tuple(v if kk == k else nom[kk] for kk in keys)
                if c not in seen:
                    seen.add(c), combos.append(c)
    else:
        raise ValueError(f"mode {mode!r}")

    def label(v):
        s = v.pwl if isinstance(v, Wave) else str(v)
        return s if len(s) <= 12 else f"w{abs(hash(s)) % 10**6:06d}"

    out = []
    for combo in combos:
        sources = dict(case.sources or {})
        params = {k: dict(v) for k, v in (case.params or {}).items()}
        tag = [case.name]
        for k, v in zip(keys, combo):
            if "." in k:                                   # instance parameter
                inst, par = k.split(".", 1)
                params.setdefault(inst, {})[par] = str(v)
            else:
                sources[k] = str(v)
            tag.append(f"{k}={label(v)}")
        name = re.sub(r"[^\w.=+-]", "_", "__".join(tag))
        out.append(Case(name=name, sources=sources or None, params=params or None, subs=case.subs,
                        tran=case.tran, analyses=case.analyses, limits=case.limits,
                        plot=case.plot, xlim=case.xlim))
    return out


def worst(results: dict, metric: str, how: str = "max", prefix: str | None = None):
    """Worst corner for a metric from run_cases() output: (case name, value).
    how: "max" or "min". prefix: only cases whose name starts with it (one corner family)."""
    vals = [(n, v["metrics"][metric]) for n, v in results.items()
            if metric in v["metrics"] and (prefix is None or n.startswith(prefix))
            and isinstance(v["metrics"][metric], (int, float, np.floating))]
    if not vals:
        return None
    return (max if how == "max" else min)(vals, key=lambda x: x[1])