s_transforms ¤

def fdomain_component(ports: tuple[str, ...]) -> Any:
    """Decorator for frequency-domain (admittance) circuit components.

    Compiles the decorated admittance function into a
    :class:`~circulax.components.base_component.CircuitComponent` subclass
    that is evaluated in the frequency domain rather than the time domain.
    The component:

    - **DC analysis** — evaluated at ``f = 0`` Hz (e.g. skin-effect reduces to
      ``R₀`` at DC).
    - **Harmonic Balance** — evaluated at each harmonic frequency ``k · f₀``,
      contributing ``Y(k · f₀) @ V_k`` directly to the frequency-domain residual.
    - **Transient simulation** — raises :exc:`RuntimeError` at setup time
      (time-domain convolution not supported).

    The decorated function must accept ``f`` as its first positional argument
    (frequency in Hz) followed by any number of keyword parameters with defaults.
    It must return a square Y-matrix of shape ``(n_ports, n_ports)`` with
    ``dtype=complex128``.

    Args:
        ports: Ordered tuple of port names matching the netlist connection keys.

    Returns:
        A decorator that accepts an admittance function and returns a
        :class:`~circulax.components.base_component.CircuitComponent` subclass.

    Example::

        @fdomain_component(ports=("p1", "p2"))
        def SkinEffectResistor(f: float, R0: float = 1.0, a: float = 0.1):
            \"\"\"Z(f) = R0 + a * sqrt(|f|)\"\"\"
            Z = R0 + a * jnp.sqrt(jnp.abs(f) + 1e-30)
            Y = 1.0 / Z
            return jnp.array([[Y, -Y], [-Y, Y]], dtype=jnp.complex128)

    """
    return lambda fn: _build_fdomain_component(fn, ports)

def fdomain_model(fn: callable, *, ports: tuple[str, ...], z0: complex = 1.0 + 1e-12j) -> callable:
    """Adapt a frequency-domain admittance function into a SAX-shaped S-parameter model.

    Bridges :func:`fdomain_component`-style physics (``fn(f, **params) ->
    Y-matrix``) into an ordinary SAX model function (``(**params) ->
    sax.SDict``) via :func:`y_to_s`. The result is meant to be passed to
    :func:`sax_component`, e.g.::

        @sax_component
        def my_rf_model(f=1e9, R0=50.0): ...  # returns a Y-matrix


        RfComp = sax_component(fdomain_model(my_rf_model, ports=("p1", "p2")))

    Once wrapped this way, the resulting component's frequency argument
    becomes an ordinary settable parameter (like ``wl`` on a photonic SAX
    model) rather than the circuit-level AC-sweep frequency: it is evaluated
    once per compile/param-update, not re-evaluated per point in a
    :meth:`~circulax.circuit.Circuit.sp` sweep. This makes it usable in the
    all-SAX, source-free :meth:`~circulax.circuit.Circuit.dc` S-parameter
    path (where the frequency becomes a batchable parameter, e.g.
    ``circuit.dc(f=jnp.linspace(...))``), but it is **not** a drop-in
    replacement for ``@fdomain_component`` inside a circuit that also has
    reactive elements (capacitors/inductors) or sources: there, the
    component's own frequency dependence is frozen at whatever value its
    parameter was last set to, decoupled from the ``freqs`` sweep argument
    driving the rest of the circuit. To explicitly restore live sweep
    semantics, wrap it as ``sax_component(fdomain_model(...),
    frequency_param="f")``; that creates a frequency-domain Circulax
    component and is consequently no longer eligible for SAX-circuit export.

    Note:
        Compiling the result with ``is_complex=True`` may emit a spurious
        "uses non-holomorphic operations: conj" warning from
        :func:`~circulax.circuit.compile_circuit`. This is a known false
        positive: ``s_to_y``/``y_to_s`` are individually ``jax.jit``-ed, so
        their default ``z0`` argument is traced as an abstract value inside
        that nested call, which trips the checker's literal-vs-traced
        heuristic. It does not indicate an actual bug — ``fn`` never
        receives port signals (only ``f`` and parameters), so its
        contribution is provably linear in the port voltages regardless.

    Args:
        fn: Admittance function with signature ``fn(f, **params) ->
            Y-matrix``, matching :func:`fdomain_component`'s contract (``f``
            first, every other parameter with a default).
        ports: Ordered tuple of port names, matching the shape of ``fn``'s
            returned Y-matrix.
        z0: Reference impedance used for the Y-to-S conversion. Defaults to
            SAX's own convention so the wrapped model composes correctly
            with other SAX models under :func:`sax_component`.

    Returns:
        A callable with the same signature as ``fn`` that returns a
        ``sax.SDict`` instead of a Y-matrix.

    Raises:
        TypeError: If ``fn``'s signature doesn't start with ``f`` or any
            other parameter lacks a default.

    """
    sig = inspect.signature(fn)
    params_list = list(sig.parameters.values())
    if not params_list or params_list[0].name != "f":
        msg = f"fdomain_model function '{fn.__name__}' must have 'f' as its first argument."
        raise TypeError(msg)
    for p in params_list[1:]:
        if p.default is inspect.Parameter.empty:
            msg = f"Parameter '{p.name}' in '{fn.__name__}' must have a default value."
            raise TypeError(msg)

    f_name = params_list[0].name

    def sax_model(*args, **kwargs) -> dict:  # noqa: ANN002, ANN003
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        arguments = dict(bound.arguments)
        f = arguments.pop(f_name)
        Y = fn(f, **arguments)
        S = y_to_s(Y, z0=z0)
        return {(a, b): S[i, j] for i, a in enumerate(ports) for j, b in enumerate(ports)}

    sax_model.__name__ = fn.__name__
    sax_model.__doc__ = fn.__doc__
    sax_model.__signature__ = sig
    # Keep enough provenance for ``sax_component(..., frequency_param="f")``
    # to recover the original admittance function.  This is deliberately
    # private metadata: without the explicit ``frequency_param`` opt-in the
    # returned callable remains a completely ordinary SAX model.
    sax_model._fdomain_admittance = fn
    sax_model._fdomain_ports = ports
    sax_model._fdomain_frequency_param = f_name
    return sax_model

