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
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
import os
from scipy.optimize import minimize
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use("Agg")
 
import sys
import glob
 
 
if __name__ == "__main__":
    ant_type='LF'
    MC=False
    run = 797
    
    RemoveOutlier=True
    
    # read pulser_directions_and_pairwise_time_delays from preprocessed feather files
    files = glob.glob(f"feather_data/pulser_directions_and_pairwise_time_delays_run*.feather")
    #files = glob.glob(f"feather_data/pulser_directions_and_pairwise_time_delays_run797_2518_7533.feather")
    pulser_directions_and_pairwise_time_delays = pl.concat([pl.read_ipc(f) for f in files])
    
    # make a correlation plot
    corr = pulser_directions_and_pairwise_time_delays["max correlation"].to_numpy()
    N_tot_pairs = len(corr)
    print("N pairs before selections = ", N_tot_pairs)
    percentage=90.0#99.8 #  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.axvline(0.9, linestyle="--", linewidth=2)
    plt.xlabel("max correlation")
    plt.ylabel("count")
    plt.legend()
    plt.tight_layout()
    os.makedirs("plot", exist_ok=True)
    plt.savefig(f"plot/{ant_type}max_correlation.png")
    
    # filter out events that doesn't have a successfully caculated incoming pulser direction
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.filter(
        ~pl.col("incoming pulse direction (Phi Theta)")
        .arr.eval(pl.element().is_nan())
        .arr.any()
    )
    
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.filter(pl.col("max correlation")> 0.90)
    pulser_directions_and_pairwise_time_delays = pulser_directions_and_pairwise_time_delays.filter(pl.col("measured time delays (ns)").is_between(-16, 16))
    
    pl.Config.set_tbl_rows(-1)   # show all rows
    '''
    print(
    pulser_directions_and_pairwise_time_delays
    .select([
        "measured time delays (ns)",
        "max correlation",
        pl.col("incoming pulse direction (Phi Theta)")
        .map_elements(
            lambda x: np.rad2deg(np.array(x.to_list())),
            return_dtype=pl.Array(pl.Float64, 2)
        )
        .alias("incoming pulse direction (deg)")
    ])
    )  
    '''
    # 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
    
    print("N total pairs after correlation seletion = ", N_pairs)
 
    relevant_antennas = (
        retrieve_relevant_antennas_from_all_pairs(
            phase_center_pairs=pulser_directions_and_pairwise_time_delays
        )
    )
 
    relevant_antennas = relevant_antennas.with_columns(
        pl.col("AntNum").cast(pl.Utf8).cast(pl.Int64).alias("AntNum")
    )
    
    print("relevant_antennas=", relevant_antennas) 
    
    axis_list_initial = ["Rho","Phi","Z", "CableDelay"] # adding random offst the each axis
 
    x0_randomguess = {}
    fixed_antenna={}
 
    # Now make an axis list which is the 1d optimization order
    # For the prerun minimization to remove the outlier
    axis_list0 = []
    for _ in range(4):  # repeat 4 times
        axis_list0.extend([["CableDelay"],["Rho"], ["Phi"], ["Z"]])
    for _ in range(4):
        axis_list0.extend([["CableDelay"],["Z"]])
    axis_list0.extend([["Rho"]])
    
    # This is for the real minimization run
    axis_list=[]
    for _ in range(4):  # repeat 4 times
        axis_list.extend([["CableDelay"],["Rho"], ["Phi"],["Z"]])
    for _ in range(4):
        axis_list.extend([["CableDelay"],["Z"]])
    axis_list.extend([["Rho"]])
 
    # Do minimization once to throw away outlier. Step1: setup the initial value 
    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) # -0.1, 0.1
                
            elif axis_label=='Phi':
                x0_randomguess[axis_label][1:] = x0_randomguess[axis_label][1:] + np.random.uniform(-0.1, 0.1, size=len(relevant_antennas)-1) # -0.05, 0.05
 
            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) #-0.1, 0.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 -0.066, 0.066
 
            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().copy()
            fixed_antenna[axis_label] = x0_randomguess[axis_label][0]
    
    print("x0_randomguess=", x0_randomguess)
    # figure to plot delta T histograms
    fig1, ax1 = plt.subplots(figsize=(8, 5))
    
    #Do minimization once to throw away outlier. Step2: minimize along the axes in axis_list0
    for index, axis_label in enumerate(axis_list0):
        AXIS = axis_label[0]  # "Rho", "Phi", "Z", or "CableDelay"
    
        x0 = x0_randomguess[AXIS]  
        
        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.0)] * (len(x0))
        if AXIS=='Phi':
            bounds = [(-3.14, 3.14)] * (len(x0))
        if AXIS == 'Z':
            bounds =[(-12, -4)] * (len(x0))
        if AXIS=='CableDelay':
            bounds =[(-1, 1)] * (len(x0)) # -0.066, 0.066
        
        result = minimize(fun=objective_function,
            x0=x0,
            args=(placeholder,
            pulser_directions_and_pairwise_time_delays, AXIS, 'l2'),
            method='Nelder-Mead',
            bounds=bounds,
            options={'maxiter': 50000})
        
        #print("result=",result.fun)
        #print("result.x=", result.x)
        
        
        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
            
 
 
        # calculate residuls distribution 
        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,
        )
        #print('df_resid=', df_resid)
        residuals = df_resid["residual (ns)"].to_numpy()
        cmap = plt.cm.copper
        color = cmap(index / (len(axis_list0) - 1)) if len(axis_list0) > 1 else cmap(1.0)
        if index>=0:
            bin_width = 0.005 #0.005
            bins = np.arange(residuals.min(), residuals.max() + bin_width, bin_width)
            ax1.hist(residuals, bins=bins, alpha=1, label=f'{index}:[{AXIS}]', color=color)
            #ax1.set_xlim(-50,50)
        
        if (index>=12 and RemoveOutlier==True):
            
            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()
            
            Nsigma=100
            
            if index>=12:
                Nsigma = 3.0
            if index>=20:
                Nsigma=2.0
                
            if Nsigma!=100:
                print(f"You're removing outlier over {Nsigma} sigma, mu=", mu, "sigma=", sigma)
            
            good_rows = (
                df_resid
                .filter((pl.col(resid_col) - mu).abs() <= Nsigma * 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")
            )
 
        
    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/{ant_type}step_timediff_RemoveOutliers.png")
    
    plt.close()
    
    # Now we've done selecting samples, we start the real optimization
    
    N_pairs = pulser_directions_and_pairwise_time_delays.height
    print("N pairs after removing outliers = ", N_pairs)
 
    x0_randomguess = {} # the initial value
    fixed_antenna={}
 
    result_rho = None
    result_z = None
    result_phi = None
    result_cabledelay = None
    tmp_trace = []
 
    for index, axis_label in enumerate(axis_list_initial):
        # apply some random variations 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() # fill in initial value 
        fixed_antenna[axis_label] = x0_randomguess[axis_label][0] # fix one antenna
 
        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.1, 0.1, size=len(relevant_antennas)-1) #-0.01, 0.01
 
        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) #-0.01, 0.01
 
        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}")
                )
        )
        
    # figure to plot delta T histograms
    fig1, ax1 = plt.subplots(figsize=(8, 5))
    
    # minimize along axis_list elements
    for index, axis_label in enumerate(axis_list):
 
        AXIS = axis_label[0]  # "Rho", "Phi", "CableDelay" or "Z"
 
        x0 = x0_randomguess[AXIS] # initial value for the AXIS
 
        x0=x0[1:] # position of the first antenna is fixed, so we start from 1
        placeholder = relevant_antennas[f"{AXIS} (placeholder)"][1:]
 
        if AXIS== 'Rho':
            bounds = [(1.0, 2.0)] * (len(x0))
        if AXIS=='Phi':
            bounds = [(-3.14, 3.14)] * (len(x0))
        if AXIS == 'Z':
            bounds =[(-12, -4)] * (len(x0))
        if AXIS=='CableDelay':
            bounds =[(-1, 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': 500_000})
        
        # save the loss function value to keep track
        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]) # place the first (the fixed position one) antenna back
        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)
        if index%3==0:
            bin_width = 0.005
            bins = np.arange(residuals.min(), residuals.max() + bin_width, bin_width)
            ax1.hist(residuals, bins=bins, alpha=0.8, label=f'{index}:[{AXIS}]', color=color)
 
 
        if AXIS == 'Rho':
            result_rho = np.insert(result.x, 0, x0_randomguess[AXIS][0])
        elif AXIS == 'Phi':
            result_phi = np.insert(result.x, 0, x0_randomguess[AXIS][0])
        elif AXIS == 'Z':
            result_z = np.insert(result.x, 0, x0_randomguess[AXIS][0])
        elif AXIS == 'CableDelay':
            result_cabledelay = np.insert(result.x, 0, x0_randomguess[AXIS][0])
 
    ax1.legend()
    fig1.savefig(f"plot/Minimization_plot.png")
    plt.close(fig1)
    print("Rho:", result_rho)
    print("Phi:", np.rad2deg(result_phi))
    print("Z:", result_z)
    print("CableDelay:", result_cabledelay)