calibration.py

Go to the documentation of this file.

r"""!
@file  calibration.py
@brief `main()` shows how to
       [minimize()](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html)
       the [objective function()](#calibration.objective_function) for antenna phase center
       calibration.
"""
import ROOT
import re
import numpy as np
import polars as pl
from pathlib import Path
from initialise import load_pueoEvent_Dataset
 
 
dirname = "simplesim_half_copy"
UPSAMPLE_FACTOR = 1
def generate_antenna_phase_center_pairs_with_placeholders(qrh_dot_dat: Path, offset_inSIM:Path | None = None, ant_type='MI') -> pl.DataFrame:
    r"""!Creates [antenna pairs](#antenna_pairs.generate_MI_antenna_pairs) with placeholder
         integers as antenna phase centers.
@ingroup CALPC
@param[in] qrh_dot_dat              Path to `qrh.dat`
@retval nominal_phase_center_pairs
 
*   Output schema:
```
$ A1_AntNum            <enum>
$ A1_Ring                <u8>
$ A1_Rho (placeholder)  <u16>
$ A1_Phi (placeholder)  <u16>
$ A1_Z (placeholder)    <u16>
$ A1_Rho                <f64>
$ A1_Phi                <f64>
$ A1_Z                  <f64>
$ A2_AntNum            <enum>
$ A2_Ring                <u8>
$ A2_Rho (placeholder)  <u16>
$ A2_Phi (placeholder)  <u16>
$ A2_Z (placeholder)    <u16>
$ A2_Rho                <f64>
$ A2_Phi                <f64>
$ A2_Z                  <f64>
```
*   Details:
    -#  Calls #antenna_attributes.read_MI_antenna_geometry() to read `qrh.dat` to obtain
        antenna face center positions
    -#  Calls #antenna_attributes.get_MI_nominal_phase_center() to convert face centers to
        nominal phase centers
    -#  Adds three columns to store placeholder values of the phase centers; these integer placeholders
        serve as dictionary keys.
```
shape: (96, 9)
┌────────┬────────┬──────┬────────────────┬────────────────┬───────────────┬──────────┬───────────┬───────┐
│ AntNum ┆ AntIdx ┆ Ring ┆ Rho            ┆ Phi            ┆ Z             ┆ Rho      ┆ Phi       ┆ Z     │
│ ---    ┆ ---    ┆ ---  ┆ (placeholder)  ┆ (placeholder)  ┆ (placeholder) ┆ ---      ┆ ---       ┆ ---   │
│ enum   ┆ u32    ┆ u8   ┆ ---            ┆ ---            ┆ ---           ┆ f64      ┆ f64       ┆ f64   │
│        ┆        ┆      ┆ u16            ┆ u16            ┆ u16           ┆          ┆           ┆       │
╞════════╪════════╪══════╪════════════════╪════════════════╪═══════════════╪══════════╪═══════════╪═══════╡
│ 101    ┆ 0      ┆ 1    ┆ 0              ┆ 96             ┆ 192           ┆ 1.188    ┆ 0.0       ┆ 5.114 │
│ 201    ┆ 1      ┆ 2    ┆ 1              ┆ 97             ┆ 193           ┆ 2.296    ┆ 0.0       ┆ 1.457 │
│ 301    ┆ 2      ┆ 3    ┆ 2              ┆ 98             ┆ 194           ┆ 2.296    ┆ 0.0       ┆ 0.728 │
│ 401    ┆ 3      ┆ 4    ┆ 3              ┆ 99             ┆ 195           ┆ 2.296    ┆ 0.0       ┆ 0.0   │
│ 102    ┆ 4      ┆ 1    ┆ 4              ┆ 100            ┆ 196           ┆ 1.076821 ┆ 0.261538  ┆ 4.476 │
│ …      ┆ …      ┆ …    ┆ …              ┆ …              ┆ …             ┆ …        ┆ …         ┆ …     │
│ 423    ┆ 91     ┆ 4    ┆ 91             ┆ 187            ┆ 283           ┆ 2.295825 ┆ -0.523638 ┆ 0.0   │
│ 124    ┆ 92     ┆ 1    ┆ 92             ┆ 188            ┆ 284           ┆ 1.07708  ┆ -0.26224  ┆ 4.476 │
│ 224    ┆ 93     ┆ 2    ┆ 93             ┆ 189            ┆ 285           ┆ 2.295502 ┆ -0.261893 ┆ 1.457 │
│ 324    ┆ 94     ┆ 3    ┆ 94             ┆ 190            ┆ 286           ┆ 2.295502 ┆ -0.261893 ┆ 0.728 │
│ 424    ┆ 95     ┆ 4    ┆ 95             ┆ 191            ┆ 287           ┆ 2.295502 ┆ -0.261893 ┆ 0.0   │
└────────┴────────┴──────┴────────────────┴────────────────┴───────────────┴──────────┴───────────┴───────┘
```
    -#  Calls #antenna_pairs.generate_MI_antenna_pairs with the above dataframe as the input.
    """
    from antenna_pairs import generate_MI_antenna_pairs, generate_LF_antenna_pairs
    from antenna_attributes import read_antenna_geometry, get_nominal_phase_center
 
    fc: pl.DataFrame = read_antenna_geometry(qrh_dot_dat=qrh_dot_dat, offset_inSIM=offset_inSIM, ant_type=ant_type)
 
    nominal_and_placeholders: pl.DataFrame = (
        get_nominal_phase_center(face_centers=fc, coordinates="cylindrical", ant_type=ant_type)
        .with_columns(
            pl.arange(pl.len()).alias("Rho (placeholder)").cast(pl.UInt16),
            pl.arange(pl.len(), pl.len() * 2).alias("Phi (placeholder)").cast(pl.UInt16),
            pl.arange(pl.len() * 2, pl.len() * 3).alias("Z (placeholder)").cast(pl.UInt16),
            pl.arange(pl.len() * 3, pl.len() * 4).alias("CableDelay (placeholder)").cast(pl.UInt16),
            pl.lit(0.0).alias("CableDelay"),   # fills whole column with 0.0
        )
        .select(
            "AntNum", "AntIdx", "Ring", "Rho (placeholder)", "Phi (placeholder)", "Z (placeholder)", "CableDelay (placeholder)",
            pl.col("Rho[m]").alias("Rho"), pl.col("Phi[rad]").alias("Phi"), pl.col("Z[m]").alias("Z"), "CableDelay"
        )
    )
    
    if ant_type=='MI':
        nominal_phase_center_pairs: pl.DataFrame = (
            generate_MI_antenna_pairs(antennas=nominal_and_placeholders)
            .select(pl.all().exclude(r"^A[12]_AntIdx$"))
        )
 
    elif ant_type=='LF':
        #print('nominal_and_placeholders=',nominal_and_placeholders)
        nominal_phase_center_pairs: pl.DataFrame = (
            generate_LF_antenna_pairs(antennas=nominal_and_placeholders)
            .select(pl.all().exclude(r"^A[12]_AntIdx$"))
        )
    else:
        print("Ant type is not supported!")
        
    return nominal_phase_center_pairs
 
 
