ac_sweep ¤

def renormalize(
    S: Array,
    z0_from: float | Array,
    z0_to: float | Array,
) -> Array:
    """Renormalize S-parameters from one reference impedance to another.

    Transforms S-parameters computed at impedance ``z0_from`` to the
    equivalent S-parameters at impedance ``z0_to``, using the standard
    power-wave renormalization formula.

    Args:
        S: S-parameter array of shape ``(N_freqs, N_ports, N_ports)``.
        z0_from: Original reference impedance.  Scalar or ``(N_ports,)``
            for frequency-independent, or ``(N_freqs, N_ports)`` for
            frequency-dependent impedance.
        z0_to: Target reference impedance, same shape options as ``z0_from``.

    Returns:
        Renormalized S-parameter array, same shape as ``S``.

    """
    n_freqs, n_ports, _ = S.shape
    z0_from = jnp.broadcast_to(jnp.atleast_1d(jnp.asarray(z0_from, dtype=jnp.complex128)), (n_freqs, n_ports))
    z0_to = jnp.broadcast_to(jnp.atleast_1d(jnp.asarray(z0_to, dtype=jnp.complex128)), (n_freqs, n_ports))

    def _renorm_one(S_f, zf, zt):
        I = jnp.eye(n_ports, dtype=jnp.complex128)
        gamma = (zt - zf) / (zt + zf)
        Gamma = jnp.diag(gamma)
        A = jnp.diag(jnp.sqrt(zt.real / zf.real) * zf / zt)
        return A @ (S_f - Gamma) @ jnp.linalg.solve(I - Gamma @ S_f, I) @ jnp.linalg.inv(A)

    return jax.vmap(_renorm_one)(S, z0_from, z0_to)

def setup_ac_response(
    groups: dict[str, Any],
    num_vars: int,
    *,
    is_complex: bool = False,
    dc_groups: dict[str, Any] | None = None,
) -> Callable[[Array, Array, Array], Array]:
    """Linearize physical response without scattering-port terminations.

    The caller supplies an independent-source RHS, leaving DC bias intact.
    Native AC charge registrations, delays and frequency-domain admittances
    use the same assembly as scattering analysis.
    """
    rows, cols, ground_idxs, _ = _build_index_arrays(groups, num_vars, is_complex=False)
    rows = jnp.asarray(rows)
    cols = jnp.asarray(cols)
    ground = jnp.asarray(ground_idxs)
    assemble = assemble_gc_complex if is_complex else assemble_gc_real

    def run_response(y_dc: Array, freqs: Array, rhs: Array) -> Array:
        """Solve one physical excitation at each frequency."""
        g_values, c_values = assemble(y_dc, groups)
        if dc_groups is not None:
            g_values, _ = assemble(y_dc, dc_groups)
        conductance = jnp.zeros((num_vars, num_vars), dtype=jnp.complex128).at[rows, cols].add(g_values)
        capacitance = jnp.zeros_like(conductance).at[rows, cols].add(c_values)

        def solve_frequency(freq: Array) -> Array:
            """Include dynamic component admittances before solving the RHS."""
            admittance = conductance + 2j * jnp.pi * freq * capacitance
            for group in groups.values():
                if getattr(group, "is_fdomain", False):
                    values = jax.vmap(functools.partial(group.physics_func, freq))(group.params)
                    admittance = admittance.at[group.jac_rows.reshape(-1), group.jac_cols.reshape(-1)].add(values.reshape(-1))
                if getattr(group, "has_delay", False):
                    values = _delay_admittance(group, y_dc, freq, is_complex=is_complex)
                    admittance = admittance.at[group.jac_rows.reshape(-1), group.jac_cols.reshape(-1)].add(values.reshape(-1))
            admittance = admittance.at[ground, ground].add(GROUND_STIFFNESS)
            return jnp.linalg.solve(admittance, rhs)

        return jax.vmap(solve_frequency)(freqs)

    return run_response

