circuit ¤

def ac(
    self,
    *,
    ports: str | Sequence[str],
    freqs: jax.Array,
    z0: float = 50.0,
    y_dc: jax.Array | None = None,
    params: dict[str, Any] | None = None,
    **param_updates: Any,
) -> jax.Array:
    """Run an AC small-signal S-parameter sweep.

    Linearises at the DC operating point and sweeps over ``freqs``,
    returning S-parameters for the given ports.

    Args:
        ports: Port name(s) to probe (e.g. ``"out"`` or ``["in", "out"]``).
        freqs: Frequency array in Hz.
        z0: Reference impedance (default 50 Ohm).
        y_dc: DC operating point. If ``None``, a DC solve is run first.
        params: Parameter updates (same format as :meth:`dc`).
            Array-valued params are **not** supported.
        **param_updates: Global parameter overrides.

    Returns:
        Complex S-parameter array of shape ``(len(freqs), n_ports, n_ports)``.

    Raises:
        ValueError: If the circuit is complex-valued (photonic). Use the
            low-level ``setup_ac_sweep`` API for complex circuits.

    """
    from circulax.solvers import setup_ac_sweep

    if self.solver.is_complex:
        msg = "Circuit.ac() currently supports real-valued circuits. Use the low-level AC API for custom paths."
        raise ValueError(msg)

    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "ac")
    groups = self._with_param_values(arrays)
    if y_dc is None:
        y_dc = self.solver.solve_dc(
            groups,
            self._zero_guess(),
            rtol=self.rtol,
            atol=self.atol,
            max_steps=self.max_steps,
        )
    port_list = [ports] if isinstance(ports, str) else list(ports)
    port_nodes = [self._resolve_port_node(port) for port in port_list]
    run_ac = setup_ac_sweep(groups, self.sys_size, port_nodes, z0=z0)
    return run_ac(y_dc, jnp.asarray(freqs))

