sys_ID_working.py

import numpy as np
from scipy.signal import lsim, TransferFunction, dlti, dlsim
from scipy.optimize import least_squares, minimize
from scipy.fft import fft, fftfreq

class TFestResult:
    """Container for estimated transfer functions (supports MIMO)."""
    def __init__(self, tf_list, Ts=0.0):
        self.tf_list = tf_list  # List of lists of TransferFunction objects
        self.Ts = Ts  # Sample time
        self.ny = len(tf_list)      # number of outputs
        self.nu = len(tf_list[0]) if self.ny > 0 else 0  # number of inputs

    def print_summary(self):
        for i, row in enumerate(self.tf_list):
            for j, tf in enumerate(row):
                print(f"TF y{i+1}/u{j+1}:")
                print(f" Num: {tf.num}")
                print(f" Den: {tf.den}\n")


def tfest(u, y, t=None, npoles=2, nzeros=None, iodelay=None, Ts=0.0):
    """
    Estimate transfer function from input u and output y.
    
    Parameters
    ----------
    u : array_like
        Input signal, shape (N,) or (N, nu)
    y : array_like
        Output signal, shape (N,) or (N, ny)
    t : array_like, optional
        Time vector. Required for continuous-time systems.
    npoles : int or array_like
        Number of poles. Can be scalar or (ny, nu) array.
    nzeros : int or array_like, optional
        Number of zeros. Defaults to npoles - 1
    iodelay : int or array_like, optional
        Input-output delay in samples. Shape (ny, nu)
    Ts : float
        Sampling time. 0 = continuous-time, >0 discrete-time
    """
    u = np.atleast_2d(u).T if u.ndim == 1 else np.array(u)
    y = np.atleast_2d(y).T if y.ndim == 1 else np.array(y)
    N, nu = u.shape
    _, ny = y.shape

    # Handle npoles/nzeros as scalars or arrays
    if np.isscalar(npoles):
        npoles_arr = np.full((ny, nu), npoles, dtype=int)
    else:
        npoles_arr = np.array(npoles)
    
    if nzeros is None:
        nzeros_arr = npoles_arr - 1
    elif np.isscalar(nzeros):
        nzeros_arr = np.full((ny, nu), nzeros, dtype=int)
    else:
        nzeros_arr = np.array(nzeros)
    
    if iodelay is None:
        iodelay_arr = np.zeros((ny, nu), dtype=int)
    else:
        iodelay_arr = np.atleast_2d(iodelay)

    # Time vector
    if t is None:
        if Ts > 0:
            t = np.arange(N) * Ts
        else:
            raise ValueError("Time vector t is required for continuous-time systems (Ts=0)")
    t = np.asarray(t)

    # Estimate each transfer function separately
    tf_list = []
    for i in range(ny):
        tf_row = []
        for j in range(nu):
            tf_ij, _ = _estimate_single_tf(
                u[:, j], y[:, i], t, 
                npoles_arr[i, j], nzeros_arr[i, j], 
                iodelay_arr[i, j], Ts
            )
            tf_row.append(tf_ij)
        tf_list.append(tf_row)
    
    return TFestResult(tf_list, Ts)


