enforce ¤

def enforce_passivity(
    model: VFModel,
    s: jnp.ndarray,
    opts: FitOptions,
    verbose: bool = True,
) -> tuple[VFModel, np.ndarray]:
    """Iteratively enforce passivity of a Y-parameter rational model.

    Port of ``RPdriver.m`` (outer loop) + ``RP_QRNNLS_Y.m`` (RP step).

    At each outer iteration:
      1. Sweep minimum eigenvalue of Re(Y(jω)) → detect violation bands.
      2. Find per-channel violation extrema → violpairs.
      3. If passive and D, E ≥ 0: stop.
      4. Solve the NNLS residue-perturbation problem.
      5. Repeat.

    Args:
        model:   Initial VFModel (pole-residue form).
        s:       ``(Ns,)`` purely imaginary sweep frequencies (s = j·ω).
        opts:    FitOptions controlling tolerances and iteration counts.
        verbose: Print per-iteration passivity summary if True.

    Returns:
        new_model: Passive (or best-effort passive) VFModel.
        gmin_final: ``(Ns,)`` minimum eigenvalue after enforcement.

    """
    # Pin everything to CPU: enforcement is a Python loop with scipy NNLS.
    # Callers may pass GPU-resident arrays (e.g. from vfdriver on CUDA);
    # mixing GPU matmuls with CPU NNLS causes device-mismatch divergence.
    cpu = jax.devices("cpu")[0]
    model = jax.device_put(model, cpu)
    s = jax.device_put(s, cpu)

    Nc = model.residues.shape[0]
    omega = np.array(jnp.imag(s), dtype=float)
    s_np = np.array(s, dtype=complex)

    ss = vfmodel_to_ss(model, Nc)
    poles = np.array(model.poles)
    cindex = compute_cindex(poles)

    _cache: dict = {}

    if verbose:
        print("--- enforce_passivity (Y-parameters) ---")

    for iter_out in range(opts.Niter_out):
        gmin, is_passive = passivity_sweep_Y(ss, s)
        gmin_np = np.array(gmin)

        d_eigs = np.linalg.eigvalsh(np.real(np.array(ss.D)))
        e_eigs = np.linalg.eigvalsh(np.real(np.array(ss.E)))
        d_ok = bool(np.all(d_eigs >= 0))
        e_ok = bool(np.all(e_eigs >= 0))

        g_min_val = float(np.min(gmin_np))
        if verbose:
            print(
                f"  iter {iter_out:2d}  gmin={g_min_val:+.3e}"
                f"  D_min={float(np.min(d_eigs)):+.3e}"
                f"  E_min={float(np.min(e_eigs)):+.3e}"
            )

        if is_passive and d_ok and e_ok:
            if verbose:
                print("  → passive ✓")
            break

        bands = find_violation_bands(gmin_np, omega, opts.TOLG)
        violpairs = find_violation_extrema(ss, omega, bands)

        if not violpairs and d_ok and e_ok:
            if verbose:
                print("  → no violation extrema found, stopping")
            break

        # RP-NNLS perturbation step.
        old_res = np.array(model.residues)
        old_D = np.array(model.D)
        old_E = np.array(model.E)
        model, ss, _cache = _rp_qrnnls_y(model, ss, s_np, violpairs, poles, cindex, opts, _cache)
        poles = np.array(model.poles)
        cindex = compute_cindex(poles)

        # Stall detection: stop if model did not change.
        if (
            np.allclose(np.array(model.residues), old_res, atol=1e-30)
            and np.allclose(np.array(model.D), old_D, atol=1e-30)
            and np.allclose(np.array(model.E), old_E, atol=1e-30)
        ):
            if verbose:
                print("  → stalled, stopping")
            break

    gmin_final, _ = passivity_sweep_Y(ss, s)
    return model, np.array(gmin_final)