polars-waveform 0.1.0__py3-none-any.whl

polars_waveform/__init__.py

@@ -0,0 +1,49 @@
+"""Waveforms on Polars: curves and curve families with circuit-style measurements.
+
+A :class:`Waveform` is a value column over one or more index columns, backed by a lazy
+``pl.LazyFrame`` and materialized only when a value is needed. The last index column is the
+sweep (time, frequency, offset, ...); any others are groups (corners, Monte Carlo runs, measured
+operating points), so one Waveform can hold a whole family of curves::
+
+    import polars as pl
+    import polars_waveform as pw
+
+    w = pw.Waveform(pl.scan_parquet("lab.parquet"), "gain", index=["dut", "temp", "freq"])
+    w.bandwidth()                  # one row per dut x temp
+    w - w.mean()                   # per-curve results broadcast back over the curves
+    w.leaf(dut="A1", temp=25.0)    # one curve
+
+Measurements follow OCEAN semantics (``xmax`` is the x at the largest y; ``cross`` counts edges
+from 1); the OCEAN and pycircuit spellings (``dB20``, ``unityGainFreq``, ``IIP3``, ...) are aliases.
+Complex data is ``Struct{re, im}``; :mod:`.cx` registers ``pl.col(...).cx`` for it.
+:func:`from_nested` builds families from nested channel columns, :meth:`Waveform.from_arrays`
+from numpy grids, and :class:`ResultSource` is the shape shared by result objects
+(``polars_psf.Result`` implements it). :class:`WaveformBase` is the contract other waveform
+kinds implement; :class:`PandasWaveform` is one for values Polars cannot compute with (sympy
+expressions, ...), and :func:`from_arrays` picks the kind the values need.
+
+Modules: :mod:`.waveform` (the class), :mod:`.functions` (calculator functions), :mod:`.cx`
+(complex helpers), :mod:`.nested`, :mod:`.source`, :mod:`.pandas_waveform`.
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from . import cx
+from .base import WaveformBase
+from .functions import *  # noqa: F403
+from .functions import __all__ as _functions
+from .nested import from_nested
+from .pandas_waveform import PandasWaveform, from_arrays
+from .source import ResultSource
+from .waveform import Waveform, either, falling, raising
+
+__all__ = [
+    "PandasWaveform", "ResultSource", "Waveform", "WaveformBase", "cx", "either", "falling", "from_arrays",
+    "from_nested", "raising", *_functions,
+]  # fmt: skip
+
+try:
+    __version__ = _version("polars-waveform")
+except PackageNotFoundError:  # running from a source tree without an install
+    __version__ = "0+unknown"

polars_waveform/base.py

@@ -0,0 +1,176 @@
+"""The contract shared by all waveform kinds.
+
+:class:`~polars_waveform.Waveform` stores numeric data in Polars (lazy, families as group
+columns). Other kinds implement the same interface with their own storage, for example a
+symbolic waveform holding sympy expressions in numpy object arrays: Polars can store Python
+objects but has no kernels to compute with them, so such data needs its own container.
+
+A subclass implements storage, construction and elementwise math (the abstract methods), and
+:meth:`WaveformBase.numeric` to turn itself into a numeric :class:`~polars_waveform.Waveform`.
+Measurements it does not implement itself (``bandwidth``, ``cross``, ``rise_time``, ...) are
+inherited from this class and run on ``numeric()``, so they work on any kind once its values are
+numbers.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+__all__ = ["WaveformBase"]
+
+# measurements that need numbers: inherited by other kinds, evaluated on numeric()
+_NUMERIC = (
+    "value", "ymax", "ymin", "xmax", "xmin", "argmax", "argmin", "average", "mean", "rms", "stddev",
+    "cross", "frequency", "period", "rise_time", "fall_time", "slew_rate", "overshoot", "settling_time",
+    "bandwidth", "unity_gain_frequency", "phase_margin", "gain_margin", "integ", "iinteg", "deriv",
+    "clip", "dft", "to_polars", "plot", "semilogx", "semilogy", "loglog", "stem", "bode",
+)  # fmt: skip
+
+
+class WaveformBase(ABC):
+    """A value over one or more index columns: groups (one curve per combination) and a sweep.
+
+    The interface every waveform kind provides: names and units, conversion from and to numpy
+    arrays (:meth:`from_arrays`, :meth:`to_arrays`), arithmetic, the elementwise functions
+    (``db20``, ``phase``, ...), and the measurements, which by default evaluate on
+    :meth:`numeric`.
+    """
+
+    # --- storage and names (abstract) -------------------------------------------------------
+    @classmethod
+    @abstractmethod
+    def from_arrays(cls, x, y, xlabels=None, ylabel=None, xunits=None, yunit=None) -> WaveformBase:
+        """A waveform from numpy arrays: one x array per sweep axis (the last one is the sweep,
+        the others become groups) and ``y`` with shape ``(len(x0), len(x1), ...)``, or ragged
+        object arrays (see :meth:`polars_waveform.Waveform.from_arrays`)."""
+
+    @abstractmethod
+    def to_arrays(self):
+        """``(xs, y)`` as numpy arrays in the layout :meth:`from_arrays` accepts."""
+
+    @abstractmethod
+    def numeric(self):
+        """This waveform as a numeric :class:`~polars_waveform.Waveform`."""
+
+    @property
+    @abstractmethod
+    def index(self) -> list[str]:
+        """Index column names: the groups followed by the sweep."""
+
+    @property
+    def groups(self) -> list[str]:
+        """Index columns other than the sweep (one curve per combination)."""
+        return [c for c in self.index if c != self.xname]
+
+    @property
+    @abstractmethod
+    def xname(self) -> str: ...
+
+    @property
+    @abstractmethod
+    def yname(self) -> str: ...
+
+    @property
+    @abstractmethod
+    def xunit(self): ...
+
+    @property
+    @abstractmethod
+    def yunit(self): ...
+
+    # --- arithmetic (abstract core, shared operators) ------------------------------------------
+    @abstractmethod
+    def _binop(self, other, op: str, *, reverse: bool = False) -> WaveformBase:
+        """``self <op> other`` (``other <op> self`` if ``reverse``) for op in + - * / < > <= >=."""
+
+    def __add__(self, other):
+        return self._binop(other, "+")
+
+    def __radd__(self, other):
+        return self._binop(other, "+", reverse=True)
+
+    def __sub__(self, other):
+        return self._binop(other, "-")
+
+    def __rsub__(self, other):
+        return self._binop(other, "-", reverse=True)
+
+    def __mul__(self, other):
+        return self._binop(other, "*")
+
+    def __rmul__(self, other):
+        return self._binop(other, "*", reverse=True)
+
+    def __truediv__(self, other):
+        return self._binop(other, "/")
+
+    def __rtruediv__(self, other):
+        return self._binop(other, "/", reverse=True)
+
+    def __lt__(self, other):
+        return self._binop(other, "<")
+
+    def __le__(self, other):
+        return self._binop(other, "<=")
+
+    def __gt__(self, other):
+        return self._binop(other, ">")
+
+    def __ge__(self, other):
+        return self._binop(other, ">=")
+
+    # --- elementwise (abstract) ----------------------------------------------------------------
+    @abstractmethod
+    def __neg__(self): ...
+
+    @abstractmethod
+    def __abs__(self): ...
+
+    @abstractmethod
+    def __pow__(self, other): ...
+
+    @abstractmethod
+    def real(self): ...
+
+    @abstractmethod
+    def imag(self): ...
+
+    @abstractmethod
+    def conj(self): ...
+
+    @abstractmethod
+    def phase(self, deg: bool = True): ...
+
+    @abstractmethod
+    def db10(self): ...
+
+    @abstractmethod
+    def db20(self): ...
+
+    def __pos__(self):
+        return self
+
+    def abs(self):
+        """Magnitude (``abs(w)``)."""
+        return abs(self)
+
+    mag = abs
+
+    def conjugate(self):
+        return self.conj()
+
+
+def _numeric_method(name: str):
+    def method(self, *args, **kwargs):
+        num = self.numeric()
+        if type(num) is type(self):  # a numeric kind must implement the measurement itself
+            raise NotImplementedError(f"{type(self).__name__}.{name}()")
+        return getattr(num, name)(*args, **kwargs)
+
+    method.__name__ = name
+    method.__doc__ = f"``{name}`` evaluated on :meth:`numeric` (see ``polars_waveform.Waveform.{name}``)."
+    return method
+
+
+for _name in _NUMERIC:
+    setattr(WaveformBase, _name, _numeric_method(_name))

