Untitled

 avatar
unknown
plain_text
9 months ago
1.4 kB
13
Indexable
def FROG(P_cur, G_cur, I_meas_f, threshold, max_iter, t_P, t_G):
    M = []
    i = 0
    
    # Initial error calculation
    S_proj = project(P_cur, G_cur, np.sqrt(I_meas_f))
    rms_error = np.sqrt(np.mean((np.abs(calc_S_t(P_cur, G_cur)) - np.abs(S_proj))**2))
    M.append(rms_error)
    
    # progress abr
    pbar = tqdm(total=max_iter, desc="Retrieval Progress", position=0)
    
    while rms_error > threshold and i < max_iter:
        S_proj = project(P_cur, G_cur, np.sqrt(I_meas_f))
        
        # Update pulse
        P_cur = next_P(G_cur, S_proj)
        P_cur /= np.max(np.abs(P_cur))
        center_idx = np.argmax(np.abs(P_cur)**2)
        P_cur = np.roll(P_cur, len(P_cur)//2 - center_idx)
        
        G_cur, t_G = Make_G(P_cur, t_P)
        
        rms_error = np.sqrt(np.mean((np.abs(calc_S_t(P_cur, G_cur)) - np.abs(S_proj))**2))
        M.append(rms_error)
        i += 1
        pbar.update(1)
    
    pbar.close()
    
    # Final outputs
    S_out = calc_S_f(calc_S_t(P_cur, G_cur))
    I_out = np.abs(S_out)**2
                     
    P_out = P_cur
    G_out = G_cur / np.max(np.abs(G_cur))
    
    # Center the probe pulse
    max_idx_G = np.argmax(np.abs(G_out)**2)
    shift_G = (len(G_out) // 2) - max_idx_G
    G_out = np.roll(G_out, shift_G)
    
    return P_out, G_out, I_out, S_out, M, t_P, t_G
Editor is loading...
Leave a Comment