@jax.jit
def s_to_y(S: jax.Array, z0: complex = 1.0 + 1e-12j) -> jax.Array:
    """Convert an S-parameter matrix to an admittance (Y) matrix.

    Kurokawa power-wave form: ``Y = (I - S) (z0 S + z0* I)^-1``. Reduces to
    ``(1/z0) (I - S) (I + S)^-1`` for real ``z0``. A small ``Im(z0)`` keeps
    the inverse well-conditioned when ``S`` has eigenvalues at -1 (e.g. an
    ideal lossless symmetric splitter).
    """
    n = S.shape[-1]
    eye = jnp.eye(n, dtype=jnp.complex128)
    Sc = S.astype(jnp.complex128)
    z0c = jnp.asarray(z0, dtype=jnp.complex128)
    if n == 2:
        # Two-port models can export to StableHLO/IREE without a host LAPACK
        # custom call. Preserve the Kurokawa normalization, including complex z0.
        a, b = z0c * Sc[..., 0, 0] + jnp.conj(z0c), z0c * Sc[..., 0, 1]
        c, d = z0c * Sc[..., 1, 0], z0c * Sc[..., 1, 1] + jnp.conj(z0c)
        det = a * d - b * c
        n00, n01 = 1 - Sc[..., 0, 0], -Sc[..., 0, 1]
        n10, n11 = -Sc[..., 1, 0], 1 - Sc[..., 1, 1]
        return jnp.stack((
            jnp.stack(((n00 * d - n01 * c) / det, (-n00 * b + n01 * a) / det), axis=-1),
            jnp.stack(((n10 * d - n11 * c) / det, (-n10 * b + n11 * a) / det), axis=-1),
        ), axis=-2)
    M = z0c * Sc + jnp.conj(z0c) * eye
    return jnp.linalg.solve(M.swapaxes(-1, -2), (eye - Sc).swapaxes(-1, -2)).swapaxes(-1, -2)

