run_LF_calibration.py

import ROOT
import re
import numpy as np
import polars as pl
from pathlib import Path
from initialise import load_pueoEvent_Dataset
 
dirname = "simplesim_lat885_lat88_lat87_snr5ish_200evt"
#dirname = "tmp"
#dirname = "simplesim_lat88_5_snr5ish"
if __name__ == "__main__":
    import os
    from scipy.optimize import minimize
    import matplotlib.pyplot as plt
    import matplotlib
    import calibration
    from calibration import generate_antenna_phase_center_pairs_with_placeholders, load_dataset_and_measure_time_delays
    from calibration import retrieve_relevant_antennas_from_all_pairs, compute_time_delay_residuals, objective_function, _plot_calibration_result
    matplotlib.use("Agg") 
    ant_type='LF'
    MC=True
    
    RemoveOutlier=True
    RemoveOutlierTimes=0
    calibration.dirname = dirname 
    calibration.UPSAMPLE_FACTOR=50
 
    # use nominal phase center as the initial guesss of the minimzation
    if ant_type=='MI':
        j25: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/qrh.dat")
        a25: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/aug25/qrh.dat")
        offset_inSIM: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/phase_center_offset_MI.dat")
    elif ant_type=='LF':
        j25: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/lf_cal.dat")#lf_cal.dat
        a25: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/aug25/lf.dat")
        offset_inSIM: Path = os.environ.get("PUEO_UTIL_INSTALL_DIR") / Path("share/pueo/geometry/jun25/phase_center_offset_LF.dat")
    else:
        print("Ant type doesn't exist!")
    
    nominal_pairs = generate_antenna_phase_center_pairs_with_placeholders(qrh_dot_dat=j25, offset_inSIM=offset_inSIM, ant_type=ant_type)
 
    # start with some run number to initialize the Dataset instance.
    ds: ROOT.pueo.Dataset = load_pueoEvent_Dataset(pueo_mc_data=Path(f"/users/PAS0654/yuchieh/pueo_2025/components/pueoAnalysisTools/SharedUtils/{dirname}/"), run_number=20)
    if ant_type=='LF':
        Icefinalpath=f'./{dirname}/'
    else:
        Icefinalpath=''
    
    pulser_directions_and_pairwise_time_delays = (
        load_dataset_and_measure_time_delays(dataset=ds, phase_center_pairs=nominal_pairs, MC=True, AntType=ant_type, IcefinalFileDir=Icefinalpath, corr_fcn='ZNCC') # option corr_fcn='ZNCC'
        )
 
    corr = pulser_directions_and_pairwise_time_delays["max correlation"].to_numpy()
    N_tot_pairs = len(corr)
    print("N_tot_pairs=", N_tot_pairs)
    percentage=99.7 #  at least 99.5 for snr=5 
    pXX = np.percentile(corr, percentage)
    print(f"{percentage}th percentile =", pXX)
    
    # histogram
    plt.figure(figsize=(8,5))
    plt.hist(corr, bins=100)
    plt.axvline(pXX, linestyle="--", linewidth=2, label=f"{percentage}% = {pXX:.3f}")
    plt.xlabel("max correlation")
    plt.ylabel("count")
    plt.legend()
    plt.tight_layout()
    os.makedirs(f"plot/{dirname}", exist_ok=True)
    plt.savefig(f"plot/{dirname}/max_correlation.png")
 
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.filter((pl.col("max correlation")> pXX) | (pl.col("max correlation")> 0.95))
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.filter(pl.col("measured time delays (ns)").is_between(-16, 16))
    
 
 
    # add index to pair in order to track them later
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.with_row_index("row_id")
 
    
 
    N_pairs = pulser_directions_and_pairwise_time_delays.height
    
    #pl.Config.set_tbl_cols(30)      # show up to 1000 columns
    #pl.Config.set_tbl_rows(30)      # optional: show more rows too
    #pl.Config.set_tbl_width_chars(200)  # optional: wider display
    print("N pairs=", N_pairs)
 
    relevant_antennas = (
        retrieve_relevant_antennas_from_all_pairs(
            phase_center_pairs=pulser_directions_and_pairwise_time_delays
        )
    )
    print("relevant_antennas=", relevant_antennas) 
    axis_list_initial = ["Rho","Phi","Z", "CableDelay"] # adding random offst the each axis
    plot_range=[0.1,0.1,10/180*3.1415926] # cm, cm, deg
    unit_scale=[100,100,180/3.1415926] # m to cm, m to cm, rad to deg
 
    x0_randomguess = {}
    fixed_antenna={}
    
    best_result_rho = None
    best_result_z = None
    best_result_phi = None
    best_result_cabledelay = None
    result_rho = None
    result_z = None
    result_phi = None
    result_cabledelay = None
    best_value = np.inf
    best_trace = []
    tmp_trace = []
    
    # repeat the minimization n times and pick the minimal loss function. This is to reduce dependency on initial value, and stuck in local minimum
    for trials in range(150): #50
        for index, axis_label in enumerate(axis_list_initial):
            if MC:
 
                # For simulation data, apply some random variation to the initial guess, or we won't initially start from the correct answer
 
                x0_randomguess[axis_label] = relevant_antennas[axis_label].to_numpy().copy()
                fixed_antenna[axis_label] = x0_randomguess[axis_label][0]
                if axis_label=='Z':
                    x0_randomguess[axis_label][1:] = x0_randomguess[axis_label][1:] + np.random.uniform(-0.1, 0.1, size=len(relevant_antennas)-1)
                    
                elif axis_label=='Phi':
                    x0_randomguess[axis_label][1:] = x0_randomguess[axis_label][1:] + np.random.uniform(-0.05, 0.05, size=len(relevant_antennas)-1)
 
                elif axis_label=='Rho':
                    x0_randomguess[axis_label][1:] = x0_randomguess[axis_label][1:] + np.random.uniform(-0.1, 0.1, size=len(relevant_antennas)-1)
                
                elif axis_label=='CableDelay':
                    x0_randomguess[axis_label][1:] = x0_randomguess[axis_label][1:] + np.random.uniform(-0.066, 0.066, size=len(relevant_antennas)-1) #ns
 
 
                pulser_directions_and_pairwise_time_delays = (
                    pulser_directions_and_pairwise_time_delays
                    .with_columns(
                    pl.col(f"A1_{axis_label} (placeholder)")
                    .replace_strict(relevant_antennas[f"{axis_label} (placeholder)"], x0_randomguess[axis_label])
                    .alias(f"A1_{axis_label}"),
 
                    pl.col(f"A2_{axis_label} (placeholder)")
                    .replace_strict(relevant_antennas[f"{axis_label} (placeholder)"], x0_randomguess[axis_label])
                    .alias(f"A2_{axis_label}")
                    )
                )
            else:
                x0_randomguess[axis_label] = (
                    relevant_antennas[axis_label].to_numpy()
                )
 
        axis_list=[] 
        for _ in range(3):  # repeat 3 times
            axis_list.extend([["CableDelay"],["Rho"], ["Phi"],["Z"]])
        for _ in range(6):
            axis_list.extend([["CableDelay"], ["Z"]])
        axis_list.extend([["Rho"], ["Phi"]])
        
        # figure to plot delta T histograms
        fig1, ax1 = plt.subplots(figsize=(8, 5))
 
        for index, axis_label in enumerate(axis_list):
            
            AXIS = axis_label[0]  # "Rho", "Phi", or "Z"
        
            #print("AXIS=",axis_label ," true:", relevant_antennas[AXIS]) 
            x0 = x0_randomguess[AXIS]  
            #print("starting points=", x0)
            
            if ant_type=='LF':
                x0=x0[1:]
                placeholder = relevant_antennas[f"{AXIS} (placeholder)"][1:]
            else:
                x0=x0[1:]
                placeholder = relevant_antennas[f"{AXIS} (placeholder)"][1:]
            
            if AXIS== 'Rho':
                bounds = [(1.0, 2.5)] * (len(x0))
            
            if AXIS=='Phi':
                bounds = [(-3.14, 3.14)] * (len(x0))
            if AXIS == 'Z':
                bounds =[(-12, -4)] * (len(x0))
            if AXIS=='CableDelay':
                bounds =[(-0.1, 0.1)] * (len(x0))
 
            result = minimize(fun=objective_function,
                x0=x0,
                args=(placeholder,
                pulser_directions_and_pairwise_time_delays, AXIS),
                method='Nelder-Mead',
                bounds=bounds,
                options={'maxiter': 500000000})
            
            print("result=",result.fun)
            tmp_trace.append(result.fun)
            
            
            pulser_directions_and_pairwise_time_delays = (
                pulser_directions_and_pairwise_time_delays
                .with_columns(
                pl.col(f"A1_{AXIS} (placeholder)")
                .replace_strict(placeholder, result.x, default=pl.col(f"A1_{AXIS}"))
                .alias(f"A1_{AXIS}"),
 
                pl.col(f"A2_{AXIS} (placeholder)")
                .replace_strict(placeholder, result.x, default = pl.col(f"A1_{AXIS}"))
                .alias(f"A2_{AXIS}")
                )
            )
 
            resultx = np.insert(result.x, 0, x0_randomguess[AXIS][0])
            x0_randomguess[AXIS] = resultx
            
 
            df_resid = compute_time_delay_residuals(
                phase_center_guess=resultx,
                phase_center_placeholders=relevant_antennas[f"{AXIS} (placeholder)"],
                pulser_directions_and_measured_time_delays=pulser_directions_and_pairwise_time_delays,
                axis=AXIS,
            )
 
            residuals = df_resid["residual (ns)"].to_numpy()
            cmap = plt.cm.copper
            color = cmap(index / (len(axis_list) - 1)) if len(axis_list) > 1 else cmap(1.0)
            ax1.hist(residuals, bins=80, alpha=0.3, label=f'{index}:[{AXIS}]', color=color)
            
            if (index>=8 and RemoveOutlier==True and RemoveOutlierTimes<5):
                RemoveOutlierTimes=RemoveOutlierTimes+1
                row_before = df_resid.height
 
                resid_col = "residual (ns)"
 
                mu = df_resid.select(pl.col(resid_col).mean()).item()
                sigma = df_resid.select(pl.col(resid_col).std()).item()
                
                print("You're removing outlier over 3 sigma, mu=", mu, "sigma=", sigma)
                good_rows = (
                    df_resid
                    .filter((pl.col(resid_col) - mu).abs() <= 3 * sigma)
                    .select("row_id")
                )
                
                row_after = good_rows.height
 
                print("Removing", row_before-row_after, "rows")
                pulser_directions_and_pairwise_time_delays = (
                    pulser_directions_and_pairwise_time_delays
                    .join(good_rows, on="row_id", how="inner")
                )
 
                #relevant_antennas = (
                #    retrieve_relevant_antennas_from_all_pairs(
                #    phase_center_pairs=pulser_directions_and_pairwise_time_delays
                #    )
                #)
 
 
            if AXIS == 'Rho':
                result_rho = result
            elif AXIS == 'Phi':
                result_phi = result
            elif AXIS == 'Z':
                result_z = result
            elif AXIS == 'CableDelay':
                result_cabledelay = result
 
            #_plot_calibration_result(axis=AXIS, before=relevant_antennas, after=resultx, anttype=ant_type)
        
 
        print("result.fun=", result.fun, "result.success=", result.success)     
        if result.fun < best_value and result.success:
            best_trace = tmp_trace.copy()
            best_value = result.fun
            best_result_rho = result_rho
            best_result_phi = result_phi
            best_result_z = result_z
            best_result_cabledelay = result_cabledelay
            print("temp best_value=", best_value)
            print("temp best result=", best_result_rho.x)
            
            ax1.set_xlabel("measured - expected (ns)")
            ax1.set_ylabel("Counts")
            ax1.set_title("Residual time delays")
            ax1.grid(True, alpha=0.3)
            ax1.legend()
            fig1.tight_layout()
            fig1.savefig(f"plot/{dirname}/{ant_type}step_timediff.png")
        
        tmp_trace.clear()
    
    resultList = [best_result_rho,  best_result_phi,  best_result_z, best_result_cabledelay]
    for i, AXIS in enumerate(["Rho", "Phi", "Z", "CableDelay"]): #["Rho", "Phi", "Z", "CableDelay"]
 
        resultx_full_axis = np.insert(resultList[i].x, 0,  fixed_antenna[AXIS]) 
        print(f"{AXIS} resultx_full_axis=", resultx_full_axis)
        _plot_calibration_result(axis=AXIS, before=relevant_antennas, after=resultx_full_axis, anttype=ant_type)
 
    print("best_value=", best_value)
    
    print("best_trace=", best_trace)
    plt.figure(figsize=(8,5)) 
    step = len(best_trace)
    if step>1:
        plt.plot(np.arange(step), np.sqrt(np.asarray(best_trace))/N_pairs)
        plt.yscale('log')
        plt.grid()
        plt.xlabel("Minimization iteration")
        plt.ylabel("Averging $\Delta$ T per pair (ns)")
        print(f"plot/{dirname}/{ant_type}best_trace.png")
        plt.savefig(f"plot/{dirname}/{ant_type}best_trace.png")