circuit ¤

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

    .. deprecated:: 0.2
        Use :meth:`sp` instead. Will be removed in 0.3.

    """
    warnings.warn("Circuit.ac() is deprecated, use Circuit.sp() instead.", DeprecationWarning, stacklevel=2)
    return self.sp(ports=ports, freqs=freqs, z0=z0, y_dc=y_dc, holomorphic=holomorphic, params=params, **param_updates)

def ac_response(
    self,
    *,
    freqs: jax.Array,
    sources: Mapping[str, complex],
    y_dc: jax.Array | None = None,
    params: dict[str, Any] | None = None,
    holomorphic: bool | str = "auto",
    **param_updates: Any,
) -> dict[str, Any]:
    """Return physical small-signal response to independent source phasors.

    ``sources`` maps source instance names to complex peak amplitudes in
    their physical source units. Unspecified sources have zero AC drive;
    DC bias and schematic waveforms remain unchanged. No matching loads
    or scattering normalization are added, and no external ports are needed.
    ``solution`` has shape ``(len(freqs), sys_size)`` with complex physical
    unknowns; ``port_map`` retains node and state aliases. Parameter batching
    is available through external JAX vmap, as for scattering analysis.

    Non-holomorphic optical response requires a quadrature/sideband result
    contract and is explicitly unsupported by this interface.

    ``holomorphic`` can only tighten, never relax, that check. ``"auto"`` and ``True`` both
    rely on the component ``holomorphic`` flags: a complex-mode circuit with any
    non-holomorphic component is refused, because the Wirtinger solve would otherwise
    return silently wrong answers and there is no cheap runtime guard. ``False`` declares
    the response non-holomorphic and refuses every complex-mode circuit. Real circuits
    ignore the argument. Other values raise ``ValueError``.
    """
    if holomorphic not in {"auto", True, False}:
        msg = f"holomorphic must be 'auto', True or False, got {holomorphic!r}"
        raise ValueError(msg)
    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "ac_response")
    groups = self._for_analysis("ac")._with_param_values(arrays)  # noqa: SLF001 -- another Circuit analysis variant
    if self.solver.is_complex and (not _infer_holomorphic(groups) or holomorphic is False):
        msg = (
            "Non-holomorphic optical AC response requires explicit quadrature/sideband observables "
            "(the holomorphic argument can only tighten the component-flag check, not relax it)"
        )
        raise ValueError(msg)
    rhs = jnp.zeros(self.sys_size, dtype=jnp.complex128)
    for name, amplitude in sources.items():
        matches = [group for group in groups.values() if name in (getattr(group, "index_map", None) or {})]
        if len(matches) != 1 or not getattr(matches[0], "ac_source_stamp", ()):
            msg = f"AC source '{name}' must identify one instance with an independent-source excitation stamp"
            raise ValueError(msg)
        value = jnp.asarray(amplitude)
        if value.ndim != 0:
            msg = f"AC source '{name}' needs a scalar phasor; use jax.vmap for batches"
            raise ValueError(msg)
        group = matches[0]
        instance = group.index_map[name]
        for local_row, coefficient in group.ac_source_stamp:
            rhs = rhs.at[group.eq_indices[instance, local_row]].add(coefficient * value)
    if y_dc is None:
        y_dc = self.dc(params=updates)
    dc_groups = (
        self._for_analysis("dc")._with_param_values(arrays)  # noqa: SLF001 -- another Circuit analysis variant
        if self._analysis_factory is not None
        else None
    )
    frequencies = jnp.asarray(freqs)
    if frequencies.ndim != 1:
        msg = "AC response frequencies must be a one-dimensional array"
        raise ValueError(msg)
    run = setup_ac_response(groups, self.sys_size, is_complex=self.solver.is_complex, dc_groups=dc_groups)
    return {"freqs": frequencies, "y_dc": y_dc, "solution": run(y_dc, frequencies, rhs), "port_map": dict(self.port_map)}

def check_sax_compatibility(self) -> bool:
    """Return True if this circuit can be represented as a pure SAX S-parameter network.

    A circuit qualifies when every component group is a SAX-wrapped
    model (:func:`sax_component`) and none is a source — i.e. a passive
    photonic (or RF) mesh with no reactive (``dQ/dt``) terms and no
    driving term. Such a circuit's DC operating point is provably
    all-zero, and :meth:`to_sax_circuit` can extract its S-parameters
    directly.
    """
    return self._is_pure_sax

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.

    """
    dc_circuit = self._for_analysis("dc")
    if dc_circuit is not self:
        return dc_circuit.dc(y_guess=y_guess, rtol=rtol, atol=atol, max_steps=max_steps, params=params, **param_updates)
    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
    if y_guess is None:
        y_guess = self._zero_guess()

    def solve_fn(groups: dict) -> jax.Array:
        return self.solver.solve_dc(groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    return self._batched_solve(params, param_updates, solve_fn)

def dc_checked(
    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,
) -> tuple[jax.Array, jax.Array]:
    """Return the DC solution and the solver's explicit convergence flag.

    Scalar parameters return a scalar boolean; array parameters use the
    same shared native registrations and vmap batching as dc, returning
    one boolean per leading batch point. Failed estimates can be finite;
    callers must use this flag rather than infer success from sample values.
    Both results remain JAX arrays and can be consumed inside jit.
    """
    dc_circuit = self._for_analysis("dc")
    if dc_circuit is not self:
        return dc_circuit.dc_checked(y_guess=y_guess, params=params, rtol=rtol, atol=atol, max_steps=max_steps, **param_updates)
    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
    if y_guess is None:
        y_guess = self._zero_guess()

    def solve_fn(groups: dict) -> tuple[jax.Array, jax.Array]:
        """Retain native convergence evidence for each batched solve."""
        return self.solver.solve_dc_checked(groups, y_guess, rtol=rtol, atol=atol, max_steps=max_steps)

    return self._batched_solve(params, param_updates, solve_fn)

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.

    Native OSDI devices without ABI state slots use transient-mode F/Q
    evaluation. Differentiating through the converged native solve is
    not yet supported.

    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

    if any(getattr(group, "num_states", 0) > 0 for group in self.groups.values()):
        msg = "Native OSDI harmonic balance with ABI state slots is not supported"
        raise NotImplementedError(msg)

    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "hb")
    # HB evaluates large-signal F/Q over one period, so use transient
    # registrations rather than the small-signal AC evaluation mode.
    groups = self._for_analysis("tran")._with_param_values(arrays)
    if y0 is None:
        y0 = self.dc(params=updates)
    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 prepare_analysis(
    self,
    analysis: str,
    *,
    ports: Sequence[str] | None = None,
    z0: float = 1.0,
) -> Circuit:
    """Prepare immutable native registrations without solving the circuit.

    Return the selected DC, AC or transient registration variant. Builtin
    circuits return themselves. The original circuit caches required
    variants so callers may retain it for later execution. DC bias setup is
    also prepared for AC, transient and HB, without calculating that bias.

    Supported names are op, dc, sp_wavelength, ac_response, sp_frequency,
    tran and hb. Native HB state-history and non-holomorphic optical
    driven-AC limitations are rejected before simulation execution.

    ``sp_wavelength`` runs on a separately compiled matched-boundary bench.
    Pass ``ports`` (and ``z0``, default 1.0) to build and cache that bench and its
    DC registrations now; later :meth:`sp_wavelength` calls with the same ports and
    ``z0`` reuse it. The returned circuit is still the selected DC variant of this one.
    """
    modes = {
        "op": "dc",
        "dc": "dc",
        "sp_wavelength": "dc",
        "ac_response": "ac",
        "sp_frequency": "ac",
        "tran": "tran",
        "hb": "tran",
    }
    if analysis not in modes:
        msg = f"Unsupported Circulax analysis {analysis!r}"
        raise ValueError(msg)

    def reject_unsupported(circuit: Circuit) -> None:
        if analysis == "hb" and any(getattr(group, "num_states", 0) > 0 for group in circuit.groups.values()):
            msg = "Native OSDI harmonic balance with ABI state slots is not supported"
            raise NotImplementedError(msg)
        if analysis == "ac_response" and circuit.solver.is_complex and not _infer_holomorphic(circuit.groups):
            msg = "Non-holomorphic optical AC response requires explicit quadrature/sideband observables"
            raise ValueError(msg)

    # Fail fast before building variants, then recheck the selected native registrations.
    reject_unsupported(self)
    self._for_analysis("dc")
    selected = self._for_analysis(modes[analysis])
    reject_unsupported(selected)
    if analysis == "sp_wavelength" and ports is not None:
        self._wavelength_bench(ports, z0).circuit.prepare_analysis("dc")
    return selected

