pole_id ¤

def identify_poles(
    f: jnp.ndarray,  # (nnn, Ns) complex — elements to fit
    s: jnp.ndarray,  # (Ns,) complex
    poles: np.ndarray,  # (N,) complex — current poles (NumPy, outside JAX)
    weight: jnp.ndarray,  # (nnn, Ns) or (1, Ns)
    opts: FitOptions,
) -> np.ndarray:
    """Identify new poles from the current poles using FRVF (Stage 1).

    Args:
        f:      (nnn, Ns) complex — frequency-domain data (upper-triangle elements).
        s:      (Ns,) complex.
        poles:  (N,) complex — current pole set (NumPy array, updated each call).
        weight: (nnn, Ns) or (1, Ns) — fitting weights.
        opts:   FitOptions.

    Returns:
        new_poles: (N,) complex NumPy array — updated pole locations.

    """
    N = len(poles)
    Ns = s.shape[0]
    nnn = f.shape[0]
    offs = opts.asymp - 1  # 0, 1, or 2

    cindex = compute_cindex(poles)
    Dk = build_dk(s, jnp.array(poles), cindex, offs)  # (Ns, N+offs)

    # --- compute scale for relaxation row ---
    w = jnp.broadcast_to(weight, (nnn, Ns))
    scale_sq = jnp.sum(jnp.array([jnp.linalg.norm(w[n] * f[n]) ** 2 for n in range(nnn)]))
    scale = float(jnp.sqrt(scale_sq) / Ns)

    # --- relaxed LS system ---
    if opts.relax:
        AA, bb = _build_pole_id_ls(f, s, Dk, weight, N, Ns, nnn, offs, scale)
        x = _solve_scaled(AA, bb)
        x_np = np.array(x)
        sigma_D = x_np[-1]
    else:
        sigma_D = 0.0  # force fallback

    # --- fallback: non-relaxed LS (degenerate sigma D) ---
    use_fallback = not opts.relax or abs(sigma_D) < _TOL_LOW or abs(sigma_D) > _TOL_HIGH
    if use_fallback:
        if sigma_D == 0.0:
            Dnew = 1.0
        elif abs(sigma_D) < _TOL_LOW:
            Dnew = float(np.sign(sigma_D) * _TOL_LOW)
        else:
            Dnew = float(np.sign(sigma_D) * _TOL_HIGH)

        AA2, bb2 = _build_pole_id_ls_norelax(f, s, Dk, weight, N, Ns, nnn, offs, Dnew)
        x2 = _solve_scaled(AA2, bb2)
        # Append Dnew at end to match expected format
        x_np = np.append(np.array(x2), Dnew)
        sigma_D = Dnew

    # --- convert sigma coefficients from real to complex ---
    sigma_C = _real_to_complex_sigma(x_np, cindex, N)

    # --- build companion matrix and compute eigenvalues ---
    ZER = _build_companion_matrix(poles, sigma_C, sigma_D, cindex)
    # Keep the eigensolve on the JAX side. This is the dense operation that
    # must be batchable when pole relocation is lifted into a fixed-shape
    # candidate kernel; the surrounding compatibility path still returns a
    # NumPy array for the existing VF driver.
    new_poles = np.array(jnp.linalg.eigvals(jnp.asarray(ZER)), copy=True)  # (N,) complex

    # --- enforce stability: flip poles with positive real part ---
    if opts.stable:
        unstable = new_poles.real > 0
        new_poles[unstable] = new_poles[unstable] - 2 * new_poles[unstable].real

    # --- sort and cleanup ---
    new_poles = sort_poles(new_poles)

    return new_poles.astype(np.complex128)