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.
    2. **Physics wrapper** — a closure is built that calls ``fn`` at runtime,
       converts the S-dict to a dense matrix via :func:`sax.sdense`, converts
       it to an admittance matrix via :func:`s_to_y`, and returns
       ``I = Y @ V`` as a port current dict.
    3. **Component registration** — the wrapper is passed to
       :func:`~circulax.components.base_component.component` with the
       discovered ports, 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)
    except Exception as exc:
        msg = f"Failed to dry-run SAX component '{cls_name}': {exc}"
        raise RuntimeError(msg) from exc

    # 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.
    port_names = tuple(_sanitize_port(p) for p in detected_ports)

    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 y_matrix; using the
        # sorted order silently scrambles coupling between non-adjacent rows.
        s_matrix, port_map = sdense(s_dict)
        y_matrix = s_to_y(s_matrix)
        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)
        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,
        ]
    )

    return component(ports=port_names)(physics_wrapper)