def __get_all_run_numbers(pueo_mc_data: Path) -> [int]:
    """! A temporary helper function of #load_dataset_and_measure_time_delays()
    @ingroup CALPC
    @param[in] pueo_mc_data  See the parameter of #initialise.load_pueoEvent_Dataset
 
    *   This function retrieves the run numbers (negative sign allowed) of the `run/` subdirectories
        in `pueo_mc_data`
 
    *   The runs are returned as a list of integers.
 
    *   The function is here only for testing purposes,
        as it makes little sense to use every single run for antenna calibration.
    """
 
    # regular expression, allowing positive and negative signs in run numbers
    run_number_pattern_ = re.compile(r'^run([+\-]?\d+)$')
    return [int(run_number_pattern_.match(run_directories_.name).group(1))
            for run_directories_ in pueo_mc_data.rglob("run*")]
 
 
def load_dataset_and_measure_time_delays(dataset: ROOT.pueo.Dataset,
        phase_center_pairs: pl.DataFrame, pulser_station: str='taylor_dome', MC=True, AntType='MI', IcefinalFileDir:str='', corr_fcn:str='') -> pl.DataFrame:
    r"""! Uses [SciPy's correlate()](https://docs.scipy.org/doc//scipy-1.16.2/reference/generated/scipy.signal.correlate.html)
    to "measure" the delays of [waveforms](#waveform_plots.load_waveforms) between antenna pairs
    @ingroup CALPC
    @param[in] dataset              The output of #initialise.load_pueoEvent_Dataset
    @param[in] phase_center_pairs   The output of #generate_antenna_phase_center_pairs_with_placeholders
    @retval    pulser_directions_and_measured_time_delays
 
    *   The input `dataset` should have been initialized with a (valid) run number,
        but that is not the only run that gets loaded.
        For now, [all runs](#__get_all_run_numbers) are used.
    *   Output schema:
```
$ measured time delays (ns)                      <f64>
$ max correlation                                <f64>
$ A1_AntNum                                     <enum>
$ A1_Ring                                         <u8>
$ A1_Rho (placeholder)                           <u16>
$ A1_Phi (placeholder)                           <u16>
$ A1_Z (placeholder)                             <u16>
$ A1_Rho                                         <f64>
$ A1_Phi                                         <f64>
$ A1_Z                                           <f64>
$ A2_AntNum                                     <enum>
$ A2_Ring                                         <u8>
$ A2_Rho (placeholder)                           <u16>
$ A2_Phi (placeholder)                           <u16>
$ A2_Z (placeholder)                             <u16>
$ A2_Rho                                         <f64>
$ A2_Phi                                         <f64>
$ A2_Z                                           <f64>
$ incoming pulse direction (Phi Theta) <array[f64, 2]>
```
    *   The `dataset` is also used to load the directions of the incoming pulses
        (\f$^\circ\phi\f$ and \f$^\circ\theta\f$ in payload coordinates).
    *   `measured time delays (ns)` is computed by the helper function
        #_scipy_cross_correlate_all_pairs_for_one_event().
    """
 
    # Note: eventually will have to get pulser direction somehow, instead of from the truth file
    from locate_signal import _get_true_direction, _get_true_pulser_direction
    from waveform_plots import load_waveforms, upsample_waveforms, load_waveforms_from_root 
 
    # initialize an empty dataframe that will be in-place updated
    pulser_directions_and_measured_time_delays = pl.DataFrame()
 
    # Note: probably doesn't make sense to use all runs, for now it's okay
    all_runs: [int] = __get_all_run_numbers(
        pueo_mc_data=Path(
            dataset.getDataDir(ROOT.pueo.Dataset.PUEO_MC_DATA)
        )
    )
    
 
    for run in all_runs:
        dataset.loadRun(run, ROOT.pueo.Dataset.PUEO_MC_DATA)
        for i in range(dataset.N()):  # loop through all events in the run
            dataset.getEntry(i)
            if AntType=='MI':
                wf: pl.DataFrame = (
                  load_waveforms(dataset=dataset)
                  .filter(pl.col("Pol") == "V")
                  .select("AntNum", "Pol", "waveforms (volts)", "step size (nanoseconds)")
                )
            elif AntType=='LF':
                IcefinalFile=IcefinalFileDir+f'run{run}/IceFinal_{run}_allTree.root'
                wf: pl.DataFrame = (
                  load_waveforms_from_root(IcefinalFile, ev=i)
                  .filter(pl.col("Pol") == "V")
                  .select("AntNum", "Pol", "waveforms (volts)", "step size (nanoseconds)")
                )
            up: pl.DataFrame = upsample_waveforms(waveforms=wf, upsample_factor=UPSAMPLE_FACTOR)
            
                
            if MC==True:
                pulser_directions_and_measured_time_delays.vstack(
                    _scipy_cross_correlate_all_pairs_for_one_event(
                        phase_center_pairs=phase_center_pairs,
                        waveforms=up, 
                        fcn=corr_fcn,
                    ).with_columns(
                        (pl.lit(_get_true_direction(dataset=dataset)) * np.pi / 180)
                        .cast(pl.Array(pl.Float64, 2))
                        .alias("incoming pulse direction (Phi Theta)")
                    ),
                    in_place=True
                )
                #print(pulser_directions_and_measured_time_delays)
            
            else:
                # for real data position calibration
 
                pulser_directions_and_measured_time_delays.vstack(
                    _scipy_cross_correlate_all_pairs_for_one_event(
                        phase_center_pairs=phase_center_pairs,
                        waveforms=up,
                        fcn=corr_fcn,
                    ).with_columns(
                        (pl.lit(_get_true_pulser_direction(dataset=dataset, pulser_station=pulser_station)) * np.pi / 180)
                        .cast(pl.Array(pl.Float64, 2))
                        .alias("incoming pulse direction (Phi Theta)")
                    ),
                    in_place=True
                )
        #print("pulser_directions_and_measured_time_delays=", pulser_directions_and_measured_time_delays)
        
    return pulser_directions_and_measured_time_delays
 
 
