aaa ¤

def aaa_driver(
    bigH: jnp.ndarray,  # (Nc, Nc, Ns) complex
    s: jnp.ndarray,  # (Ns,) complex
    opts: FitOptions,  # reuses asymp, weightparam; N is ignored
    tol: float = 1e-10,  # AAA convergence tolerance
    mmax: int = 100,  # max support points per element
    verbose: bool = True,
) -> tuple[VFModel, SSModel, float, jnp.ndarray]:
    """Fit bigH(s) using AAA rational approximation for pole finding.

    Replaces the VF pole-finding phase with AAA applied to each diagonal
    element. The poles found are then used with the standard residue
    identification step.

    Args:
        bigH:    (Nc, Nc, Ns) complex — frequency-domain matrix data.
        s:       (Ns,) complex — frequency points (= j*omega).
        opts:    FitOptions — asymp and weightparam are used; N is ignored.
        tol:     AAA convergence tolerance (per diagonal element).
        mmax:    maximum AAA support points per diagonal element.
        verbose: print progress.

    Returns:
        model:   VFModel — fitted pole-residue model.
        ss:      SSModel — state-space form.
        rmserr:  float   — final RMS fitting error.
        bigHfit: (Nc, Nc, Ns) — reconstructed frequency response.

    """
    Nc = bigH.shape[0]
    Ns = s.shape[0]

    bigH_np = np.asarray(bigH, dtype=np.complex128)
    s_np = np.asarray(s, dtype=np.complex128)

    # -------------------------------------------------------------------------
    # Step 1-2: Run AAA on each diagonal element and extract poles
    # -------------------------------------------------------------------------
    all_poles = []

    for i in range(Nc):
        f_diag = bigH_np[i, i, :]  # (Ns,) complex diagonal element

        w, zj, fj = aaa_scalar(f_diag, s_np, tol=tol, mmax=mmax)

        if len(zj) > 1:
            pols = _aaa_poles(w, zj)
        else:
            pols = np.array([], dtype=np.complex128)

        all_poles.append(pols)

        if verbose:
            print(f"  AAA diagonal [{i},{i}]: {len(zj)} support pts, {len(pols)} raw poles")

    # -------------------------------------------------------------------------
    # Step 3: Collect and clean all poles
    # -------------------------------------------------------------------------
    poles_np = _collect_poles(all_poles)

    if verbose:
        print(f"AAA driver: {len(poles_np)} poles after post-processing")

    if len(poles_np) == 0:
        raise RuntimeError("AAA found no poles. Try increasing mmax or reducing tol.")

    # -------------------------------------------------------------------------
    # Steps 5-7: Residue identification using standard VF machinery
    # -------------------------------------------------------------------------
    f_full = stack_upper_triangle(bigH)  # (nnn, Ns)
    w_full = compute_weights(bigH, opts.weightparam)  # (nnn, Ns) or (1, Ns)

    C_flat, D_vec, E_vec = identify_residues(f_full, s, poles_np, w_full, opts)
    # C_flat: (nnn, N) complex
    # D_vec:  (nnn,) float
    # E_vec:  (nnn,) float

    # -------------------------------------------------------------------------
    # Step 8: Reconstruct VFModel
    # -------------------------------------------------------------------------
    N = len(poles_np)
    nnn = f_full.shape[0]
    idx = _upper_triangle_indices(Nc)

    residues = jnp.zeros((Nc, Nc, N), dtype=jnp.complex128)
    D_mat = jnp.zeros((Nc, Nc), dtype=jnp.float64)
    E_mat = jnp.zeros((Nc, Nc), dtype=jnp.float64)

    for k, (r, c) in enumerate(idx):
        residues = residues.at[r, c, :].set(C_flat[k])
        D_mat = D_mat.at[r, c].set(float(D_vec[k]))
        E_mat = E_mat.at[r, c].set(float(E_vec[k]))
        if r != c:
            # Y-matrices are symmetric (not Hermitian): H[c,r] = H[r,c], not conj.
            residues = residues.at[c, r, :].set(C_flat[k])
            D_mat = D_mat.at[c, r].set(float(D_vec[k]))
            E_mat = E_mat.at[c, r].set(float(E_vec[k]))

    model = VFModel(
        poles=jnp.array(poles_np),
        residues=residues,
        D=D_mat,
        E=E_mat,
    )

    # -------------------------------------------------------------------------
    # Step 9: Build SSModel, evaluate, compute error
    # -------------------------------------------------------------------------
    ss = vfmodel_to_ss(model, Nc)
    bigHfit_stacked = eval_model(s, ss)  # (Ns, Nc, Nc)
    bigHfit = jnp.moveaxis(bigHfit_stacked, 0, -1)  # (Nc, Nc, Ns)

    H_data = jnp.moveaxis(bigH, -1, 0)  # (Ns, Nc, Nc)
    rmserr = compute_rmserr(H_data, bigHfit_stacked)

    return model, ss, rmserr, bigHfit

def aaa_scalar(f, z, tol=1e-10, mmax=100, *, backend="numpy") -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Discover AAA support with either backend, returning compact NumPy arrays.

    Use ``aaa_scalar_jax`` directly for padded device arrays and batching.
    This wrapper is the explicit host boundary before SciPy pole extraction.
    """
    if backend == "numpy":
        return aaa_scalar_numpy(f, z, tol=tol, mmax=mmax)
    if backend != "jax":
        raise ValueError("AAA backend must be 'numpy' or 'jax'")
    from .aaa_jax import aaa_scalar_jax

    result = aaa_scalar_jax(f, z, tol=tol, mmax=mmax)
    count = int(result.count)
    return tuple(np.asarray(value)[:count] for value in result[:3])