ovf/main.py

from core.orthonormal_basis import generate_laguerre_basis
from core.sk_iter import generate_starting_poles
import numpy as np
import skrf as rf
import matplotlib.pyplot as plt
import sympy as sp
from scipy.linalg import null_space
from plotly.subplots import make_subplots
import plotly.graph_objs


network = rf.Network("/tmp/paramer/simulation/4000/4000.s2p")
freqs = network.f
s = freqs * 2j * np.pi
vf = rf.VectorFitting(network)
vf.vector_fit(n_poles_real=2, n_poles_cmplx=1,parameter_type="y")

poles = vf.poles
residues = vf.residues
y = vf.network.y

# fig, ax = plt.subplots(2, 2)
# fig.set_size_inches(12, 8)
# vf.plot("mag",0,0,freqs,ax=ax[0][0],parameter="y")
# rms_error11 = vf.get_rms_error(0,0,"y")
# print("rms_error11",rms_error11)
# vf.plot("mag",1,0,freqs,ax=ax[1][0],parameter="y")
# rms_error21 = vf.get_rms_error(1,0,"y")
# print("rms_error21",rms_error21)
# vf.plot("mag",0,1,freqs,ax=ax[0][1],parameter="y")
# rms_error12 = vf.get_rms_error(0,1,"y")
# print("rms_error12",rms_error12)
# vf.plot("mag",1,1,freqs,ax=ax[1][1],parameter="y")
# rms_error22 = vf.get_rms_error(1,1,"y")
# print("rms_error22",rms_error22)
# fig.tight_layout()
# plt.show()
# plt.savefig(f"img.png")