def sax_component(
    fn: callable,
    *,
    name: str | None = None,
    frequency_param: str | None = None,
) -> callable:
    """Decorator to convert a SAX model function into a circulax component.

    Inspects ``fn`` at decoration time to discover its port interface via a
    dry run, then wraps its S-matrix output in an admittance-based physics
    function compatible with the circulax nodal solver.

    The conversion proceeds in three stages:

    1. **Discovery** — ``fn`` is called once with its default (or dummy)
       parameter values and :func:`sax.get_ports` extracts the sorted port
       names from the resulting S-parameter dict. The same dry-run matrix is
       checked via :func:`_needs_wave_stamp` to decide, once per wrapped
       model, which physics stamp this component type gets.
    2. **Physics wrapper** — a closure is built that calls ``fn`` at runtime
       and converts the S-dict to a dense matrix via :func:`sax.sdense`.
       Components that passed the dry-run conditioning check use
       :func:`s_to_y` directly and return ``I = Y @ V``. Components flagged
       as near-lossless instead get one auxiliary "incident wave" state per
       port (``wave_<port>``) and are stamped with ``b = S @ a``,
       ``I = a - b``, ``a + b - V = 0`` — algebraically equivalent to the
       Y-matrix form (for z0=1, solving the constraint recovers
       ``I = (I-S)(I+S)^-1 V``) but never inverts ``S`` directly, so the
       ill-conditioning is resolved by the global sparse solve instead of a
       local matrix inverse.
    3. **Component registration** — the wrapper is passed to
       :func:`~circulax.components.base_component.component` with the
       discovered ports (and, for wave-stamped components, the auxiliary
       states), producing a :class:`~circulax.components.base_component.CircuitComponent`
       subclass.

    ``fn`` may be a plain function or a :class:`functools.partial` wrapping
    one (SAX PDKs typically use partials to bind fab-specific defaults).
    Partials are unwrapped to the innermost callable for ``__name__`` /
    ``__doc__`` recovery; :func:`inspect.signature` handles the parameter
    reduction.

    Args:
        fn: A SAX model function whose keyword arguments are scalar
            parameters and whose return value is a SAX S-parameter dict.
            All parameters must have defaults, or will be substituted with
            ``1.0`` during the dry run.
        name: Optional override for the resulting class name. Useful when
            wrapping :class:`functools.partial` objects where several
            partials share the same underlying ``__name__`` — e.g.
            ``{key: sax_component(val, name=key) for key, val in pdk.items()}``.
        frequency_param: Optional SAX parameter that is driven by
            :meth:`~circulax.circuit.Circuit.sp`'s frequency at every sweep
            point. This creates a Circulax frequency-domain component rather
            than a normal static SAX wrapper, so it cannot be exported via
            :meth:`~circulax.circuit.Circuit.to_sax_circuit`. Use this for an
            explicit round trip such as
            ``sax_component(fdomain_model(y_fn, ports=("p1", "p2")),
            frequency_param="f")``. For models made by
            :func:`fdomain_model`, the original Y function is recovered
            directly; arbitrary SAX models are converted from S to Y at each
            sweep point.

    Returns:
        A :class:`~circulax.components.base_component.CircuitComponent`
        subclass named after ``name`` (if given) else after the unwrapped
        function.

    Raises:
        RuntimeError: If the dry run fails for any reason.

    """
    if frequency_param is not None:
        return _build_frequency_sax_component(fn, frequency_param=frequency_param, name=name)

    sig = inspect.signature(fn)
    base_fn = _unwrap(fn)
    cls_name = name if name is not None else getattr(base_fn, "__name__", "SaxComponent")
    defaults = {
        param.name: param.default if param.default is not inspect.Parameter.empty else 1.0 for param in sig.parameters.values()
    }

    try:
        dummy_s_dict = fn(**defaults)
        detected_ports = get_ports(dummy_s_dict)
        dummy_s_matrix, _ = sdense(dummy_s_dict)
    except Exception as exc:
        msg = f"Failed to dry-run SAX component '{cls_name}': {exc}"
        raise RuntimeError(msg) from exc

    use_wave_stamp = _needs_wave_stamp(dummy_s_matrix)

    # Component variables use attribute access, which requires every port
    # name to be a valid Python identifier. Some SAX PDKs label
    # ports numerically ('1', '2'); coerce those to identifiers while keeping
    # the index ordering.
    raw_to_sanitized = {str(p): _sanitize_port(p) for p in detected_ports}
    sanitized_to_raw: dict[str, list[str]] = {}
    for raw, sanitized in raw_to_sanitized.items():
        sanitized_to_raw.setdefault(sanitized, []).append(raw)
    port_names = tuple(raw_to_sanitized[str(p)] for p in detected_ports)
    aux_state_names = tuple(f"wave_{p}" for p in port_names) if use_wave_stamp else ()

    def physics_wrapper(signals: Signals, **kwargs) -> tuple[dict, dict]:  # noqa: ANN003
        s_dict = fn(**kwargs)
        # `sdense` returns the S-matrix *and* a port_map keyed in dict-insertion
        # order (dict[raw_port, matrix_row_index]). `get_ports` returns ports
        # sorted alphabetically — the two can differ (e.g. an MMI2x2 SDict
        # built {(o1,o1):…,(o3,o3):…,(o4,o4):…,(o2,o2):…}). We must use the
        # port_map to line up v_vec with the rows/cols of the matrix; using the
        # sorted order silently scrambles coupling between non-adjacent rows.
        s_matrix, port_map = sdense(s_dict)
        matrix_order = sorted(port_map, key=port_map.get)  # raw names, matrix-row order
        sanitized_in_order = [_sanitize_port(p) for p in matrix_order]
        v_vec = jnp.array([getattr(signals, p) for p in sanitized_in_order], dtype=jnp.complex128)

        if use_wave_stamp:
            a_vec = jnp.array([getattr(signals, f"wave_{p}") for p in sanitized_in_order], dtype=jnp.complex128)
            b_vec = s_matrix.astype(jnp.complex128) @ a_vec
            i_vec = a_vec - b_vec
            constraints = a_vec + b_vec - v_vec
            f_dict = {p: i_vec[k] for k, p in enumerate(sanitized_in_order)}
            f_dict.update({f"wave_{p}": constraints[k] for k, p in enumerate(sanitized_in_order)})
            return f_dict, {}

        y_matrix = s_to_y(s_matrix)
        i_vec = y_matrix @ v_vec
        return {p: i_vec[k] for k, p in enumerate(sanitized_in_order)}, {}

    physics_wrapper.__name__ = cls_name
    physics_wrapper.__doc__ = getattr(base_fn, "__doc__", None)

    # Synthesise a signature that base_component._build_component can consume:
    # it must begin with the reserved ``signals`` arg and expose every SAX
    # parameter as a keyword-only entry with a default. The wrapper's runtime
    # body still accepts them via **kwargs.
    _sax_params = [
        inspect.Parameter(
            p.name,
            inspect.Parameter.KEYWORD_ONLY,
            default=defaults[p.name],
            annotation=p.annotation if p.annotation is not inspect.Parameter.empty else inspect.Parameter.empty,
        )
        for p in sig.parameters.values()
    ]
    physics_wrapper.__signature__ = inspect.Signature(
        parameters=[
            inspect.Parameter("signals", inspect.Parameter.POSITIONAL_OR_KEYWORD),
            *_sax_params,
        ]
    )

    cls = component(ports=port_names, states=aux_state_names, holomorphic=True)(physics_wrapper)
    cls._is_sax_wrapped = True
    cls._raw_to_sanitized_ports = raw_to_sanitized
    cls._sanitized_to_raw_ports = {sanitized: tuple(raws) for sanitized, raws in sanitized_to_raw.items()}
    return cls

@jax.jit
def y_to_s(Y: jax.Array, z0: complex = 1.0 + 1e-12j) -> jax.Array:
    """Convert an admittance (Y) matrix to an S-parameter matrix.

    Exact inverse of :func:`s_to_y`: ``S = (I - z0 Y) (I + z0* Y)^-1``.
    Reduces to ``(I - z0 Y) (I + z0 Y)^-1`` for real ``z0``. ``(I + z0* Y)``
    and ``(I - z0 Y)`` are both polynomials in ``Y`` and so commute, making
    this well-defined regardless of multiplication order.
    """
    n = Y.shape[-1]
    eye = jnp.eye(n, dtype=jnp.complex128)
    Yc = Y.astype(jnp.complex128)
    z0c = jnp.asarray(z0, dtype=jnp.complex128)
    M = eye + jnp.conj(z0c) * Yc
    return jnp.linalg.solve(M.swapaxes(-1, -2), (eye - z0c * Yc).swapaxes(-1, -2)).swapaxes(-1, -2)