polars_waveform/cx.py

@@ -0,0 +1,145 @@
+"""Complex-number helpers for Struct{re, im} columns.
+
+Registers a ``cx`` expression namespace on import::
+
+    pl.col("out").cx.abs()      # magnitude
+    pl.col("out").cx.db20()     # 20*log10(|z|)
+    pl.col("out").cx.phase()    # degrees (deg=False for radians)
+
+The same functions are available as ``pw.cx.abs("out")`` etc. (accept a column name or expression).
+Arithmetic on two complex columns: ``cx.mul("a", "b")``, ``cx.div``, ``cx.add``, ``cx.sub`` and
+``cx.pow(x, n)`` for a real exponent; ``cx.from_real(x)`` lifts a real column to ``Struct{re, im}``.
+"""
+
+import math
+
+import polars as pl
+
+__all__ = [
+    "abs",
+    "add",
+    "complex",
+    "conj",
+    "db10",
+    "db20",
+    "div",
+    "from_real",
+    "im",
+    "mul",
+    "phase",
+    "pow",
+    "re",
+    "sub",
+]
+
+
+def _e(x) -> pl.Expr:
+    return pl.col(x) if isinstance(x, str) else x
+
+
+def re(x) -> pl.Expr:
+    return _e(x).struct.field("re")
+
+
+def im(x) -> pl.Expr:
+    return _e(x).struct.field("im")
+
+
+def abs(x) -> pl.Expr:
+    return (re(x) ** 2 + im(x) ** 2).sqrt()
+
+
+def db20(x) -> pl.Expr:
+    return 20 * abs(x).log10()
+
+
+def db10(x) -> pl.Expr:
+    """10*log10(|z|), for power quantities."""
+    return 10 * abs(x).log10()
+
+
+def phase(x, deg: bool = True) -> pl.Expr:
+    p = pl.arctan2(im(x), re(x))
+    return p * (180 / math.pi) if deg else p
+
+
+def conj(x) -> pl.Expr:
+    return pl.struct(re(x).alias("re"), (-im(x)).alias("im"))
+
+
+def complex(re_expr, im_expr) -> pl.Expr:
+    """Build a Struct{re, im} column."""
+    return pl.struct(_e(re_expr).alias("re"), _e(im_expr).alias("im"))
+
+
+def from_real(x) -> pl.Expr:
+    """A real column as ``Struct{re, im}`` with ``im = 0``."""
+    return complex(x, _e(x) * 0.0)
+
+
+def add(a, b) -> pl.Expr:
+    return complex(re(a) + re(b), im(a) + im(b))
+
+
+def sub(a, b) -> pl.Expr:
+    return complex(re(a) - re(b), im(a) - im(b))
+
+
+def mul(a, b) -> pl.Expr:
+    ar, ai, br, bi = re(a), im(a), re(b), im(b)
+    return complex(ar * br - ai * bi, ar * bi + ai * br)
+
+
+def div(a, b) -> pl.Expr:
+    ar, ai, br, bi = re(a), im(a), re(b), im(b)
+    den = br * br + bi * bi
+    return complex((ar * br + ai * bi) / den, (ai * br - ar * bi) / den)
+
+
+def pow(x, n: float) -> pl.Expr:
+    """``x ** n`` for a real exponent (principal branch)."""
+    mag = abs(x) ** n
+    th = pl.arctan2(im(x), re(x)) * n
+    return complex(mag * th.cos(), mag * th.sin())
+
+
+@pl.api.register_expr_namespace("cx")
+class _CxNamespace:
+    def __init__(self, expr: pl.Expr):
+        self._e = expr
+
+    def re(self) -> pl.Expr:
+        return re(self._e)
+
+    def im(self) -> pl.Expr:
+        return im(self._e)
+
+    def abs(self) -> pl.Expr:
+        return abs(self._e)
+
+    def db20(self) -> pl.Expr:
+        return db20(self._e)
+
+    def db10(self) -> pl.Expr:
+        return db10(self._e)
+
+    def phase(self, deg: bool = True) -> pl.Expr:
+        return phase(self._e, deg)
+
+    def conj(self) -> pl.Expr:
+        return conj(self._e)
+
+    def mul(self, other) -> pl.Expr:
+        return mul(self._e, other)
+
+    def div(self, other) -> pl.Expr:
+        return div(self._e, other)
+
+    def add(self, other) -> pl.Expr:
+        return add(self._e, other)
+
+    def sub(self, other) -> pl.Expr:
+        return sub(self._e, other)
+
+    def pow(self, n: float) -> pl.Expr:
+        return pow(self._e, n)