def set_simparams(self, simparams: Mapping[str, float] | None = None, **settings: float) -> Circuit:
    """Return an updated circuit with numeric Verilog-A simulator settings.

    Accept a mapping, keyword arguments, or both (keywords take precedence).
    Updates merge into the existing settings of every native descriptor.
    Native setup and analysis variants are rebuilt; existing registrations
    are reused for matching settings. The current circuit is unchanged so
    JIT callables already compiled for it remain valid.

    Call outside JIT, then rerun DC, AC or transient on the returned circuit.
    This updates model $simparam queries, not the host solver's options.
    """
    if simparams is not None and not isinstance(simparams, Mapping):
        msg = "simparams must be a mapping of names to finite numbers"
        raise TypeError(msg)
    if self._simparams_factory is None:
        msg = "set_simparams requires a native OSDI circuit created by compile_circuit"
        raise ValueError(msg)
    updates = dict(simparams or {})
    updates.update(settings)
    return self._simparams_factory(updates)

def sp(
    self,
    *,
    ports: str | Sequence[str],
    freqs: jax.Array,
    z0: float | jax.Array = 50.0,
    y_dc: jax.Array | None = None,
    holomorphic: bool | str = "auto",
    params: dict[str, Any] | None = None,
    **param_updates: Any,
) -> jax.Array:
    """Run a 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 in ohms.  Scalar (uniform) or array of
            shape ``(N_ports,)`` (per-port).  Default 50 Ohm.  For
            frequency-dependent impedance, solve at a fixed z0 and use
            :func:`~circulax.solvers.renormalize` afterwards.
        y_dc: DC operating point. If ``None``, a DC solve is run first.
        holomorphic: If ``"auto"`` (default), infer from component flags —
            if any component in the circuit has ``holomorphic=False``, the
            full 2N×2N real-block system is used.  Pass ``True`` to force
            the N×N Wirtinger system, or ``False`` to force the 2N×2N
            system.
        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)``.

    """
    from circulax.solvers import setup_ac_sweep

    updates = self._coerce_param_updates(params, param_updates)
    arrays = self._require_scalar_params(updates, "sp")
    groups = self._for_analysis("ac")._with_param_values(arrays)  # noqa: SLF001 -- another Circuit analysis variant
    dc_groups = (
        self._for_analysis("dc")._with_param_values(arrays)  # noqa: SLF001 -- another Circuit analysis variant
        if self._analysis_factory is not None
        else None
    )
    if holomorphic == "auto":
        holomorphic = _infer_holomorphic(groups)
    if y_dc is None:
        y_dc = self.dc(params=updates)
    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,
        is_complex=self.solver.is_complex,
        holomorphic=holomorphic,
        dc_groups=dc_groups,
    )
    return run_ac(y_dc, jnp.asarray(freqs))

