wavelength ¤

def build_wavelength_bench(circuit: Any, compile_fn: Callable, selected: tuple[str, ...], z0: float) -> WavelengthBench:
    """Clone the source topology, attach matched-port fixtures and compile the bench.

    The bench depends only on the circuit's source topology and models, the selected
    ports and ``z0``; parameter values are runtime inputs of the solves.
    """
    _validate_boundary(circuit, selected)
    _validate_reference(z0)
    device = (
        kfnl.Netlist.from_dict(circuit.source_netlist.to_dict())
        if isinstance(circuit.source_netlist, kfnl.Netlist)
        else deepcopy(circuit.source_netlist)
    )
    instances = device.instances if isinstance(device, kfnl.Netlist) else device["instances"]
    models = dict(circuit.source_models)
    model_key = "__circulax_matched_port__"
    while model_key in models:
        model_key += "_"
    models[model_key] = _matched_port
    fixtures = {}
    for index, port in enumerate(selected):
        name = f"__circulax_analysis_{index}"
        while name in instances:
            name += "_"
        fixtures[port] = {"name": name, "component": model_key, "settings": {"z0": z0}}
    # Topology is static: compile eagerly even when the caller is tracing so index arrays stay concrete.
    with jax.ensure_compile_time_eval():
        bench = compile_fn(
            attach_testbench(device, sources=fixtures),
            models,
            is_complex=True,
            g_leak=getattr(circuit.solver, "g_leak", 1e-9),
            rtol=circuit.rtol,
            atol=circuit.atol,
            max_steps=circuit.max_steps,
        )
    return WavelengthBench(bench, {port: fixture["name"] for port, fixture in fixtures.items()})

def wavelength_scattering(
    circuit: Any,
    bench_for: Callable[[tuple[str, ...], float], WavelengthBench],
    *,
    ports: Sequence[str],
    wavelengths: jax.Array,
    amplitude: float,
    amplitude_unit: str,
    params: Mapping[str, Any] | None,
    z0: float,
    wl_param: str,
) -> dict[str, Any]:
    """Run independent port excitations with wavelength continuation.

    Optical discipline and wave normalization are declared by the caller; no
    classification is inferred from solver dtype. For SAX optical models the
    reference is z0=1 and amplitude_unit identifies sqrt(W) or sqrt(mW).
    ``bench_for(selected, z0)`` supplies the (cached) compiled matched-boundary bench.
    """
    selected = tuple(ports)
    _validate_boundary(circuit, selected)
    coordinates = _validate_coordinates(wavelengths, amplitude, amplitude_unit, z0)
    bench, fixtures = bench_for(selected, z0)
    count = len(selected)
    excitations = jnp.concatenate([jnp.zeros((1, count)), amplitude * jnp.eye(count)])
    guesses = jnp.zeros((count + 1, 2 * bench.sys_size))

    def solve(guess: jax.Array, incident: jax.Array, wavelength: jax.Array) -> tuple[jax.Array, jax.Array]:
        """Solve the DUT with fixed bias and this incident-wave column."""
        updates = {**(params or {}), wl_param: wavelength}
        updates.update({f"{fixtures[port]}.incident": incident[index] for index, port in enumerate(selected)})
        return bench.dc_checked(guess, params=updates)

    def sweep_point(previous: jax.Array, wavelength: jax.Array) -> tuple[jax.Array, jax.Array]:
        """Maintain separate continuation guesses for baseline and each column."""
        solved, converged = jax.vmap(solve, in_axes=(0, 0, None))(previous, excitations, wavelength)
        return solved, (solved, converged)

    _, (solved, converged) = jax.lax.scan(sweep_point, guesses, coordinates)
    size = bench.sys_size
    physical = solved[..., :size] + 1j * solved[..., size:]
    baseline = physical[:, 0]
    excited = physical[:, 1:]
    if not isinstance(solved, jax.core.Tracer):
        # Concrete results are validated here; under jit callers inspect ``converged`` instead.
        if not bool(jnp.all(converged)):
            msg = "Wavelength scattering DC solve did not converge for at least one wavelength or excitation"
            raise ValueError(msg)
        if not bool(jnp.all(jnp.isfinite(solved))):
            msg = "Wavelength scattering produced nonfinite baseline or excited solutions"
            raise ValueError(msg)
        # One quadrature probe at the first wavelength detects responses that are not complex-linear.
        probe, probe_ok = solve(solved[0, 1], 1j * excitations[1], coordinates[0])
        if not bool(probe_ok):
            msg = "Wavelength scattering quadrature probe did not converge"
            raise ValueError(msg)
        _check_complex_linear(excited[0, 0] - baseline[0], probe[:size] + 1j * probe[size:] - baseline[0])
    port_nodes = jnp.asarray([bench.port_map[port] for port in selected])
    outgoing = (excited[..., port_nodes] - baseline[:, None, port_nodes]) / jnp.sqrt(z0) / amplitude
    scattering = jnp.swapaxes(outgoing, 1, 2) - jnp.eye(count)
    return {
        "wavelengths": coordinates,
        "ports": selected,
        "s": scattering,
        "amplitude": amplitude,
        "amplitude_unit": amplitude_unit,
        "converged": converged,
        "baseline_solution": baseline,
        "solution": excited,
        "port_map": dict(bench.port_map),
    }