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)

@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)
    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) -> 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()}``.

    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.

    """
    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)

    # base_component builds a namedtuple over the port tuple, 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, s: States, **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(s, 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, s) args 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),
            inspect.Parameter("s", inspect.Parameter.POSITIONAL_OR_KEYWORD),
            *_sax_params,
        ]
    )

    cls = component(ports=port_names, states=aux_state_names)(physics_wrapper)
    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