def formula_67(s,y):
    diag_values = [y[i][0][0] for i in range(len(y))]
    H = np.diag(diag_values)
    P = 2
    start_poles = generate_starting_poles(P,beta_min=1e8,beta_max=freqs[-1])
    basis = generate_laguerre_basis(start_poles,s).T
    print("start_poles",start_poles)
    print("basis",basis)

    # first step iteration
    # A*x = b

    A11 = np.real(basis)
    print("A11 shape:",A11.shape)
    A12 = np.real(- H @ basis[:,1:])
    print("A12 shape:",A12.shape)
    # print("A11",A11)
    A21 = np.imag(basis)
    print("A21 shape:",A21.shape)
    A22 = np.imag(- H @ basis[:,1:])
    print("A22 shape:",A22.shape)
    A1 = np.hstack([A11,A12])
    A2 = np.hstack([A21,A22])
    A = np.vstack([A1,A2])
    print ("A shape:",A.shape)

    b1 = np.real(H @ basis[:,0])
    b2 = np.imag(H @ basis[:,0])
    b = np.hstack([b1,b2])
    Q, R = np.linalg.qr(A, mode='reduced')
    print("Q :",Q)
    print("R :",R)
    print("b shape:",b.shape)
    # x = np.linalg.solve(R, Q.T @ b)
    x_star, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)

    # print("x_qr",x_star)
    print("residuals",residuals)
    # print("rank",rank)
    # print("s",s)

    x_ne = np.linalg.inv(A.T @ A) @ A.T @ b
    print("x_ne",x_ne)

    # sk iteration
    target_vec = np.array([1+0j] + list(x_ne[len(x_ne)//2+1:]))
    D_t = basis @ target_vec
    print("D_t",D_t)
    K = 25

    for i in range(K):
        print(f"Iteration {i+1}/{K}")
        A11 = np.real(basis / D_t[:, np.newaxis])
        A12 = np.real(- H @ basis[:,1:] / D_t[:, np.newaxis])
        A21 = np.imag(basis / D_t[:, np.newaxis])
        A22 = np.imag(- H @ basis[:,1:] / D_t[:, np.newaxis])
        A1 = np.hstack([A11,A12])
        A2 = np.hstack([A21,A22])
        A = np.vstack([A1,A2])
        b1 = np.real(H @ basis[:,0])
        b2 = np.imag(H @ basis[:,0])
        b = np.hstack([b1,b2])
        x_star, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)
        # print("x_lstsq",x_star)
        # print("residuals",residuals)
        # print("rank",rank)
        # print("s",s)
        x_ne = np.linalg.inv(A.T @ A) @ A.T @ b
        # print("x_ne",x_ne)
        target_vec = np.array([1+0j] + list(x_ne[len(x_ne)//2+1:]))
        D_t_pre = D_t
        D_t = basis @ target_vec
        print("Dt", D_t)
        # print("D_t/D_t_pre",D_t/D_t_pre)
        # print("D_t",D_t)
        n_target_vec = np.array(list(x_ne[:len(x_ne)//2+1]))
        N_t = basis @ n_target_vec
        # print("H = N_t / D_t",np.abs(N_t / D_t))

class formula_70:
    """
    VF-(70) with final pole-residue model.
    After fit():
        self.poles : (P,) complex
        self.res   : (P, M) complex  # residues per response (columns)
        self.h     : (M,) complex    # optional constants
        self.g     : (M,) complex    # optional proportional terms
    """

    # -------- internals --------

    def __init__(self):
        self.cond = []
        self.rel = []
    @staticmethod
    def _psi(s, z):
        s = np.asarray(s, np.complex128).reshape(-1)
        z = np.asarray(z, np.complex128).reshape(-1)
        return 1.0 / (s[:, None] + z[None, :])

    @staticmethod
    def _lhp(z):
        z = np.asarray(z, np.complex128).reshape(-1).copy()
        z[z.real > 0] = -np.conj(z[z.real > 0])
        return z

    def _build_70(self, s, H_list, z_ref, d0):
        Hs = [np.asarray(h, np.complex128).reshape(-1) for h in H_list]
        K, P, M = len(s), len(z_ref), len(Hs)
        Psi = self._psi(s, z_ref)
        Z = np.zeros((K, P))
        rows, rhs = [], []
        for m, H in enumerate(Hs):
            Hp = H[:, None]
            L_re = np.real(Hp * Psi);  L_im = np.imag(Hp * Psi)   # acts on c
            R_re = -np.real(Psi);      R_im = -np.imag(Psi)       # acts on r^(m)
            rows.append(np.hstack([L_re] + [R_re if j == m else Z for j in range(M)]))
            rhs.append(-np.real(d0 * H))
            rows.append(np.hstack([L_im] + [R_im if j == m else Z for j in range(M)]))
            rhs.append(-np.imag(d0 * H))
        A = np.vstack(rows)
        b = np.concatenate(rhs)
        return A, b, M, P

    def _step_70(self, s, H_list, z_ref, d0=1.0, scale=True):
        A, b, M, P = self._build_70(s, H_list, z_ref, d0)
        if scale:
            coln = np.maximum(np.linalg.norm(A, axis=0), 1e-12)
            x, *_ = np.linalg.lstsq(A / coln, b, rcond=None)
            x = x / coln
            cond = np.linalg.cond(A)
        else:
            x, *_ = np.linalg.lstsq(A, b, rcond=None)
            cond = np.linalg.cond(A)

        c = x[:P]
        res_ratio = np.empty((P, M), np.complex128)
        off = P
        for m in range(M):
            res_ratio[:, m] = x[off:off+P]; off += P
        # relocate poles with test matrix
        Sigma = np.diag(-np.asarray(z_ref, np.complex128))
        T = Sigma - (np.ones((P, 1), np.complex128) @ (c.reshape(1, -1) / d0))
        z_new = -np.linalg.eigvals(T)
        z_new = self._lhp(z_new)
        # cond = np.linalg.cond(T)
        return z_new, c, cond, res_ratio

    # -------- public API --------
    def fit(self, s, H, z0, n_iter=20, d0=1.0, include_const=True, include_linear=False, verbose=True):
        """
        s : (K,) complex samples (j*2π*f_sel)
        H : (K,) complex  or (K,M) or list of M vectors
        z0: initial poles (P,) complex (LHP + conjugate pairs recommended)
        """
        # normalize responses -> list
        if isinstance(H, (list, tuple)):
            H_list = [np.asarray(h, np.complex128).reshape(-1) for h in H]
        else:
            H_arr = np.asarray(H, np.complex128)
            if H_arr.ndim == 1:
                H_list = [H_arr]
            elif H_arr.ndim == 2 and H_arr.shape[0] == len(s):
                H_list = [H_arr[:, i].copy() for i in range(H_arr.shape[1])]
            else:
                raise ValueError("H must be (K,), list of (K,), or (K,M) with M responses.")
        M = len(H_list)

        z = self._lhp(np.asarray(z0, np.complex128))
        # SK/VF relocations
        for it in range(n_iter):
            z_next, c_last, cond, _ = self._step_70(s, H_list, z, d0=d0, scale=True)
            rel = np.linalg.norm(z_next) / max(1.0, np.linalg.norm(z))
            if verbose:
                print(f"[VF-70] iter {it+1:02d}/{n_iter:02d}  Δz_rel={rel} cond(z)={cond}")
            self.cond.append(cond)
            self.rel.append(rel)
            z = z_next

        # ---- Finalize: refit residues with fixed poles z (pole–residue model) ----
        Phi = self._psi(s, z)                   # K x P
        # Build design matrix for extras
        extras = []
        if include_const:  extras.append(np.ones(len(s), np.complex128))
        if include_linear: extras.append(s.astype(np.complex128))
        if extras:
            X_base = np.column_stack([Phi] + extras)   # K x (P + E)
        else:
            X_base = Phi

        res = np.empty((len(z), M), np.complex128)
        h = np.zeros(M, np.complex128)
        g = np.zeros(M, np.complex128)

        for m, Hm in enumerate(H_list):
            theta, *_ = np.linalg.lstsq(X_base, Hm, rcond=None)  # complex LS
            res[:, m] = theta[:len(z)]
            e = len(theta) - len(z)
            if e >= 1: h[m] = theta[len(z)]
            if e >= 2: g[m] = theta[len(z)+1]

        # store the rational function
        self.poles = z              # (P,)
        self.res   = res            # (P, M)
        self.h     = h              # (M,)
        self.g     = g              # (M,)
        return self

    # Evaluate the stored **rational** model on any grid
    def evaluate(self, s_eval, m=None):
        if not hasattr(self, "poles"):
            raise RuntimeError("Model not fitted. Call fit(...) first.")
        s_eval = np.asarray(s_eval, np.complex128).reshape(-1)
        Phi = self._psi(s_eval, self.poles)         # K_eval x P
        if m is None:
            H = Phi @ self.res
            if np.any(self.h): H += self.h
            if np.any(self.g): H += s_eval[:, None] * self.g
            return H                                  # (K_eval, M) or (K_eval,1)
        m = int(m)
        H = Phi @ self.res[:, m]
        H += self.h[m]
        H += s_eval * self.g[m]
        return H                                      # (K_eval,)
    
    # def plot_rel_and_cond(self):
    #     fig, (ax1, ax2) = plt.subplots(2,1,figsize=(15,20))
    #     ax1.plot(np.log(self.rel), 'g-', label='rel')
    #     ax2.plot(np.log(self.cond), 'b-', label='cond')
    #     ax1.set_xlabel('Iteration')
    #     ax1.set_ylabel('Relative Change', color='g')
    #     ax2.set_ylabel('Condition Number', color='b')
    #     ax1.tick_params(axis='y', labelcolor='g')
    #     ax2.tick_params(axis='y', labelcolor='b')
    #     fig.tight_layout()
    #     plt.title('Relative Change and Condition Number per Iteration')
    #     # plt.show()
    #     plt.savefig(f"img_rel_cond.png")
    def plot_rel_and_cond(self):
        fig = make_subplots(rows=2, cols=1, subplot_titles=('Relative Change', 'Condition Number'))
        fig.add_trace(plotly.graph_objs.Scatter(y=self.rel, mode='lines+markers', name='rel'), row=1, col=1)
        fig.add_trace(plotly.graph_objs.Scatter(y=self.cond, mode='lines+markers', name='cond'), row=2, col=1)
        fig.update_xaxes(title_text='Iteration', row=1, col=1)
        fig.update_yaxes(title_text='Relative Change', row=1, col=1)
        fig.update_yaxes(title_text='Condition Number', row=2, col=1)
        fig.update_layout(height=800, width=800, title_text='Relative Change and Condition Number per Iteration')
        fig.write_image("img_rel_cond.png")
        # fig.show()

    def evaluate_on_freq(self, freq_eval, m=None):
        return self.evaluate(1j * 2*np.pi * np.asarray(freq_eval, float), m=m)

    # Optional: save/load the final rational model
    def save(self, path):
        if not hasattr(self, "poles"):
            raise RuntimeError("Nothing to save; call fit(...) first.")
        np.savez(path, poles=self.poles, res=self.res, h=self.h, g=self.g)

    @classmethod
    def load(cls, path):
        d = np.load(path, allow_pickle=False)
        obj = cls()
        obj.poles = d["poles"]; obj.res = d["res"]; obj.h = d["h"]; obj.g = d["g"]
        return obj
    
def auto_select(H, freq,
                n_baseline=64,          # log-spaced backbone points
                peak_prominence=0.05,   # fraction of |H| dB dynamic range for peak detection
                peak_window=5,          # take ±peak_window samples around each peak
                topgrad_q=0.98,         # keep top 2% largest slope/phase-change points
                max_points=25,         # final cap on selected samples (None = no cap)
                ensure_ends=True):
    """
    Select several significant sample points for vector fitting.

    Strategy:
      1) Always keep endpoints (optional).
      2) Add a log-spaced baseline over the band.
      3) Detect resonance peaks in |H| (on a log scale) and keep small windows around them.
      4) Add points with the largest magnitude slope and phase-change (w.r.t log-f).
      5) De-duplicate, sort, and optionally thin to 'max_points' with priority
         to endpoints and detected peaks.

    Parameters
    ----------
    H : (N,) complex array
        Frequency response samples.
    freq : (N,) float array
        Frequency axis [Hz], strictly increasing.
    n_baseline : int
        Count of log-spaced baseline samples across the band.
    peak_prominence : float
        Peak prominence threshold as a fraction of the dynamic range in log|H|.
        0.05 ≈ keep peaks ≥ 5% of the range.
    peak_window : int
        Number of neighbor indices to include on each side of every detected peak.
    topgrad_q : float in (0,1)
        Quantile for selecting strong slope/phase points.
        0.98 ⇒ keep the top 2% largest derivatives.
    max_points : int or None
        If not None, cap the total number of selected indices to this value.
    ensure_ends : bool
        Always include the first and last samples.

    Returns
    -------
    H_sel : (K,) complex array
    freq_sel : (K,) float array
    """
    H = np.asarray(H).reshape(-1)
    f = np.asarray(freq).reshape(-1)
    if H.size != f.size:
        raise ValueError("H and freq must have the same length.")
    N = f.size
    if N < 4:
        return H.copy(), f.copy()

    eps = 1e-16
    mag = np.abs(H)
    logmag = np.log10(mag + eps)
    phase = np.unwrap(np.angle(H))

    # log-frequency axis (scale-invariant derivatives)
    # keep it linear if any non-positive freq sneaks in
    if np.all(f > 0):
        lf = np.log(f)
    else:
        lf = f.copy()

    dlf = np.gradient(lf)
    d_logmag = np.gradient(logmag) / (dlf + 1e-16)
    d_phase  = np.gradient(phase)  / (dlf + 1e-16)

    idx = set()
    if ensure_ends:
        idx.update([0, N-1])

    # 1) log-spaced baseline
    if n_baseline > 0:
        # map a log grid to nearest indices
        grid = np.linspace(lf.min(), lf.max(), n_baseline)
        base_idx = np.clip(np.searchsorted(lf, grid), 0, N-1)
        idx.update(np.unique(base_idx).tolist())

    # 2) peaks in |H|
    try:
        from scipy.signal import find_peaks
        dyn = logmag.max() - logmag.min()
        prom = peak_prominence * (dyn + 1e-12)
        peaks, _ = find_peaks(logmag, prominence=prom)
    except Exception:
        # simple fallback: strict local maxima
        peaks = np.where((mag[1:-1] > mag[:-2]) & (mag[1:-1] > mag[2:]))[0] + 1

    for p in peaks:
        lo = max(0, p - peak_window)
        hi = min(N, p + peak_window + 1)
        idx.update(range(lo, hi))

    # 3) strongest slope / phase-change points
    thr_slope = np.quantile(np.abs(d_logmag), topgrad_q)
    thr_phase = np.quantile(np.abs(d_phase),  topgrad_q)
    idx.update(np.where(np.abs(d_logmag) >= thr_slope)[0].tolist())
    idx.update(np.where(np.abs(d_phase)  >= thr_phase)[0].tolist())

    # 4) finalize set
    sel = np.array(sorted(idx), dtype=int)

    # 5) optional thinning with priority to endpoints and peaks
    if max_points is not None and sel.size > max_points:
        priority = np.zeros(sel.size, dtype=int)
        if ensure_ends:
            priority[(sel == 0) | (sel == N-1)] = 3
        if peaks.size:
            priority[np.isin(sel, peaks)] = np.maximum(priority[np.isin(sel, peaks)], 2)

        keep = []
        budget = max_points
        # keep highest-priority first
        for lev in (3, 2, 1, 0):
            cand = sel[priority == lev]
            if cand.size == 0:
                continue
            if cand.size <= budget:
                keep.extend(cand.tolist())
                budget -= cand.size
            else:
                step = max(1, int(np.ceil(cand.size / budget)))
                keep.extend(cand[::step][:budget].tolist())
                budget = 0
            if budget == 0:
                break
        sel = np.array(sorted(set(keep)), dtype=int)

    return H[sel], f[sel]


if __name__ == "__main__":
    # formula_67(s,y)
    H11 = np.array([y[i,0,0] for i in range(len(y))])
    H11_slice,freqs_slice = auto_select(H11,freqs,max_points=20)
    s_slice = freqs_slice * 2j * np.pi
    P_pairs = 1
    z0 = generate_starting_poles(P_pairs, beta_min=1e8, beta_max=freqs_slice[-1])
    z0 = np.array(z0, dtype=np.complex128)

    f70 = formula_70()
    K = 10
    model = f70.fit(s_slice, H11_slice, z0, n_iter=K, d0=1.0, verbose=True)
    model.save("vf70_model.npz")
    model.plot_rel_and_cond()

    Hfit_dense = model.evaluate_on_freq(freqs)

    fig, axes = plt.subplots(2, 1, figsize=(10, 8), sharex=False)

    # (1) Magnitude plot like your screenshot
    ax0 = axes[0]
    ax0.plot(freqs, np.abs(H11), 'o', ms=4, color='red', label='Samples')
    ax0.plot(freqs, np.abs(Hfit_dense), '-', lw=2, color='k', label='Fit')
    ax0.plot(freqs_slice, np.abs(H11_slice), 'x', ms=4, color='blue', label='Input Samples')
    ax0.set_title("Response i=0, j=0")
    ax0.set_ylabel("Magnitude")
    ax0.legend(loc="best")

    # (2) RMS error vs iteration
    # ax1 = axes[1]
    # its = np.arange(1, K+1)
    # ax1.plot(its, hist["rms_rel"], '-o', lw=2)
    # ax1.set_xlabel("Iteration")
    # ax1.set_ylabel("RMS error (relative)")
    # ax1.grid(True, alpha=0.3)
    # ax1.set_title(f"RMS(final) = {hist['rms_rel'][-1]:.3e}")

    fig.tight_layout()
    plt.savefig(f"img_formula_70.png")