polars_waveform/functions.py

@@ -0,0 +1,325 @@
+"""Calculator functions: ``db20(w)``, ``bandwidth(w, 3, "low")``, ``rise_time(w)``.
+
+Each forwards to the :class:`Waveform` method of the same name, so it works per curve of a
+family (corners, Monte Carlo) and returns one row per curve there. The semantics follow OCEAN:
+``xmax`` is the x at the largest y, and ``cross`` counts edges from 1 with
+``"rising"``/``"falling"``/``"either"``. The OCEAN names (``dB20``, ``unityGainFreq``,
+``riseTime``, ``leafValue``, ...) and pycircuit names (``unityGainFrequency``, ``IIP3``, ...) are
+aliases.
+
+Arguments follow OCEAN's order where it has positional ones; OCEAN's ``initType``/``finalType``
+flags (initial value given as an x instead of a y) are not supported: pass y values.
+"""
+
+from __future__ import annotations
+
+import math
+
+from .base import WaveformBase
+from .plot import compression_plot, set_plot_backend
+from .waveform import Waveform
+
+__all__ = [
+    # calculator functions
+    "db10", "db20", "mag", "phase", "phase_deg", "phase_rad", "real", "imag", "conjugate",
+    "value", "ymax", "ymin", "xmax", "xmin", "average", "rms", "stddev",
+    "cross", "deriv", "integ", "iinteg", "clip", "dft", "leaf_value", "log10", "exp", "sqrt",
+    "bandwidth", "unity_gain_frequency", "phase_margin", "gain_margin",
+    "rise_time", "fall_time", "slew_rate", "overshoot", "settling_time", "delay", "frequency", "period",
+    "im2", "im3", "iip2", "iip3", "calc_extrapolation_line", "compression_point", "compression_plot",
+    "set_plot_backend",
+    # OCEAN aliases
+    "dB10", "dB20", "phaseDeg", "phaseRad", "leafValue", "unityGainFreq", "phaseMargin", "gainMargin",
+    "riseTime", "fallTime", "slewRate", "settlingTime",
+    # pycircuit aliases
+    "unityGainFrequency", "IM2", "IM3", "IIP2", "IIP3",
+]
+
+
+# --- elementwise: also on plain numbers and numpy arrays --------------------------------------
+def _num(w):
+    """True if ``w`` is a plain number or array (not a waveform)."""
+    return not isinstance(w, WaveformBase)
+
+
+def db20(w):
+    """``20*log10(|w|)``; works on waveforms, numbers and numpy arrays."""
+    if _num(w):
+        import numpy as np
+
+        return 20.0 * np.log10(np.abs(w))
+    return w.db20()
+
+
+def db10(w):
+    """``10*log10(|w|)``, for power quantities."""
+    if _num(w):
+        import numpy as np
+
+        return 10.0 * np.log10(np.abs(w))
+    return w.db10()
+
+
+def mag(w):
+    if _num(w):
+        import numpy as np
+
+        return np.abs(w)
+    return w.mag()
+
+
+def phase(w, deg: bool = True):
+    """Phase in degrees (radians with ``deg=False``)."""
+    if _num(w):
+        import numpy as np
+
+        return np.angle(w, deg=deg)
+    return w.phase(deg)
+
+
+def phase_deg(w):
+    return phase(w, deg=True)
+
+
+def phase_rad(w):
+    return phase(w, deg=False)
+
+
+def real(w):
+    if _num(w):
+        import numpy as np
+
+        return np.real(w)
+    return w.real()
+
+
+def imag(w):
+    if _num(w):
+        import numpy as np
+
+        return np.imag(w)
+    return w.imag()
+
+
+def conjugate(w):
+    if _num(w):
+        import numpy as np
+
+        return np.conjugate(w)
+    return w.conj()
+
+
+def log10(w):
+    if _num(w):
+        import numpy as np
+
+        return np.log10(w)
+    return w.log10()
+
+
+def exp(w):
+    if _num(w):
+        import numpy as np
+
+        return np.exp(w)
+    return w.exp()
+
+
+def sqrt(w):
+    if _num(w):
+        import numpy as np
+
+        return np.sqrt(w)
+    return w.sqrt()
+
+
+def deriv(w: Waveform) -> Waveform:
+    return w.deriv()
+
+
+def iinteg(w: Waveform) -> Waveform:
+    return w.iinteg()
+
+
+def clip(w: Waveform, xfrom: float, xto: float | None = None) -> Waveform:
+    return w.clip(xfrom, xto)
+
+
+def dft(w: Waveform) -> Waveform:
+    return w.dft()
+
+
+def leaf_value(w: Waveform, **values) -> Waveform:
+    """One leaf of a family: ``leaf_value(w, temp=27, rval=1000)``."""
+    return w.leaf(**values)
+
+
+# --- measurements (scalar per curve) ----------------------------------------------------------
+def value(w: Waveform, x):
+    return w.value(x)
+
+
+def ymax(w: Waveform, axis=-1):
+    return w.ymax(axis)
+
+
+def ymin(w: Waveform, axis=-1):
+    return w.ymin(axis)
+
+
+def xmax(w: Waveform):
+    """x where y is largest."""
+    return w.xmax()
+
+
+def xmin(w: Waveform):
+    """x where y is smallest."""
+    return w.xmin()
+
+
+def average(w: Waveform, axis=-1):
+    return w.average(axis)
+
+
+def rms(w: Waveform, axis=-1):
+    return w.rms(axis)
+
+
+def stddev(w: Waveform, axis=-1):
+    return w.stddev(axis)
+
+
+def integ(w: Waveform, xfrom: float | None = None, xto: float | None = None):
+    return w.integ(xfrom, xto)
+
+
+def cross(w: Waveform, threshold: float = 0.0, edge: int = 1, type: str = "either"):
+    """x of the ``edge``-th crossing of ``threshold`` (from 1; negative counts from the end)."""
+    return w.cross(threshold, edge, type)
+
+
+def bandwidth(w: Waveform, db: float = 3.0, type: str = "low"):
+    return w.bandwidth(db, type)
+
+
+def unity_gain_frequency(w: Waveform):
+    return w.unity_gain_frequency()
+
+
+def phase_margin(w: Waveform):
+    return w.phase_margin()
+
+
+def gain_margin(w: Waveform):
+    return w.gain_margin()
+
+
+def rise_time(w: Waveform, initial=None, final=None, theta1: float = 10.0, theta2: float = 90.0):
+    """Rise time (OCEAN ``riseTime(wave initVal nil finalVal nil theta1 theta2)``)."""
+    return w.rise_time(theta1, theta2, initial, final)
+
+
+def fall_time(w: Waveform, initial=None, final=None, theta1: float = 10.0, theta2: float = 90.0):
+    return w.fall_time(theta1, theta2, initial, final)
+
+
+def slew_rate(w: Waveform, initial=None, final=None, theta1: float = 10.0, theta2: float = 90.0):
+    return w.slew_rate(theta1, theta2, initial, final)
+
+
+def overshoot(w: Waveform, initial=None, final=None):
+    return w.overshoot(initial, final)
+
+
+def settling_time(w: Waveform, initial=None, final=None, tolerance: float = 1.0):
+    return w.settling_time(tolerance, initial, final)
+
+
+def delay(
+    wf1: Waveform,
+    wf2: Waveform,
+    value1: float = 0.0,
+    value2: float | None = None,
+    edge1: str = "either",
+    nth1: int = 1,
+    edge2: str | None = None,
+    nth2: int = 1,
+):
+    """x of ``wf2``'s ``nth2`` crossing of ``value2`` minus that of ``wf1`` (OCEAN ``delay``
+    keyword names)."""
+    return wf1.delay(wf2, value1, value2, nth1, nth2, edge1, edge2)
+
+
+def frequency(w: Waveform, threshold: float | None = None):
+    return w.frequency(threshold)
+
+
+def period(w: Waveform, threshold: float | None = None):
+    return w.period(threshold)
+
+
+# --- RF (pycircuit) ---------------------------------------------------------------------------
+def im3(w: Waveform, fund1: float, fund2: float, fund0=None) -> float:
+    """Third-order intermodulation tone at ``fund1 + 2*fund2`` of an output spectrum.
+
+    ``fund0`` (the LO of a mixer) is accepted for pycircuit compatibility; output tones are
+    measured at IF, so it does not shift them (only :func:`iip3` uses it, for the input)."""
+    return abs(w).value(fund1 + 2 * fund2)
+
+
+def im2(w: Waveform, fund1: float, fund2: float, fund0=None) -> float:
+    """Second-order intermodulation tone at ``fund1 + fund2`` of an output spectrum (``fund0``:
+    see :func:`im3`)."""
+    return abs(w).value(fund1 + fund2)
+
+
+def iip3(output: Waveform, input: Waveform, fund1: float, fund2: float, fund0=None) -> float:
+    """Input-referred third-order intercept point."""
+    s = abs(output)
+    if fund0 is None:
+        gain = (s / abs(input)).value(fund1)
+    else:
+        gain = s.value(abs(fund1)) / abs(input).value(abs(abs(fund1) + fund0))
+    return math.sqrt(s.value(abs(fund1)) * s.value(abs(fund2)) ** 2 / s.value(fund1 + 2 * fund2)) / gain
+
+
+def iip2(output: Waveform, input: Waveform, fund1: float, fund2: float, fund0=None) -> float:
+    """Input-referred second-order intercept point."""
+    s = abs(output)
+    if fund0 is None:
+        gain = (s / abs(input)).value(fund1)
+    else:
+        gain = s.value(abs(fund1)) / abs(input).value(abs(abs(fund1) + fund0))
+    return s.value(abs(fund1)) * s.value(abs(fund2)) / s.value(fund1 + fund2) / gain
+
+
+def calc_extrapolation_line(
+    w_db: Waveform, slope: float = 1.0, extrapolation_point: float | None = None, **_
+) -> Waveform:
+    """Linear (in dB) extrapolation of ``w_db`` through its low end (smallest x)."""
+    if extrapolation_point is None:
+        extrapolation_point = w_db.x.min()
+    m = w_db.value(extrapolation_point) - slope * extrapolation_point
+    return m + slope * w_db.xval()
+
+
+def compression_point(
+    w_db: Waveform, slope: float = 1.0, compression: float = 1.0, extrapolation_point=None, **_
+) -> float:
+    """Input-referred compression point: where the response falls ``compression`` dB below the
+    extrapolated small-signal line."""
+    line = calc_extrapolation_line(w_db, slope, extrapolation_point)
+    return (line - w_db).cross(compression)
+
+
+# --- aliases ----------------------------------------------------------------------------------
+# OCEAN
+dB20, dB10 = db20, db10
+phaseDeg, phaseRad = phase_deg, phase_rad
+leafValue = leaf_value
+unityGainFreq = unity_gain_frequency
+phaseMargin, gainMargin = phase_margin, gain_margin
+riseTime, fallTime, slewRate, settlingTime = rise_time, fall_time, slew_rate, settling_time
+# pycircuit
+unityGainFrequency = unity_gain_frequency
+IM2, IM3, IIP2, IIP3 = im2, im3, iip2, iip3