def _estimate_single_tf(u, y, t, npoles, nzeros, iodelay, Ts):
    """Estimate a single SISO transfer function with robust initialization."""
    
    # Apply IO delay
    if iodelay > 0:
        u_delayed = np.zeros_like(u)
        u_delayed[iodelay:] = u[:-iodelay]
    else:
        u_delayed = u.copy()
    
    # Helper: convert parameter vector to TF
    def build_tf(x):
        num = x[:nzeros + 1]
        den = np.r_[1, x[nzeros + 1:]]  # leading 1 in denominator
        if Ts > 0:
            return dlti(num, den, dt=Ts)
        else:
            return TransferFunction(num, den)

    # Residuals for least_squares with regularization
    def residuals(x, lambda_reg=1e-6):
        try:
            tf = build_tf(x)
            if Ts == 0:  # continuous-time
                _, y_pred, _ = lsim(tf, u_delayed, t)
            else:  # discrete-time
                _, y_pred = dlsim(tf, u_delayed, t=t)
            
            err = y_pred.flatten() - y
            
            # Add small regularization to prevent extreme coefficients
            reg = np.sqrt(lambda_reg) * x
            
            return np.concatenate([err, reg])
        except:
            # Return large error if simulation fails
            return np.ones(len(y) + len(x)) * 1e6

    # **SMART INITIALIZATION**
    # Estimate DC gain
    u_std = np.std(u_delayed)
    y_std = np.std(y)
    dc_gain = y_std / (u_std + 1e-10)
    
    # Estimate dominant frequency from output
    Y = fft(y)
    freqs = fftfreq(len(t), t[1] - t[0])
    positive_freqs = freqs[1:len(freqs)//2]
    positive_Y = np.abs(Y[1:len(Y)//2])
    
    if len(positive_Y) > 0 and np.max(positive_Y) > 0:
        peak_idx = np.argmax(positive_Y)
        peak_freq = abs(positive_freqs[peak_idx])
        wn = 2 * np.pi * peak_freq if peak_freq > 0 else 5.0
    else:
        wn = 5.0
    
    # Initialize based on system order
    if npoles == 2 and nzeros == 0:
        # Second-order system without zeros: gain / (s^2 + 2*zeta*wn*s + wn^2)
        # For your system: 26.5 / (s^2 + s + 26.5)
        # This means: wn^2 = 26.5, 2*zeta*wn = 1
        
        # Use identified natural frequency
        wn_init = np.sqrt(wn**2) if wn > 0 else 5.0
        zeta_init = 0.1  # light damping
        
        x0_num = [dc_gain * wn_init**2]  # Numerator: K*wn^2 for DC gain = K
        x0_den = [2*zeta_init*wn_init, wn_init**2]  # Denominator: [2*zeta*wn, wn^2]
        
    elif npoles == 2 and nzeros == 1:
        wn_init = wn if wn > 0 else 5.0
        zeta_init = 0.7
        x0_num = [dc_gain, dc_gain * wn_init]
        x0_den = [2*zeta_init*wn_init, wn_init**2]
        
    elif npoles == 1 and nzeros == 0:
        # First-order: K / (tau*s + 1)
        tau_init = 1.0
        x0_num = [dc_gain / tau_init]
        x0_den = [1.0 / tau_init]
        
    else:
        # Generic initialization
        x0_num = np.ones(nzeros + 1) * dc_gain / (nzeros + 1)
        x0_den = np.ones(npoles)
    
    x0 = np.concatenate([x0_num, x0_den])
    
    # Multi-start optimization with different strategies
    best_res = None
    best_cost = np.inf
    
    strategies = [
        ('initial', x0, 1e-6),
        ('scaled', x0 * 0.5, 1e-6),
        ('freq_based', None, 1e-5),
    ]
    
    for strategy_name, x_init, lambda_reg in strategies:
        if x_init is None:
            # Frequency-based alternative initialization
            x_init = np.concatenate([
                [dc_gain * wn**2],
                [1.0, wn**2]
            ])
        
        try:
            res = least_squares(
                lambda x: residuals(x, lambda_reg), 
                x_init, 
                method='lm',
                max_nfev=2000,
                ftol=1e-10,
                xtol=1e-10,
                gtol=1e-10
            )
            
            # Calculate actual fit cost (without regularization)
            tf_test = build_tf(res.x)
            if Ts == 0:
                _, y_pred, _ = lsim(tf_test, u_delayed, t)
            else:
                _, y_pred = dlsim(tf_test, u_delayed, t=t)
            
            actual_cost = np.sum((y_pred.flatten() - y)**2)
            
            if actual_cost < best_cost and actual_cost < np.inf:
                best_cost = actual_cost
                best_res = res
                
        except Exception as e:
            continue
    
    # If all failed, try simple gradient descent
    if best_res is None or best_cost > np.var(y) * len(y):
        print("Warning: Least squares failed, trying alternative optimization...")
        
        def cost_function(x):
            try:
                tf = build_tf(x)
                if Ts == 0:
                    _, y_pred, _ = lsim(tf, u_delayed, t)
                else:
                    _, y_pred = dlsim(tf, u_delayed, t=t)
                return np.sum((y_pred.flatten() - y)**2)
            except:
                return 1e10
        
        res_nm = minimize(cost_function, x0, method='Nelder-Mead', 
                         options={'maxiter': 5000, 'xatol': 1e-8, 'fatol': 1e-8})
        
        if res_nm.fun < best_cost:
            best_res = type('obj', (object,), {'x': res_nm.x, 'cost': res_nm.fun})()
            best_cost = res_nm.fun
    
    # Final fallback
    if best_res is None:
        print("Warning: All optimization attempts failed, using initial guess")
        best_res = type('obj', (object,), {'x': x0, 'cost': np.inf})()
    
    # Build final transfer function
    tf_final = build_tf(best_res.x)
    return tf_final, best_res


# -------------------------
# Example usage
# -------------------------
if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from scipy.signal import lti, chirp

    # --- Define SISO system ---
    num = [26.5]
    den = [1.0, 1, 26.5]
    sys_true = lti(num, den)

    # --- Generate input/output data ---
    T = 30.0             # total duration [s]
    fs = 500.0  
    t = np.arange(0, T, 1/fs)
    u = chirp(t, f0=0.5, f1=5.0, t1=T, method='linear')  # Chirp signal
    
    _, y, _ = lsim(sys_true, u, t)
    
    # Add some noise
    np.random.seed(42)  # For reproducibility
    y += np.random.normal(0, 0.01, y.shape)

    # --- Estimate transfer function with TRUE order ---
    print("Estimating transfer function...")
    result = tfest(u, y, t=t, npoles=2, nzeros=0, Ts=0.0)

    # --- Print results ---
    print("\nTrue numerator:", num)
    print("True denominator:", den)
    print("\nEstimated transfer function:")
    result.print_summary()
    
    # --- Compare outputs ---
    tf_est = result.tf_list[0][0]
    _, y_est, _ = lsim(tf_est, u, t)
    
    # Calculate fit percentage
    fit_percent = 100 * (1 - np.linalg.norm(y - y_est) / np.linalg.norm(y - np.mean(y)))
    print(f"Fit: {fit_percent:.2f}%")
    
    # Calculate normalized error
    nrmse = np.sqrt(np.mean((y - y_est)**2)) / (np.max(y) - np.min(y))
    print(f"NRMSE: {nrmse:.4f}")
    
    plt.figure(figsize=(12, 8))
    
    plt.subplot(2, 1, 1)
    plt.plot(t, y, label='True output', alpha=0.7, linewidth=1.5)
    plt.plot(t, y_est, '--', label='Estimated output', alpha=0.7, linewidth=1.5)
    plt.xlabel('Time [s]')
    plt.ylabel('Output')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.title(f'Transfer Function Estimation (Fit: {fit_percent:.1f}%)')
    
    plt.subplot(2, 1, 2)
    plt.plot(t, y - y_est, label='Error', color='red', alpha=0.7)
    plt.xlabel('Time [s]')
    plt.ylabel('Error')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.title('Prediction Error')
    
    plt.tight_layout()
    plt.show()