def _scipy_cross_correlate_all_pairs_for_one_event(phase_center_pairs: pl.DataFrame,
                                                   waveforms: pl.DataFrame,
                                                   fcn: str='') -> pl.DataFrame:
    r"""! A helper function of #load_dataset_and_measure_time_delays()
    @ingroup CALPC
    @param[in] phase_center_pairs   The output of #generate_antenna_phase_center_pairs_with_placeholders
    @param[in] waveforms            The output of #waveform_plots.load_waveforms
    @retval    measured_time_delays
 
    *   Output schema:
```
$ measured time delays (ns)  <f64>
$ max correlation            <f64>
$ A1_AntNum                 <enum>
$ A1_Ring                     <u8>
$ A1_Rho (placeholder)       <u16>
$ A1_Phi (placeholder)       <u16>
$ A1_Z (placeholder)         <u16>
$ A1_Rho                     <f64>
$ A1_Phi                     <f64>
$ A1_Z                       <f64>
$ A2_AntNum                 <enum>
$ A2_Ring                     <u8>
$ A2_Rho (placeholder)       <u16>
$ A2_Phi (placeholder)       <u16>
$ A2_Z (placeholder)         <u16>
$ A2_Rho                     <f64>
$ A2_Phi                     <f64>
$ A2_Z                       <f64>
```
    *   Details:
        -#  [Load the waveforms](#waveform_plots.load_waveforms) of both antennas in the antenna pair
        -#  Cross correlate the waveform pair
        -#  Figure out by how much one needs to shift the second waveform against the first
            in order to achieve the maximum correlation.
        -#  The above is called `measured time delays (ns)` in the output dataframe.
        -#  For more details, see [Explanation](#scipy_corr_expl).
    """
    from scipy.signal import correlate, correlation_lags
 
    # the fact that the waveforms are in the same column guarantees they have the same length,
    # which is probably an integer multiple of 1024 (if upsampled)
    signal_length = len(waveforms["waveforms (volts)"][0])
    CORRLAGS: np.ndarray = correlation_lags(signal_length, signal_length)
 
    STEP_SIZE = waveforms["step size (nanoseconds)"][0]
    # For peace of mind, check that all signals were sampled at the same step size :)
    assert np.mean(waveforms["step size (nanoseconds)"].to_numpy()) - STEP_SIZE < 1e-10
    
    if fcn == 'ZNCC':
        corr_func = zncc_pairwise_correlation
    else:
        corr_func = pairwise_correlation
 
    measured_time_delays = (
        phase_center_pairs
        .join(waveforms, left_on="A1_AntNum", right_on="AntNum")
        .join(waveforms, left_on=["A2_AntNum", "Pol"], right_on=["AntNum", "Pol"])
        .with_columns(
            pl.struct("waveforms (volts)", "waveforms (volts)_right")
            .map_batches(corr_func).alias("correlation"),
        )
        .with_columns(
            pl.col("correlation").arr.max().alias("max correlation"),
            pl.col("correlation").arr.arg_max().alias("max correlation idx"),
        )
        .with_columns(  # Find the time where the maximum correlation occurs and convert to [ns]
            pl.col("max correlation idx").map_batches(lambda idx: CORRLAGS[idx] * STEP_SIZE)
            .alias("measured time delays (ns)")
        )
        .select(
            "measured time delays (ns)", "max correlation",
            r"^A[12]_(?:AntNum|Ring|Rho|Phi|Z|CableDelay)(?: \‍(placeholder\‍))?$",
        )
 
    )
    return measured_time_delays
 