def setup_ac_sweep(
    groups: dict[str, Any],
    num_vars: int,
    port_nodes: list[int],
    *,
    z0: float | Array = 50.0,
    is_complex: bool = False,
    holomorphic: bool = False,
    dc_groups: dict[str, Any] | None = None,
) -> Callable[[Array, Array], Array]:
    """Configure and return a callable for AC small-signal S-parameter sweep.

    Linearises the circuit DAE at the DC operating point and sweeps over an
    array of frequencies, returning the complex S-parameter matrix at each
    frequency.  The returned callable is compatible with :func:`jax.jit` and
    :func:`jax.vmap`.

    The analysis solves ``Y(jω) · V = RHS`` at each frequency, where::

        Y(jω) = G + jωC + Y_fdomain(f) + port_terminations + ground_penalty

    ``G = ∂F/∂y`` and ``C = ∂Q/∂y`` are extracted once at the DC operating
    point.  ``Y_fdomain(f)`` is the admittance contribution from
    frequency-domain components, re-evaluated at each frequency.

    **S-parameter convention** — matched-load verification:

    - Matched load (Z_circuit = Z0) → S11 = 0
    - Open circuit (Z_circuit → ∞) → S11 = +1
    - Short circuit (Z_circuit = 0) → S11 = −1

    Args:
        groups: Compiled component groups from :func:`~circulax.compile_netlist`.
        dc_groups: Optional groups registered for DC, providing conductance while
            ``groups`` provides AC charge derivatives. Both must share a scatter layout.
        num_vars: Total number of scalar unknowns (second return value of
            :func:`~circulax.compile_netlist`).
        port_nodes: Global node indices for each circuit port, in the desired
            port ordering.  Obtain from the port-to-node map returned by
            :func:`~circulax.compile_netlist`::

                _, _, pmap = compile_netlist(net_dict, models)
                port_nodes = [pmap["R1,p1"], pmap["C1,p1"]]

        z0: Reference impedance in ohms.  Accepts:

            - **scalar** — same impedance for all ports (default: 50.0).
            - **array of shape** ``(N_ports,)`` — per-port impedance.

            For frequency-dependent impedance, solve at a fixed z0 and use
            :func:`renormalize` afterwards.

        is_complex: If ``True``, use complex-valued assembly for photonic
            circuits.  The DC operating point ``y_dc`` is expected in unrolled
            block format (shape ``(2 * num_vars,)``).
        holomorphic: If ``True``, use the N×N Wirtinger system for complex
            circuits.  Defaults to ``False`` so the full 2N×2N real-block
            system is used safely when any component uses
            non-holomorphic operations (e.g. ``jnp.real()``, ``jnp.abs()``),
            which couple the field and its conjugate.  This option has no
            effect for real-valued circuits.

    Returns:
        A callable ``run_ac(y_dc, freqs) -> S`` where:

        - **y_dc** — DC operating point, shape ``(num_vars,)`` for real
        circuits or ``(2 * num_vars,)`` for complex circuits.
        - **freqs** — frequencies in Hz, shape ``(N_freqs,)``.
        - **S** — complex128 S-parameter matrix with shape
        ``(N_freqs, N_ports, N_ports)``.

        Compatible with :func:`jax.jit` and :func:`jax.vmap` over ``y_dc``.

    """
    if 0 in port_nodes:
        msg = "Port node cannot be the ground node (index 0)."
        raise ValueError(msg)

    if is_complex and not holomorphic:
        return _setup_ac_sweep_2n(groups, num_vars, port_nodes, z0=z0, dc_groups=dc_groups)

    # --- Pre-compute static COO index arrays (captured in closure) -----------
    static_rows, static_cols, ground_idxs, _ = _build_index_arrays(groups, num_vars, is_complex=False)
    static_rows_jax = jnp.array(static_rows)
    static_cols_jax = jnp.array(static_cols)
    ground_indices = jnp.array(ground_idxs)

    N_ports = len(port_nodes)
    port_nodes_arr = jnp.array(port_nodes, dtype=jnp.int32)

    fdomain_scatter: dict[str, tuple[Array, Array]] = {
        gk: (
            jnp.array(groups[gk].jac_rows).reshape(-1),
            jnp.array(groups[gk].jac_cols).reshape(-1),
        )
        for gk in sorted(groups)
        if groups[gk].is_fdomain
    }
    delay_scatter: dict[str, tuple[Array, Array]] = {
        gk: (
            jnp.array(groups[gk].jac_rows).reshape(-1),
            jnp.array(groups[gk].jac_cols).reshape(-1),
        )
        for gk in sorted(groups)
        if getattr(groups[gk], "has_delay", False)
    }

    gc_assemble = assemble_gc_complex if is_complex else assemble_gc_real

    z0_arr = _normalize_z0(z0, N_ports)

    # -------------------------------------------------------------------------
    def run_ac(y_dc: Array, freqs: Array) -> Array:
        G_vals, C_vals = gc_assemble(y_dc, groups)
        if dc_groups is not None:
            G_vals, _ = gc_assemble(y_dc, dc_groups)

        G_mat = jnp.zeros((num_vars, num_vars), dtype=jnp.complex128)
        G_mat = G_mat.at[static_rows_jax, static_cols_jax].add(G_vals)

        C_mat = jnp.zeros((num_vars, num_vars), dtype=jnp.complex128)
        C_mat = C_mat.at[static_rows_jax, static_cols_jax].add(C_vals)

        RHS = jnp.zeros((num_vars, N_ports), dtype=jnp.complex128)
        RHS = RHS.at[port_nodes_arr, jnp.arange(N_ports)].set(2.0 / z0_arr)

        def _solve_one_freq(f: Array) -> Array:
            omega = 2.0 * jnp.pi * f
            Y = G_mat + 1j * omega * C_mat

            for gk, (rows_fd, cols_fd) in fdomain_scatter.items():
                group_fd = groups[gk]
                Y_mats = jax.vmap(functools.partial(group_fd.physics_func, f))(group_fd.params)
                Y = Y.at[rows_fd, cols_fd].add(Y_mats.reshape(-1))

            for gk, (rows_delay, cols_delay) in delay_scatter.items():
                Y_delay = _delay_admittance(groups[gk], y_dc, f, is_complex=is_complex)
                Y = Y.at[rows_delay, cols_delay].add(Y_delay.reshape(-1))

            Y = Y.at[port_nodes_arr, port_nodes_arr].add(1.0 / z0_arr)
            Y = Y.at[ground_indices, ground_indices].add(GROUND_STIFFNESS)

            V = jnp.linalg.solve(Y, RHS)
            V_ports = V[port_nodes_arr, :]
            return V_ports - jnp.eye(N_ports, dtype=jnp.complex128)

        return jax.vmap(_solve_one_freq)(freqs)

    return run_ac