def dc(
    self,
    y_guess: jax.Array | None = None,
    *,
    params: dict[str, Any] | None = None,
    rtol: float | None = None,
    atol: float | None = None,
    max_steps: int | None = None,
    **param_updates: Any,
) -> jax.Array:
    """Solve the DC operating point for the given parameters.

    Scalar params produce a single solve returning shape ``(n,)``.
    Array params trigger ``jax.vmap`` over the leading dimension,
    returning shape ``(batch, n)``. All array params must share the
    same leading dimension size.

    Args:
        y_guess: Initial guess for the Newton solver. Defaults to zeros.
        rtol: Relative tolerance override. Defaults to value from
            :func:`compile_circuit` (``1e-6``).
        atol: Absolute tolerance override. Defaults to value from
            :func:`compile_circuit` (``1e-6``).
        max_steps: Max Newton iterations override. Defaults to value
            from :func:`compile_circuit` (``100``).
        params: Optional mapping of parameter updates. Keys without a dot
            are broadcast to every component group declaring that parameter.
            Keys like ``"R1.R"`` update one instance.
        **param_updates: Convenience form for global parameter values, e.g.
            ``wavelength_nm=1310.0`` or
            ``wavelength_nm=jnp.linspace(1260, 1360, 1000)``.

    Returns:
        Flat solution vector of shape ``(n,)`` or ``(batch, n)``.

    Raises:
        ValueError: If multiple array params have different leading dims.

    """
    rtol = self.rtol if rtol is None else rtol
    atol = self.atol if atol is None else atol
    max_steps = self.max_steps if max_steps is None else max_steps

    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._as_arrays(updates)
    batch_keys = [k for k, v in arrays.items() if v.ndim > 0]

    if y_guess is None:
        y_guess = self._zero_guess()

    if not batch_keys:
        return self.solver.solve_dc(self._with_param_values(arrays), y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    batch_sizes = {k: arrays[k].shape[0] for k in batch_keys}
    if len(set(batch_sizes.values())) > 1:
        msg = f"All batched params must share the same leading dim. Got: {batch_sizes}"
        raise ValueError(msg)

    scalar_params = {k: v for k, v in arrays.items() if k not in batch_keys}

    def solve_single(*batch_vals: jax.Array) -> jax.Array:
        kw = dict(zip(batch_keys, batch_vals, strict=True))
        kw.update(scalar_params)
        return self.solver.solve_dc(self._with_param_values(kw), y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    return jax.vmap(solve_single)(*[arrays[k] for k in batch_keys])

def get_port_field(self, y: jax.Array, port: str) -> jax.Array:
    """Extract the (possibly complex) field at a named port.

    Args:
        y: Solution array of shape ``(n,)`` or ``(batch, n)``.
        port: Port key as ``"InstanceName,port_name"``.

    Returns:
        Real or complex scalar/array. For complex circuits, reconstructs
        the field from the unrolled block format.

    """
    idx = self._resolve_port_node(port)
    if self.solver.is_complex:
        return y[..., idx] + 1j * y[..., idx + self.sys_size]
    return y[..., idx]

def hb(
    self,
    *,
    freq: float,
    harmonics: int = 5,
    y0: jax.Array | None = None,
    params: dict[str, Any] | None = None,
    y_flat_init: jax.Array | None = None,
    rtol: float | None = None,
    atol: float | None = None,
    max_steps: int | None = None,
    osc_node: int | str | None = None,
    amplitude_tries: jax.Array | None = None,
    g_leak: float | None = None,
    **param_updates: Any,
) -> tuple[jax.Array, jax.Array]:
    """Run harmonic balance to find the periodic steady state.

    Args:
        freq: Fundamental frequency in Hz.
        harmonics: Number of harmonics (K = 2*harmonics + 1 time samples).
        y0: DC operating point. If ``None``, a DC solve is run first.
        params: Parameter updates (same format as :meth:`dc`).
            Array-valued params are **not** supported.
        y_flat_init: Flat initial waveform, shape ``(K * sys_size,)``.
            Overrides the automatic multi-start when ``osc_node`` is set.
        rtol: Relative tolerance override.
        atol: Absolute tolerance override.
        max_steps: Max Newton iterations override.
        osc_node: Port name or state index for oscillator multi-start.
        amplitude_tries: Amplitudes to try when ``osc_node`` is set.
        g_leak: Override the solver's regularisation conductance.
        **param_updates: Global parameter overrides.

    Returns:
        ``(y_time, y_freq)`` — time samples ``(K, n)`` and normalised
        Fourier coefficients ``(harmonics+1, n)`` complex.

    """
    from circulax.solvers import setup_harmonic_balance

    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "hb")
    groups = self._with_param_values(arrays)
    if y0 is None:
        y0 = self.solver.solve_dc(
            groups,
            self._zero_guess(),
            rtol=self.rtol,
            atol=self.atol,
            max_steps=self.max_steps,
        )
    osc_idx = self._resolve_port_node(osc_node) if isinstance(osc_node, str) else osc_node
    run_hb = setup_harmonic_balance(
        groups,
        self.sys_size,
        freq=freq,
        num_harmonics=harmonics,
        is_complex=self.solver.is_complex,
        g_leak=getattr(self.solver, "g_leak", 1e-9) if g_leak is None else g_leak,
        osc_node=osc_idx,
        amplitude_tries=amplitude_tries,
    )
    rtol = self.rtol if rtol is None else rtol
    atol = self.atol if atol is None else atol
    max_steps = self.max_steps if max_steps is None else max_steps
    return run_hb(y0, y_flat_init=y_flat_init, max_steps=max_steps, rtol=rtol, atol=atol)

def port(self, y: jax.Array, port: str) -> jax.Array:
    """Alias for :meth:`get_port_field`."""
    return self.get_port_field(y, port)

def transient(
    self,
    *,
    t0: float,
    t1: float,
    dt0: float,
    y0: jax.Array | None = None,
    saveat: Any = None,
    params: dict[str, Any] | None = None,
    transient_solver: Any = None,
    **kwargs: Any,
) -> Any:
    """Run transient (time-domain) analysis.

    Args:
        t0: Start time.
        t1: End time.
        dt0: Initial time step.
        y0: Initial state vector. If ``None``, a DC solve is run first.
        saveat: Times at which to save the solution. Accepts an array of
            timestamps or a ``diffrax.SaveAt`` object.
        params: Parameter updates (same format as :meth:`dc`).
            Array-valued params are **not** supported — raises ``ValueError``.
        transient_solver: Override the Diffrax solver (default BDF2).
        **kwargs: Forwarded to ``setup_transient`` / Diffrax (e.g.
            ``max_steps``).

    Returns:
        A ``diffrax.Solution`` with ``.ts`` and ``.ys`` attributes.

    """
    from diffrax import SaveAt

    from circulax.solvers import setup_transient

    param_updates = kwargs.pop("param_updates", {})
    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "transient")
    groups = self._with_param_values(arrays)
    if y0 is None:
        y0 = self.solver.solve_dc(
            groups,
            self._zero_guess(),
            rtol=self.rtol,
            atol=self.atol,
            max_steps=self.max_steps,
        )
    saveat_obj = SaveAt(ts=saveat) if saveat is not None and not isinstance(saveat, SaveAt) else saveat
    run_transient = setup_transient(groups=groups, linear_strategy=self.solver, transient_solver=transient_solver)
    return run_transient(t0=t0, t1=t1, dt0=dt0, y0=y0, saveat=saveat_obj, **kwargs)

def with_groups(self, groups: dict) -> Circuit:
    """Return a new Circuit with replaced component groups.

    Use together with :func:`~circulax.utils.update_params_dict` for
    instance-specific parameter changes before solving.

    Args:
        groups: New groups dict.

    Returns:
        A new :class:`Circuit` with the updated groups.

    """
    return Circuit(
        self.solver,
        groups,
        self.sys_size,
        self.port_map,
        rtol=self.rtol,
        atol=self.atol,
        max_steps=self.max_steps,
    )