def zncc_pairwise_correlation(s):  # zero-center usefrontlobeized cross-correlation (ZNCC), bounded b/w [-1, 1]
    from scipy.signal import correlate
    return np.array([
        correlate(
            (wf1 - np.mean(wf1)) / np.std(wf1),
            (wf2 - np.mean(wf2)) / np.std(wf2),
            method='fft'
        ) / len(wf1)
        for wf2, wf1 in s.to_numpy()
    ])
 
def pairwise_correlation(s):  # cross-correlation 
    from scipy.signal import correlate
    return np.array([
        correlate(
            wf1,
            wf2,
            method='fft'
        ) / len(wf1)
        for wf2, wf1 in s.to_numpy()
    ])
def retrieve_relevant_antennas_from_all_pairs(phase_center_pairs: pl.DataFrame):
    r"""! Retrieves the (unique) antennas, given a bunch of antenna pairs.
    @ingroup CALPC
    @param[in] phase_center_pairs   The output of #load_dataset_and_measure_time_delays
    @retval    relevant_antennas
 
    *   This step is not trivial, as explained below.
    *   It is conceivable that not all antennas receive signals clear enough for calibration purposes
        (resulting in weak correlations).
    *   As such, the weakly correlated antenna pairs should be filtered out.
        For example,
```python3
    pulser_directions_and_pairwise_time_delays = (
        load_dataset_and_measure_time_delays(dataset=ds, phase_center_pairs=nominal_pairs)
        .filter(pl.col("max correlation") > 0.8)
    )
```
    *   It is possible that some antennas end up getting filtered out.
    *   Sample output (note that only 76 out of 96 MI antennas survived the filter in this case):
```
shape: (76, 8)
┌────────┬──────┬───────────────────┬───────────────────┬─────────────────┬──────────┬───────────┬───────┐
│ AntNum ┆ Ring ┆ Rho (placeholder) ┆ Phi (placeholder) ┆ Z (placeholder) ┆ Rho      ┆ Phi       ┆ Z     │
│ ---    ┆ ---  ┆ ---               ┆ ---               ┆ ---             ┆ ---      ┆ ---       ┆ ---   │
│ enum   ┆ u8   ┆ u16               ┆ u16               ┆ u16             ┆ f64      ┆ f64       ┆ f64   │
╞════════╪══════╪═══════════════════╪═══════════════════╪═════════════════╪══════════╪═══════════╪═══════╡
│ 101    ┆ 1    ┆ 0                 ┆ 96                ┆ 192             ┆ 1.188    ┆ 0.0       ┆ 5.114 │
│ 201    ┆ 2    ┆ 1                 ┆ 97                ┆ 193             ┆ 2.296    ┆ 0.0       ┆ 1.457 │
│ 301    ┆ 3    ┆ 2                 ┆ 98                ┆ 194             ┆ 2.296    ┆ 0.0       ┆ 0.728 │
│ 401    ┆ 4    ┆ 3                 ┆ 99                ┆ 195             ┆ 2.296    ┆ 0.0       ┆ 0.0   │
│ 102    ┆ 1    ┆ 4                 ┆ 100               ┆ 196             ┆ 1.076821 ┆ 0.261538  ┆ 4.476 │
│ …      ┆ …    ┆ …                 ┆ …                 ┆ …               ┆ …        ┆ …         ┆ …     │
│ 423    ┆ 4    ┆ 91                ┆ 187               ┆ 283             ┆ 2.295825 ┆ -0.523638 ┆ 0.0   │
│ 124    ┆ 1    ┆ 92                ┆ 188               ┆ 284             ┆ 1.07708  ┆ -0.26224  ┆ 4.476 │
│ 224    ┆ 2    ┆ 93                ┆ 189               ┆ 285             ┆ 2.295502 ┆ -0.261893 ┆ 1.457 │
│ 324    ┆ 3    ┆ 94                ┆ 190               ┆ 286             ┆ 2.295502 ┆ -0.261893 ┆ 0.728 │
│ 424    ┆ 4    ┆ 95                ┆ 191               ┆ 287             ┆ 2.295502 ┆ -0.261893 ┆ 0.0   │
└────────┴──────┴───────────────────┴───────────────────┴─────────────────┴──────────┴───────────┴───────┘
```
    *   If we still pass these filtered antennas to the #objective_function,
        SciPy's minimizer would act as if these antennas are relevant and try to still modify their
        phase centers during the minization.
    """
    a1 = (
        phase_center_pairs
        .select(pl.col(r"^A1_(?:AntNum|Ring|Rho|Phi|Z|CableDelay)(?: \‍(placeholder\‍))?$"))
        .rename(lambda name: re.sub(r"^A1_(.*)$", r"\1", name))
    )
    a2 = (
        phase_center_pairs
        .select(pl.col(r"^A2_(?:AntNum|Ring|Rho|Phi|Z|CableDelay)(?: \‍(placeholder\‍))?$"))
        .rename(lambda name: re.sub(r"^A2_(.*)$", r"\1", name))
    )
 
    relevant_antennas = (
        a1.vstack(a2).unique().sort("AntNum")
    )
    return relevant_antennas
 
 