def sp_wavelength(
    self,
    *,
    ports: Sequence[str],
    wavelengths: jax.Array,
    amplitude_unit: str,
    amplitude: float = 1e-4,
    params: Mapping[str, Any] | None = None,
    z0: float = 1.0,
    wl_param: str = "wl",
) -> dict[str, Any]:
    """Characterize declared optical ports with normalized low-power DC waves.

    The caller declares optical discipline and field power normalization;
    z0=1 is the SAX optical reference. Temporary matching and excitation
    preserve physical DUT sources, loads and electrical bias. Independent
    input columns and their zero-incident baseline use scan continuation
    over wavelength. ``s`` has shape (wavelength, output_port, input_port).

    ``converged`` has shape (wavelength, 1 + len(ports)): the baseline and each
    excitation. Concrete results raise if any solve failed to converge or is
    not finite, and if a quadrature probe at the first wavelength shows the
    response is not complex-linear (a non-holomorphic model, or an amplitude
    outside the linear regime). Under jit these checks are skipped; inspect
    ``converged`` instead. The bench inherits this circuit's ``g_leak``. It is compiled once
    per (ports, z0) and cached on this circuit; :meth:`prepare_analysis` can build it ahead.
    """
    return wavelength_scattering(
        self,
        self._wavelength_bench,
        ports=ports,
        wavelengths=wavelengths,
        amplitude=amplitude,
        amplitude_unit=amplitude_unit,
        params=params,
        z0=z0,
        wl_param=wl_param,
    )