def objective_function(phase_center_guess: pl.Series,
                       phase_center_placeholders: pl.Series,
                       pulser_directions_and_measured_time_delays: pl.DataFrame,
                       axis: str) -> int:
    r"""!Computes the expected time delays and compares them with the measured time delays,
    designed to be run by SciPy's minimizer.
    @ingroup CALPC
    @param[in] phase_center_guess               The output of #retrieve_relevant_antennas_from_all_pairs
    @param[in] phase_center_placeholders        The output of #retrieve_relevant_antennas_from_all_pairs
    @param[in] pulser_directions_and_measured_time_delays  The output of #load_dataset_and_measure_time_delays
    @param[in] axis                             Axis to fit: \f$\rho\f$, \f$\phi\f$, or \f$z\f$
    @retval    total_difference_bw_expected_time_delay_and_measured_time_delay  A scalar
 
    * In #load_dataset_and_measure_time_delays(), the `measured_time_delay` has already been computed.
    * The job of this function is to
        -#  Given `phase_center_guess`, replace the placeholder antenna phase centers in
            `pulser_directions_and_pairwise_time_delays`
        -#  Compute the pairwise antenna displacement vectors (call these \f$\bf{d}\f$),
            using the updated antenna positions.
        -#  Compute the `expected time delay` based on the dot product between \f$\bf{d}\f$ and
            `incoming pulse direction (Phi Theta)`,
            see [this figure](@ref timedelaybydotproduct) for an illustration.
 ```
 ┌───────────────────────────┬───────────────────────────┐
 │ measured time delays (ns) ┆ expected time delays (ns) │
 │ ---                       ┆ ---                       │
 │ f64                       ┆ f64                       │
 ╞═══════════════════════════╪═══════════════════════════╡
 │ -0.866667                 ┆ -0.853384                 │
 │ -0.866667                 ┆ -0.853384                 │
 │ -7.466667                 ┆ -7.487111                 │
 │ 0.6                       ┆ 0.611216                  │
 │ -9.466667                 ┆ -9.48138                  │
 │ …                         ┆ …                         │
 │ 8.333333                  ┆ 8.33698                   │
 │ 0.2                       ┆ 0.15054                   │
 │ -0.866667                 ┆ -0.855732                 │
 │ 0.8                       ┆ 0.760589                  │
 │ 1.466667                  ┆ 1.463433                  │
 └───────────────────────────┴───────────────────────────┘
 ```
        -#  Quantify the difference between the expectation and the measurement; this is done by
            summing the difference over all rows above.
    * SciPy's minimizer tries to minimize the total difference by updating `phase_center_guess`.
    """
    from scipy.constants import c as speed_of_light
 
    phi_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(0)
    theta_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(1)
    x_p = phi_p.cos() * theta_p.sin()
    y_p = phi_p.sin() * theta_p.sin()
    z_p = theta_p.cos()
 
    r_1 = pl.col("A1_Rho")
    r_2 = pl.col("A2_Rho")
    z_1 = pl.col("A1_Z")
    z_2 = pl.col("A2_Z")
    phi_1 = pl.col("A1_Phi")
    phi_2 = pl.col("A2_Phi")
    cabledelay_1 = pl.col("A1_CableDelay")
    cabledelay_2 = pl.col("A2_CableDelay")
 
    A1minusA2_dotProduct_signalDir = (
        x_p * (r_1 * phi_1.cos() - r_2 * phi_2.cos()) +
        y_p * (r_1 * phi_1.sin() - r_2 * phi_2.sin()) + z_p * (z_1 - z_2)
    )
    
    
    total_difference_bw_expected_time_delay_and_measured_time_delay = (
        pulser_directions_and_measured_time_delays
        .with_columns(
            pl.col(f"A1_{axis} (placeholder)").replace_strict(phase_center_placeholders, phase_center_guess, default=pl.col(f"A1_{axis}")).alias(f"A1_{axis}"),
            pl.col(f"A2_{axis} (placeholder)").replace_strict(phase_center_placeholders, phase_center_guess, default=pl.col(f"A2_{axis}")).alias(f"A2_{axis}")
        )
        .with_columns(
            (A1minusA2_dotProduct_signalDir / speed_of_light * 1e9 + cabledelay_1-cabledelay_2).alias("expected time delays (ns)")
        )
        .select(
            ((pl.col("measured time delays (ns)") - pl.col("expected time delays (ns)"))**2)
            .sum().alias("time delay difference")
        )
    )
    return total_difference_bw_expected_time_delay_and_measured_time_delay.item()
 
def objective_function_MultipleAxes(phase_center_guess: np.ndarray,
                       phase_center_placeholders_list: list[pl.Series],
                       pulser_directions_and_measured_time_delays: pl.DataFrame,
                       axis_list:list[str]) -> int:
    
    from scipy.constants import c as speed_of_light
    N_ant_involved = len(phase_center_placeholders_list[0])
 
    phi_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(0)
    theta_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(1)
    x_p = phi_p.cos() * theta_p.sin()
    y_p = phi_p.sin() * theta_p.sin()
    z_p = theta_p.cos()
 
    r_1 = pl.col("A1_Rho")
    r_2 = pl.col("A2_Rho")
    z_1 = pl.col("A1_Z")
    z_2 = pl.col("A2_Z")
    phi_1 = pl.col("A1_Phi")
    phi_2 = pl.col("A2_Phi")
 
    A1minusA2_dotProduct_signalDir = (
        x_p * (r_1 * phi_1.cos() - r_2 * phi_2.cos()) +
        y_p * (r_1 * phi_1.sin() - r_2 * phi_2.sin()) + z_p * (z_1 - z_2)
    )
    
    start_idx = 0
 
    # Start with the original DataFrame
    df = pulser_directions_and_measured_time_delays
 
    for i, AXIS in enumerate(axis_list):
        placeholders = phase_center_placeholders_list[i]
        end_idx = start_idx + len(placeholders)
        
        # Replace both A1 and A2 columns for this axis
        df = df.with_columns(
            pl.col(f"A1_{AXIS} (placeholder)")
              .replace_strict(placeholders, phase_center_guess[start_idx:end_idx])
              .alias(f"A1_{AXIS}"),
            
            pl.col(f"A2_{AXIS} (placeholder)")
              .replace_strict(placeholders, phase_center_guess[start_idx:end_idx])
              .alias(f"A2_{AXIS}")
        )
        
        start_idx = end_idx  # move to next slice of phase_center_guess
 
    # Compute expected time delays
    df = df.with_columns(
        (A1minusA2_dotProduct_signalDir / speed_of_light * 1e9).alias("expected time delays (ns)")
    )
 
    # Compute total difference
    total_difference_bw_expected_time_delay_and_measured_time_delay = df.select(
        ((pl.col("measured time delays (ns)") - pl.col("expected time delays (ns)"))**2)
        .sum()
        .alias("time delay difference")
    )
 
    return total_difference_bw_expected_time_delay_and_measured_time_delay.item()
 