def to_sax_circuit(self) -> Callable[..., dict[tuple[str, str], jax.Array]]:
    """Convert this circuit into a native SAX model function.

    Only valid when :meth:`check_sax_compatibility` is True. Returns a
    callable ``model(**params) -> sax.SDict`` matching the calling
    convention of :func:`sax.circuit`'s own returned model: scalar
    params produce one ``SDict``; array-valued params trigger
    ``jax.vmap`` and return an ``SDict`` whose values are stacked
    arrays, e.g. ``model(wl=jnp.linspace(1500, 1600, 101))``.

    Raises:
        ValueError: If this circuit is not all-SAX/source-free, or has
            no known external ports (e.g. built via :meth:`with_groups`).

    """
    if not self._is_pure_sax:
        msg = (
            "Circuit.to_sax_circuit() requires an all-SAX, source-free circuit "
            "(check_sax_compatibility() is False for this circuit)."
        )
        raise ValueError(msg)
    if not self.ports:
        msg = (
            "This all-SAX, source-free circuit has no known external ports "
            "(e.g. built via with_groups(), which drops the source netlist), "
            "so it cannot be converted to a SAX circuit."
        )
        raise ValueError(msg)

    def model(*, params: dict[str, Any] | None = None, **param_updates: Any) -> dict[tuple[str, str], jax.Array]:
        return self._batched_solve(params, param_updates, self._solve_sparams_sdict)

    return model

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,
    tstep: float | None = 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``, solve the time-zero operating
            point using transient source values.
        tstep: SPICE source-default time step; defaults to ``dt0``.
        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")
    transient_circuit = self._for_analysis("tran")
    groups = transient_circuit._with_param_values(arrays)  # noqa: SLF001 -- another Circuit analysis variant
    if y0 is None:
        from circulax.solvers.source_settings import waveform_groups

        dc_circuit = self._for_analysis("dc")
        initial_groups = waveform_groups(
            dc_circuit._with_param_values(arrays),  # noqa: SLF001 -- another Circuit analysis variant
            tstep=dt0 if tstep is None else tstep,
            tstop=t1,
        )
        y0 = dc_circuit.solver.solve_dc(
            initial_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=transient_circuit.solver, transient_solver=transient_solver)
    return run_transient(t0=t0, t1=t1, dt0=dt0, y0=y0, saveat=saveat_obj, tstep=tstep, **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.

    """
    is_pure_sax = bool(groups) and all(getattr(g, "is_sax_wrapped", False) for g in groups.values())
    return Circuit(
        self.solver,
        groups,
        self.sys_size,
        self.port_map,
        rtol=self.rtol,
        atol=self.atol,
        max_steps=self.max_steps,
        _is_pure_sax=is_pure_sax,
    )