def compute_time_delay_residuals(
    phase_center_guess: np.ndarray,
    phase_center_placeholders: pl.Series,
    pulser_directions_and_measured_time_delays: pl.DataFrame,
    axis: str,
):
    from scipy.constants import c as speed_of_light
    phi_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(0)
    theta_p = pl.col("incoming pulse direction (Phi Theta)").arr.get(1)
 
    x_p = phi_p.cos() * theta_p.sin()
    y_p = phi_p.sin() * theta_p.sin()
    z_p = theta_p.cos()
 
    r_1 = pl.col("A1_Rho")
    r_2 = pl.col("A2_Rho")
    z_1 = pl.col("A1_Z")
    z_2 = pl.col("A2_Z")
    phi_1 = pl.col("A1_Phi")
    phi_2 = pl.col("A2_Phi")
    cabledelay_1 = pl.col("A1_CableDelay")
    cabledelay_2 = pl.col("A2_CableDelay")
 
    A1minusA2_dotProduct_signalDir = (
        x_p * (r_1 * phi_1.cos() - r_2 * phi_2.cos()) +
        y_p * (r_1 * phi_1.sin() - r_2 * phi_2.sin()) +
        z_p * (z_1 - z_2)
    )
 
    df_resid = (
        pulser_directions_and_measured_time_delays
        .with_columns(
            pl.col(f"A1_{axis} (placeholder)")
              .replace_strict(phase_center_placeholders, phase_center_guess, default=pl.col(f"A1_{axis}"))
              .alias(f"A1_{axis}"),
 
            pl.col(f"A2_{axis} (placeholder)")
              .replace_strict(phase_center_placeholders, phase_center_guess, default=pl.col(f"A2_{axis}"))
              .alias(f"A2_{axis}")
        )
        .with_columns(
            (A1minusA2_dotProduct_signalDir / speed_of_light * 1e9 + cabledelay_1 - cabledelay_2)
            .alias("expected time delays (ns)")
        )
        .with_columns(
            (pl.col("measured time delays (ns)") - pl.col("expected time delays (ns)"))
            .alias("residual (ns)")
        )
    )
 
    return df_resid
 
def _plot_calibration_result(axis: str, before: pl.DataFrame, after: np.ndarray, anttype='MI'):
    """!
    @param[in] axis     Axis to fit: \f$\rho\f$, \f$\phi\f$, or \f$z\f$
    @param[in] before   The output of of #retrieve_relevant_antennas_from_all_pairs
    @param[in] after    The `restul.x` of running `scipy.optimize.minimize` on #objective_function
    @ingroup CALPC
    Nothing crazy here!
 
    ![](example_phase_center_fit_result.png)
    """
    import matplotlib.pyplot as plt
    import matplotlib
    matplotlib.use("Agg")
 
    plt.rcParams.update({'font.size': 13})
    fig1, ax1 = plt.subplots(nrows=4, figsize=(10, 15), tight_layout=True)
 
    before_and_after = (
        before.with_columns(
            pl.Series(after).alias(f"{axis} minimization result")
        )
    )
    
    before_and_after = before_and_after.with_columns(
      pl.col("AntNum").cast(pl.Int32)
      )
    for i in range(before['Ring'].unique().len()):
        if anttype=='MI':
            ring = before_and_after.filter(pl.col("Ring") == (i + 1))
        elif anttype=='LF':
            ring = before_and_after.filter(pl.col("Ring") == (i + 5))
 
        #ax1[i].scatter(ring['AntNum'], ring[f"{axis}"] * 1e2,
        #               label=f"Simulated geometry")
        #ax1[i].scatter(ring['AntNum'], ring[f"{axis} minimization result"] * 1e2,
        #               label=f"minimization result")
        if axis=='Rho' or axis=='Z':
            ax1[i].scatter(ring['AntNum'], ring[f"{axis} minimization result"] * 1e2 - ring[f"{axis}"] * 1e2,
                       label=f"Minimized - Simulated geometry", s=60)
        elif axis=='Phi':
            diff = (ring[f"{axis} minimization result"] - ring[f"{axis}"]) * 180/np.pi
            diff = (diff + 180) % 360 - 180
            ax1[i].scatter(ring['AntNum'], diff,
                       label=f"Minimized - Simulated geometry", s=60)
        elif axis=='CableDelay':
            ax1[i].scatter(ring['AntNum'], ring[f"{axis} minimization result"] - ring[f"{axis}"],
                       label=f"Minimized - Simulated Cable Delay", s=60)
 
        x_ticks_to_show = [val for val in sorted(set(ring['AntNum']))]
        ax1[i].set_xticks(x_ticks_to_show)
        ax1[i].set_xticklabels([str(int(val)) for val in x_ticks_to_show])
 
        if anttype=='LF':
            ax1[i].set_title(f"Ring {i + 5}", fontsize=20)
            #ax1[i].set_ylim(150,165)
        else:
            ax1[i].set_title(f"Ring {i + 1}", fontsize=20)
        ax1[i].tick_params(axis='x', labelrotation=45, labelright=True, labelsize=20)
        ax1[i].tick_params(axis='y', labelright=True, labelsize=16)
        ax1[i].grid(alpha=.3)
    
    if anttype=='LF':
        ax1[0].set_xlim(500,525)
        ax1[1].set_xlim(600,625)
        ax1[2].set_xlim(700, 725)
        ax1[3].set_xlim(800,825)
      
    ax1[0].legend(fontsize=20)
    
    if axis=='Rho' or axis=='Z':
        fig1.supylabel(f"$\Delta$ {axis} [cm]", fontsize=20)
    elif axis=='Phi':
        fig1.supylabel(f"$\Delta$ {axis} [deg]", fontsize=20)
    elif axis=='CableDelay':
        fig1.supylabel(f"$\Delta$ {axis} [ns]", fontsize=20)
 
 
    fig1.supxlabel('Antenna Label', fontsize=20)
    fig1.suptitle('Antenna Phase Centers (Payload Coordinates)\n', fontsize=24)
    fig1.savefig(f"plot/{dirname}/{anttype}_phase_center_fit_result_{axis}_tmp.png")
 
 
if __name__ == "__main__":
    import os
    from scipy.optimize import minimize
    import matplotlib.pyplot as plt
    import matplotlib
    matplotlib.use("Agg") 
    ant_type='LF'
    MC=True
    dirnametest="test"
    # 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'
        )
    
    relevant_antennas = (
        retrieve_relevant_antennas_from_all_pairs(
            phase_center_pairs=pulser_directions_and_pairwise_time_delays
        )
    )
 
    AXIS = "Rho"  # "Rho", "Phi", or "Z"
    result = minimize(fun=objective_function,
                      x0=relevant_antennas[AXIS],
                      args=(relevant_antennas[f"{AXIS} (placeholder)"],
                            pulser_directions_and_pairwise_time_delays, AXIS),
                      method='Nelder-Mead',
                      options={'maxiter': 500_000})
 
    print(result)
    _plot_calibration_result(axis=AXIS, before=relevant_antennas, after=result.x)